#include <sys/reg.h>
#include <sys/ptrace.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <errno.h>
#define Mapfile "/proc/%d/maps"
//用于获取plt内容
union pltval{
    size_t val;
    unsigned char chars[sizeof(size_t)];
};
void usage(char** argv){
    printf("USAGE: \n --- %s　pid plt_offset \n", argv[0]);
}
void printerror(){
    printf("Status: %s\n", strerror(errno));
}
void mod_handle(pid_t tracee, void* addr1, void* addr2)
{
    union pltval buf;
    buf.val = ptrace(PTRACE_PEEKDATA, tracee, addr1, 0);
    printf("mod_handle: ");
    printerror();
    memcpy(buf.chars, "hooked", 6);
    buf.chars[6] = 0;

    ptrace(PTRACE_POKEDATA, tracee, addr1, buf.val);
    printf("hook: ");
    printerror();
    buf.val = ptrace(PTRACE_PEEKDATA, tracee, addr2, 0);
    printf("mod_handle: ");
    printerror();
    memcpy(buf.chars, "/hooked", 7);
    buf.chars[7] = 0;

    ptrace(PTRACE_POKEDATA, tracee, addr2, buf.val);
    printf("/hooked: ");
    printerror();
}

size_t getCodebase(pid_t pid)
{
    size_t addr;
    char buf[2 * sizeof(size_t)];
    char* end;
    char* mapfile[0x18];
    sprintf(mapfile, Mapfile, pid);
    int fd = open(mapfile, O_RDONLY);
    if(fd == -1)
    {
        printf("open maps error!");
        exit(1);
    }
    read(fd, buf, 2 * sizeof(size_t));
    end = strchr(buf, '-');
    addr = strtol(buf, &end, 16);
    printf("The codebase is: 0x%lx\n", addr);

    close(fd);
    return addr;
}

int main(int argc, char* argv[]){
    pid_t tracee; 
    union pltval plt;
    struct user_regs_struct regs;
    siginfo_t si;

    int status;
    size_t plt_offset, plt_addr, bak;

    if(argc < 2){
        usage(argv);
        exit(1);
    }

    tracee = atoi(argv[1]);
    plt_offset = atoi(argv[2]);
    //获取codebase
    plt_addr = plt_offset + getCodebase(tracee);
    printf("plt_addr ==> %lx\n", plt_addr);

    //attach the process
    ptrace(PTRACE_ATTACH, tracee, 0, 0);
    printf("Attach: ");
    printerror();

    wait(&status);

    //获取目标的plt值，保存，修改，写入，　继续运行
    plt.val = ptrace(PTRACE_PEEKDATA, tracee, plt_addr, 0);
    bak = plt.val;
    plt.chars[0] = 0xcc;        //breakpoint
    ptrace(PTRACE_POKEDATA, tracee, plt_addr, plt.val);
    ptrace(PTRACE_CONT, tracee, 0, 0);

    //监视有没有触发断点
    while(1){
        printf("Wait....\n");
        wait(&status);
        printf("Done!\n");

        if(WIFEXITED(status)) break;

        //获取regs和sig信息，判断是否到达plt
        ptrace(PTRACE_GETSIGINFO, tracee, 0, &si);
        ptrace(PTRACE_GETREGS, tracee, 0, &regs);

        if((si.si_signo != SIGTRAP) || (regs.rip != (size_t)plt_addr + 1)){
            ptrace(PTRACE_GETREGS, tracee, 0, &regs);
            ptrace(PTRACE_CONT, tracee, 0, 0);
            continue;
        }

        //hook & modify
        mod_handle(tracee, (void*)argv[0], (void*)argv[1]);

        //修改回原值
        plt.val = bak;
        ptrace(PTRACE_POKEDATA, tracee, plt_addr, plt.val);
        //返回0xcc前
        regs.rip -= 1;
        ptrace(PTRACE_SETREGS, tracee, 0, &regs);
        ptrace(PTRACE_SINGLESTEP, tracee, 0, 0);
        wait(0);

        ptrace(PTRACE_GETREGS, tracee, 0, &regs);
        plt.chars[0] = 0xcc;
        ptrace(PTRACE_POKEDATA, tracee, plt_addr, plt.val);
        ptrace(PTRACE_CONT, tracee, 0, 0);

    }
    return 0;
}