diff --git a/.gitignore b/.gitignore index 05840ee..27096bb 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,11 @@ DerivedData/ .swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata .netrc .vscode + +# Swift build artifacts that might be created in the wrong location +*.swiftdeps +*.o +*.d +*.swiftmodule +*.swiftsourceinfo +*.swiftdoc diff --git a/Package.swift b/Package.swift index ad13d20..ceda0de 100644 --- a/Package.swift +++ b/Package.swift @@ -31,6 +31,7 @@ let package = Package( .product(name: "GRPCNIOTransportHTTP2", package: "grpc-swift-nio-transport"), .product(name: "EdgeAgentGRPC", package: "edge-agent-common"), .target(name: "EdgeCLI"), + .target(name: "Imager"), ], resources: [ .copy("Resources") @@ -55,5 +56,10 @@ let package = Package( .product(name: "Crypto", package: "swift-crypto"), ] ), + + /// Tools to put EdgeOS images onto drives. + .target( + name: "Imager" + ), ] ) diff --git a/Sources/Imager/Drive.swift b/Sources/Imager/Drive.swift new file mode 100644 index 0000000..25d3f7b --- /dev/null +++ b/Sources/Imager/Drive.swift @@ -0,0 +1,35 @@ +import Foundation +/// A structure representing a drive. +public struct Drive { + /// The path to the drive. + public var path: String + /// The total capacity of the drive in bytes. + public var capacity: Int64 + /// The available free space on the drive in bytes, if known. + public var available: Int64? + /// The name of the drive. + public var name: String? + + /// The total capacity of the drive in a human-readable format. + public var capacityHumanReadable: String { + let formatter = ByteCountFormatter() + formatter.allowedUnits = [.useAll] + formatter.countStyle = .file + formatter.includesUnit = true + formatter.isAdaptive = true + return formatter.string(fromByteCount: capacity) + } + + /// The available space on the drive in a human-readable format, or "N/A" if unknown. + public var availableHumanReadable: String { + guard let availableSpace = available else { + return "N/A" + } + let formatter = ByteCountFormatter() + formatter.allowedUnits = [.useAll] + formatter.countStyle = .file + formatter.includesUnit = true + formatter.isAdaptive = true + return formatter.string(fromByteCount: availableSpace) + } +} diff --git a/Sources/Imager/ImagerError.swift b/Sources/Imager/ImagerError.swift new file mode 100644 index 0000000..2a8c03b --- /dev/null +++ b/Sources/Imager/ImagerError.swift @@ -0,0 +1,102 @@ +/// Error types that can occur during the imaging process. +/// +/// These errors represent various failure scenarios that might occur when +/// attempting to write an image file to a drive. +public enum ImagerError: Error, Sendable { + /// The provided image file is invalid. + /// + /// This error occurs when the image file cannot be read or is in an unsupported format. + /// Only ZIP, TAR, and IMG formats are supported. + case invalidImageFile(reason: String) + + /// The selected drive is invalid. + /// + /// This error occurs when the specified drive cannot be used as a target. + /// The associated value provides the operating system's specific reason for the failure. + case invalidDrive(reason: String) + + /// Permission to access the drive was denied. + /// + /// This error occurs when the application does not have sufficient privileges + /// to write to the selected drive. The associated value provides the operating + /// system's specific reason for the permission denial. + case permissionDenied(reason: String) + + /// The image file is too large for the selected drive. + /// + /// This error occurs when the size of the image file exceeds the available + /// space on the target drive. + case imageTooLargeForDrive(imageSize: UInt64, driveSize: UInt64) + + /// The imaging process was interrupted. + /// + /// This error occurs when the imaging process is unexpectedly terminated + /// before completion. The associated value provides information about the + /// cause of the interruption. + case processingInterrupted(reason: String) + + /// Error detecting or enumerating drives. + /// + /// This error occurs when there is an issue detecting or enumerating drives. + /// The associated value provides information about the cause of the error. + case driveDetectionError(reason: String) + + /// Error related to disk watching functionality. + /// + /// This error occurs when setting up or managing DiskArbitration callbacks fails. + case diskWatchError(reason: String) + + /// Functionality not implemented. + /// + /// This error occurs when a feature is not yet implemented for a particular platform. + case notImplemented(reason: String) + + /// An unknown error occurred. + /// + /// This error represents any failure that doesn't fall into the other categories. + /// It should be used as a last resort when the specific error type cannot be determined. + case unknown(reason: String?) + +} + +// MARK: - CustomStringConvertible + +extension ImagerError: CustomStringConvertible { + /// A user-friendly description of the error. + public var description: String { + switch self { + case .invalidImageFile(let reason): + return "Invalid Image File: \(reason). Only ZIP, TAR, and IMG formats are supported." + + case .invalidDrive(let reason): + return "Invalid drive: \(reason)" + + case .permissionDenied(let reason): + return "Permission denied: \(reason)" + + case .imageTooLargeForDrive(let imageSize, let driveSize): + let imageSizeMB = Double(imageSize) / 1_048_576 + let driveSizeMB = Double(driveSize) / 1_048_576 + return "Image too large for drive: Image size is \(String(format: "%.1f", imageSizeMB)) MB, but drive size is only \(String(format: "%.1f", driveSizeMB)) MB" + + case .processingInterrupted(let reason): + return "Imaging process was interrupted: \(reason)" + + case .driveDetectionError(let reason): + return "Error detecting drives: \(reason)" + + case .diskWatchError(let reason): + return "Disk Watch Error: \(reason)" + + case .notImplemented(let reason): + return "Functionality not implemented: \(reason)" + + case .unknown(let reason): + if let reason = reason, !reason.isEmpty { + return "An unknown error occurred: \(reason)" + } else { + return "An unknown error occurred" + } + } + } +} \ No newline at end of file diff --git a/Sources/Imager/ImagerFactory.swift b/Sources/Imager/ImagerFactory.swift new file mode 100644 index 0000000..8aafb8e --- /dev/null +++ b/Sources/Imager/ImagerFactory.swift @@ -0,0 +1,34 @@ +import Foundation + +/// Provides factory methods for creating platform-specific Imager and DiskLister instances. +public struct ImagerFactory { + + /// Creates a platform-specific object conforming to the Imager protocol. + /// - Parameters: + /// - imageFilePath: The path to the source image file. + /// - drivePath: The path to the target drive (e.g., "disk4"). + /// - Returns: An object conforming to Imager. + public static func createImager(imageFilePath: String, drivePath: String) -> Imager { + #if os(macOS) + return MacOSImager(imageFilePath: imageFilePath, drivePath: drivePath) + #elseif os(Linux) + // return LinuxImager(imageFilePath: imageFilePath, drivePath: drivePath) + fatalError("Linux implementation not yet available.") + #else + fatalError("Unsupported operating system.") + #endif + } + + /// Creates a platform-specific object conforming to the DiskLister protocol. + /// - Returns: An object conforming to DiskLister. + public static func createDiskLister() -> DiskLister { + #if os(macOS) + return MacOSDiskLister() + #elseif os(Linux) + // return LinuxDiskLister() + fatalError("Linux implementation not yet available.") + #else + fatalError("Unsupported operating system.") + #endif + } +} diff --git a/Sources/Imager/ImagerState.swift b/Sources/Imager/ImagerState.swift new file mode 100644 index 0000000..f275206 --- /dev/null +++ b/Sources/Imager/ImagerState.swift @@ -0,0 +1,52 @@ +/// Represents the current status of the imaging process. +/// +/// This enum tracks the lifecycle of an imaging operation, from initialization through +/// completion or failure. +public enum ImagerState: Sendable { + /// The imager is initialized but has not started imaging. + /// + /// This is the default state before `startImaging()` is called. + case idle + + /// The imaging process has started but has not yet completed or failed. + /// + /// This state indicates that the operation is active. + case imaging + + /// The imaging process is being cancelled. + case cancelling + + /// The imaging process completed successfully. + /// + /// This state indicates that the image file was written to the drive without errors. + case completed + + /// The imaging process failed. + /// + /// The associated value `ImagerError` provides details about the failure reason. + case failed(ImagerError) +} + +// MARK: - CustomStringConvertible + +extension ImagerState: CustomStringConvertible { + /// A user-friendly description of the current imaging state. + public var description: String { + switch self { + case .idle: + return "Ready to begin imaging" + + case .imaging: + return "Imaging in progress" + + case .cancelling: + return "Cancelling imaging process" + + case .completed: + return "Imaging completed successfully" + + case .failed(let error): + return "Imaging failed: \(error.localizedDescription)" + } + } +} \ No newline at end of file diff --git a/Sources/Imager/MacOS/MacOSDiskLister.swift b/Sources/Imager/MacOS/MacOSDiskLister.swift new file mode 100644 index 0000000..77ae6b6 --- /dev/null +++ b/Sources/Imager/MacOS/MacOSDiskLister.swift @@ -0,0 +1,156 @@ +#if os(macOS) +import Foundation +import DiskArbitration +import IOKit + +/// Concrete implementation of the Imager protocol for macOS. +public struct MacOSDiskLister: DiskLister, @unchecked Sendable { + + public init() {} + + public func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] { + // Get whole disk identifiers using diskutil list -plist + let listTask = Process() + listTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + listTask.arguments = ["list", "-plist"] + let listPipe = Pipe() + listTask.standardOutput = listPipe + + var wholeDiskIdentifiers: [String] = [] + var wholeDiskToPartitionMap: [String: String] = [:] + + do { + try listTask.run() + listTask.waitUntilExit() + let listData = listPipe.fileHandleForReading.readDataToEndOfFile() + + if listTask.terminationStatus != 0 { + throw ImagerError.driveDetectionError(reason: "diskutil list failed with status \(listTask.terminationStatus)") + } + guard let listPlist = try? PropertyListSerialization.propertyList(from: listData, options: [], format: nil) as? [String: Any], + let parsedWholeDisks = listPlist["WholeDisks"] as? [String], + let allDisksAndPartitions = listPlist["AllDisksAndPartitions"] as? [[String: Any]] else { + throw ImagerError.driveDetectionError(reason: "Error running or parsing diskutil list: Invalid output format or missing keys") + } + wholeDiskIdentifiers = parsedWholeDisks + + // Create a map from whole disk identifier (e.g., "disk28") to its first mountable partition identifier (e.g., "disk28s1") + for diskOrPartData in allDisksAndPartitions { + if let deviceID = diskOrPartData["DeviceIdentifier"] as? String, + wholeDiskIdentifiers.contains(deviceID), + let partitions = diskOrPartData["Partitions"] as? [[String: Any]], + let firstPartition = partitions.first?["DeviceIdentifier"] as? String + { + wholeDiskToPartitionMap[deviceID] = firstPartition + } + } + + } catch { + throw ImagerError.driveDetectionError(reason: "Failed to execute or process diskutil list: \(error.localizedDescription)") + } + + var drives: [Drive] = [] + let fileManager = FileManager.default + let currentDASession = DASessionCreate(kCFAllocatorDefault) + guard let currentDASession = currentDASession else { + throw ImagerError.driveDetectionError(reason: "Failed to create temporary DiskArbitration session") + } + + for wholeDiskIdentifier in wholeDiskIdentifiers { + let infoTask = Process() + infoTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + infoTask.arguments = ["info", "-plist", wholeDiskIdentifier] + let infoPipe = Pipe() + infoTask.standardOutput = infoPipe + + do { + try infoTask.run() + infoTask.waitUntilExit() + let infoData = infoPipe.fileHandleForReading.readDataToEndOfFile() + + if infoTask.terminationStatus != 0 { + continue + } + + guard let diskInfoPlist = try? PropertyListSerialization.propertyList(from: infoData, options: [], format: nil) as? [String: Any] else { + continue + } + + if onlyExternalDrives, !isSuitableExternal(diskInfo: diskInfoPlist) { + continue + } + + var volumeName: String? = "(No Name)" + var availableBytes: Int64? = nil + let bsdName = "/dev/\(wholeDiskIdentifier)" + + guard let partitionIdentifier = wholeDiskToPartitionMap[wholeDiskIdentifier] else { + volumeName = diskInfoPlist["VolumeName"] as? String ?? "(No Name)" + let capacity = diskInfoPlist["TotalSize"] as? Int64 ?? 0 + let drive = Drive( + path: bsdName, + capacity: capacity, + available: nil, + name: volumeName + ) + drives.append(drive) + continue + } + + let partitionBsdName = "/dev/\(partitionIdentifier)" + if let partitionDisk = DADiskCreateFromBSDName(kCFAllocatorDefault, currentDASession, partitionBsdName) { + if let partitionInfo = DADiskCopyDescription(partitionDisk) as? [String: Any] { + volumeName = partitionInfo[kDADiskDescriptionVolumeNameKey as String] as? String ?? volumeName + + if let mountPath = partitionInfo[kDADiskDescriptionVolumePathKey as String] as? URL { + do { + let attributes = try fileManager.attributesOfFileSystem(forPath: mountPath.path) + availableBytes = attributes[.systemFreeSize] as? Int64 + } catch { + + } + } else { + // print("Partition \(partitionIdentifier) is not mounted.") + } + } else { + // print("Failed to get description for partition \(partitionIdentifier).") + } + } else { + // print("Failed to create DADiskRef for partition \(partitionIdentifier).") + } + + let capacity = diskInfoPlist["TotalSize"] as? Int64 ?? 0 + let drive = Drive( + path: bsdName, + capacity: capacity, + available: availableBytes, + name: volumeName + ) + drives.append(drive) + + } catch { + continue + } + } + return drives.sorted { $0.path.localizedStandardCompare($1.path) == .orderedAscending } + } + + // ... isSuitableExternal (no changes needed) ... + private func isSuitableExternal(diskInfo: [String: Any]) -> Bool { + guard let isInternal = diskInfo["Internal"] as? Bool, !isInternal else { + return false + } + if diskInfo["VirtualOrPhysical"] as? String == "Virtual" { + return false + } + if diskInfo["SystemImage"] as? Bool == true { + return false + } + if (diskInfo["VolumeName"] as? String)?.contains("Recovery") == true { + // return false // Maybe allow imaging recovery? + } + return true + } +} + +#endif // os(macOS) \ No newline at end of file diff --git a/Sources/Imager/MacOS/MacOSImager.swift b/Sources/Imager/MacOS/MacOSImager.swift new file mode 100644 index 0000000..f8a1dcf --- /dev/null +++ b/Sources/Imager/MacOS/MacOSImager.swift @@ -0,0 +1,291 @@ +#if os(macOS) + + import Foundation + import DiskArbitration + import Darwin + + public struct MacOSImager: Imager, @unchecked Sendable { + public var imageFilePath: String + public var drivePath: String + private var fullDrivePath: String { + return "/dev/\(drivePath)" + } + + public init(imageFilePath: String, drivePath: String) { + self.imageFilePath = imageFilePath + self.drivePath = drivePath + } + + public func startImaging( + handler: @Sendable @escaping (_ progress: Foundation.Progress, _ error: ImagerError?) -> + Void + ) { + + // 1. check that the image file exists + guard FileManager.default.fileExists(atPath: imageFilePath) else { + handler(.init(), ImagerError.invalidImageFile(reason: "Image file does not exist")) + return + } + + // 2. check that the drive exists + guard FileManager.default.fileExists(atPath: fullDrivePath) else { + handler(.init(), ImagerError.invalidDrive(reason: "Drive does not exist or is not accessible")) + return + } + + let isMounted = self.isMounted(drivePath: drivePath) + // 3. If the drive is mounted, try to unmount it + if isMounted { + print("Drive is mounted, unmounting...") + // Try to unmount the drive + let unmountProcess = Process() + unmountProcess.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + unmountProcess.arguments = ["unmountDisk", drivePath] + + do { + try unmountProcess.run() + unmountProcess.waitUntilExit() + + // Check if unmount was successful + if unmountProcess.terminationStatus != 0 { + handler( + .init(), + ImagerError.permissionDenied(reason: "Failed to unmount drive") + ) + return + } + + // Double-check that the drive is actually unmounted now + if self.isMounted(drivePath: drivePath) { + handler( + .init(), + ImagerError.permissionDenied(reason: "Drive is still mounted after unmount attempt") + ) + return + } + + print("Drive successfully unmounted.") + } catch { + handler( + .init(), + ImagerError.permissionDenied( + reason: "Failed to unmount drive: \(error.localizedDescription)" + ) + ) + return + } + } + + // 4. Check permissions on the drive + do { + // We don't need to store the attributes, just check if we can access them + _ = try FileManager.default.attributesOfItem(atPath: fullDrivePath) + } catch { + handler( + .init(), + ImagerError.permissionDenied( + reason: "Cannot get attributes of drive: \(error.localizedDescription)" + ) + ) + return + } + + // 5. Check if we're running with sufficient privileges + if geteuid() != 0 { + handler( + .init(), + ImagerError.permissionDenied( + reason: "This operation requires administrative privileges. Please run with sudo." + ) + ) + return + } + + // 6. Get the image file size for progress tracking + let imageFileAttributes: [FileAttributeKey: Any] + do { + imageFileAttributes = try FileManager.default.attributesOfItem(atPath: imageFilePath) + } catch { + handler( + .init(), + ImagerError.invalidImageFile( + reason: "Cannot get attributes of image file: \(error.localizedDescription)" + ) + ) + return + } + + guard let imageFileSize = imageFileAttributes[.size] as? Int64, imageFileSize > 0 else { + handler( + .init(), + ImagerError.invalidImageFile(reason: "Cannot determine image file size") + ) + return + } + + // 7. start imaging + // Create a progress object to track and report imaging progress + let progress = Progress(totalUnitCount: imageFileSize) + progress.kind = .file + progress.fileOperationKind = .copying + + // Setup the dd command to perform the imaging + let ddProcess = Process() + ddProcess.executableURL = URL(fileURLWithPath: "/bin/dd") + ddProcess.arguments = [ + "if=\(imageFilePath)", + "of=\(fullDrivePath)", + "bs=1m", + "status=progress", + ] + + // Set up pipes for stdout and stderr + let outputPipe = Pipe() + let errorPipe = Pipe() + ddProcess.standardOutput = outputPipe + ddProcess.standardError = errorPipe + + // Monitor output to track progress + outputPipe.fileHandleForReading.readabilityHandler = { fileHandle in + let data = fileHandle.availableData + if !data.isEmpty, let output = String(data: data, encoding: .utf8) { + // Parse dd output for progress information + if let bytesTransferred = self.parseProgress(from: output) { + DispatchQueue.main.async { + // Update progress with the bytes transferred + progress.completedUnitCount = bytesTransferred + handler(progress, nil) + } + } + + // Print raw output for debugging (to standard error) + fputs("DD Output: \(output)\n", stderr) + } + } + + // Monitor for errors + errorPipe.fileHandleForReading.readabilityHandler = { fileHandle in + let data = fileHandle.availableData + if !data.isEmpty, let errorOutput = String(data: data, encoding: .utf8) { + if !errorOutput.isEmpty { + print("Error output: \(errorOutput)") + } + } + } + + // Handle process termination + ddProcess.terminationHandler = { process in + // Clean up file handles + outputPipe.fileHandleForReading.readabilityHandler = nil + errorPipe.fileHandleForReading.readabilityHandler = nil + + if process.terminationStatus == 0 { + progress.completedUnitCount = progress.totalUnitCount + handler(progress, nil) + } else { + handler( + progress, + ImagerError.processingInterrupted( + reason: "dd process exited with status \(process.terminationStatus)" + ) + ) + } + } + + do { + try ddProcess.run() + // Initial progress update + handler(progress, nil) + } catch { + handler( + progress, + ImagerError.processingInterrupted( + reason: "Failed to start dd process: \(error.localizedDescription)" + ) + ) + } + + } + + private func parseProgress(from output: String) -> Int64? { + // dd on macOS typically outputs progress in formats like: + // "123456789 bytes transferred in 0.012345 secs (123456789 bytes/sec)" + // or "1.2 GB copied, 30.1 s, 40.2 MB/s" + + // Look for patterns in dd output + let bytesPattern = "(\\d+)\\s+bytes\\s+(transferred|copied)" + let mbPattern = "(\\d+(\\.\\d+)?)\\s+[MG]B\\s+(transferred|copied)" + let gbPattern = "(\\d+(\\.\\d+)?)\\s+GB\\s+(transferred|copied)" + + // Try to match bytes first (most precise) + if let bytesMatch = output.range(of: bytesPattern, options: .regularExpression), + let bytesStr = output[bytesMatch].firstMatch(of: /(\d+)/), + let bytes = Int64(bytesStr.1) + { + return bytes + } + // Try to match MB + else if let mbMatch = output.range(of: mbPattern, options: .regularExpression), + let mbStr = output[mbMatch].firstMatch(of: /(\d+(\.\d+)?)/), + let mb = Double(mbStr.1) + { + // Convert MB to bytes (1 MB = 1,048,576 bytes) + return Int64(mb * 1_048_576) + } + // Try to match GB + else if let gbMatch = output.range(of: gbPattern, options: .regularExpression), + let gbStr = output[gbMatch].firstMatch(of: /(\d+(\.\d+)?)/), + let gb = Double(gbStr.1) + { + // Convert GB to bytes (1 GB = 1,073,741,824 bytes) + return Int64(gb * 1_073_741_824) + } + + // Alternative pattern for macOS dd status=progress output + // Example: "1234+0 records in\n1234+0 records out\n1234567890 bytes transferred in 123.456789 secs (123456789 bytes/sec)" + if let recordsOutMatch = output.range(of: "(\\d+)\\+\\d+\\s+records\\s+out", options: .regularExpression), + let recordsOutStr = output[recordsOutMatch].firstMatch(of: /(\d+)/), + let recordsOut = Int64(recordsOutStr.1) + { + // Each record is typically 512 bytes or the specified block size (bs) + // We're using bs=1m (1 MiB) in our dd command + return recordsOut * 1_048_576 + } + + return nil + } + + // MARK: - Helper Functions + + private func isMounted(drivePath: String) -> Bool { + guard let session = DASessionCreate(kCFAllocatorDefault) else { + fputs("Failed to create Disk Arbitration session.\n", stderr) + return false + } + + // We expect drivePath to be just the disk identifier (e.g., "disk28") + // No need to extract from URL, just use it directly + let bsdName = drivePath + + guard !bsdName.isEmpty, + let disk = DADiskCreateFromBSDName(kCFAllocatorDefault, session, bsdName) + else { + // Disk might not exist or path is invalid + print("Debug: Could not create disk reference for \(bsdName)") + return false + } + + // DADiskCopyDescription returns a CFDictionary?. Cast to Swift dictionary. + guard let description = DADiskCopyDescription(disk) as? [String: AnyObject] else { + // Failed to get description + print("Debug: Could not get disk description for \(bsdName)") + return false + } + + // Check if the volume path key exists. If it does, the volume is mounted. + // The value associated with kDADiskDescriptionVolumePathKey is a CFURLRef. + return description[kDADiskDescriptionVolumePathKey as String] != nil + } + } + +#endif // os(macOS) diff --git a/Sources/Imager/Protocols/DiskLister.swift b/Sources/Imager/Protocols/DiskLister.swift new file mode 100644 index 0000000..bee5431 --- /dev/null +++ b/Sources/Imager/Protocols/DiskLister.swift @@ -0,0 +1,15 @@ +import Foundation + +/// Defines the interface for listing disk drives available for imaging. +public protocol DiskLister: Sendable { + /// Retrieves a list of drives that can potentially be used as imaging targets. + /// + /// This function should identify physical drives connected to the system. + /// + /// - Parameter onlyExternalDrives: If `true`, only external drives (like USB drives) + /// are returned. If `false`, internal drives might also + /// be included. + /// - Returns: An array of `Drive` objects representing the available drives. + /// - Throws: An error if the drive list cannot be retrieved (e.g., permission issues). + func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] +} \ No newline at end of file diff --git a/Sources/Imager/Protocols/Imager.swift b/Sources/Imager/Protocols/Imager.swift new file mode 100644 index 0000000..59430f5 --- /dev/null +++ b/Sources/Imager/Protocols/Imager.swift @@ -0,0 +1,26 @@ +import Foundation + +/// Defines the interface for an object capable of writing an image file to a drive. +public protocol Imager: Sendable { + /// The path to the source image file. + var imageFilePath: String { get } + + /// The path to the target drive (e.g., "disk4" on macOS). + var drivePath: String { get } + + /// Initializes the Imager with the necessary paths. + /// - Parameters: + /// - imageFilePath: The full path to the source image file. + /// - drivePath: The identifier for the target drive (e.g., "disk4"). + init(imageFilePath: String, drivePath: String) + + /// Starts the imaging process. + /// + /// This method performs the actual writing of the image file to the target drive. + /// It should handle necessary pre-checks (like drive mounting status) and execute + /// the underlying imaging command (e.g., `dd`). + /// + /// - Parameter handler: A closure that is called periodically with progress updates + /// and upon completion or failure. + func startImaging(handler: @Sendable @escaping (_ progress: Progress, _ error: ImagerError?) -> Void) +} diff --git a/Sources/edge/Commands/Imager/ImagerCommand.swift b/Sources/edge/Commands/Imager/ImagerCommand.swift new file mode 100644 index 0000000..a625b34 --- /dev/null +++ b/Sources/edge/Commands/Imager/ImagerCommand.swift @@ -0,0 +1,14 @@ +import ArgumentParser +import Imager + +struct ImagerCommand: AsyncParsableCommand { + + static let configuration = CommandConfiguration( + commandName: "imager", + abstract: "Commands for listing available disks and writing OS images to drives.", + subcommands: [ + ImagerListDisksSubCommand.self, + ImagerImageDiskSubCommand.self + ] + ) +} diff --git a/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift b/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift new file mode 100644 index 0000000..04e2b08 --- /dev/null +++ b/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift @@ -0,0 +1,106 @@ +import ArgumentParser +import Foundation +import Imager + +struct ImagerImageDiskSubCommand: AsyncParsableCommand { + static let configuration = CommandConfiguration( + commandName: "image", + abstract: "Write an image file to a specified drive." + ) + + @Option(name: [.long, .customShort("s")], help: "The path to the source image file (.img, .zip, .tar).", completion: .file(extensions: ["img", "zip", "tar"])) + var source: String + + @Option(name: [.long, .customShort("t")], help: "The target drive path (e.g., disk4). Use 'list-disks' to find available drives.") + var target: String + + func run() async throws { + // Basic path validation + guard !source.isEmpty, !target.isEmpty else { + print("Error: Both source image path (--source) and target drive path (--target) must be provided.") + throw ExitCode.failure + } + + // Ensure target doesn't contain /dev/ prefix as MacOSImager expects raw disk name (e.g., disk4) + let drivePath = target.replacingOccurrences(of: "/dev/", with: "") + + print("Starting imaging process...") + print(" Source Image: \(source)") + print(" Target Drive: /dev/\(drivePath)") + print(" This may take some time...") + + let imager = ImagerFactory.createImager(imageFilePath: source, drivePath: drivePath) + + // Setup for progress display + print("\nProgress: ") + + // Create a task to handle the imaging process + let task = Task { + // Create an actor to safely handle the lastUpdateTime + actor ProgressState { + private var lastUpdateTime = Date() + + func shouldUpdate(now: Date, isFinished: Bool) -> Bool { + if isFinished || now.timeIntervalSince(lastUpdateTime) >= 0.5 { + lastUpdateTime = now + return true + } + return false + } + } + + let progressState = ProgressState() + + // Create a continuation to handle the async callback + return await withCheckedContinuation { continuation in + // The startImaging method uses a closure for progress/error handling + imager.startImaging { [progressState] progress, error in + if let error = error { + // Clear the current line before showing the error + print("\r\u{1B}[2K", terminator: "") + print("Error during imaging: \(error.description)") + continuation.resume(returning: false) + } else { + // Handle progress updates + Task { + let now = Date() + // Check if we should update the display + let shouldUpdate = await progressState.shouldUpdate(now: now, isFinished: progress.isFinished) + + if shouldUpdate { + // Calculate percentage + let percentage = Int(progress.fractionCompleted * 100) + + // Create progress bar + let barWidth = 50 + let completedWidth = Int(Double(barWidth) * progress.fractionCompleted) + let bar = String(repeating: "=", count: completedWidth) + + String(repeating: " ", count: barWidth - completedWidth) + + // Display progress + print("\r[\(bar)] \(percentage)%", terminator: "") + fflush(stdout) + + // Check for completion + if progress.isFinished { + print("\nImaging process completed successfully!") + continuation.resume(returning: true) + } + } + } + } + } + } + } + + print("\nImaging process initiated. Press Ctrl+C to cancel.") + + // Wait for the imaging process to complete + let success = await task.value + + // Exit with appropriate code + if !success { + throw ExitCode.failure + } + } +} \ No newline at end of file diff --git a/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift new file mode 100644 index 0000000..c5957aa --- /dev/null +++ b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift @@ -0,0 +1,59 @@ +import ArgumentParser +import Foundation +import Imager + +struct ImagerListDisksSubCommand: AsyncParsableCommand { + + static let configuration = CommandConfiguration( + commandName: "list-disks", + abstract: "Lists available drives that can be targeted for imaging." + ) + + @Flag(name: .long, help: "Include internal drives in the list.") + var all: Bool = false + + // Note: Timeout functionality is not implemented yet. + @Option(name: .long, help: "Timeout in seconds (functionality not yet implemented).") + var timeout: Int? + + func run() async throws { + let diskLister = ImagerFactory.createDiskLister() + try await MainActor.run { + try listAndPrintDisks(lister: diskLister, listAll: self.all) + } + } + + @MainActor + private func listAndPrintDisks(lister: DiskLister, listAll: Bool) throws { + let drives = try lister.availableDrivesToImage(onlyExternalDrives: !listAll) + + if drives.isEmpty { + print("No suitable drives found.") + fflush(stdout) + return + } + + let nameHeader = "Name" + let pathHeader = "Path" + let availableHeader = "Available" + let capacityHeader = "Capacity" + + let maxNameWidth = drives.map { $0.name?.count ?? "(No Name)".count }.max() ?? nameHeader.count + let maxPathWidth = drives.map { $0.path.count }.max() ?? pathHeader.count + let maxAvailableWidth = drives.map { $0.availableHumanReadable.count }.max() ?? availableHeader.count + let maxCapacityWidth = drives.map { $0.capacityHumanReadable.count }.max() ?? capacityHeader.count + + let header = "\(nameHeader.padding(toLength: maxNameWidth, withPad: " ", startingAt: 0)) \(pathHeader.padding(toLength: maxPathWidth, withPad: " ", startingAt: 0)) \(availableHeader.padding(toLength: maxAvailableWidth, withPad: " ", startingAt: 0)) \(capacityHeader.padding(toLength: maxCapacityWidth, withPad: " ", startingAt: 0))" + print(header) + print(String(repeating: "-", count: header.count)) + + for drive in drives { + let name = (drive.name ?? "(No Name)").padding(toLength: maxNameWidth, withPad: " ", startingAt: 0) + let path = drive.path.padding(toLength: maxPathWidth, withPad: " ", startingAt: 0) + let available = drive.availableHumanReadable.padding(toLength: maxAvailableWidth, withPad: " ", startingAt: 0) + let capacity = drive.capacityHumanReadable.padding(toLength: maxCapacityWidth, withPad: " ", startingAt: 0) + print("\(name) \(path) \(available) \(capacity)") + } + fflush(stdout) // Flush after printing the table + } +} \ No newline at end of file diff --git a/Sources/edge/Commands/ProgressBar/ProgressBar.swift b/Sources/edge/Commands/ProgressBar/ProgressBar.swift new file mode 100644 index 0000000..9249c8f --- /dev/null +++ b/Sources/edge/Commands/ProgressBar/ProgressBar.swift @@ -0,0 +1,100 @@ +import Foundation + +/// A simple CLI progress bar that displays progress using ASCII block characters. +public struct ProgressBar { + /// The width of the progress bar in characters. + private let width: Int + + /// The character used for filled portions of the progress bar. + private let filledChar: Character + + /// The character used for empty portions of the progress bar. + private let emptyChar: Character + + /// The left bracket character. + private let leftBracket: Character + + /// The right bracket character. + private let rightBracket: Character + + /// Creates a new progress bar with the specified configuration. + /// - Parameters: + /// - width: The width of the progress bar in characters. Default is 10. + /// - filledChar: The character used for filled portions. Default is "█". + /// - emptyChar: The character used for empty portions. Default is "░". + /// - leftBracket: The left bracket character. Default is "[". + /// - rightBracket: The right bracket character. Default is "]". + public init( + width: Int = 10, + filledChar: Character = "█", + emptyChar: Character = "░", + leftBracket: Character = "[", + rightBracket: Character = "]" + ) { + self.width = width + self.filledChar = filledChar + self.emptyChar = emptyChar + self.leftBracket = leftBracket + self.rightBracket = rightBracket + } + + /// Renders the progress bar as a string based on the given progress value. + /// - Parameter progress: A value between 0.0 and 1.0 representing the progress. + /// - Returns: A string representation of the progress bar. + public func render(progress: Double, additionalText: String? = nil) -> String { + let clampedProgress = min(1.0, max(0.0, progress)) + let filledWidth = Int(Double(width) * clampedProgress) + let emptyWidth = width - filledWidth + + let filledPart = String(repeating: filledChar, count: filledWidth) + let emptyPart = String(repeating: emptyChar, count: emptyWidth) + + return "\(leftBracket)\(filledPart)\(emptyPart)\(rightBracket) \(additionalText ?? "")" + } + + /// Renders the progress bar with a percentage value. + /// - Parameter progress: A value between 0.0 and 1.0 representing the progress. + /// - Returns: A string representation of the progress bar with percentage. + public func renderWithPercentage(progress: Double, additionalText: String? = nil) -> String { + let percentage = Int(min(1.0, max(0.0, progress)) * 100) + return "\(render(progress: progress)) \(percentage)% \(additionalText ?? "")" + } + + /// Updates the progress bar in place by clearing the current line and printing the new progress. + /// - Parameters: + /// - progress: A value between 0.0 and 1.0 representing the progress. + /// - showPercentage: Whether to show the percentage value. Default is true. + /// - stream: The output stream to write to. Default is stdout. + public func update( + progress: Double, + showPercentage: Bool = true, + stream: UnsafeMutablePointer = stdout, + additionalText: String? = nil + ) { + // Clear the current line and move cursor to beginning + fputs("\r\u{1B}[2K", stream) + + // Print the progress bar + let progressBar = showPercentage ? renderWithPercentage(progress: progress, additionalText: additionalText) : render(progress: progress, additionalText: additionalText) + fputs("\r\(progressBar)", stream) + fflush(stream) + } + + /// Completes the progress bar by setting it to 100% and adding a newline. + /// - Parameters: + /// - message: An optional message to display after completing the progress. + /// - stream: The output stream to write to. Default is stdout. + public func complete( + message: String? = nil, + stream: UnsafeMutablePointer = stdout, + additionalText: String? = nil + ) { + update(progress: 1.0, stream: stream, additionalText: additionalText) + + if let message = message { + fputs("\n\(message)\n", stream) + } else { + fputs("\n", stream) + } + } +} diff --git a/Sources/edge/Commands/ProgressBar/ProgressBarExample.swift b/Sources/edge/Commands/ProgressBar/ProgressBarExample.swift new file mode 100644 index 0000000..425ed95 --- /dev/null +++ b/Sources/edge/Commands/ProgressBar/ProgressBarExample.swift @@ -0,0 +1,82 @@ +import Foundation +import ArgumentParser + +/// Example functions demonstrating how to use the ProgressBar +public enum ProgressBarExample { + + /// Demonstrates a simple progress bar that updates at a fixed rate + public static func simpleExample(additionalText: String?) { + let progressBar = ProgressBar(width: 20) + + print("Simple progress bar example:") + if let text = additionalText { + print("Additional text: \(text)") + } + + // Simulate progress + for i in 0...10 { + let progress = Double(i) / 10.0 + progressBar.update(progress: progress, additionalText: additionalText) + Thread.sleep(forTimeInterval: 0.2) + } + + progressBar.complete(message: "Completed!") + } + + /// Demonstrates a progress bar with custom characters + public static func customExample(additionalText: String?) { + let progressBar = ProgressBar( + width: 20, + filledChar: "#", + emptyChar: "-", + leftBracket: "(", + rightBracket: ")" + ) + + print("Custom progress bar example:") + + // Simulate progress + for i in 0...20 { + let progress = Double(i) / 20.0 + progressBar.update(progress: progress, additionalText: additionalText) + Thread.sleep(forTimeInterval: 0.1) + } + + progressBar.complete(message: "Completed with custom style!") + } + + /// Demonstrates how to use the progress bar with a real task + public static func realTaskExample(additionalText: String?) { + let progressBar = ProgressBar(width: 30) + let totalItems = 100 + + print("Processing items:") + + // Simulate processing items with variable timing + for i in 1...totalItems { + // Simulate work with random duration + let workDuration = Double.random(in: 0.01...0.05) + Thread.sleep(forTimeInterval: workDuration) + + // Update progress + let progress = Double(i) / Double(totalItems) + progressBar.update(progress: progress, additionalText: additionalText) + } + + progressBar.complete(message: "All items processed successfully!") + } +} + +public struct ProgressBarExampleCommand: ParsableCommand { + public init() {} + + public static let configuration = CommandConfiguration( + commandName: "progress-bar-example", + abstract: "Example of how to use the ProgressBar" + ) + + public func run() { + // Just run the simple example without additional text + ProgressBarExample.simpleExample(additionalText: "Progress bar example") + } +} \ No newline at end of file diff --git a/Sources/edge/EdgeCLI.swift b/Sources/edge/EdgeCLI.swift index 706e58f..e309bea 100644 --- a/Sources/edge/EdgeCLI.swift +++ b/Sources/edge/EdgeCLI.swift @@ -7,7 +7,9 @@ struct EdgeCLI: AsyncParsableCommand { commandName: "edge", abstract: "Edge CLI", subcommands: [ - RunCommand.self + RunCommand.self, + ImagerCommand.self, + ProgressBarExampleCommand.self ] ) } diff --git a/imager_command_features.md b/imager_command_features.md new file mode 100644 index 0000000..30cc056 --- /dev/null +++ b/imager_command_features.md @@ -0,0 +1,15 @@ +# Imager Command Features + +## List Disk Subcommand + +1. `edge-cli imager list-disks` should list available drives to image. The timeout option will default to 3 seconds. +2. `edge-cli imager list-disks --watch` should list available drives to image and watch for changes. The user should able to press `q`, `esc` or `ctrl-c` to exit the watch mode. +3. `edge-cli imager list-disks --timeout ` should list available drives to image and prompt the user to select a drive to image. The command should exit after the specified timeout. +4. `edge-cli imager list-disks --all` should list all available drives, including internal ones. This can be used with `--watch` and `--timeout` options. + +## Image Subcommand + +1. `edge-cli imager image --source --target ` should image the specified drive with the specified image file. It will print progress to the console. If there is an error, it will print the `ImagerError` from the `Imager` package's description and exit with a non-zero code. +2. The progress bar should be displayed in the console with ASCII characters using █ for filled and ░ for empty. + +Example: [█████░░░░░] to represent 50% progress.