#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/user.h>
#include <sys/uio.h>
#include <elf.h>
#include <errno.h>

#define LIB_PATH "/data/local/tmp/libcamera_hook.so"
#define ARM64_BRK 0xd4200000

static int ptrace_write_data(long pid, unsigned long addr, const void *buf, size_t len) {
    for (size_t i = 0; i < len; i += sizeof(long)) {
        long val = 0;
        size_t cpy = (len - i < sizeof(long)) ? len - i : sizeof(long);
        memcpy(&val, (const char*)buf + i, cpy);
        if (ptrace(PTRACE_POKETEXT, pid, addr + i, val) < 0) {
            fprintf(stderr, "POKETEXT failed at 0x%lx\n", addr + i);
            return -1;
        }
    }
    return 0;
}

static unsigned long find_linker_rx_base(long pid, unsigned long *file_offset_out) {
    char path[256];
    snprintf(path, sizeof(path), "/proc/%ld/maps", pid);
    FILE *f = fopen(path, "r");
    if (!f) { perror("fopen maps"); return 0; }

    char line[1024];
    while (fgets(line, sizeof(line), f)) {
        if (strstr(line, "/apex/com.android.runtime/bin/linker64")) {
            unsigned long start, end, offset;
            char perm[8];
            sscanf(line, "%lx-%lx %4s %lx", &start, &end, perm, &offset);
            if (strstr(perm, "r-xp")) {
                fclose(f);
                fprintf(stderr, "linker64 r-xp: 0x%lx (file offset 0x%lx)\n", start, offset);
                if (file_offset_out) *file_offset_out = offset;
                return start;
            }
        }
    }
    fclose(f);
    fprintf(stderr, "linker64 r-xp not found\n");
    return 0;
}

int main(int argc, char *argv[]) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <pid>\n", argv[0]);
        return 1;
    }

    long pid = atol(argv[1]);

    unsigned long rx_offset = 0;
    unsigned long rx_base = find_linker_rx_base(pid, &rx_offset);
    if (!rx_base) {
        fprintf(stderr, "Failed to find linker64 r-xp\n");
        return 1;
    }

    unsigned long dlopen_offset = 0x86fa0;
    unsigned long dlopen_addr = rx_base + (dlopen_offset - rx_offset);
    fprintf(stderr, "dlopen address: 0x%lx (rx_base=0x%lx, sym_offset=0x%lx, rx_offset=0x%lx)\n",
            dlopen_addr, rx_base, dlopen_offset, rx_offset);

    fprintf(stderr, "Attaching to PID %ld...\n", pid);
    if (ptrace(PTRACE_ATTACH, pid, 0, 0) < 0) {
        perror("PTRACE_ATTACH");
        return 1;
    }
    int status;
    waitpid(pid, &status, WUNTRACED);
    if (!WIFSTOPPED(status)) {
        fprintf(stderr, "Target did not stop after attach\n");
        ptrace(PTRACE_DETACH, pid, 0, 0);
        return 1;
    }
    fprintf(stderr, "Attached, target stopped (signal=%d)\n", WSTOPSIG(status));

    struct user_pt_regs saved_regs;
    struct iovec iov_save = { &saved_regs, sizeof(saved_regs) };
    if (ptrace(PTRACE_GETREGSET, pid, NT_PRSTATUS, &iov_save) < 0) {
        perror("PTRACE_GETREGSET (save)");
        ptrace(PTRACE_DETACH, pid, 0, 0);
        return 1;
    }
    fprintf(stderr, "Saved: PC=0x%lx SP=0x%lx\n", saved_regs.pc, saved_regs.sp);

    unsigned long data_addr = saved_regs.sp - 0x8000;
    fprintf(stderr, "Data area at 0x%lx\n", data_addr);

    size_t path_len = strlen(LIB_PATH) + 1;
    if (ptrace_write_data(pid, data_addr, LIB_PATH, path_len) < 0) {
        fprintf(stderr, "Failed to write lib path\n");
        ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov_save);
        ptrace(PTRACE_DETACH, pid, 0, 0);
        return 1;
    }
    fprintf(stderr, "Wrote '%s' to 0x%lx\n", LIB_PATH, data_addr);

    unsigned long ret_addr = data_addr + 0x200;
    errno = 0;
    long orig_at_ret = ptrace(PTRACE_PEEKTEXT, pid, ret_addr, 0);
    if (errno) {
        perror("PEEKTEXT at return address");
        ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov_save);
        ptrace(PTRACE_DETACH, pid, 0, 0);
        return 1;
    }
    if (ptrace(PTRACE_POKETEXT, pid, ret_addr, (void*)(long)ARM64_BRK) < 0) {
        perror("POKETEXT breakpoint");
        ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov_save);
        ptrace(PTRACE_DETACH, pid, 0, 0);
        return 1;
    }
    fprintf(stderr, "Breakpoint at 0x%lx (orig=0x%lx)\n", ret_addr, orig_at_ret);

    struct user_pt_regs call_regs = saved_regs;
    call_regs.regs[0] = data_addr;
    call_regs.regs[1] = 2;
    call_regs.regs[2] = saved_regs.regs[30];
    call_regs.regs[3] = 0;
    call_regs.regs[30] = ret_addr;
    call_regs.pc = dlopen_addr;

    struct iovec iov_call = { &call_regs, sizeof(call_regs) };
    if (ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov_call) < 0) {
        perror("PTRACE_SETREGSET (call)");
        ptrace(PTRACE_POKETEXT, pid, ret_addr, (void*)(long)orig_at_ret);
        ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov_save);
        ptrace(PTRACE_DETACH, pid, 0, 0);
        return 1;
    }

    fprintf(stderr, "Calling dlopen(\"%s\", 2) at 0x%lx...\n", LIB_PATH, dlopen_addr);
    ptrace(PTRACE_CONT, pid, 0, 0);
    waitpid(pid, &status, WUNTRACED);

    if (WIFSTOPPED(status)) {
        fprintf(stderr, "Target stopped (signal=%d)\n", WSTOPSIG(status));
    } else if (WIFEXITED(status)) {
        fprintf(stderr, "Target exited! (code=%d)\n", WEXITSTATUS(status));
    } else if (WIFSIGNALED(status)) {
        fprintf(stderr, "Target killed! (signal=%d)\n", WTERMSIG(status));
    }

    ptrace(PTRACE_POKETEXT, pid, ret_addr, (void*)(long)orig_at_ret);

    struct user_pt_regs result_regs;
    struct iovec iov_result = { &result_regs, sizeof(result_regs) };
    if (ptrace(PTRACE_GETREGSET, pid, NT_PRSTATUS, &iov_result) < 0) {
        perror("PTRACE_GETREGSET (result)");
    } else {
        fprintf(stderr, "dlopen returned: 0x%llx\n", (unsigned long long)result_regs.regs[0]);
        if (result_regs.regs[0] != 0) {
            fprintf(stderr, "SUCCESS! Library loaded.\n");
        } else {
            fprintf(stderr, "FAILED: dlopen returned NULL\n");
        }
    }

    ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &iov_save);
    ptrace(PTRACE_DETACH, pid, 0, 0);
    fprintf(stderr, "Detached from PID %ld\n", pid);
    return 0;
}
