#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/user.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <fcntl.h>

#define r_eax 0
#define r_ecx 1
#define r_edx 2
#define r_ebx 3
#define r_esp 4
#define r_ebp 5
#define r_esi 6       
#define r_edi 7
#define r_eip 8
#define r_eflags 9
#define dr0 252
#define dr1 256
#define dr2 260
#define dr3 264
#define dr4 268
#define dr5 272
#define dr6 276
#define dr7 280
 
pid_t traced_process;
int fd_msr_r, fd_msr_w;

unsigned int getRegister(unsigned int reg);

void init(){
        fd_msr_r = open("/dev/cpu/0/msr", O_RDONLY);
        fd_msr_w = open("/dev/cpu/0/msr", O_WRONLY);
        if(fd_msr_r <0 || fd_msr_w < 0){
                perror("open");
        }
}

int enableSingleStep(){
        int status;
        int ret = 0;


        if (ptrace(PTRACE_SINGLESTEP, traced_process, 0, 0) != 0){
                perror("ptrace single step");
                ret = -1;
        }

        wait(&status);
        if(!WIFSTOPPED(status)){
                wait(&status);
        }
        return ret;
}

enableBranchStep(){
        if(wrmsr(0x1d9, "00:00:00:00:00:00:00:02") < 0){
                fprintf(stderr, "Couldn't set msr\n");
                exit(-1);
        }
}

void cleanup(){
        close(fd_msr_r);
        close(fd_msr_w);
}

void rdmsr(unsigned long addy){
    unsigned char buf[8];
    int i;

    lseek(fd_msr_r, addy, SEEK_SET);
    read(fd_msr_r, buf, 8);

    printf("MSR register 0x%lx => ", addy);
    for (i = 7; i > 0; i--)
        printf("%2.2x:", buf[i]);
    printf("%2.2x\n", buf[i]);
}

int wrmsr(unsigned long addy, char *p){
    unsigned char buf[8];
    int i;

    for (i = 7; i > 0; i--) {
        buf[i] = strtol(p, &p, 16);
        p++;
    }
    buf[i] = strtol(p, &p, 16);
    if(lseek(fd_msr_w, addy, SEEK_SET)<0){
        perror("lseek");
    }
    if (write(fd_msr_w, buf, 8) < 0)
    {
        perror("write");
        return -1;
    }
    return 0;
}

unsigned int stepToNextBranch(){

        enableBranchStep();
        if(enableSingleStep() < 0){
                return -1;
        }

        return getRegister(r_eip);
}

void attach(int process){
        traced_process = process;
	init();
        int status;
        if(ptrace(PTRACE_ATTACH, traced_process, NULL, NULL)<0){
                fprintf(stderr, "Couldn't attach\n");
		perror("attach");
                exit(-1);
        }
        wait(&status);
        if(!WIFSTOPPED(status)){
                printf("Wasn't stopped by trace\n");
        }
}

void detach(){
        if(ptrace(PTRACE_DETACH, traced_process, NULL, NULL)<0){
                fprintf(stderr, "Couldn't detach\n");
                exit(-1);
        }
	cleanup();
}

continueExecution(){
        int status;
        ptrace(PTRACE_CONT, traced_process, NULL, NULL);
        wait(&status);
        if(!WIFSTOPPED(status)){
                printf("Wasn't stopped by trace\n");
        }
}

unsigned int getRegister(unsigned int reg){
        struct user_regs_struct regs;
        unsigned int val;
        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
        switch(reg){
                case r_eax:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.eax;
                        break;
                case r_ecx:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.ecx;
                        break;
                case r_edx:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.edx;
                        break;
                case r_ebx:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.ebx;
                        break;
                case r_esp:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.esp;
                case r_ebp:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.ebp;
                        break;
                case r_esi:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.esi;
                        break;
                case r_edi:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.edi;
                        break;
                case r_eip:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.eip;
                        break;
                case r_eflags:
                        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
                        val = regs.eflags;
                        break;
                case dr0:
                case dr1:
                case dr2:
                case dr3:
                case dr4:
                case dr5:
                case dr6:
                case dr7:
                        val = ptrace(PTRACE_PEEKUSER, traced_process, reg, NULL);
                        break;
                default:
                        fprintf(stderr, "Unknown register\n");
                        exit(-1);
        }
        return val;
}

void setRegValue(unsigned int reg, unsigned int value){
        struct user_regs_struct regs;
        unsigned int val;
        ptrace(PTRACE_GETREGS, traced_process, NULL, &regs);
        switch(reg){
                case r_eax:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.eax, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_ecx:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.ecx, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_edx:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.edx, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_ebx:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.ebx, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_esp:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.esp, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_ebp:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.ebp, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_esi:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.esi, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_edi:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.edi, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_eip:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.eip, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case r_eflags:
                        val = ptrace(PTRACE_POKETEXT, traced_process, regs.eflags, value);
                        ptrace(PTRACE_SETREGS, traced_process, NULL, &regs);
                        break;
                case dr0:
                case dr1:
                case dr2:
                case dr3:
                case dr4:
                case dr5:
                case dr6:
                case dr7:
                        ptrace(PTRACE_POKEUSER, traced_process, reg, value);
                        break;
                default:
                        fprintf(stderr, "Unknown register\n");
                        exit(-1);
        }
}


int main(int argc, int argv[]){
	int i;
	attach(atoi(argv[1]));
	for(i=0; i<10; i++){
		printf("0x%08x\n", stepToNextBranch());
	}
	detach();
}