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

#include "patchfinder64.h"
#include "kernel_utils.h"
#include "sha1.h"
#include "sha256.h"

#import <Foundation/Foundation.h>

//#define LOG(str, args...) do { NSLog(@"[*] " str "\n", ##args); } while(0)
#define LOG(str, args...)

struct trust_mem {
    uint64_t next; //struct trust_mem *next;
    unsigned char uuid[16];
    unsigned int count;
    //unsigned char data[];
} __attribute__((packed));


uint32_t swap_uint32( uint32_t val ) {
    val = ((val << 8) & 0xFF00FF00 ) | ((val >> 8) & 0xFF00FF );
    return (val << 16) | (val >> 16);
}

uint32_t read_magic(FILE* file, off_t offset) {
    uint32_t magic;
    fseek(file, offset, SEEK_SET);
    fread(&magic, sizeof(uint32_t), 1, file);
    return magic;
}

void *load_bytes(FILE *file, off_t offset, size_t size) {
    void *buf = calloc(1, size);
    fseek(file, offset, SEEK_SET);
    fread(buf, size, 1, file);
    return buf;
}

uint8_t *get_code_directory(const char* name) {
    
    FILE* fd = fopen(name, "r");
    
    uint32_t magic;
    fread(&magic, sizeof(magic), 1, fd);
    fseek(fd, 0, SEEK_SET);
    
    long off, file_off = 0;
    int ncmds;
    int foundarm64 = 0;
    
    if (magic == MH_MAGIC_64) { // 0xFEEDFACF
        struct mach_header_64 mh64;
        fread(&mh64, sizeof(mh64), 1, fd);
        off = sizeof(mh64);
        ncmds = mh64.ncmds;
    }
    else if (magic == MH_MAGIC) {
        printf("[-] %s is 32bit. What are you doing here?\n", name);
        fclose(fd);
        return NULL;
    }
    else if (magic == 0xBEBAFECA) { //FAT binary magic
        
        size_t header_size = sizeof(struct fat_header);
        size_t arch_size = sizeof(struct fat_arch);
        size_t arch_off = header_size;
        
        struct fat_header *fat = (struct fat_header*)load_bytes(fd, 0, header_size);
        struct fat_arch *arch = (struct fat_arch *)load_bytes(fd, arch_off, arch_size);
        
        int n = swap_uint32(fat->nfat_arch);
        printf("[*] Binary is FAT with %d architectures\n", n);
        
        while (n-- > 0) {
            magic = read_magic(fd, swap_uint32(arch->offset));
            
            if (magic == 0xFEEDFACF) {
                printf("[*] Found arm64\n");
                foundarm64 = 1;
                struct mach_header_64* mh64 = (struct mach_header_64*)load_bytes(fd, swap_uint32(arch->offset), sizeof(struct mach_header_64));
                file_off = swap_uint32(arch->offset);
                off = swap_uint32(arch->offset) + sizeof(struct mach_header_64);
                ncmds = mh64->ncmds;
                break;
            }
            
            arch_off += arch_size;
            arch = load_bytes(fd, arch_off, arch_size);
        }
        
        if (!foundarm64) { // by the end of the day there's no arm64 found
            printf("[-] No arm64? RIP\n");
            fclose(fd);
            return NULL;
        }
    }
    else {
        printf("[-] %s is not a macho! (or has foreign endianness?) (magic: %x)\n", name, magic);
        fclose(fd);
        return NULL;
    }
    
    for (int i = 0; i < ncmds; i++) {
        struct load_command cmd;
        fseek(fd, off, SEEK_SET);
        fread(&cmd, sizeof(struct load_command), 1, fd);
        if (cmd.cmd == LC_CODE_SIGNATURE) {
            uint32_t off_cs;
            fread(&off_cs, sizeof(uint32_t), 1, fd);
            uint32_t size_cs;
            fread(&size_cs, sizeof(uint32_t), 1, fd);
            
            uint8_t *cd = malloc(size_cs);
            fseek(fd, off_cs + file_off, SEEK_SET);
            fread(cd, size_cs, 1, fd);
            fclose(fd);
            return cd;
        } else {
            off += cmd.cmdsize;
        }
    }
    fclose(fd);
    return NULL;
}

void get_sha256_hash(const uint8_t* data, uint32_t datasize, uint8_t *out) {
  SHA256_CTX ctx;
	sha256_init(&ctx);
	sha256_update(&ctx, data, datasize);
	sha256_final(&ctx, out);
}

void get_sha1_hash(const uint8_t* data, uint32_t datasize, uint8_t *out) {
  SHA1_CTX ctx;
  SHA1Init(&ctx);
  SHA1Update(&ctx, data, datasize);
  SHA1Final(out, &ctx);
}

int trust_bin(const char* filepath) 
{
  uint64_t trust_chain = Find_trustcache();
  if (!trust_chain) {
    trust_chain = Find_trustcache10_3_2();
  }
  LOG("trust %p\n", (void*)trust_chain);

  struct trust_mem fake_chain;
  fake_chain.next = KernelRead_64bits(trust_chain);
  *(uint64_t *)&fake_chain.uuid[0] = 0xabadbabeabadbabe;
  *(uint64_t *)&fake_chain.uuid[8] = 0xabadbabeabadbabe;
  LOG("trust_chain %p\n", (void*)fake_chain.next);

  uint8_t *cd = get_code_directory(filepath);
  if (!cd) {
    return -1;
  }

  uint32_t* code_dir_int = (uint32_t*)cd;
  uint32_t realsize = 0;
  for (int j = 0; j < 10; j++) {
      if (swap_uint32(code_dir_int[j]) == 0xfade0c02) {
          realsize = swap_uint32(code_dir_int[j+1]);
          cd += 4*j;
      }
  }

  uint8_t *hash;
  size_t hash_size;
  // iOS 11 uses sha-256
  if (kCFCoreFoundationVersionNumber >= 1443.00) {
    hash_size = 32;
    hash = malloc(hash_size);
    get_sha256_hash(cd, realsize, hash);
  } else {
    hash_size = 20;
    hash = malloc(hash_size);
    get_sha1_hash(cd, realsize, hash);
  }

  fake_chain.count = 1;
  size_t length = (sizeof(fake_chain) + hash_size + 0xFFFF) & ~0xFFFF;
  uint64_t kernel_trust = Kernel_alloc(length);
  LOG("[*] allocated: 0x%zx => 0x%llx\n", length, kernel_trust);

  KernelWrite(kernel_trust, &fake_chain, sizeof(fake_chain));
  KernelWrite(kernel_trust + sizeof(fake_chain), hash, hash_size);
  KernelWrite_64bits(trust_chain, kernel_trust);

  return 0;
}