Add support for taking screenshots with the scripts

StikJIT/JSSupport/RunJSView.swift

@@ -8,6 +8,9 @@
 import SwiftUI
 import JavaScriptCore
 
+typealias RemoteServerHandle = OpaquePointer
+typealias ScreenshotClientHandle = OpaquePointer
+
 class RunJSViewModel: ObservableObject {
     var context: JSContext?
     @Published var logs: [String] = []
@@ -15,13 +18,15 @@ class RunJSViewModel: ObservableObject {
     @Published var executionInterrupted = false
     var pid: Int
     var debugProxy: OpaquePointer?
+    var remoteServer: OpaquePointer?
     var semaphore: dispatch_semaphore_t?
     private var progressTimer: DispatchSourceTimer?
     private var reportedProgress: Double = 0
     
-    init(pid: Int, debugProxy: OpaquePointer?, semaphore: dispatch_semaphore_t?) {
+    init(pid: Int, debugProxy: OpaquePointer?, remoteServer: OpaquePointer?, semaphore: dispatch_semaphore_t?) {
         self.pid = pid
         self.debugProxy = debugProxy
+        self.remoteServer = remoteServer
         self.semaphore = semaphore
     }
     
@@ -61,6 +66,10 @@ class RunJSViewModel: ObservableObject {
             return handleJITPageWrite(self.context, startAddr, regionSize, self.debugProxy) ?? ""
         }
         
+        let takeScreenshotFunction: @convention(block) (String?) -> String? = { fileName in
+            return self.captureScreenshot(named: fileName)
+        }
+        
         let hasTXMFunction: @convention(block) () -> Bool = {
             return ProcessInfo.processInfo.hasTXM
         }
@@ -70,6 +79,7 @@ class RunJSViewModel: ObservableObject {
         context?.setObject(getPidFunction, forKeyedSubscript: "get_pid" as NSString)
         context?.setObject(sendCommandFunction, forKeyedSubscript: "send_command" as NSString)
         context?.setObject(prepareMemoryRegionFunction, forKeyedSubscript: "prepare_memory_region" as NSString)
+        context?.setObject(takeScreenshotFunction, forKeyedSubscript: "take_screenshot" as NSString)
         context?.setObject(logFunction, forKeyedSubscript: "log" as NSString)
         
         context?.evaluateScript(scriptContent)
@@ -88,6 +98,115 @@ class RunJSViewModel: ObservableObject {
         }
     }
     
+    private func captureScreenshot(named preferredName: String?) -> String {
+        if executionInterrupted {
+            raiseException("Script execution is interrupted by StikDebug.")
+            return ""
+        }
+        guard let remoteServer else {
+            raiseException("Screenshot capture is unavailable in the current session.")
+            return ""
+        }
+        
+        var screenshotClient: ScreenshotClientHandle?
+        let creationError = screenshot_client_new(remoteServer, &screenshotClient)
+        if let creationError {
+            let message = describeIdeviceError(creationError)
+            idevice_error_free(creationError)
+            raiseException("Failed to create screenshot client: \(message)")
+            return ""
+        }
+        guard let screenshotClient else {
+            raiseException("Failed to allocate screenshot client.")
+            return ""
+        }
+        defer { screenshot_client_free(screenshotClient) }
+        
+        var buffer: UnsafeMutablePointer<UInt8>?
+        var length: UInt = 0
+        let captureError = screenshot_client_take_screenshot(screenshotClient, &buffer, &length)
+        if let captureError {
+            let message = describeIdeviceError(captureError)
+            idevice_error_free(captureError)
+            raiseException("Failed to take screenshot: \(message)")
+            return ""
+        }
+        guard let buffer else {
+            raiseException("Device returned empty screenshot data.")
+            return ""
+        }
+        defer { idevice_data_free(buffer, length) }
+        
+        let data = Data(bytes: buffer, count: Int(length))
+        do {
+            let fileURL = try screenshotFileURL(preferredName: preferredName)
+            try data.write(to: fileURL, options: .atomic)
+            return fileURL.path
+        } catch {
+            raiseException("Failed to save screenshot: \(error.localizedDescription)")
+            return ""
+        }
+    }
+    
+    private func screenshotFileURL(preferredName: String?) throws -> URL {
+        let directory = URL.documentsDirectory.appendingPathComponent("screenshots", isDirectory: true)
+        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
+        let fileManager = FileManager.default
+        let initialName = sanitizedScreenshotName(from: preferredName)
+        var targetURL = directory.appendingPathComponent(initialName)
+        guard fileManager.fileExists(atPath: targetURL.path) else {
+            return targetURL
+        }
+        
+        let baseName = targetURL.deletingPathExtension().lastPathComponent
+        let ext = targetURL.pathExtension.isEmpty ? "png" : targetURL.pathExtension
+        var counter = 1
+        repeat {
+            let candidate = "\(baseName)-\(counter).\(ext)"
+            targetURL = directory.appendingPathComponent(candidate)
+            counter += 1
+        } while fileManager.fileExists(atPath: targetURL.path)
+        return targetURL
+    }
+    
+    private func sanitizedScreenshotName(from preferredName: String?) -> String {
+        let defaultName = "screenshot-\(Int(Date().timeIntervalSince1970))"
+        guard var candidate = preferredName?.trimmingCharacters(in: .whitespacesAndNewlines),
+              !candidate.isEmpty else {
+            return "\(defaultName).png"
+        }
+        
+        let allowed = CharacterSet.alphanumerics.union(CharacterSet(charactersIn: "-_."))
+        var sanitized = ""
+        sanitized.reserveCapacity(candidate.count)
+        for scalar in candidate.unicodeScalars {
+            if allowed.contains(scalar) {
+                sanitized.append(Character(scalar))
+            } else {
+                sanitized.append("_")
+            }
+        }
+        if sanitized.isEmpty {
+            sanitized = defaultName
+        }
+        if !sanitized.lowercased().hasSuffix(".png") {
+            sanitized += ".png"
+        }
+        return sanitized
+    }
+    
+    private func describeIdeviceError(_ error: UnsafeMutablePointer<IdeviceFfiError>) -> String {
+        if let messagePointer = error.pointee.message {
+            return "[\(error.pointee.code)] \(String(cString: messagePointer))"
+        }
+        return "[\(error.pointee.code)] Unknown error"
+    }
+    
+    private func raiseException(_ message: String) {
+        guard let context else { return }
+        context.exception = JSValue(object: message, in: context)
+    }
+    
     private func startContinuedProcessing(withTitle title: String) {
         guard ContinuedProcessingManager.shared.isSupported,
               UserDefaults.standard.bool(forKey: UserDefaults.Keys.enableContinuedProcessing) else { return }

StikJIT/JSSupport/ScriptListView.swift

@@ -26,6 +26,11 @@ struct ScriptListView: View {
 
     @State private var showDeleteConfirmation = false
     @State private var pendingDelete: URL? = nil
+    @State private var scriptToRun: URL? = nil
+    @State private var showRunConfiguration = false
+    @State private var runConfiguration = ScriptRunConfiguration()
+    @State private var runnerModel: RunJSViewModel?
+    @State private var showRunnerLogs = false
 
     var onSelectScript: ((URL?) -> Void)? = nil
 
@@ -136,6 +141,42 @@ struct ScriptListView: View {
                 case .failure(let error): presentError(title: "Import Failed", message: error.localizedDescription)
                 }
             }
+            .sheet(isPresented: $showRunConfiguration) {
+                if let scriptToRun {
+                    RunScriptConfigurationView(
+                        scriptName: scriptToRun.lastPathComponent,
+                        configuration: $runConfiguration,
+                        onCancel: { showRunConfiguration = false },
+                        onRun: {
+                            let configuration = runConfiguration
+                            showRunConfiguration = false
+                            startRunningScript(scriptToRun, configuration: configuration)
+                        }
+                    )
+                }
+            }
+            .sheet(isPresented: $showRunnerLogs, onDismiss: { runnerModel = nil }) {
+                NavigationView {
+                    if let runnerModel {
+                        RunJSView(model: runnerModel)
+                            .toolbar {
+                                ToolbarItem(placement: .topBarTrailing) {
+                                    Button("Done") { showRunnerLogs = false }
+                                }
+                            }
+                            .navigationTitle(runnerModel.scriptName)
+                            .navigationBarTitleDisplayMode(.inline)
+                    } else {
+                        Text("No script output available.")
+                            .navigationTitle("Scripts")
+                            .toolbar {
+                                ToolbarItem(placement: .topBarTrailing) {
+                                    Button("Done") { showRunnerLogs = false }
+                                }
+                            }
+                    }
+                }
+            }
         }
         .preferredColorScheme(preferredScheme)
     }
@@ -212,6 +253,16 @@ struct ScriptListView: View {
                     .foregroundColor(.yellow)
             }
 
+            if !isPickerMode {
+                Button {
+                    prepareToRun(script)
+                } label: {
+                    Image(systemName: "play.circle.fill")
+                        .foregroundColor(.green)
+                }
+                .buttonStyle(.borderless)
+            }
+
             if showDelete {
                 Button(role: .destructive) {
                     pendingDelete = script
@@ -307,19 +358,28 @@ struct ScriptListView: View {
             }
             if !exists || !isDir.boolValue {
                 try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
-                if let bundleURL = Bundle.main.url(forResource: "attachDetach", withExtension: "js") {
-                    let dest = dir.appendingPathComponent("attachDetach.js")
-                    if !FileManager.default.fileExists(atPath: dest.path) {
-                        try FileManager.default.copyItem(at: bundleURL, to: dest)
-                    }
-                }
             }
+            try ensureDefaultScripts(in: dir)
         } catch {
             presentError(title: "Unable to Create Scripts Folder", message: error.localizedDescription)
         }
         return dir
     }
 
+    private func ensureDefaultScripts(in directory: URL) throws {
+        let fm = FileManager.default
+        if let bundleURL = Bundle.main.url(forResource: "attachDetach", withExtension: "js") {
+            let dest = directory.appendingPathComponent("attachDetach.js")
+            if !fm.fileExists(atPath: dest.path) {
+                try fm.copyItem(at: bundleURL, to: dest)
+            }
+        }
+        let screenshotURL = directory.appendingPathComponent("screenshot-demo.js")
+        if !fm.fileExists(atPath: screenshotURL.path) {
+            try screenshotDemoScript.write(to: screenshotURL, atomically: true, encoding: .utf8)
+        }
+    }
+
     private func loadScripts() {
         let dir = scriptsDirectory()
         scripts = (try? FileManager.default.contentsOfDirectory(at: dir, includingPropertiesForKeys: nil))?
@@ -385,6 +445,90 @@ struct ScriptListView: View {
         }
     }
 
+    // MARK: - Script Execution
+
+    private func prepareToRun(_ script: URL) {
+        guard ProcessInfo.processInfo.hasTXM else {
+            presentError(title: "Scripts Not Supported",
+                         message: "This device does not have a Trusted Execution Monitor. Enable the override in Settings if you want to force script execution.")
+            return
+        }
+        scriptToRun = script
+        showRunConfiguration = true
+    }
+
+    private func startRunningScript(_ script: URL, configuration: ScriptRunConfiguration) {
+        guard ProcessInfo.processInfo.hasTXM else {
+            presentError(title: "Scripts Not Supported",
+                         message: "This device does not have a Trusted Execution Monitor.")
+            return
+        }
+        guard configuration.isValid else {
+            presentError(title: "Invalid Input", message: "Enter a bundle identifier or PID to continue.")
+            return
+        }
+        let scriptData: Data
+        do {
+            scriptData = try Data(contentsOf: script)
+        } catch {
+            presentError(title: "Failed to Read Script", message: error.localizedDescription)
+            return
+        }
+
+        scriptToRun = nil
+        isBusy = true
+        let scriptName = script.lastPathComponent
+        let callback = makeJsCallback(scriptData: scriptData, scriptName: scriptName)
+        let logger: LogFunc = { message in
+            if let message {
+                LogManager.shared.addInfoLog(message)
+            }
+        }
+
+        DispatchQueue.global(qos: .userInitiated).async {
+            let success: Bool
+            switch configuration.mode {
+            case .bundleID:
+                let trimmed = configuration.bundleID.trimmingCharacters(in: .whitespacesAndNewlines)
+                success = JITEnableContext.shared.debugApp(withBundleID: trimmed, logger: logger, jsCallback: callback)
+            case .pid:
+                let trimmed = configuration.pid.trimmingCharacters(in: .whitespacesAndNewlines)
+                let pidValue = Int32(trimmed) ?? 0
+                success = JITEnableContext.shared.debugApp(withPID: pidValue, logger: logger, jsCallback: callback)
+            }
+            DispatchQueue.main.async {
+                self.isBusy = false
+                if success {
+                    self.presentSuccess(title: "Script Running", message: scriptName)
+                } else {
+                    self.presentError(title: "Failed to Start Script", message: "Ensure the provider is connected and try again.")
+                }
+            }
+        }
+    }
+
+    private func makeJsCallback(scriptData: Data, scriptName: String) -> DebugAppCallback {
+        return { pid, debugProxyHandle, remoteServerHandle, semaphore in
+            let model = RunJSViewModel(pid: Int(pid),
+                                       debugProxy: debugProxyHandle,
+                                       remoteServer: remoteServerHandle,
+                                       semaphore: semaphore)
+            DispatchQueue.main.async {
+                self.runnerModel = model
+                self.showRunnerLogs = true
+            }
+            DispatchQueue.global(qos: .background).async {
+                do {
+                    try model.runScript(data: scriptData, name: scriptName)
+                } catch {
+                    DispatchQueue.main.async {
+                        self.presentError(title: "Script Execution Failed", message: error.localizedDescription)
+                    }
+                }
+            }
+        }
+    }
+
     // MARK: - Feedback helpers
 
     private func presentError(title: String, message: String) {
@@ -449,3 +593,118 @@ private struct WideGlassyButton: View {
         .buttonStyle(.plain)
     }
 }
+
+private struct RunScriptConfigurationView: View {
+    let scriptName: String
+    @Binding var configuration: ScriptRunConfiguration
+    let onCancel: () -> Void
+    let onRun: () -> Void
+
+    var body: some View {
+        NavigationStack {
+            Form {
+                Section(header: Text("Script")) {
+                    Text(scriptName)
+                        .font(.body.weight(.semibold))
+                        .textSelection(.enabled)
+                    Picker("Target", selection: $configuration.mode) {
+                        ForEach(ScriptRunMode.allCases) { mode in
+                            Text(mode.label).tag(mode)
+                        }
+                    }
+                    .pickerStyle(.segmented)
+                }
+
+                Section(header: Text(configuration.mode == .bundleID ? "Bundle Identifier" : "PID")) {
+                    if configuration.mode == .bundleID {
+                        TextField("com.example.app", text: $configuration.bundleID)
+                            .textInputAutocapitalization(.none)
+                            .autocorrectionDisabled()
+                    } else {
+                        TextField("1234", text: $configuration.pid)
+                            .keyboardType(.numberPad)
+                    }
+                }
+
+                Section {
+                    Text("The app will be attached, the script will run, and then it will detach automatically.")
+                        .font(.footnote)
+                        .foregroundColor(.secondary)
+                }
+            }
+            .navigationTitle("Run Script")
+            .toolbar {
+                ToolbarItem(placement: .cancellationAction) {
+                    Button("Cancel", action: onCancel)
+                }
+                ToolbarItem(placement: .confirmationAction) {
+                    Button("Run", action: onRun)
+                        .disabled(!configuration.isValid)
+                }
+            }
+        }
+    }
+}
+
+private enum ScriptRunMode: String, CaseIterable, Identifiable {
+    case bundleID
+    case pid
+
+    var id: String { rawValue }
+
+    var label: String {
+        switch self {
+        case .bundleID: return "Bundle ID"
+        case .pid: return "PID"
+        }
+    }
+}
+
+private struct ScriptRunConfiguration {
+    var mode: ScriptRunMode = .bundleID
+    var bundleID: String = ""
+    var pid: String = ""
+
+    var isValid: Bool {
+        switch mode {
+        case .bundleID:
+            return !bundleID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
+        case .pid:
+            guard let value = Int(pid.trimmingCharacters(in: .whitespacesAndNewlines)), value > 0 else {
+                return false
+            }
+            return true
+        }
+    }
+}
+
+private let screenshotDemoScript = """
+// Screenshot Demo Script
+// Attaches to the target, captures a PNG screenshot, and detaches.
+
+function takeScreenshotDemo() {
+    log("[ScreenshotDemo] Starting demo");
+
+    const pid = get_pid();
+    log(`[ScreenshotDemo] Target PID: ${pid}`);
+
+    const attachResponse = send_command(`vAttach;${pid.toString(16)}`);
+    log(`[ScreenshotDemo] attach_response = ${attachResponse}`);
+
+    const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
+    const fileName = `screenshot-${timestamp}.png`;
+    const savedPath = take_screenshot(fileName);
+
+    if (savedPath && savedPath.length > 0) {
+        log(`[ScreenshotDemo] Screenshot saved to ${savedPath}`);
+    } else {
+        log("[ScreenshotDemo] Device did not report a saved path.");
+    }
+
+    const detachResponse = send_command("D");
+    log(`[ScreenshotDemo] detach_response = ${detachResponse}`);
+    log("[ScreenshotDemo] Demo complete.");
+}
+
+takeScreenshotDemo();
+"""

StikJIT/Views/HomeView.swift

@@ -1361,8 +1361,11 @@ struct HomeView: View {
     }
     
     private func getJsCallback(_ script: Data, name: String? = nil) -> DebugAppCallback {
-        return { pid, debugProxyHandle, semaphore in
-            jsModel = RunJSViewModel(pid: Int(pid), debugProxy: debugProxyHandle, semaphore: semaphore)
+        return { pid, debugProxyHandle, remoteServerHandle, semaphore in
+            jsModel = RunJSViewModel(pid: Int(pid),
+                                     debugProxy: debugProxyHandle,
+                                     remoteServer: remoteServerHandle,
+                                     semaphore: semaphore)
             scriptViewShow = true
             DispatchQueue.global(qos: .background).async {
                 do { try jsModel?.runScript(data: script, name: name) }

StikJIT/idevice/jit.c

@@ -18,7 +18,11 @@
 
 #include "jit.h"
 
-void runDebugServerCommand(int pid, DebugProxyHandle* debug_proxy, LogFuncC logger, DebugAppCallback callback) {
+void runDebugServerCommand(int pid,
+                           DebugProxyHandle* debug_proxy,
+                           RemoteServerHandle* remote_server,
+                           LogFuncC logger,
+                           DebugAppCallback callback) {
     // enable QStartNoAckMode
     char *disableResponse = NULL;
     debug_proxy_send_ack(debug_proxy);
@@ -32,7 +36,7 @@ void runDebugServerCommand(int pid, DebugProxyHandle* debug_proxy, LogFuncC logg
     
     if(callback) {
         dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
-        callback(pid, debug_proxy, semaphore);
+        callback(pid, debug_proxy, remote_server, semaphore);
         dispatch_semaphore_wait(semaphore, DISPATCH_TIME_FOREVER);
         err = debug_proxy_send_raw(debug_proxy, "\x03", 1);
         usleep(500);
@@ -193,7 +197,7 @@ int debug_app(IdeviceProviderHandle* tcp_provider, const char *bundle_id, LogFun
       return 1;
     }
     
-    runDebugServerCommand((int)pid, debug_proxy, logger, callback);
+    runDebugServerCommand((int)pid, debug_proxy, remote_server, logger, callback);
     
     /*****************************************************************
      * Cleanup
@@ -296,7 +300,7 @@ int debug_app_pid(IdeviceProviderHandle* tcp_provider, int pid, LogFuncC logger,
     }
     
     
-    runDebugServerCommand(pid, debug_proxy, logger, callback);
+    runDebugServerCommand(pid, debug_proxy, remote_server, logger, callback);
     
     /*****************************************************************
      * Cleanup

StikJIT/idevice/jit.h

@@ -11,7 +11,10 @@
 #include "idevice.h"
 
 typedef void (^LogFuncC)(const char* message, ...);
-typedef void (^DebugAppCallback)(int pid, struct DebugProxyHandle* debug_proxy, dispatch_semaphore_t semaphore);
+typedef void (^DebugAppCallback)(int pid,
+                                 struct DebugProxyHandle* debug_proxy,
+                                 struct RemoteServerHandle* remote_server,
+                                 dispatch_semaphore_t semaphore);
 int debug_app(IdeviceProviderHandle* tcp_provider, const char *bundle_id, LogFuncC logger, DebugAppCallback callback);
 int debug_app_pid(IdeviceProviderHandle* tcp_provider, int pid, LogFuncC logger, DebugAppCallback callback);
 int launch_app_via_proxy(IdeviceProviderHandle* tcp_provider, const char *bundle_id, LogFuncC logger);