From 64c7a0dc2380fe2bd2a75ff4444ee7d68b9656a3 Mon Sep 17 00:00:00 2001 From: Zamderax Date: Thu, 17 Apr 2025 12:57:35 -0700 Subject: [PATCH 1/4] Add a MacOSImager --- .gitignore | 8 + Package.swift | 5 + Sources/Imager/Drive.swift | 42 ++++ Sources/Imager/Imager.swift | 47 ++++ Sources/Imager/ImagerError.swift | 85 +++++++ Sources/Imager/ImagerState.swift | 60 +++++ Sources/Imager/MacOSImager.swift | 381 +++++++++++++++++++++++++++++++ Sources/Imager/Progress.swift | 12 + 8 files changed, 640 insertions(+) create mode 100644 Sources/Imager/Drive.swift create mode 100644 Sources/Imager/Imager.swift create mode 100644 Sources/Imager/ImagerError.swift create mode 100644 Sources/Imager/ImagerState.swift create mode 100644 Sources/Imager/MacOSImager.swift create mode 100644 Sources/Imager/Progress.swift 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..a6e9890 100644 --- a/Package.swift +++ b/Package.swift @@ -55,5 +55,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..e1bc95d --- /dev/null +++ b/Sources/Imager/Drive.swift @@ -0,0 +1,42 @@ +#if canImport(FoundationEssentials) + import FoundationEssentials +#else + import Foundation +#endif + +struct Drive { + /// The path to the drive. + var path: String + /// The size of the drive in bytes. + var size: Int64 + /// The name of the drive. + var name: String? + + /// The size of the drive in a human-readable format. + var sizeHumanReadable: String { +#if canImport(Foundation) || canImport(FoundationEssentials) + let formatter = ByteCountFormatter() + formatter.allowedUnits = [.useAll] + formatter.countStyle = .file + formatter.includesUnit = true + formatter.isAdaptive = true + return formatter.string(fromByteCount: size) +#else + // Fallback implementation when ByteCountFormatter is not available + let kb = Double(size) / 1024.0 + if kb < 1024 { + return String(format: "%.1f KB", kb) + } + let mb = kb / 1024.0 + if mb < 1024 { + return String(format: "%.1f MB", mb) + } + let gb = mb / 1024.0 + if gb < 1024 { + return String(format: "%.1f GB", gb) + } + let tb = gb / 1024.0 + return String(format: "%.1f TB", tb) +#endif + } +} diff --git a/Sources/Imager/Imager.swift b/Sources/Imager/Imager.swift new file mode 100644 index 0000000..c1a79a2 --- /dev/null +++ b/Sources/Imager/Imager.swift @@ -0,0 +1,47 @@ +import Foundation + +/// Protocol defining the interface for an image writer. +/// +/// This protocol defines the requirements for a component that can write +/// an image file to a drive. +protocol Imager { + /// Path to the image file to be imaged. + var imageFilePath: String { get } + + /// Path to the drive to be imaged. + var drivePath: String { get } + + /// The current state of the imaging process. + var state: ImagerState { get } + + /// Get a list of available drives that can be imaged in alphabetical order. + /// + /// - Returns: An array of drive identifiers that can be used for imaging. + /// - Throws: An `ImagerError` if the drives cannot be enumerated. + func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [String] + + /// Begin imaging the drive. + /// + /// This method starts the process of writing the image file to the drive. + /// It may throw an `ImagerError` if the operation cannot be completed. + func startImaging() throws + + /// Stop the imaging process. + /// + /// This method stops the imaging process if it is currently running. + /// It may throw an `ImagerError` if the operation cannot be stopped. + func stopImaging() throws + + /// Register a handler to receive progress updates. + /// + /// - Parameter handler: A closure that will be called with progress updates. + /// The closure receives the current progress and an optional error if one occurred. + func progress(_ handler: @escaping (Progress, ImagerError?) -> Void) + + /// Initialize a new imager with the specified image file and drive paths. + /// + /// - Parameters: + /// - imageFilePath: The path to the image file to be written. + /// - drivePath: The path to the drive where the image will be written. + init(imageFilePath: String, drivePath: String) +} \ No newline at end of file diff --git a/Sources/Imager/ImagerError.swift b/Sources/Imager/ImagerError.swift new file mode 100644 index 0000000..866bc81 --- /dev/null +++ b/Sources/Imager/ImagerError.swift @@ -0,0 +1,85 @@ +/// 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. +enum ImagerError: Error { + /// 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) + + /// 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 .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/ImagerState.swift b/Sources/Imager/ImagerState.swift new file mode 100644 index 0000000..d1c7f57 --- /dev/null +++ b/Sources/Imager/ImagerState.swift @@ -0,0 +1,60 @@ +/// Represents the current state of the imaging process. +/// +/// This enum tracks the lifecycle of an imaging operation, from initialization through +/// completion or failure, including progress updates during the process. +enum ImagerState { + /// The imager is initialized but has not started imaging. + /// + /// This is the initial state before the imaging process begins. + case idle + + /// The imaging process is actively running. + /// + /// This state indicates that data is being written to the target drive. + case imaging + + /// The imaging process has successfully completed. + /// + /// This state indicates that all data has been successfully written to the target drive. + case completed + + /// The imaging process has failed. + /// + /// This state includes an associated `ImagerError` value that provides + /// detailed information about the cause of the failure. + case failed(ImagerError) + + /// A progress update for the ongoing imaging process. + /// + /// This state includes a `Progress` value that contains information about + /// the current progress of the imaging operation, including the total bytes, + /// completed bytes, and percentage completion. + case progress(Progress) +} + +// 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 .completed: + return "Imaging completed successfully" + + case .failed(let error): + return "Imaging failed: \(error)" + + case .progress(let progress): + let percentage = Int(progress.percentage * 100) + let completedMB = Double(progress.completedBytes) / 1_048_576 + let totalMB = Double(progress.totalBytes) / 1_048_576 + return "Progress: \(percentage)% (\(String(format: "%.1f", completedMB)) MB of \(String(format: "%.1f", totalMB)) MB)" + } + } +} \ No newline at end of file diff --git a/Sources/Imager/MacOSImager.swift b/Sources/Imager/MacOSImager.swift new file mode 100644 index 0000000..e374bf4 --- /dev/null +++ b/Sources/Imager/MacOSImager.swift @@ -0,0 +1,381 @@ +#if os(macOS) +import Foundation +import DiskArbitration + +/// A macOS implementation of the `Imager` protocol. +/// +/// This class provides functionality for writing image files to drives on macOS systems +/// using native APIs. +@available(macOS 10.15, *) +class MacOSImager: Imager, @unchecked Sendable { + /// Path to the image file to be imaged. + let imageFilePath: String + + /// Path to the drive to be imaged. + let drivePath: String + + /// The current state of the imaging process. + private(set) var state: ImagerState = .idle + + /// Progress handler closure. + private var progressHandler: ((Progress, ImagerError?) -> Void)? + + /// File handle for the image file. + private var imageFileHandle: FileHandle? + + /// File handle for the drive. + private var driveFileHandle: FileHandle? + + /// Buffer size for reading/writing operations. + private let bufferSize = 1024 * 1024 // 1MB buffer + + /// Get a list of available drives that can be imaged in alphabetical order. + /// + /// - Parameter onlyExternalDrives: If true, only external drives will be returned. + /// - Returns: An array of drive identifiers that can be used for imaging. + /// - Throws: An `ImagerError` if the drives cannot be enumerated. + func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [String] { + // Get list of all physical drives from the system + let task = Process() + task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + task.arguments = ["list", "-plist", "physical"] + + let pipe = Pipe() + task.standardOutput = pipe + + do { + try task.run() + let data = pipe.fileHandleForReading.readDataToEndOfFile() + task.waitUntilExit() + + if task.terminationStatus == 0, let plist = try PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], + let allDisks = plist["WholeDisks"] as? [String] { + + if onlyExternalDrives { + // Filter to only include external drives + var externalDisks: [String] = [] + var errors: [String] = [] + + for disk in allDisks { + do { + if try isExternalDrive(disk) { + externalDisks.append(disk) + } + } catch let error as ImagerError { + // Collect errors but continue processing other disks + errors.append("\(disk): \(error.description)") + } catch { + errors.append("\(disk): \(error.localizedDescription)") + } + } + + // If we couldn't process any disks due to errors, throw an error + if externalDisks.isEmpty && !errors.isEmpty { + throw ImagerError.driveDetectionError(reason: "Failed to detect external drives: \(errors.joined(separator: "; "))") + } + + return externalDisks.sorted() + } else { + // Return all drives + return allDisks.sorted() + } + } else { + throw ImagerError.driveDetectionError(reason: "Failed to parse disk information") + } + } catch let error as ImagerError { + throw error + } catch { + throw ImagerError.driveDetectionError(reason: "Error getting available drives: \(error.localizedDescription)") + } + } + + /// Determines if a disk is an external drive. + /// + /// - Parameter diskName: The name of the disk to check. + /// - Returns: `true` if the disk is an external drive, `false` otherwise. + /// - Throws: An `ImagerError` if there is an error checking the drive. + private func isExternalDrive(_ diskName: String) throws -> Bool { + let task = Process() + task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + task.arguments = ["info", "-plist", diskName] + + let pipe = Pipe() + task.standardOutput = pipe + + do { + try task.run() + let data = pipe.fileHandleForReading.readDataToEndOfFile() + task.waitUntilExit() + + if task.terminationStatus == 0, let plist = try PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any] { + // Check if it's external + if let isExternal = plist["External"] as? Bool, isExternal { + // Check if it's ejectable (typically means removable) + if let isEjectable = plist["Ejectable"] as? Bool, isEjectable { + // Check if it's not the system disk + if let isSystemDisk = plist["SystemImage"] as? Bool, !isSystemDisk { + // Additional safety check: make sure it's not the boot drive + if let volumeName = plist["VolumeName"] as? String, volumeName != "Macintosh HD" { + return true + } + } + } + } + + // Not an external drive based on our criteria + return false + } else { + throw ImagerError.driveDetectionError(reason: "Failed to get disk information for \(diskName)") + } + } catch let error as ImagerError { + throw error + } catch { + throw ImagerError.driveDetectionError(reason: "Error checking if \(diskName) is external: \(error.localizedDescription)") + } + } + + /// Initialize a new imager with the specified image file and drive paths. + /// + /// - Parameters: + /// - imageFilePath: The path to the image file to be written. + /// - drivePath: The path to the drive where the image will be written. + required init(imageFilePath: String, drivePath: String) { + self.imageFilePath = imageFilePath + self.drivePath = drivePath + } + + /// Begin imaging the drive. + /// + /// This method starts the process of writing the image file to the drive. + /// It may throw an `ImagerError` if the operation cannot be completed. + func startImaging() throws { + // Validate image file + guard FileManager.default.fileExists(atPath: imageFilePath) else { + state = .failed(.invalidImageFile(reason: "Image file does not exist")) + throw ImagerError.invalidImageFile(reason: "Image file does not exist") + } + + // Check file extension + let fileExtension = URL(fileURLWithPath: imageFilePath).pathExtension.lowercased() + guard ["zip", "tar", "img"].contains(fileExtension) else { + let reason = "Unsupported file format: \(fileExtension). Only zip, tar, and img formats are supported." + state = .failed(.invalidImageFile(reason: reason)) + throw ImagerError.invalidImageFile(reason: reason) + } + + // Validate drive + let driveURL = URL(fileURLWithPath: "/dev/\(drivePath)") + let drivePath = driveURL.path + + guard FileManager.default.fileExists(atPath: drivePath) else { + let reason = "Drive does not exist at path: \(drivePath)" + state = .failed(.invalidDrive(reason: reason)) + throw ImagerError.invalidDrive(reason: reason) + } + + // Verify it's an external drive for safety + do { + guard try isExternalDrive(self.drivePath) else { + let reason = "Selected drive is not an external drive. For safety, only external drives can be imaged." + state = .failed(.invalidDrive(reason: reason)) + throw ImagerError.invalidDrive(reason: reason) + } + } catch { + if let imagerError = error as? ImagerError { + state = .failed(imagerError) + throw imagerError + } else { + let imagerError = ImagerError.driveDetectionError(reason: "Failed to verify if drive is external: \(error.localizedDescription)") + state = .failed(imagerError) + throw imagerError + } + } + + // Get image file size + let imageAttributes = try FileManager.default.attributesOfItem(atPath: imageFilePath) + guard let imageSize = imageAttributes[.size] as? Int64 else { + let reason = "Could not determine image file size" + state = .failed(.invalidImageFile(reason: reason)) + throw ImagerError.invalidImageFile(reason: reason) + } + + // Get drive size + let driveSize = try getDriveSize(drive: drivePath) + + // Check if image fits on drive + if imageSize > driveSize { + state = .failed(.imageTooLargeForDrive(imageSize: UInt64(imageSize), driveSize: UInt64(driveSize))) + throw ImagerError.imageTooLargeForDrive(imageSize: UInt64(imageSize), driveSize: UInt64(driveSize)) + } + + // Open file handles + do { + imageFileHandle = try FileHandle(forReadingFrom: URL(fileURLWithPath: imageFilePath)) + driveFileHandle = try FileHandle(forWritingTo: driveURL) + } catch { + let reason = "Failed to open file handles: \(error.localizedDescription)" + state = .failed(.permissionDenied(reason: reason)) + throw ImagerError.permissionDenied(reason: reason) + } + + // Start imaging process + state = .imaging + + // Create a background task for imaging + Task { + await performImaging(imageSize: imageSize) + } + } + + /// Stop the imaging process. + /// + /// This method stops the imaging process if it is currently running. + /// It may throw an `ImagerError` if the operation cannot be stopped. + func stopImaging() throws { + if case .imaging = state { + state = .failed(.processingInterrupted(reason: "Imaging process was stopped by user")) + + // Close file handles + if let handle = imageFileHandle { + try handle.close() + imageFileHandle = nil + } + + if let handle = driveFileHandle { + try handle.close() + driveFileHandle = nil + } + } else { + throw ImagerError.unknown(reason: "Cannot stop imaging: No imaging process is currently running") + } + } + + /// Register a handler to receive progress updates. + /// + /// - Parameter handler: A closure that will be called with progress updates. + /// The closure receives the current progress and an optional error if one occurred. + func progress(_ handler: @escaping (Progress, ImagerError?) -> Void) { + progressHandler = handler + } + + // MARK: - Private Methods + + /// Get the size of a drive in bytes. + /// + /// - Parameter drive: The path to the drive. + /// - Returns: The size of the drive in bytes. + /// - Throws: An `ImagerError` if the drive size cannot be determined. + private func getDriveSize(drive: String) throws -> Int64 { + let task = Process() + task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + task.arguments = ["info", "-plist", drive] + + let pipe = Pipe() + task.standardOutput = pipe + + do { + try task.run() + let data = pipe.fileHandleForReading.readDataToEndOfFile() + task.waitUntilExit() + + if task.terminationStatus == 0, let plist = try PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], + let size = plist["TotalSize"] as? Int64 { + return size + } else { + throw ImagerError.invalidDrive(reason: "Could not determine drive size") + } + } catch let error as ImagerError { + throw error + } catch { + throw ImagerError.invalidDrive(reason: "Error getting drive size: \(error.localizedDescription)") + } + } + + /// Perform the actual imaging process. + /// + /// - Parameter imageSize: The size of the image file in bytes. + private func performImaging(imageSize: Int64) async { + guard let imageFileHandle = imageFileHandle, let driveFileHandle = driveFileHandle else { + let error = ImagerError.unknown(reason: "File handles not initialized") + state = .failed(error) + progressHandler?(Progress(totalBytes: imageSize, completedBytes: 0), error) + return + } + + var bytesWritten: Int64 = 0 + var lastProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) + + do { + // Reset file positions + try imageFileHandle.seek(toOffset: 0) + try driveFileHandle.seek(toOffset: 0) + + // Read and write in chunks + while bytesWritten < imageSize { + // Check if we should continue + if case .imaging = state { + // Continue imaging + } else { + break + } + + // Read a chunk from the image file + let data = imageFileHandle.readData(ofLength: bufferSize) + if data.isEmpty { + break + } + + // Write the chunk to the drive + try driveFileHandle.write(contentsOf: data) + + // Update progress + bytesWritten += Int64(data.count) + let currentProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) + + // Only update if progress has changed significantly (1% or more) + if currentProgress.percentage - lastProgress.percentage >= 0.01 { + lastProgress = currentProgress + state = .progress(currentProgress) + progressHandler?(currentProgress, nil) + } + } + + // Ensure all data is written to disk + try driveFileHandle.synchronize() + + // Update final progress + let finalProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) + state = .progress(finalProgress) + progressHandler?(finalProgress, nil) + + // Set state to completed + state = .completed + + } catch { + let imagerError = ImagerError.processingInterrupted(reason: error.localizedDescription) + state = .failed(imagerError) + progressHandler?(lastProgress, imagerError) + } + + // Close file handles + try? imageFileHandle.close() + try? driveFileHandle.close() + + // Reset file handles + self.imageFileHandle = nil + self.driveFileHandle = nil + } + + /// Cancel the imaging process. + @available(*, deprecated, message: "Use stopImaging() instead") + func cancel() { + try? stopImaging() + } + + deinit { + // Ensure file handles are closed + try? imageFileHandle?.close() + try? driveFileHandle?.close() + } +} +#endif \ No newline at end of file diff --git a/Sources/Imager/Progress.swift b/Sources/Imager/Progress.swift new file mode 100644 index 0000000..653ab63 --- /dev/null +++ b/Sources/Imager/Progress.swift @@ -0,0 +1,12 @@ +struct Progress { + /// The total number of bytes to be written. + var totalBytes: Int64 + + /// The number of bytes that have been written. + var completedBytes: Int64 + + /// The percentage of the image that has been written. + var percentage: Double { + return Double(completedBytes) / Double(totalBytes) + } +} \ No newline at end of file From e2d31dac89377da758583f76b34a02c83f0cdcdf Mon Sep 17 00:00:00 2001 From: Zamderax Date: Thu, 17 Apr 2025 16:26:54 -0700 Subject: [PATCH 2/4] ability to list disks --- Package.swift | 1 + Sources/Imager/Drive.swift | 54 ++- Sources/Imager/Imager.swift | 3 +- Sources/Imager/ImagerError.swift | 4 +- Sources/Imager/ImagerState.swift | 32 +- Sources/Imager/MacOSImager.swift | 336 ++++++++++-------- Sources/Imager/Progress.swift | 8 +- .../edge/Commands/Imager/ImagerCommand.swift | 14 + .../Imager/ImagerImageDiskSubCommand.swift | 50 +++ .../Imager/ImagerListDiskSubCommand.swift | 52 +++ Sources/edge/EdgeCLI.swift | 3 +- imager_command_features.md | 8 + 12 files changed, 383 insertions(+), 182 deletions(-) create mode 100644 Sources/edge/Commands/Imager/ImagerCommand.swift create mode 100644 Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift create mode 100644 Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift create mode 100644 imager_command_features.md diff --git a/Package.swift b/Package.swift index a6e9890..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") diff --git a/Sources/Imager/Drive.swift b/Sources/Imager/Drive.swift index e1bc95d..95c2485 100644 --- a/Sources/Imager/Drive.swift +++ b/Sources/Imager/Drive.swift @@ -4,26 +4,60 @@ import Foundation #endif -struct Drive { +/// A structure representing a drive. +public struct Drive { /// The path to the drive. - var path: String - /// The size of the drive in bytes. - var size: Int64 + 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. - var name: String? + public var name: String? - /// The size of the drive in a human-readable format. - var sizeHumanReadable: String { + /// The total capacity of the drive in a human-readable format. + public var capacityHumanReadable: String { #if canImport(Foundation) || canImport(FoundationEssentials) let formatter = ByteCountFormatter() formatter.allowedUnits = [.useAll] formatter.countStyle = .file formatter.includesUnit = true formatter.isAdaptive = true - return formatter.string(fromByteCount: size) + return formatter.string(fromByteCount: capacity) #else - // Fallback implementation when ByteCountFormatter is not available - let kb = Double(size) / 1024.0 + // Basic fallback for non-Apple platforms + let kb = Double(capacity) / 1024.0 + if kb < 1024 { + return String(format: "%.1f KB", kb) + } + let mb = kb / 1024.0 + if mb < 1024 { + return String(format: "%.1f MB", mb) + } + let gb = mb / 1024.0 + if gb < 1024 { + return String(format: "%.1f GB", gb) + } + let tb = gb / 1024.0 + return String(format: "%.1f TB", tb) +#endif + } + + /// 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" + } +#if canImport(Foundation) || canImport(FoundationEssentials) + let formatter = ByteCountFormatter() + formatter.allowedUnits = [.useAll] + formatter.countStyle = .file + formatter.includesUnit = true + formatter.isAdaptive = true + return formatter.string(fromByteCount: availableSpace) +#else + // Basic fallback for non-Apple platforms + let kb = Double(availableSpace) / 1024.0 if kb < 1024 { return String(format: "%.1f KB", kb) } diff --git a/Sources/Imager/Imager.swift b/Sources/Imager/Imager.swift index c1a79a2..9452a2b 100644 --- a/Sources/Imager/Imager.swift +++ b/Sources/Imager/Imager.swift @@ -16,9 +16,10 @@ protocol Imager { /// Get a list of available drives that can be imaged in alphabetical order. /// + /// - Parameter onlyExternalDrives: If true, only external drives will be returned. /// - Returns: An array of drive identifiers that can be used for imaging. /// - Throws: An `ImagerError` if the drives cannot be enumerated. - func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [String] + func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] /// Begin imaging the drive. /// diff --git a/Sources/Imager/ImagerError.swift b/Sources/Imager/ImagerError.swift index 866bc81..770da1f 100644 --- a/Sources/Imager/ImagerError.swift +++ b/Sources/Imager/ImagerError.swift @@ -2,7 +2,7 @@ /// /// These errors represent various failure scenarios that might occur when /// attempting to write an image file to a drive. -enum ImagerError: Error { +public enum ImagerError: Error { /// The provided image file is invalid. /// /// This error occurs when the image file cannot be read or is in an unsupported format. @@ -55,7 +55,7 @@ extension ImagerError: CustomStringConvertible { public var description: String { switch self { case .invalidImageFile(let reason): - return "Invalid image file: \(reason). Only ZIP, TAR, and IMG formats are supported." + return "Invalid Image File: \(reason). Only ZIP, TAR, and IMG formats are supported." case .invalidDrive(let reason): return "Invalid drive: \(reason)" diff --git a/Sources/Imager/ImagerState.swift b/Sources/Imager/ImagerState.swift index d1c7f57..3f59961 100644 --- a/Sources/Imager/ImagerState.swift +++ b/Sources/Imager/ImagerState.swift @@ -1,35 +1,33 @@ -/// Represents the current state of the imaging process. +/// Represents the current status of the imaging process. /// /// This enum tracks the lifecycle of an imaging operation, from initialization through /// completion or failure, including progress updates during the process. -enum ImagerState { +public enum ImagerState { /// The imager is initialized but has not started imaging. /// - /// This is the initial state before the imaging process begins. + /// This is the default state before `startImaging()` is called. case idle - /// The imaging process is actively running. + /// The imaging process is currently in progress. /// - /// This state indicates that data is being written to the target drive. + /// The associated value `Progress` provides details about the current progress, + /// such as bytes written and total bytes. + case progress(Progress) + + /// 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 has successfully completed. + /// The imaging process completed successfully. /// - /// This state indicates that all data has been successfully written to the target drive. + /// This state indicates that the image file was written to the drive without errors. case completed - /// The imaging process has failed. + /// The imaging process failed. /// - /// This state includes an associated `ImagerError` value that provides - /// detailed information about the cause of the failure. + /// The associated value `ImagerError` provides details about the failure reason. case failed(ImagerError) - - /// A progress update for the ongoing imaging process. - /// - /// This state includes a `Progress` value that contains information about - /// the current progress of the imaging operation, including the total bytes, - /// completed bytes, and percentage completion. - case progress(Progress) } // MARK: - CustomStringConvertible diff --git a/Sources/Imager/MacOSImager.swift b/Sources/Imager/MacOSImager.swift index e374bf4..a3ce85e 100644 --- a/Sources/Imager/MacOSImager.swift +++ b/Sources/Imager/MacOSImager.swift @@ -7,15 +7,15 @@ import DiskArbitration /// This class provides functionality for writing image files to drives on macOS systems /// using native APIs. @available(macOS 10.15, *) -class MacOSImager: Imager, @unchecked Sendable { +public class MacOSImager: Imager, @unchecked Sendable { /// Path to the image file to be imaged. - let imageFilePath: String + public let imageFilePath: String /// Path to the drive to be imaged. - let drivePath: String + public let drivePath: String /// The current state of the imaging process. - private(set) var state: ImagerState = .idle + public private(set) var state: ImagerState = .idle /// Progress handler closure. private var progressHandler: ((Progress, ImagerError?) -> Void)? @@ -31,130 +31,157 @@ class MacOSImager: Imager, @unchecked Sendable { /// Get a list of available drives that can be imaged in alphabetical order. /// - /// - Parameter onlyExternalDrives: If true, only external drives will be returned. - /// - Returns: An array of drive identifiers that can be used for imaging. - /// - Throws: An `ImagerError` if the drives cannot be enumerated. - func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [String] { - // Get list of all physical drives from the system - let task = Process() - task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") - task.arguments = ["list", "-plist", "physical"] - - let pipe = Pipe() - task.standardOutput = pipe - + /// - Parameter onlyExternalDrives: If `true`, only external drives are returned. + /// - Returns: An array of `Drive` objects representing the available drives. + /// - Throws: An `ImagerError` if there is an error getting or parsing drive information. + public func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] { + // Step 1: Get the list of whole disk identifiers and map whole disks to their mountable partitions + let listTask = Process() + listTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + listTask.arguments = ["list", "-plist"] + let listPipe = Pipe() + listTask.standardOutput = listPipe + + var wholeDiskIdentifiers: [String] = [] do { - try task.run() - let data = pipe.fileHandleForReading.readDataToEndOfFile() - task.waitUntilExit() - - if task.terminationStatus == 0, let plist = try PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], - let allDisks = plist["WholeDisks"] as? [String] { - - if onlyExternalDrives { - // Filter to only include external drives - var externalDisks: [String] = [] - var errors: [String] = [] - - for disk in allDisks { - do { - if try isExternalDrive(disk) { - externalDisks.append(disk) - } - } catch let error as ImagerError { - // Collect errors but continue processing other disks - errors.append("\(disk): \(error.description)") - } catch { - errors.append("\(disk): \(error.localizedDescription)") + try listTask.run() + let listData = listPipe.fileHandleForReading.readDataToEndOfFile() + listTask.waitUntilExit() + guard listTask.terminationStatus == 0 else { + 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 wholeDisksIdentifiers = 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 = wholeDisksIdentifiers + + // Create a map from whole disk identifier (e.g., "disk28") to its first mountable partition identifier (e.g., "disk28s1") + var wholeDiskToPartitionMap: [String: String] = [:] + for diskEntry in allDisksAndPartitions { + if let wholeDiskId = diskEntry["DeviceIdentifier"] as? String, + let partitions = diskEntry["Partitions"] as? [[String: Any]] { + for partition in partitions { + if let partitionId = partition["DeviceIdentifier"] as? String, + let mountPoint = partition["MountPoint"] as? String, + !mountPoint.isEmpty { + wholeDiskToPartitionMap[wholeDiskId] = partitionId + break // Found the first mountable partition for this disk } } - - // If we couldn't process any disks due to errors, throw an error - if externalDisks.isEmpty && !errors.isEmpty { - throw ImagerError.driveDetectionError(reason: "Failed to detect external drives: \(errors.joined(separator: "; "))") - } - - return externalDisks.sorted() - } else { - // Return all drives - return allDisks.sorted() } - } else { - throw ImagerError.driveDetectionError(reason: "Failed to parse disk information") } - } catch let error as ImagerError { - throw error - } catch { - throw ImagerError.driveDetectionError(reason: "Error getting available drives: \(error.localizedDescription)") - } - } - - /// Determines if a disk is an external drive. - /// - /// - Parameter diskName: The name of the disk to check. - /// - Returns: `true` if the disk is an external drive, `false` otherwise. - /// - Throws: An `ImagerError` if there is an error checking the drive. - private func isExternalDrive(_ diskName: String) throws -> Bool { - let task = Process() - task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") - task.arguments = ["info", "-plist", diskName] - - let pipe = Pipe() - task.standardOutput = pipe - - do { - try task.run() - let data = pipe.fileHandleForReading.readDataToEndOfFile() - task.waitUntilExit() - - if task.terminationStatus == 0, let plist = try PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any] { - // Check if it's external - if let isExternal = plist["External"] as? Bool, isExternal { - // Check if it's ejectable (typically means removable) - if let isEjectable = plist["Ejectable"] as? Bool, isEjectable { - // Check if it's not the system disk - if let isSystemDisk = plist["SystemImage"] as? Bool, !isSystemDisk { - // Additional safety check: make sure it's not the boot drive - if let volumeName = plist["VolumeName"] as? String, volumeName != "Macintosh HD" { - return true - } - } + + // Step 2: Get detailed info, filter, and find available space + var availableDrives: [Drive] = [] + guard let session = DASessionCreate(kCFAllocatorDefault) else { + throw ImagerError.driveDetectionError(reason: "Failed to create DiskArbitration session") + } + + for diskIdentifier in wholeDiskIdentifiers { + let path = "/dev/\(diskIdentifier)" + let infoTask = Process() + infoTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + infoTask.arguments = ["info", "-plist", path] + let infoPipe = Pipe() + infoTask.standardOutput = infoPipe + + do { + try infoTask.run() + let infoData = infoPipe.fileHandleForReading.readDataToEndOfFile() + infoTask.waitUntilExit() + + guard infoTask.terminationStatus == 0, + let infoPlist = try? PropertyListSerialization.propertyList(from: infoData, options: [], format: nil) as? [String: Any] else { + continue } + + // --- Get Base Info & Filter using diskutil info on Whole Disk (diskX) --- + let isInternal = infoPlist["Internal"] as? Bool ?? false + let isSystemImage = infoPlist["SystemImage"] as? Bool ?? false + let virtualOrPhysical = infoPlist["VirtualOrPhysical"] as? String ?? "Physical" + let isVirtualDevice = (virtualOrPhysical == "Virtual") + + // Apply filtering first + let isSuitableExternal = !isInternal && !isSystemImage && !isVirtualDevice + if onlyExternalDrives && !isSuitableExternal { + continue // Skip this whole disk if filtering for external and it doesn't match + } + + // --- Get Mount Point, Available Space, and Name (via Partition diskXsY) --- + var available: Int64? = nil + var volumeNameFromPartition: String? = nil // Store name found from partition + let capacity = infoPlist["TotalSize"] as? Int64 ?? 0 // Capacity comes from whole disk + + // Find the partition identifier associated with this whole disk + if let partitionIdentifier = wholeDiskToPartitionMap[diskIdentifier] { + let partitionPath = "/dev/\(partitionIdentifier)" + + if let disk = DADiskCreateFromBSDName(kCFAllocatorDefault, session, partitionPath) { + if let diskInfoDA = DADiskCopyDescription(disk) as? [String: Any] { + // Attempt to get name from partition's DA description + volumeNameFromPartition = diskInfoDA[kDADiskDescriptionVolumeNameKey as String] as? String + + if let mountURL = diskInfoDA[kDADiskDescriptionVolumePathKey as String] as? URL { + let mountPoint = mountURL.path + if !mountPoint.isEmpty { + // Extra check: Don't include root if filtering external + if !(onlyExternalDrives && mountPoint == "/") { + do { + let attributes = try FileManager.default.attributesOfFileSystem(forPath: mountPoint) + available = attributes[.systemFreeSize] as? Int64 + } catch { + available = nil // Failed FileManager + } + } // else: Skip root disk if filtering + } // else: DA gave empty mount point for partition + } // else: DA description for partition missing mount point + } // else: Failed DA CopyDescription for partition + } // else: Failed DA Create Disk for partition + } // else: No mountable partition found for this whole disk in the map + + // Only add the drive if it passed the initial filter (or we aren't filtering) + // Use the name found from the partition if available + let drive = Drive(path: path, capacity: capacity, available: available, name: volumeNameFromPartition) + availableDrives.append(drive) + + } catch { + // Error running diskutil info process for this disk + continue } - - // Not an external drive based on our criteria - return false - } else { - throw ImagerError.driveDetectionError(reason: "Failed to get disk information for \(diskName)") } + + // Sort drives by path for consistent ordering + return availableDrives.sorted { $0.path < $1.path } } catch let error as ImagerError { throw error } catch { - throw ImagerError.driveDetectionError(reason: "Error checking if \(diskName) is external: \(error.localizedDescription)") + throw ImagerError.driveDetectionError(reason: "Error running or parsing diskutil list: \(error.localizedDescription)") } } - + /// Initialize a new imager with the specified image file and drive paths. /// /// - Parameters: /// - imageFilePath: The path to the image file to be written. /// - drivePath: The path to the drive where the image will be written. - required init(imageFilePath: String, drivePath: String) { + public required init(imageFilePath: String, drivePath: String) { self.imageFilePath = imageFilePath self.drivePath = drivePath } - + /// Begin imaging the drive. /// /// This method starts the process of writing the image file to the drive. /// It may throw an `ImagerError` if the operation cannot be completed. - func startImaging() throws { + public func startImaging() throws { // Validate image file guard FileManager.default.fileExists(atPath: imageFilePath) else { state = .failed(.invalidImageFile(reason: "Image file does not exist")) throw ImagerError.invalidImageFile(reason: "Image file does not exist") } - + // Check file extension let fileExtension = URL(fileURLWithPath: imageFilePath).pathExtension.lowercased() guard ["zip", "tar", "img"].contains(fileExtension) else { @@ -162,35 +189,50 @@ class MacOSImager: Imager, @unchecked Sendable { state = .failed(.invalidImageFile(reason: reason)) throw ImagerError.invalidImageFile(reason: reason) } - + // Validate drive let driveURL = URL(fileURLWithPath: "/dev/\(drivePath)") let drivePath = driveURL.path - + guard FileManager.default.fileExists(atPath: drivePath) else { let reason = "Drive does not exist at path: \(drivePath)" state = .failed(.invalidDrive(reason: reason)) throw ImagerError.invalidDrive(reason: reason) } - + // Verify it's an external drive for safety do { - guard try isExternalDrive(self.drivePath) else { - let reason = "Selected drive is not an external drive. For safety, only external drives can be imaged." - state = .failed(.invalidDrive(reason: reason)) - throw ImagerError.invalidDrive(reason: reason) - } - } catch { - if let imagerError = error as? ImagerError { - state = .failed(imagerError) - throw imagerError + // Get list of all physical drives from the system + let task = Process() + task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + task.arguments = ["list", "-plist", "physical"] + + let pipe = Pipe() + task.standardOutput = pipe + + try task.run() + let data = pipe.fileHandleForReading.readDataToEndOfFile() + task.waitUntilExit() + + if task.terminationStatus == 0, let plist = try? PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], + let allDisks = plist["WholeDisks"] as? [String] { + + if allDisks.contains(drivePath) { + // It's a physical drive + } else { + let reason = "Selected drive is not a physical drive. For safety, only physical drives can be imaged." + state = .failed(.invalidDrive(reason: reason)) + throw ImagerError.invalidDrive(reason: reason) + } } else { - let imagerError = ImagerError.driveDetectionError(reason: "Failed to verify if drive is external: \(error.localizedDescription)") - state = .failed(imagerError) - throw imagerError + throw ImagerError.driveDetectionError(reason: "Failed to parse disk information") } + } catch let error as ImagerError { + throw error + } catch { + throw ImagerError.driveDetectionError(reason: "Failed to verify if drive is physical: \(error.localizedDescription)") } - + // Get image file size let imageAttributes = try FileManager.default.attributesOfItem(atPath: imageFilePath) guard let imageSize = imageAttributes[.size] as? Int64 else { @@ -198,16 +240,16 @@ class MacOSImager: Imager, @unchecked Sendable { state = .failed(.invalidImageFile(reason: reason)) throw ImagerError.invalidImageFile(reason: reason) } - + // Get drive size let driveSize = try getDriveSize(drive: drivePath) - + // Check if image fits on drive if imageSize > driveSize { state = .failed(.imageTooLargeForDrive(imageSize: UInt64(imageSize), driveSize: UInt64(driveSize))) throw ImagerError.imageTooLargeForDrive(imageSize: UInt64(imageSize), driveSize: UInt64(driveSize)) } - + // Open file handles do { imageFileHandle = try FileHandle(forReadingFrom: URL(fileURLWithPath: imageFilePath)) @@ -217,30 +259,30 @@ class MacOSImager: Imager, @unchecked Sendable { state = .failed(.permissionDenied(reason: reason)) throw ImagerError.permissionDenied(reason: reason) } - + // Start imaging process state = .imaging - + // Create a background task for imaging Task { await performImaging(imageSize: imageSize) } } - + /// Stop the imaging process. /// /// This method stops the imaging process if it is currently running. /// It may throw an `ImagerError` if the operation cannot be stopped. - func stopImaging() throws { + public func stopImaging() throws { if case .imaging = state { state = .failed(.processingInterrupted(reason: "Imaging process was stopped by user")) - + // Close file handles if let handle = imageFileHandle { try handle.close() imageFileHandle = nil } - + if let handle = driveFileHandle { try handle.close() driveFileHandle = nil @@ -249,17 +291,17 @@ class MacOSImager: Imager, @unchecked Sendable { throw ImagerError.unknown(reason: "Cannot stop imaging: No imaging process is currently running") } } - + /// Register a handler to receive progress updates. /// /// - Parameter handler: A closure that will be called with progress updates. /// The closure receives the current progress and an optional error if one occurred. - func progress(_ handler: @escaping (Progress, ImagerError?) -> Void) { + public func progress(_ handler: @escaping (Progress, ImagerError?) -> Void) { progressHandler = handler } - + // MARK: - Private Methods - + /// Get the size of a drive in bytes. /// /// - Parameter drive: The path to the drive. @@ -269,16 +311,16 @@ class MacOSImager: Imager, @unchecked Sendable { let task = Process() task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") task.arguments = ["info", "-plist", drive] - + let pipe = Pipe() task.standardOutput = pipe - + do { try task.run() let data = pipe.fileHandleForReading.readDataToEndOfFile() task.waitUntilExit() - - if task.terminationStatus == 0, let plist = try PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], + + if task.terminationStatus == 0, let plist = try? PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], let size = plist["TotalSize"] as? Int64 { return size } else { @@ -290,7 +332,7 @@ class MacOSImager: Imager, @unchecked Sendable { throw ImagerError.invalidDrive(reason: "Error getting drive size: \(error.localizedDescription)") } } - + /// Perform the actual imaging process. /// /// - Parameter imageSize: The size of the image file in bytes. @@ -301,15 +343,15 @@ class MacOSImager: Imager, @unchecked Sendable { progressHandler?(Progress(totalBytes: imageSize, completedBytes: 0), error) return } - + var bytesWritten: Int64 = 0 var lastProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) - + do { // Reset file positions try imageFileHandle.seek(toOffset: 0) try driveFileHandle.seek(toOffset: 0) - + // Read and write in chunks while bytesWritten < imageSize { // Check if we should continue @@ -318,20 +360,20 @@ class MacOSImager: Imager, @unchecked Sendable { } else { break } - + // Read a chunk from the image file let data = imageFileHandle.readData(ofLength: bufferSize) if data.isEmpty { break } - + // Write the chunk to the drive try driveFileHandle.write(contentsOf: data) - + // Update progress bytesWritten += Int64(data.count) let currentProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) - + // Only update if progress has changed significantly (1% or more) if currentProgress.percentage - lastProgress.percentage >= 0.01 { lastProgress = currentProgress @@ -339,39 +381,39 @@ class MacOSImager: Imager, @unchecked Sendable { progressHandler?(currentProgress, nil) } } - + // Ensure all data is written to disk try driveFileHandle.synchronize() - + // Update final progress let finalProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) state = .progress(finalProgress) progressHandler?(finalProgress, nil) - + // Set state to completed state = .completed - + } catch { let imagerError = ImagerError.processingInterrupted(reason: error.localizedDescription) state = .failed(imagerError) progressHandler?(lastProgress, imagerError) } - + // Close file handles try? imageFileHandle.close() try? driveFileHandle.close() - + // Reset file handles self.imageFileHandle = nil self.driveFileHandle = nil } - + /// Cancel the imaging process. @available(*, deprecated, message: "Use stopImaging() instead") func cancel() { try? stopImaging() } - + deinit { // Ensure file handles are closed try? imageFileHandle?.close() diff --git a/Sources/Imager/Progress.swift b/Sources/Imager/Progress.swift index 653ab63..70eb753 100644 --- a/Sources/Imager/Progress.swift +++ b/Sources/Imager/Progress.swift @@ -1,12 +1,12 @@ -struct Progress { +public struct Progress { /// The total number of bytes to be written. - var totalBytes: Int64 + public var totalBytes: Int64 /// The number of bytes that have been written. - var completedBytes: Int64 + public var completedBytes: Int64 /// The percentage of the image that has been written. - var percentage: Double { + public var percentage: Double { return Double(completedBytes) / Double(totalBytes) } } \ No newline at end of file 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..f888c6f --- /dev/null +++ b/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift @@ -0,0 +1,50 @@ +import Imager +import ArgumentParser +import Foundation + +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 = MacOSImager(imageFilePath: source, drivePath: drivePath) + + do { + // startImaging is synchronous, but handles launching the process + try imager.startImaging() + // If startImaging returns without throwing, the process was initiated. + // Actual success/failure/progress is handled internally by MacOSImager (potentially). + // A more robust solution would involve async streams for progress. + print("\nImaging process initiated successfully.") + print("Monitor system activity or logs for completion.") + } catch let error as ImagerError { + print("\nError during imaging: \(error.description)") + throw ExitCode.failure + } catch { + print("\nAn unexpected error occurred: \(error.localizedDescription)") + 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..d1e7598 --- /dev/null +++ b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift @@ -0,0 +1,52 @@ +import Imager +import ArgumentParser +import Foundation + +struct ImagerListDisksSubCommand: AsyncParsableCommand { + static let configuration = CommandConfiguration( + commandName: "list-disks", + abstract: "List available drives to image." + ) + + @Flag(name: .long, help: "Shows both internal and external drives.") + var all: Bool = false + + @Option(name: [.long, .customShort("t")], help: "The timeout in seconds.") + var timeout: Int = 3 + + func run() async throws { + let imager = MacOSImager(imageFilePath: "", drivePath: "") + let drives = try imager.availableDrivesToImage(onlyExternalDrives: !all) + + if drives.isEmpty { + print(all ? "No drives found." : "No suitable external drives found. Use --all to list internal drives.") + return + } + + // Headers + let nameHeader = "Name" + let pathHeader = "Path" + let availableHeader = "Available" + let capacityHeader = "Capacity" + + // Calculate column widths + 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 + + // Print header + 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)) + + // Print rows + 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)") + } + } +} \ No newline at end of file diff --git a/Sources/edge/EdgeCLI.swift b/Sources/edge/EdgeCLI.swift index 706e58f..c804253 100644 --- a/Sources/edge/EdgeCLI.swift +++ b/Sources/edge/EdgeCLI.swift @@ -7,7 +7,8 @@ struct EdgeCLI: AsyncParsableCommand { commandName: "edge", abstract: "Edge CLI", subcommands: [ - RunCommand.self + RunCommand.self, + ImagerCommand.self ] ) } diff --git a/imager_command_features.md b/imager_command_features.md new file mode 100644 index 0000000..dd1919b --- /dev/null +++ b/imager_command_features.md @@ -0,0 +1,8 @@ +# Imager Command Features + + +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 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. +5. `edge-cli imager list-disks --all` should list all available drives, including internal ones. This can be used with `--watch` and `--timeout` options. \ No newline at end of file From d81d698ea158225759442b6d63423c4c9756890f Mon Sep 17 00:00:00 2001 From: Zamderax Date: Thu, 17 Apr 2025 17:30:09 -0700 Subject: [PATCH 3/4] Refactor imager to use platform-specific implementations and dd for disk writing --- Sources/Imager/Imager.swift | 7 +- Sources/Imager/ImagerError.swift | 17 + Sources/Imager/ImagerLib.swift | 20 + Sources/Imager/ImagerState.swift | 22 +- Sources/Imager/LinuxImager.swift | 42 ++ Sources/Imager/MacOSImager.swift | 631 ++++++++---------- Sources/Imager/Progress.swift | 12 - .../Imager/ImagerListDiskSubCommand.swift | 34 +- imager_command_features.md | 11 +- 9 files changed, 414 insertions(+), 382 deletions(-) create mode 100644 Sources/Imager/ImagerLib.swift create mode 100644 Sources/Imager/LinuxImager.swift delete mode 100644 Sources/Imager/Progress.swift diff --git a/Sources/Imager/Imager.swift b/Sources/Imager/Imager.swift index 9452a2b..69d291f 100644 --- a/Sources/Imager/Imager.swift +++ b/Sources/Imager/Imager.swift @@ -1,10 +1,13 @@ import Foundation +/// Public typealias for disk update notifications, usable by consumers +public typealias DiskUpdateCallback = () -> Void + /// Protocol defining the interface for an image writer. /// /// This protocol defines the requirements for a component that can write /// an image file to a drive. -protocol Imager { +public protocol Imager { /// Path to the image file to be imaged. var imageFilePath: String { get } @@ -37,7 +40,7 @@ protocol Imager { /// /// - Parameter handler: A closure that will be called with progress updates. /// The closure receives the current progress and an optional error if one occurred. - func progress(_ handler: @escaping (Progress, ImagerError?) -> Void) + func progress(_ handler: @escaping (Foundation.Progress, ImagerError?) -> Void) /// Initialize a new imager with the specified image file and drive paths. /// diff --git a/Sources/Imager/ImagerError.swift b/Sources/Imager/ImagerError.swift index 770da1f..ef86ef8 100644 --- a/Sources/Imager/ImagerError.swift +++ b/Sources/Imager/ImagerError.swift @@ -41,11 +41,22 @@ public enum ImagerError: Error { /// 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 @@ -74,6 +85,12 @@ extension ImagerError: CustomStringConvertible { 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)" diff --git a/Sources/Imager/ImagerLib.swift b/Sources/Imager/ImagerLib.swift new file mode 100644 index 0000000..fbaf607 --- /dev/null +++ b/Sources/Imager/ImagerLib.swift @@ -0,0 +1,20 @@ +import Foundation + +/// Factory methods for creating the appropriate Imager implementation for the current platform. +public enum ImagerFactory { + /// Creates an Imager instance appropriate for the current platform. + /// + /// - Parameters: + /// - imageFilePath: The path to the image file to be written. + /// - drivePath: The path to the drive where the image will be written. + /// - Returns: An Imager instance for the current platform. + 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) + #else + fatalError("Unsupported platform") + #endif + } +} \ No newline at end of file diff --git a/Sources/Imager/ImagerState.swift b/Sources/Imager/ImagerState.swift index 3f59961..aa31736 100644 --- a/Sources/Imager/ImagerState.swift +++ b/Sources/Imager/ImagerState.swift @@ -1,24 +1,21 @@ /// Represents the current status of the imaging process. /// /// This enum tracks the lifecycle of an imaging operation, from initialization through -/// completion or failure, including progress updates during the process. +/// completion or failure. public enum ImagerState { /// The imager is initialized but has not started imaging. /// /// This is the default state before `startImaging()` is called. case idle - /// The imaging process is currently in progress. - /// - /// The associated value `Progress` provides details about the current progress, - /// such as bytes written and total bytes. - case progress(Progress) - /// 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. @@ -42,17 +39,14 @@ extension ImagerState: CustomStringConvertible { 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)" - - case .progress(let progress): - let percentage = Int(progress.percentage * 100) - let completedMB = Double(progress.completedBytes) / 1_048_576 - let totalMB = Double(progress.totalBytes) / 1_048_576 - return "Progress: \(percentage)% (\(String(format: "%.1f", completedMB)) MB of \(String(format: "%.1f", totalMB)) MB)" + return "Imaging failed: \(error.localizedDescription)" } } } \ No newline at end of file diff --git a/Sources/Imager/LinuxImager.swift b/Sources/Imager/LinuxImager.swift new file mode 100644 index 0000000..94fb88f --- /dev/null +++ b/Sources/Imager/LinuxImager.swift @@ -0,0 +1,42 @@ +#if os(Linux) +import Foundation + +/// Concrete implementation of the Imager protocol for Linux. +public class LinuxImager: Imager, @unchecked Sendable { + public let imageFilePath: String + public let drivePath: String + public var state: ImagerState = .idle + private var progress: Foundation.Progress? + private var progressHandler: ((Foundation.Progress, ImagerError?) -> Void)? + + private var ddProcess: Process? + private var outputPipe: Pipe? + private var errorPipe: Pipe? + + public required init(imageFilePath: String = "", drivePath: String = "") { + self.imageFilePath = imageFilePath + self.drivePath = drivePath + self.progress = Progress(totalUnitCount: 0) + } + + public func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] { + // Implementation for Linux would use commands like lsblk to list block devices + // For now, return an empty array + return [] + } + + public func startImaging() throws { + // Implementation would use dd command similar to MacOSImager + throw ImagerError.notImplemented(reason: "Linux imaging not yet implemented") + } + + public func stopImaging() throws { + // Implementation would stop the dd process + throw ImagerError.notImplemented(reason: "Linux imaging not yet implemented") + } + + public func progress(_ handler: @escaping (Foundation.Progress, ImagerError?) -> Void) { + self.progressHandler = handler + } +} +#endif // os(Linux) \ No newline at end of file diff --git a/Sources/Imager/MacOSImager.swift b/Sources/Imager/MacOSImager.swift index a3ce85e..a6d2320 100644 --- a/Sources/Imager/MacOSImager.swift +++ b/Sources/Imager/MacOSImager.swift @@ -1,41 +1,34 @@ #if os(macOS) import Foundation import DiskArbitration +import IOKit -/// A macOS implementation of the `Imager` protocol. -/// -/// This class provides functionality for writing image files to drives on macOS systems -/// using native APIs. -@available(macOS 10.15, *) -public class MacOSImager: Imager, @unchecked Sendable { - /// Path to the image file to be imaged. +/// Concrete implementation of the Imager protocol for macOS. +public class MacOSImager: Imager, @unchecked Sendable { public let imageFilePath: String - - /// Path to the drive to be imaged. public let drivePath: String - - /// The current state of the imaging process. - public private(set) var state: ImagerState = .idle - - /// Progress handler closure. - private var progressHandler: ((Progress, ImagerError?) -> Void)? - - /// File handle for the image file. + public var state: ImagerState = .idle + var progress: Foundation.Progress? + var progressObservation: NSKeyValueObservation? + var progressHandler: ((Foundation.Progress, ImagerError?) -> Void)? + private var imageFileHandle: FileHandle? - - /// File handle for the drive. private var driveFileHandle: FileHandle? - - /// Buffer size for reading/writing operations. - private let bufferSize = 1024 * 1024 // 1MB buffer - - /// Get a list of available drives that can be imaged in alphabetical order. - /// - /// - Parameter onlyExternalDrives: If `true`, only external drives are returned. - /// - Returns: An array of `Drive` objects representing the available drives. - /// - Throws: An `ImagerError` if there is an error getting or parsing drive information. + private var ddProcess: Process? + private var outputPipe: Pipe? + private var errorPipe: Pipe? + private let progressQueue = DispatchQueue(label: "com.apache-edge.imager.progress", qos: .utility) + private var lastProgressUpdate = Date() + + public required init(imageFilePath: String = "", drivePath: String = "") { + self.imageFilePath = imageFilePath + self.drivePath = drivePath + self.progress = Foundation.Progress(totalUnitCount: 0) + } + + // ... availableDrivesToImage (no changes needed here for this fix) ... public func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] { - // Step 1: Get the list of whole disk identifiers and map whole disks to their mountable partitions + // Get whole disk identifiers using diskutil list -plist let listTask = Process() listTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") listTask.arguments = ["list", "-plist"] @@ -43,381 +36,339 @@ public class MacOSImager: Imager, @unchecked Sendable { listTask.standardOutput = listPipe var wholeDiskIdentifiers: [String] = [] + var wholeDiskToPartitionMap: [String: String] = [:] + do { try listTask.run() - let listData = listPipe.fileHandleForReading.readDataToEndOfFile() listTask.waitUntilExit() - guard listTask.terminationStatus == 0 else { + 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 wholeDisksIdentifiers = listPlist["WholeDisks"] as? [String], + 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 = wholeDisksIdentifiers + wholeDiskIdentifiers = parsedWholeDisks // Create a map from whole disk identifier (e.g., "disk28") to its first mountable partition identifier (e.g., "disk28s1") - var wholeDiskToPartitionMap: [String: String] = [:] - for diskEntry in allDisksAndPartitions { - if let wholeDiskId = diskEntry["DeviceIdentifier"] as? String, - let partitions = diskEntry["Partitions"] as? [[String: Any]] { - for partition in partitions { - if let partitionId = partition["DeviceIdentifier"] as? String, - let mountPoint = partition["MountPoint"] as? String, - !mountPoint.isEmpty { - wholeDiskToPartitionMap[wholeDiskId] = partitionId - break // Found the first mountable partition for this disk - } - } + 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 } } - // Step 2: Get detailed info, filter, and find available space - var availableDrives: [Drive] = [] - guard let session = DASessionCreate(kCFAllocatorDefault) else { - throw ImagerError.driveDetectionError(reason: "Failed to create DiskArbitration session") - } + } catch { + throw ImagerError.driveDetectionError(reason: "Failed to execute or process diskutil list: \(error.localizedDescription)") + } - for diskIdentifier in wholeDiskIdentifiers { - let path = "/dev/\(diskIdentifier)" - let infoTask = Process() - infoTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") - infoTask.arguments = ["info", "-plist", path] - let infoPipe = Pipe() - infoTask.standardOutput = infoPipe - - do { - try infoTask.run() - let infoData = infoPipe.fileHandleForReading.readDataToEndOfFile() - infoTask.waitUntilExit() - - guard infoTask.terminationStatus == 0, - let infoPlist = try? PropertyListSerialization.propertyList(from: infoData, options: [], format: nil) as? [String: Any] else { - continue - } + 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") + } - // --- Get Base Info & Filter using diskutil info on Whole Disk (diskX) --- - let isInternal = infoPlist["Internal"] as? Bool ?? false - let isSystemImage = infoPlist["SystemImage"] as? Bool ?? false - let virtualOrPhysical = infoPlist["VirtualOrPhysical"] as? String ?? "Physical" - let isVirtualDevice = (virtualOrPhysical == "Virtual") + 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 - // Apply filtering first - let isSuitableExternal = !isInternal && !isSystemImage && !isVirtualDevice - if onlyExternalDrives && !isSuitableExternal { - continue // Skip this whole disk if filtering for external and it doesn't match - } + do { + try infoTask.run() + infoTask.waitUntilExit() + let infoData = infoPipe.fileHandleForReading.readDataToEndOfFile() - // --- Get Mount Point, Available Space, and Name (via Partition diskXsY) --- - var available: Int64? = nil - var volumeNameFromPartition: String? = nil // Store name found from partition - let capacity = infoPlist["TotalSize"] as? Int64 ?? 0 // Capacity comes from whole disk - - // Find the partition identifier associated with this whole disk - if let partitionIdentifier = wholeDiskToPartitionMap[diskIdentifier] { - let partitionPath = "/dev/\(partitionIdentifier)" - - if let disk = DADiskCreateFromBSDName(kCFAllocatorDefault, session, partitionPath) { - if let diskInfoDA = DADiskCopyDescription(disk) as? [String: Any] { - // Attempt to get name from partition's DA description - volumeNameFromPartition = diskInfoDA[kDADiskDescriptionVolumeNameKey as String] as? String - - if let mountURL = diskInfoDA[kDADiskDescriptionVolumePathKey as String] as? URL { - let mountPoint = mountURL.path - if !mountPoint.isEmpty { - // Extra check: Don't include root if filtering external - if !(onlyExternalDrives && mountPoint == "/") { - do { - let attributes = try FileManager.default.attributesOfFileSystem(forPath: mountPoint) - available = attributes[.systemFreeSize] as? Int64 - } catch { - available = nil // Failed FileManager - } - } // else: Skip root disk if filtering - } // else: DA gave empty mount point for partition - } // else: DA description for partition missing mount point - } // else: Failed DA CopyDescription for partition - } // else: Failed DA Create Disk for partition - } // else: No mountable partition found for this whole disk in the map - - // Only add the drive if it passed the initial filter (or we aren't filtering) - // Use the name found from the partition if available - let drive = Drive(path: path, capacity: capacity, available: available, name: volumeNameFromPartition) - availableDrives.append(drive) - - } catch { - // Error running diskutil info process for this disk + if infoTask.terminationStatus != 0 { + continue + } + + guard let diskInfoPlist = try? PropertyListSerialization.propertyList(from: infoData, options: [], format: nil) as? [String: Any] else { continue } - } - // Sort drives by path for consistent ordering - return availableDrives.sorted { $0.path < $1.path } - } catch let error as ImagerError { - throw error - } catch { - throw ImagerError.driveDetectionError(reason: "Error running or parsing diskutil list: \(error.localizedDescription)") + 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 } } - /// Initialize a new imager with the specified image file and drive paths. - /// - /// - Parameters: - /// - imageFilePath: The path to the image file to be written. - /// - drivePath: The path to the drive where the image will be written. - public required init(imageFilePath: String, drivePath: String) { - self.imageFilePath = imageFilePath - self.drivePath = drivePath + // ... 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 } - /// Begin imaging the drive. - /// - /// This method starts the process of writing the image file to the drive. - /// It may throw an `ImagerError` if the operation cannot be completed. + + // ... startImaging, stopImaging ... public func startImaging() throws { - // Validate image file - guard FileManager.default.fileExists(atPath: imageFilePath) else { - state = .failed(.invalidImageFile(reason: "Image file does not exist")) - throw ImagerError.invalidImageFile(reason: "Image file does not exist") + guard imageFilePath != "", drivePath != "" else { + throw ImagerError.invalidDrive(reason: "Image file path or drive path not set.") } - - // Check file extension - let fileExtension = URL(fileURLWithPath: imageFilePath).pathExtension.lowercased() - guard ["zip", "tar", "img"].contains(fileExtension) else { - let reason = "Unsupported file format: \(fileExtension). Only zip, tar, and img formats are supported." - state = .failed(.invalidImageFile(reason: reason)) - throw ImagerError.invalidImageFile(reason: reason) + guard case .idle = state else { + throw ImagerError.processingInterrupted(reason: "Imaging process already running or not idle.") } - // Validate drive - let driveURL = URL(fileURLWithPath: "/dev/\(drivePath)") - let drivePath = driveURL.path - - guard FileManager.default.fileExists(atPath: drivePath) else { - let reason = "Drive does not exist at path: \(drivePath)" - state = .failed(.invalidDrive(reason: reason)) - throw ImagerError.invalidDrive(reason: reason) + let fileManager = FileManager.default + guard fileManager.fileExists(atPath: imageFilePath) else { + throw ImagerError.invalidImageFile(reason: "Image file not found at \(imageFilePath)") } - // Verify it's an external drive for safety + guard fileManager.fileExists(atPath: drivePath) else { + throw ImagerError.invalidDrive(reason: "Drive path not found at \(drivePath)") + } + let rawDrivePath = drivePath.replacingOccurrences(of: "/dev/disk", with: "/dev/rdisk") + guard fileManager.fileExists(atPath: rawDrivePath) else { + throw ImagerError.invalidDrive(reason: "Raw drive device not found at \(rawDrivePath)") + } + + print("Unmounting drive \(drivePath)...") + let unmountTask = Process() + unmountTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") + unmountTask.arguments = ["unmountDisk", "force", drivePath] do { - // Get list of all physical drives from the system - let task = Process() - task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") - task.arguments = ["list", "-plist", "physical"] - - let pipe = Pipe() - task.standardOutput = pipe - - try task.run() - let data = pipe.fileHandleForReading.readDataToEndOfFile() - task.waitUntilExit() - - if task.terminationStatus == 0, let plist = try? PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], - let allDisks = plist["WholeDisks"] as? [String] { - - if allDisks.contains(drivePath) { - // It's a physical drive - } else { - let reason = "Selected drive is not a physical drive. For safety, only physical drives can be imaged." - state = .failed(.invalidDrive(reason: reason)) - throw ImagerError.invalidDrive(reason: reason) - } - } else { - throw ImagerError.driveDetectionError(reason: "Failed to parse disk information") + try unmountTask.run() + unmountTask.waitUntilExit() + if unmountTask.terminationStatus != 0 { + print("Warning: diskutil unmountDisk failed (status \(unmountTask.terminationStatus)). Proceeding anyway...") } - } catch let error as ImagerError { - throw error } catch { - throw ImagerError.driveDetectionError(reason: "Failed to verify if drive is physical: \(error.localizedDescription)") + throw ImagerError.processingInterrupted(reason: "Failed to run diskutil unmountDisk: \(error.localizedDescription)") } - // Get image file size - let imageAttributes = try FileManager.default.attributesOfItem(atPath: imageFilePath) - guard let imageSize = imageAttributes[.size] as? Int64 else { - let reason = "Could not determine image file size" - state = .failed(.invalidImageFile(reason: reason)) - throw ImagerError.invalidImageFile(reason: reason) + let imageSize: UInt64 + do { + let imageAttrs = try fileManager.attributesOfItem(atPath: imageFilePath) + guard let size = imageAttrs[.size] as? UInt64 else { + throw ImagerError.invalidImageFile(reason: "Could not read image file size attribute.") + } + imageSize = size + } catch { + throw ImagerError.invalidImageFile(reason: "Could not read image file attributes: \(error.localizedDescription)") } + progress?.totalUnitCount = Int64(imageSize) + progress?.completedUnitCount = 0 - // Get drive size - let driveSize = try getDriveSize(drive: drivePath) + print("Starting imaging process using dd...") + print("Source: \(imageFilePath)") + print("Target: \(rawDrivePath)") - // Check if image fits on drive - if imageSize > driveSize { - state = .failed(.imageTooLargeForDrive(imageSize: UInt64(imageSize), driveSize: UInt64(driveSize))) - throw ImagerError.imageTooLargeForDrive(imageSize: UInt64(imageSize), driveSize: UInt64(driveSize)) - } + ddProcess = Process() + outputPipe = Pipe() + errorPipe = Pipe() - // Open file handles - do { - imageFileHandle = try FileHandle(forReadingFrom: URL(fileURLWithPath: imageFilePath)) - driveFileHandle = try FileHandle(forWritingTo: driveURL) - } catch { - let reason = "Failed to open file handles: \(error.localizedDescription)" - state = .failed(.permissionDenied(reason: reason)) - throw ImagerError.permissionDenied(reason: reason) + guard let ddProcess = ddProcess, let errorPipe = errorPipe else { + throw ImagerError.processingInterrupted(reason: "Failed to create process or pipes for dd.") } - // Start imaging process - state = .imaging + ddProcess.executableURL = URL(fileURLWithPath: "/bin/dd") + ddProcess.arguments = ["if=\(imageFilePath)", "of=\(rawDrivePath)", "bs=4m", "status=progress"] - // Create a background task for imaging - Task { - await performImaging(imageSize: imageSize) - } - } + ddProcess.standardError = errorPipe - /// Stop the imaging process. - /// - /// This method stops the imaging process if it is currently running. - /// It may throw an `ImagerError` if the operation cannot be stopped. - public func stopImaging() throws { - if case .imaging = state { - state = .failed(.processingInterrupted(reason: "Imaging process was stopped by user")) - - // Close file handles - if let handle = imageFileHandle { - try handle.close() - imageFileHandle = nil + ddProcess.terminationHandler = { [weak self] process in + self?.progressQueue.async { + self?.handleImagingCompletion(process: process) } + } - if let handle = driveFileHandle { - try handle.close() - driveFileHandle = nil + errorPipe.fileHandleForReading.readabilityHandler = { [weak self] fileHandle in + let data = fileHandle.availableData + if data.isEmpty { + self?.progressQueue.async { + errorPipe.fileHandleForReading.readabilityHandler = nil + } + } else { + self?.progressQueue.async { + if let output = String(data: data, encoding: .utf8) { + self?.parseDdProgress(output) + } + } } - } else { - throw ImagerError.unknown(reason: "Cannot stop imaging: No imaging process is currently running") } - } - - /// Register a handler to receive progress updates. - /// - /// - Parameter handler: A closure that will be called with progress updates. - /// The closure receives the current progress and an optional error if one occurred. - public func progress(_ handler: @escaping (Progress, ImagerError?) -> Void) { - progressHandler = handler - } - - // MARK: - Private Methods - - /// Get the size of a drive in bytes. - /// - /// - Parameter drive: The path to the drive. - /// - Returns: The size of the drive in bytes. - /// - Throws: An `ImagerError` if the drive size cannot be determined. - private func getDriveSize(drive: String) throws -> Int64 { - let task = Process() - task.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") - task.arguments = ["info", "-plist", drive] - - let pipe = Pipe() - task.standardOutput = pipe do { - try task.run() - let data = pipe.fileHandleForReading.readDataToEndOfFile() - task.waitUntilExit() - - if task.terminationStatus == 0, let plist = try? PropertyListSerialization.propertyList(from: data, options: [], format: nil) as? [String: Any], - let size = plist["TotalSize"] as? Int64 { - return size - } else { - throw ImagerError.invalidDrive(reason: "Could not determine drive size") - } - } catch let error as ImagerError { - throw error + try ddProcess.run() + state = .imaging + print("dd process started (PID: \(ddProcess.processIdentifier)).") } catch { - throw ImagerError.invalidDrive(reason: "Error getting drive size: \(error.localizedDescription)") + state = .idle + cleanupResources() + throw ImagerError.processingInterrupted(reason: "Failed to launch dd process: \(error.localizedDescription)") } } - /// Perform the actual imaging process. - /// - /// - Parameter imageSize: The size of the image file in bytes. - private func performImaging(imageSize: Int64) async { - guard let imageFileHandle = imageFileHandle, let driveFileHandle = driveFileHandle else { - let error = ImagerError.unknown(reason: "File handles not initialized") - state = .failed(error) - progressHandler?(Progress(totalBytes: imageSize, completedBytes: 0), error) + public func stopImaging() { + guard let ddProcess = ddProcess, ddProcess.isRunning else { + print("Imaging process not running.") return } - var bytesWritten: Int64 = 0 - var lastProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) - - do { - // Reset file positions - try imageFileHandle.seek(toOffset: 0) - try driveFileHandle.seek(toOffset: 0) - - // Read and write in chunks - while bytesWritten < imageSize { - // Check if we should continue - if case .imaging = state { - // Continue imaging - } else { - break - } + state = .cancelling + print("Stopping imaging process (PID: \(ddProcess.processIdentifier))...") + ddProcess.interrupt() - // Read a chunk from the image file - let data = imageFileHandle.readData(ofLength: bufferSize) - if data.isEmpty { - break - } - - // Write the chunk to the drive - try driveFileHandle.write(contentsOf: data) - - // Update progress - bytesWritten += Int64(data.count) - let currentProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) - - // Only update if progress has changed significantly (1% or more) - if currentProgress.percentage - lastProgress.percentage >= 0.01 { - lastProgress = currentProgress - state = .progress(currentProgress) - progressHandler?(currentProgress, nil) - } + DispatchQueue.global(qos: .userInitiated).asyncAfter(deadline: .now() + 2.0) { [weak self] in + if self?.ddProcess?.isRunning == true { + print("Process still running, sending SIGKILL...") + self?.ddProcess?.terminate() } - - // Ensure all data is written to disk - try driveFileHandle.synchronize() - - // Update final progress - let finalProgress = Progress(totalBytes: imageSize, completedBytes: bytesWritten) - state = .progress(finalProgress) - progressHandler?(finalProgress, nil) - - // Set state to completed - state = .completed - - } catch { - let imagerError = ImagerError.processingInterrupted(reason: error.localizedDescription) - state = .failed(imagerError) - progressHandler?(lastProgress, imagerError) } - - // Close file handles - try? imageFileHandle.close() - try? driveFileHandle.close() - - // Reset file handles - self.imageFileHandle = nil - self.driveFileHandle = nil } - /// Cancel the imaging process. - @available(*, deprecated, message: "Use stopImaging() instead") - func cancel() { - try? stopImaging() + private func handleImagingCompletion(process: Process) { + let exitCode = process.terminationStatus + let reason = process.terminationReason + + errorPipe?.fileHandleForReading.readabilityHandler = nil + + if case .cancelling = state { + print("Imaging process cancelled.") + progress?.cancel() + state = .idle + } else if exitCode == 0 { + print("Imaging process completed successfully.") + progress?.completedUnitCount = progress?.totalUnitCount ?? 0 + state = .completed + } else { + print("Imaging process failed. Exit code: \(exitCode), Reason: \(reason).") + let errorData = errorPipe?.fileHandleForReading.readDataToEndOfFile() ?? Data() + if let errorString = String(data: errorData, encoding: .utf8), !errorString.isEmpty { + print("dd stderr: \(errorString)") + } + state = .failed(ImagerError.processingInterrupted(reason: "dd process failed with exit code \(exitCode)")) + } + + cleanupResources() + let workItem = DispatchWorkItem { + if case .completed = self.state { + self.state = .idle + } else if case .cancelling = self.state { + self.state = .idle + } else if case .failed = self.state { + self.state = .idle + } + } + DispatchQueue.main.asyncAfter(deadline: .now() + 0.5, execute: workItem) + } + + private func parseDdProgress(_ output: String) { + let lines = output.split(separator: "\r") + for line in lines { + if let bytesString = line.split(separator: " ").first, + let bytesCopied = Int64(bytesString) { + + if Date().timeIntervalSince(lastProgressUpdate) > 0.1 { + DispatchQueue.main.async { + self.progress?.completedUnitCount = bytesCopied + self.lastProgressUpdate = Date() + } + } + } + } + } + + private func cleanupResources() { + errorPipe?.fileHandleForReading.readabilityHandler = nil + + try? imageFileHandle?.close() + try? driveFileHandle?.close() + try? outputPipe?.fileHandleForReading.close() + try? outputPipe?.fileHandleForWriting.close() + try? errorPipe?.fileHandleForReading.close() + try? errorPipe?.fileHandleForWriting.close() + + imageFileHandle = nil + driveFileHandle = nil + outputPipe = nil + errorPipe = nil + ddProcess = nil + progressObservation = nil + progressHandler = nil + } + + // Deinitializer to ensure cleanup + deinit { + cleanupResources() + progressObservation = nil } - deinit { - // Ensure file handles are closed - try? imageFileHandle?.close() - try? driveFileHandle?.close() + // --- Imager Protocol Implementation --- + public func progress(_ handler: @escaping (Foundation.Progress, ImagerError?) -> Void) { + self.progressHandler = handler } } -#endif \ No newline at end of file + +#endif // os(macOS) \ No newline at end of file diff --git a/Sources/Imager/Progress.swift b/Sources/Imager/Progress.swift deleted file mode 100644 index 70eb753..0000000 --- a/Sources/Imager/Progress.swift +++ /dev/null @@ -1,12 +0,0 @@ -public struct Progress { - /// The total number of bytes to be written. - public var totalBytes: Int64 - - /// The number of bytes that have been written. - public var completedBytes: Int64 - - /// The percentage of the image that has been written. - public var percentage: Double { - return Double(completedBytes) / Double(totalBytes) - } -} \ No newline at end of file diff --git a/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift index d1e7598..653e1c0 100644 --- a/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift +++ b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift @@ -1,46 +1,55 @@ -import Imager import ArgumentParser import Foundation +import Imager +#if canImport(Darwin) +import Darwin // For signal handling +#endif struct ImagerListDisksSubCommand: AsyncParsableCommand { + static let configuration = CommandConfiguration( commandName: "list-disks", - abstract: "List available drives to image." + abstract: "Lists available drives that can be targeted for imaging." ) - @Flag(name: .long, help: "Shows both internal and external drives.") + @Flag(name: .long, help: "Include internal drives in the list.") var all: Bool = false - @Option(name: [.long, .customShort("t")], help: "The timeout in seconds.") - var timeout: Int = 3 + // 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 imager = MacOSImager(imageFilePath: "", drivePath: "") - let drives = try imager.availableDrivesToImage(onlyExternalDrives: !all) + let imager = ImagerFactory.createImager() + try await MainActor.run { + try listAndPrintDisks(imager: imager, listAll: self.all) + } + } + + @MainActor + private func listAndPrintDisks(imager: Imager, listAll: Bool) throws { + let drives = try imager.availableDrivesToImage(onlyExternalDrives: !listAll) if drives.isEmpty { - print(all ? "No drives found." : "No suitable external drives found. Use --all to list internal drives.") + print(listAll ? "No drives found." : "No suitable external drives found. Use --all to list internal drives.") + fflush(stdout) return } - // Headers let nameHeader = "Name" let pathHeader = "Path" let availableHeader = "Available" let capacityHeader = "Capacity" - // Calculate column widths 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 - // Print header 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)) - // Print rows 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) @@ -48,5 +57,6 @@ struct ImagerListDisksSubCommand: AsyncParsableCommand { 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/imager_command_features.md b/imager_command_features.md index dd1919b..30cc056 100644 --- a/imager_command_features.md +++ b/imager_command_features.md @@ -1,8 +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 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. -5. `edge-cli imager list-disks --all` should list all available drives, including internal ones. This can be used with `--watch` and `--timeout` options. \ No newline at end of file +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. From 514bdf86c1b8cbc5ae431e6b8b7b238222eea6ae Mon Sep 17 00:00:00 2001 From: Zamderax Date: Thu, 17 Apr 2025 22:54:44 -0700 Subject: [PATCH 4/4] Refactor imager to use factory pattern and add progress bar functionality --- Sources/Imager/Drive.swift | 43 +- Sources/Imager/Imager.swift | 51 --- Sources/Imager/ImagerError.swift | 2 +- Sources/Imager/ImagerFactory.swift | 34 ++ Sources/Imager/ImagerLib.swift | 20 - Sources/Imager/ImagerState.swift | 2 +- Sources/Imager/LinuxImager.swift | 42 -- Sources/Imager/MacOS/MacOSDiskLister.swift | 156 ++++++++ Sources/Imager/MacOS/MacOSImager.swift | 291 ++++++++++++++ Sources/Imager/MacOSImager.swift | 374 ------------------ Sources/Imager/Protocols/DiskLister.swift | 15 + Sources/Imager/Protocols/Imager.swift | 26 ++ .../Imager/ImagerImageDiskSubCommand.swift | 88 ++++- .../Imager/ImagerListDiskSubCommand.swift | 13 +- .../Commands/ProgressBar/ProgressBar.swift | 100 +++++ .../ProgressBar/ProgressBarExample.swift | 82 ++++ Sources/edge/EdgeCLI.swift | 3 +- 17 files changed, 786 insertions(+), 556 deletions(-) delete mode 100644 Sources/Imager/Imager.swift create mode 100644 Sources/Imager/ImagerFactory.swift delete mode 100644 Sources/Imager/ImagerLib.swift delete mode 100644 Sources/Imager/LinuxImager.swift create mode 100644 Sources/Imager/MacOS/MacOSDiskLister.swift create mode 100644 Sources/Imager/MacOS/MacOSImager.swift delete mode 100644 Sources/Imager/MacOSImager.swift create mode 100644 Sources/Imager/Protocols/DiskLister.swift create mode 100644 Sources/Imager/Protocols/Imager.swift create mode 100644 Sources/edge/Commands/ProgressBar/ProgressBar.swift create mode 100644 Sources/edge/Commands/ProgressBar/ProgressBarExample.swift diff --git a/Sources/Imager/Drive.swift b/Sources/Imager/Drive.swift index 95c2485..25d3f7b 100644 --- a/Sources/Imager/Drive.swift +++ b/Sources/Imager/Drive.swift @@ -1,9 +1,4 @@ -#if canImport(FoundationEssentials) - import FoundationEssentials -#else - import Foundation -#endif - +import Foundation /// A structure representing a drive. public struct Drive { /// The path to the drive. @@ -17,30 +12,12 @@ public struct Drive { /// The total capacity of the drive in a human-readable format. public var capacityHumanReadable: String { -#if canImport(Foundation) || canImport(FoundationEssentials) let formatter = ByteCountFormatter() formatter.allowedUnits = [.useAll] formatter.countStyle = .file formatter.includesUnit = true formatter.isAdaptive = true return formatter.string(fromByteCount: capacity) -#else - // Basic fallback for non-Apple platforms - let kb = Double(capacity) / 1024.0 - if kb < 1024 { - return String(format: "%.1f KB", kb) - } - let mb = kb / 1024.0 - if mb < 1024 { - return String(format: "%.1f MB", mb) - } - let gb = mb / 1024.0 - if gb < 1024 { - return String(format: "%.1f GB", gb) - } - let tb = gb / 1024.0 - return String(format: "%.1f TB", tb) -#endif } /// The available space on the drive in a human-readable format, or "N/A" if unknown. @@ -48,29 +25,11 @@ public struct Drive { guard let availableSpace = available else { return "N/A" } -#if canImport(Foundation) || canImport(FoundationEssentials) let formatter = ByteCountFormatter() formatter.allowedUnits = [.useAll] formatter.countStyle = .file formatter.includesUnit = true formatter.isAdaptive = true return formatter.string(fromByteCount: availableSpace) -#else - // Basic fallback for non-Apple platforms - let kb = Double(availableSpace) / 1024.0 - if kb < 1024 { - return String(format: "%.1f KB", kb) - } - let mb = kb / 1024.0 - if mb < 1024 { - return String(format: "%.1f MB", mb) - } - let gb = mb / 1024.0 - if gb < 1024 { - return String(format: "%.1f GB", gb) - } - let tb = gb / 1024.0 - return String(format: "%.1f TB", tb) -#endif } } diff --git a/Sources/Imager/Imager.swift b/Sources/Imager/Imager.swift deleted file mode 100644 index 69d291f..0000000 --- a/Sources/Imager/Imager.swift +++ /dev/null @@ -1,51 +0,0 @@ -import Foundation - -/// Public typealias for disk update notifications, usable by consumers -public typealias DiskUpdateCallback = () -> Void - -/// Protocol defining the interface for an image writer. -/// -/// This protocol defines the requirements for a component that can write -/// an image file to a drive. -public protocol Imager { - /// Path to the image file to be imaged. - var imageFilePath: String { get } - - /// Path to the drive to be imaged. - var drivePath: String { get } - - /// The current state of the imaging process. - var state: ImagerState { get } - - /// Get a list of available drives that can be imaged in alphabetical order. - /// - /// - Parameter onlyExternalDrives: If true, only external drives will be returned. - /// - Returns: An array of drive identifiers that can be used for imaging. - /// - Throws: An `ImagerError` if the drives cannot be enumerated. - func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] - - /// Begin imaging the drive. - /// - /// This method starts the process of writing the image file to the drive. - /// It may throw an `ImagerError` if the operation cannot be completed. - func startImaging() throws - - /// Stop the imaging process. - /// - /// This method stops the imaging process if it is currently running. - /// It may throw an `ImagerError` if the operation cannot be stopped. - func stopImaging() throws - - /// Register a handler to receive progress updates. - /// - /// - Parameter handler: A closure that will be called with progress updates. - /// The closure receives the current progress and an optional error if one occurred. - func progress(_ handler: @escaping (Foundation.Progress, ImagerError?) -> Void) - - /// Initialize a new imager with the specified image file and drive paths. - /// - /// - Parameters: - /// - imageFilePath: The path to the image file to be written. - /// - drivePath: The path to the drive where the image will be written. - init(imageFilePath: String, drivePath: String) -} \ No newline at end of file diff --git a/Sources/Imager/ImagerError.swift b/Sources/Imager/ImagerError.swift index ef86ef8..2a8c03b 100644 --- a/Sources/Imager/ImagerError.swift +++ b/Sources/Imager/ImagerError.swift @@ -2,7 +2,7 @@ /// /// These errors represent various failure scenarios that might occur when /// attempting to write an image file to a drive. -public enum ImagerError: Error { +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. 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/ImagerLib.swift b/Sources/Imager/ImagerLib.swift deleted file mode 100644 index fbaf607..0000000 --- a/Sources/Imager/ImagerLib.swift +++ /dev/null @@ -1,20 +0,0 @@ -import Foundation - -/// Factory methods for creating the appropriate Imager implementation for the current platform. -public enum ImagerFactory { - /// Creates an Imager instance appropriate for the current platform. - /// - /// - Parameters: - /// - imageFilePath: The path to the image file to be written. - /// - drivePath: The path to the drive where the image will be written. - /// - Returns: An Imager instance for the current platform. - 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) - #else - fatalError("Unsupported platform") - #endif - } -} \ No newline at end of file diff --git a/Sources/Imager/ImagerState.swift b/Sources/Imager/ImagerState.swift index aa31736..f275206 100644 --- a/Sources/Imager/ImagerState.swift +++ b/Sources/Imager/ImagerState.swift @@ -2,7 +2,7 @@ /// /// This enum tracks the lifecycle of an imaging operation, from initialization through /// completion or failure. -public enum ImagerState { +public enum ImagerState: Sendable { /// The imager is initialized but has not started imaging. /// /// This is the default state before `startImaging()` is called. diff --git a/Sources/Imager/LinuxImager.swift b/Sources/Imager/LinuxImager.swift deleted file mode 100644 index 94fb88f..0000000 --- a/Sources/Imager/LinuxImager.swift +++ /dev/null @@ -1,42 +0,0 @@ -#if os(Linux) -import Foundation - -/// Concrete implementation of the Imager protocol for Linux. -public class LinuxImager: Imager, @unchecked Sendable { - public let imageFilePath: String - public let drivePath: String - public var state: ImagerState = .idle - private var progress: Foundation.Progress? - private var progressHandler: ((Foundation.Progress, ImagerError?) -> Void)? - - private var ddProcess: Process? - private var outputPipe: Pipe? - private var errorPipe: Pipe? - - public required init(imageFilePath: String = "", drivePath: String = "") { - self.imageFilePath = imageFilePath - self.drivePath = drivePath - self.progress = Progress(totalUnitCount: 0) - } - - public func availableDrivesToImage(onlyExternalDrives: Bool) throws -> [Drive] { - // Implementation for Linux would use commands like lsblk to list block devices - // For now, return an empty array - return [] - } - - public func startImaging() throws { - // Implementation would use dd command similar to MacOSImager - throw ImagerError.notImplemented(reason: "Linux imaging not yet implemented") - } - - public func stopImaging() throws { - // Implementation would stop the dd process - throw ImagerError.notImplemented(reason: "Linux imaging not yet implemented") - } - - public func progress(_ handler: @escaping (Foundation.Progress, ImagerError?) -> Void) { - self.progressHandler = handler - } -} -#endif // os(Linux) \ 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/MacOSImager.swift b/Sources/Imager/MacOSImager.swift deleted file mode 100644 index a6d2320..0000000 --- a/Sources/Imager/MacOSImager.swift +++ /dev/null @@ -1,374 +0,0 @@ -#if os(macOS) -import Foundation -import DiskArbitration -import IOKit - -/// Concrete implementation of the Imager protocol for macOS. -public class MacOSImager: Imager, @unchecked Sendable { - public let imageFilePath: String - public let drivePath: String - public var state: ImagerState = .idle - var progress: Foundation.Progress? - var progressObservation: NSKeyValueObservation? - var progressHandler: ((Foundation.Progress, ImagerError?) -> Void)? - - private var imageFileHandle: FileHandle? - private var driveFileHandle: FileHandle? - private var ddProcess: Process? - private var outputPipe: Pipe? - private var errorPipe: Pipe? - private let progressQueue = DispatchQueue(label: "com.apache-edge.imager.progress", qos: .utility) - private var lastProgressUpdate = Date() - - public required init(imageFilePath: String = "", drivePath: String = "") { - self.imageFilePath = imageFilePath - self.drivePath = drivePath - self.progress = Foundation.Progress(totalUnitCount: 0) - } - - // ... availableDrivesToImage (no changes needed here for this fix) ... - 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 - } - - - // ... startImaging, stopImaging ... - public func startImaging() throws { - guard imageFilePath != "", drivePath != "" else { - throw ImagerError.invalidDrive(reason: "Image file path or drive path not set.") - } - guard case .idle = state else { - throw ImagerError.processingInterrupted(reason: "Imaging process already running or not idle.") - } - - let fileManager = FileManager.default - guard fileManager.fileExists(atPath: imageFilePath) else { - throw ImagerError.invalidImageFile(reason: "Image file not found at \(imageFilePath)") - } - - guard fileManager.fileExists(atPath: drivePath) else { - throw ImagerError.invalidDrive(reason: "Drive path not found at \(drivePath)") - } - let rawDrivePath = drivePath.replacingOccurrences(of: "/dev/disk", with: "/dev/rdisk") - guard fileManager.fileExists(atPath: rawDrivePath) else { - throw ImagerError.invalidDrive(reason: "Raw drive device not found at \(rawDrivePath)") - } - - print("Unmounting drive \(drivePath)...") - let unmountTask = Process() - unmountTask.executableURL = URL(fileURLWithPath: "/usr/sbin/diskutil") - unmountTask.arguments = ["unmountDisk", "force", drivePath] - do { - try unmountTask.run() - unmountTask.waitUntilExit() - if unmountTask.terminationStatus != 0 { - print("Warning: diskutil unmountDisk failed (status \(unmountTask.terminationStatus)). Proceeding anyway...") - } - } catch { - throw ImagerError.processingInterrupted(reason: "Failed to run diskutil unmountDisk: \(error.localizedDescription)") - } - - let imageSize: UInt64 - do { - let imageAttrs = try fileManager.attributesOfItem(atPath: imageFilePath) - guard let size = imageAttrs[.size] as? UInt64 else { - throw ImagerError.invalidImageFile(reason: "Could not read image file size attribute.") - } - imageSize = size - } catch { - throw ImagerError.invalidImageFile(reason: "Could not read image file attributes: \(error.localizedDescription)") - } - progress?.totalUnitCount = Int64(imageSize) - progress?.completedUnitCount = 0 - - print("Starting imaging process using dd...") - print("Source: \(imageFilePath)") - print("Target: \(rawDrivePath)") - - ddProcess = Process() - outputPipe = Pipe() - errorPipe = Pipe() - - guard let ddProcess = ddProcess, let errorPipe = errorPipe else { - throw ImagerError.processingInterrupted(reason: "Failed to create process or pipes for dd.") - } - - ddProcess.executableURL = URL(fileURLWithPath: "/bin/dd") - ddProcess.arguments = ["if=\(imageFilePath)", "of=\(rawDrivePath)", "bs=4m", "status=progress"] - - ddProcess.standardError = errorPipe - - ddProcess.terminationHandler = { [weak self] process in - self?.progressQueue.async { - self?.handleImagingCompletion(process: process) - } - } - - errorPipe.fileHandleForReading.readabilityHandler = { [weak self] fileHandle in - let data = fileHandle.availableData - if data.isEmpty { - self?.progressQueue.async { - errorPipe.fileHandleForReading.readabilityHandler = nil - } - } else { - self?.progressQueue.async { - if let output = String(data: data, encoding: .utf8) { - self?.parseDdProgress(output) - } - } - } - } - - do { - try ddProcess.run() - state = .imaging - print("dd process started (PID: \(ddProcess.processIdentifier)).") - } catch { - state = .idle - cleanupResources() - throw ImagerError.processingInterrupted(reason: "Failed to launch dd process: \(error.localizedDescription)") - } - } - - public func stopImaging() { - guard let ddProcess = ddProcess, ddProcess.isRunning else { - print("Imaging process not running.") - return - } - - state = .cancelling - print("Stopping imaging process (PID: \(ddProcess.processIdentifier))...") - ddProcess.interrupt() - - DispatchQueue.global(qos: .userInitiated).asyncAfter(deadline: .now() + 2.0) { [weak self] in - if self?.ddProcess?.isRunning == true { - print("Process still running, sending SIGKILL...") - self?.ddProcess?.terminate() - } - } - } - - private func handleImagingCompletion(process: Process) { - let exitCode = process.terminationStatus - let reason = process.terminationReason - - errorPipe?.fileHandleForReading.readabilityHandler = nil - - if case .cancelling = state { - print("Imaging process cancelled.") - progress?.cancel() - state = .idle - } else if exitCode == 0 { - print("Imaging process completed successfully.") - progress?.completedUnitCount = progress?.totalUnitCount ?? 0 - state = .completed - } else { - print("Imaging process failed. Exit code: \(exitCode), Reason: \(reason).") - let errorData = errorPipe?.fileHandleForReading.readDataToEndOfFile() ?? Data() - if let errorString = String(data: errorData, encoding: .utf8), !errorString.isEmpty { - print("dd stderr: \(errorString)") - } - state = .failed(ImagerError.processingInterrupted(reason: "dd process failed with exit code \(exitCode)")) - } - - cleanupResources() - let workItem = DispatchWorkItem { - if case .completed = self.state { - self.state = .idle - } else if case .cancelling = self.state { - self.state = .idle - } else if case .failed = self.state { - self.state = .idle - } - } - DispatchQueue.main.asyncAfter(deadline: .now() + 0.5, execute: workItem) - } - - private func parseDdProgress(_ output: String) { - let lines = output.split(separator: "\r") - for line in lines { - if let bytesString = line.split(separator: " ").first, - let bytesCopied = Int64(bytesString) { - - if Date().timeIntervalSince(lastProgressUpdate) > 0.1 { - DispatchQueue.main.async { - self.progress?.completedUnitCount = bytesCopied - self.lastProgressUpdate = Date() - } - } - } - } - } - - private func cleanupResources() { - errorPipe?.fileHandleForReading.readabilityHandler = nil - - try? imageFileHandle?.close() - try? driveFileHandle?.close() - try? outputPipe?.fileHandleForReading.close() - try? outputPipe?.fileHandleForWriting.close() - try? errorPipe?.fileHandleForReading.close() - try? errorPipe?.fileHandleForWriting.close() - - imageFileHandle = nil - driveFileHandle = nil - outputPipe = nil - errorPipe = nil - ddProcess = nil - progressObservation = nil - progressHandler = nil - } - - // Deinitializer to ensure cleanup - deinit { - cleanupResources() - progressObservation = nil - } - - // --- Imager Protocol Implementation --- - public func progress(_ handler: @escaping (Foundation.Progress, ImagerError?) -> Void) { - self.progressHandler = handler - } -} - -#endif // os(macOS) \ No newline at end of file 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/ImagerImageDiskSubCommand.swift b/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift index f888c6f..04e2b08 100644 --- a/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift +++ b/Sources/edge/Commands/Imager/ImagerImageDiskSubCommand.swift @@ -1,6 +1,6 @@ -import Imager import ArgumentParser import Foundation +import Imager struct ImagerImageDiskSubCommand: AsyncParsableCommand { static let configuration = CommandConfiguration( @@ -29,21 +29,77 @@ struct ImagerImageDiskSubCommand: AsyncParsableCommand { print(" Target Drive: /dev/\(drivePath)") print(" This may take some time...") - let imager = MacOSImager(imageFilePath: source, drivePath: drivePath) - - do { - // startImaging is synchronous, but handles launching the process - try imager.startImaging() - // If startImaging returns without throwing, the process was initiated. - // Actual success/failure/progress is handled internally by MacOSImager (potentially). - // A more robust solution would involve async streams for progress. - print("\nImaging process initiated successfully.") - print("Monitor system activity or logs for completion.") - } catch let error as ImagerError { - print("\nError during imaging: \(error.description)") - throw ExitCode.failure - } catch { - print("\nAn unexpected error occurred: \(error.localizedDescription)") + 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 } } diff --git a/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift index 653e1c0..c5957aa 100644 --- a/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift +++ b/Sources/edge/Commands/Imager/ImagerListDiskSubCommand.swift @@ -1,9 +1,6 @@ import ArgumentParser import Foundation import Imager -#if canImport(Darwin) -import Darwin // For signal handling -#endif struct ImagerListDisksSubCommand: AsyncParsableCommand { @@ -20,18 +17,18 @@ struct ImagerListDisksSubCommand: AsyncParsableCommand { var timeout: Int? func run() async throws { - let imager = ImagerFactory.createImager() + let diskLister = ImagerFactory.createDiskLister() try await MainActor.run { - try listAndPrintDisks(imager: imager, listAll: self.all) + try listAndPrintDisks(lister: diskLister, listAll: self.all) } } @MainActor - private func listAndPrintDisks(imager: Imager, listAll: Bool) throws { - let drives = try imager.availableDrivesToImage(onlyExternalDrives: !listAll) + private func listAndPrintDisks(lister: DiskLister, listAll: Bool) throws { + let drives = try lister.availableDrivesToImage(onlyExternalDrives: !listAll) if drives.isEmpty { - print(listAll ? "No drives found." : "No suitable external drives found. Use --all to list internal drives.") + print("No suitable drives found.") fflush(stdout) return } 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 c804253..e309bea 100644 --- a/Sources/edge/EdgeCLI.swift +++ b/Sources/edge/EdgeCLI.swift @@ -8,7 +8,8 @@ struct EdgeCLI: AsyncParsableCommand { abstract: "Edge CLI", subcommands: [ RunCommand.self, - ImagerCommand.self + ImagerCommand.self, + ProgressBarExampleCommand.self ] ) }