#include <mach-o/dyld.h>
#include <mach/mach.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

// #include <Foundation/Foundation.h>
#include <objc/objc.h>

#include "out/restore.h"

int custom_strcmp(const char *p1, const char *p2) {
  const unsigned char *s1 = (const unsigned char *)p1;
  const unsigned char *s2 = (const unsigned char *)p2;
  unsigned char c1, c2;
  do {
    c1 = (unsigned char)*s1++;
    c2 = (unsigned char)*s2++;
    if (c1 == '\0')
      return c1 - c2;
  } while (c1 == c2);
  return c1 - c2;
}

int custom_strncmp(const char *s1, const char *s2, register size_t n) {
  register unsigned char u1, u2;

  while (n-- > 0) {
    u1 = (unsigned char)*s1++;
    u2 = (unsigned char)*s2++;
    if (u1 != u2)
      return u1 - u2;
    if (u1 == '\0')
      return 0;
  }
  return 0;
}

const uint32_t magic64 = 0xfeedfacf;
const uint32_t magic32 = 0xfeedface;

struct ProgramVars {
  void *mh; // mach_header or mach_header64
  int *NXArgcPtr;
  const char ***NXArgvPtr;
  const char ***environPtr;
  const char **__prognamePtr;
};

void restore_strings(void* main);

__attribute__((constructor)) static void
bruh(int argc, const char *const argv[], const char *const envp[],
     const char *const apple[], const struct ProgramVars *vars) {
  void* main = (void *)(vars->mh);
  restore_strings(main);
}

/// strings in __TEXT,__cstring has been removed and this
/// function tries to recover those strings. Using either
/// these methods below.
///
/// 1. Recover __TEXT,__cstring
/// 2. Build a new segment with section for strings
///
/// (1) might seem reasonable at first, but requires __TEXT
/// segment to be writable. Although we can make that, but
/// we are not sure if the modification is allowed by Apple.
///
/// (2) actually require a little bit more work, by defining
/// a new segment with a section inside. This segment is
/// mounted readable/writable. Not only that, all string
/// references must also be updated.
/// In code, ARMv8, the sequence `adrp` `add` referencing
/// string must now be updated with new parameters as the
/// address/offset has now been changed.
/// In ARMv8, every instruction is 8 bytes, so looping
/// through all the code and change the instruction is easy.
///
/// It can be seen that opting for method (2) is safer,
/// as Apple allows for arbitrary segment. This option
/// requires that there is enough space left for a new segment.
/// Calculated, it should be around 152 bytes.
///
/// 4 + 4 + 16 + 8*4 + 4*4 + 16 + 16 + 8*2 + 4*8
/// ^~~~^   ^~~~~~~~~~~~~^   ^~~~~~~~~~~~~~~~~~^
///   1           2                  3
///
/// 1: load command header
/// 2: segment data
/// 3: section
///
/// However, if we can expand the old section, and remove the
/// old section entry, then we only need 72 bytes. Because,
/// we only move the section entry.
///
/// Remember to update the command count in macho header (+1).
void restore_strings(void* main) {
  printf("=== rebuilding the strings ===\n");

  void* header = main;
  const uint32_t magic = *(uint32_t *)header;
  char *ptr = (char *)header;
  if (magic == magic64) {
    ptr += 0x20;
  } else {
    ptr += 0x20 - 0x4;
  }
  const uint32_t ncmds = *((uint32_t *)header + 4);

  uint32_t slide = 0;

  char* secrets = 0;

  for (int i = 0; i < ncmds; i++) {
    const uint32_t cmd = *((uint32_t *)ptr + 0);
    const uint32_t cmdsize = *((uint32_t *)ptr + 1);
    if (cmd == LC_SEGMENT_64) {
      char *name = (char *)((uint64_t *)ptr + 1);
      uint64_t vmaddr = *((uint64_t *)ptr + 3);
      uint64_t fileoffset = *((uint64_t *)ptr + 5);
      // this assumes that __TEXT comes before __DATA_CONST
      if (custom_strcmp(name, "__TEXT") == 0) {
        slide = (uint64_t)header - vmaddr;

      } else if (custom_strcmp(name, "__BSHIELD") == 0) {
        printf("found __BSHIELD segment at %p\n", ptr);

        uint64_t nsect = *((uint32_t *)ptr + 8 * 2);
        char *sections_ptr = (char *)((uint32_t *)ptr + 18);
        for (int sec = 0; sec < nsect; sec++) {
          char *secname = sections_ptr;
          if (custom_strncmp(secname, "__secrets", 16) == 0) {
            uint64_t addr = *((uint64_t *)sections_ptr + 4);
            uint64_t size = *((uint64_t *)sections_ptr + 5);

            secrets = (char*)(addr + slide);
          }
          sections_ptr += 16 * 2 + 8 * 2 + 4 * 8;
        }
      }
    }
    ptr += cmdsize;
  }

  secrets[0] = 'F';
  secrets[1] = 'R';
  secrets[2] = 'E';
  secrets[3] = 'E';
  secrets[4] = ' ';
  secrets[5] = 'S';
  secrets[6] = 'P';
  secrets[7] = 'A';
  secrets[8] = 'C';
  secrets[9] = 'E';
  secrets[10] = '\n';
  secrets[11] = 0;
}