Merge pull request #51 from Automattic/update/vm-commands

Update Commands to use libhostmgr

Author: jkmassel

Package.swift

@@ -27,7 +27,6 @@ let package = Package(
             name: "hostmgr",
             dependencies: [
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
-                .product(name: "prlctl", package: "prlctl"),
                 .product(name: "Logging", package: "swift-log"),
                 .product(name: "kcpassword", package: "kcpassword-swift"),
                 .target(name: "libhostmgr"),

Sources/hostmgr/VMCommand.swift

@@ -1,89 +1,18 @@
 import Foundation
 import ArgumentParser
-import prlctl
 
 struct VMCommand: AsyncParsableCommand {
 
     static let configuration = CommandConfiguration(
         commandName: "vm",
         abstract: "Allows working with VMs",
         subcommands: [
+            VMCleanCommand.self,
+            VMExistsCommand.self,
+            VMFetchCommand.self,
             VMListCommand.self,
             VMStartCommand.self,
-            VMStopCommand.self,
-            VMDetailsCommand.self,
-            VMCleanCommand.self
+            VMStopCommand.self
         ]
     )
 }
-
-// Allow passing VM objects directly
-extension VM: ExpressibleByArgument {
-
-    public init?(argument: String) {
-        guard let foundVM = try? Parallels().lookupVM(named: argument) else {
-            return nil
-        }
-
-        self = foundVM
-    }
-
-    public static var allValueStrings: [String] {
-        guard let vms = try? Parallels().lookupAllVMs() else {
-            return []
-        }
-
-        let names = vms.map { $0.name }
-        let uuids = vms.map { $0.uuid }
-
-        return names + uuids
-    }
-}
-
-extension StoppedVM: ExpressibleByArgument {
-    public init?(argument: String) {
-        guard let foundVM = try? Parallels()
-            .lookupStoppedVMs()
-            .first(where: { $0.name == argument || $0.uuid == argument })
-        else {
-            return nil
-        }
-
-        self = foundVM
-    }
-
-    public static var allValueStrings: [String] {
-        guard let vms = try? Parallels().lookupStoppedVMs() else {
-            return []
-        }
-
-        let names = vms.map { $0.name }
-        let uuids = vms.map { $0.uuid }
-
-        return names + uuids
-    }
-}
-
-extension RunningVM: ExpressibleByArgument {
-    public init?(argument: String) {
-        guard let foundVM = try? Parallels()
-            .lookupRunningVMs()
-            .first(where: { $0.name == argument || $0.uuid == argument })
-        else {
-            return nil
-        }
-
-        self = foundVM
-    }
-
-    public static var allValueStrings: [String] {
-        guard let vms = try? Parallels().lookupRunningVMs() else {
-            return []
-        }
-
-        let names = vms.map { $0.name }
-        let uuids = vms.map { $0.uuid }
-
-        return names + uuids
-    }
-}

Sources/hostmgr/commands/benchmark/DiskBenchmark.swift

@@ -1,6 +1,5 @@
 import Foundation
 import ArgumentParser
-import prlctl
 import Logging
 
 private let startDate = Date()

Sources/hostmgr/commands/benchmark/NetworkBenchmark.swift

@@ -1,6 +1,5 @@
 import Foundation
 import ArgumentParser
-import prlctl
 import Logging
 import libhostmgr
 
@@ -13,30 +12,19 @@ struct NetworkBenchmark: AsyncParsableCommand {
         abstract: "Test Network Speed"
     )
 
-    private static let limiter = Limiter(policy: .throttle, operationsPerSecond: 1)
-
     func run() async throws {
         let remoteImages = try await RemoteVMRepository().listImages(sortedBy: .size)
 
-        guard let file = remoteImages.last else {
-            throw CleanExit.message("Unable to find a remote image to use as a network benchmark")
+        guard !remoteImages.isEmpty else {
+            Console.error("Unable to find a remote image to use as a network benchmark")
+            throw ExitCode(rawValue: -1)
         }
 
-        let manager = S3Manager(
-            bucket: Configuration.shared.vmImagesBucket,
-            region: Configuration.shared.vmImagesRegion
-        )
-
-        let progressBar = Console.startFileDownload(file.imageObject)
+        Console.heading("Starting Benchmark")
 
-        try await manager.download(
-            object: file.imageObject,
-            to: FileManager.default.temporaryFilePath(),
-            progressCallback: progressBar.update
-        )
-    }
-
-    private func imageSizeSort(_ lhs: RemoteVMImage, _ rhs: RemoteVMImage) -> Bool {
-        lhs.imageObject.size < rhs.imageObject.size
+        for remoteImage in remoteImages {
+            let path = try await libhostmgr.downloadRemoteImage(remoteImage)
+            try FileManager.default.removeItem(at: path)
+        }
     }
 }

Sources/hostmgr/commands/generate/GenerateBuildkiteJobScript.swift

@@ -1,6 +1,5 @@
 import Foundation
 import ArgumentParser
-import prlctl
 import libhostmgr
 
 struct GenerateBuildkiteJobScript: ParsableCommand {

Sources/hostmgr/commands/sync/SyncAuthorizedKeysCommand.swift

@@ -4,35 +4,15 @@ import libhostmgr
 
 struct SyncAuthorizedKeysCommand: AsyncParsableCommand, FollowsCommandPolicies {
 
+    enum Constants {
+        static let s3Key = "authorized_keys"
+    }
+
     static let configuration = CommandConfiguration(
         commandName: Configuration.SchedulableSyncCommand.authorizedKeys.rawValue,
         abstract: "Set this machine's authorized_keys file"
     )
 
-    @Option(
-        name: .shortAndLong,
-        help: "The S3 bucket containing the `authorized_keys` file"
-    )
-    var bucket: String = Configuration.shared.authorizedKeysBucket
-
-    @Option(
-        name: .shortAndLong,
-        help: "The S3 region for the bucket"
-    )
-    var region: String = Configuration.shared.authorizedKeysRegion
-
-    @Option(
-        name: .shortAndLong,
-        help: "The S3 path to the authorized_keys file"
-    )
-    var key: String = "authorized_keys"
-
-    @Option(
-        name: .shortAndLong,
-        help: "The path to your authorized_keys file on disk (defaults to ~/.ssh/authorized_keys)"
-    )
-    var destination: String = Paths.authorizedKeysFilePath.path
-
     @OptionGroup
     var options: SharedSyncOptions
 
@@ -44,25 +24,36 @@ struct SyncAuthorizedKeysCommand: AsyncParsableCommand, FollowsCommandPolicies {
     ]
 
     func run() async throws {
+        let destination = Paths.authorizedKeysFilePath
+
         try to(evaluateCommandPolicies(), unless: options.force)
-        logger.debug("Job schedule allows for running")
 
-        logger.info("Downloading file from s3://\(bucket)/\(key) in \(region) to \(destination)")
+        Console.heading("Syncing Authorized Keys")
 
-        let s3Manager = S3Manager(bucket: self.bucket, region: self.region)
+        let s3Manager = S3Manager(
+            bucket: Configuration.shared.authorizedKeysBucket,
+            region: Configuration.shared.authorizedKeysRegion
+        )
 
-        guard let object = try await s3Manager.lookupObject(atPath: key) else {
-            logger.error("Unable to locate authorized_keys file – exiting")
+        guard let object = try await s3Manager.lookupObject(atPath: Constants.s3Key) else {
+            Console.error("Unable to locate authorized_keys file – exiting")
             throw ExitCode(rawValue: 1)
         }
 
-        let url = URL(fileURLWithPath: self.destination)
-        try await s3Manager.download(object: object, to: url, progressCallback: nil)
+        let progressBar = Console.startFileDownload(object)
+
+        try await s3Manager.download(
+            object: object,
+            to: destination,
+            progressCallback: progressBar.update
+        )
 
         /// Fix the permissions on the file, if needed
+        Console.info("Setting file permissions on \(destination)")
         try FileManager.default.setAttributes([
             .posixPermissions: 0o600
-        ], ofItemAtPath: self.destination)
+        ], ofItemAt: destination)
+        Console.success("Authorized Key Sync Complete")
 
         try recordLastRun()
     }

Sources/hostmgr/commands/vm/VMClean.swift

@@ -1,21 +1,19 @@
 import Foundation
 import ArgumentParser
-import prlctl
+import libhostmgr
 
-struct VMCleanCommand: ParsableCommand {
+struct VMCleanCommand: AsyncParsableCommand {
 
     static let configuration = CommandConfiguration(
         commandName: "clean",
-        abstract: "Cleans up a VM prior to it being reused"
+        abstract: "Clean up the VM environment prior to running another job"
     )
 
-    @Option(
-        name: .shortAndLong,
-        help: "The VM to clean"
-    )
-    var virtualMachine: StoppedVM
+    func run() async throws {
+        try libhostmgr.resetVMStorage()
 
-    func run() throws {
-        try virtualMachine.clean()
+        // Clean up no-longer-needed local images
+        let deleteList = try await libhostmgr.listLocalImagesToDelete()
+        try libhostmgr.deleteLocalImages(list: deleteList)
     }
 }

Sources/hostmgr/commands/vm/VMDetails.swift

@@ -0,0 +1,21 @@
+import Foundation
+import ArgumentParser
+import libhostmgr
+
+struct VMExistsCommand: ParsableCommand {
+    static let configuration = CommandConfiguration(
+        commandName: "exists",
+        abstract: "Exits with code 0 if the named VM exists. Otherwise exits with code 1"
+    )
+
+    @Argument(help: "The exact name of the VM")
+    var name: String
+
+    func run() throws {
+        guard let localVM = try LocalVMRepository().lookupVM(withName: self.name) else {
+            Console.crash(message: "There is no local VM named \(self.name)", reason: .fileNotFound)
+        }
+
+        Console.success("VM \(localVM.basename) exists")
+    }
+}

Sources/hostmgr/commands/vm/VMExists.swift

@@ -0,0 +1,33 @@
+import ArgumentParser
+import libhostmgr
+
+struct VMFetchCommand: AsyncParsableCommand {
+
+    static let configuration = CommandConfiguration(
+        commandName: "fetch",
+        abstract: "Download a given image if it's not already present"
+    )
+
+    @Argument(
+        help: "The name of the image you would like to download"
+    )
+    var name: String
+
+    func run() async throws {
+
+        if let localVM = try LocalVMRepository().lookupVM(withName: name) {
+            if localVM.state == .packaged {
+                try await libhostmgr.unpackVM(name: localVM.basename)
+                return
+            } else {
+                Console.exit(
+                    message: "VM is present locally",
+                    style: .success
+                )
+            }
+
+        }
+
+        try await libhostmgr.fetchRemoteImage(name: self.name)
+    }
+}