Skip to content

Commit

Permalink
DefaultAgent with a few interface improvements
Browse files Browse the repository at this point in the history
- Adding state as the agent can be unavailable.
- Differentiate when Settings do not exist (nil), from where there are issues
reading them.
  • Loading branch information
Carlos Cabanero committed Jun 24, 2024
1 parent 33aec3c commit b82e66b
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 35 deletions.
13 changes: 8 additions & 5 deletions Blink/Commands/ssh/SSHAgentAdd.swift
Original file line number Diff line number Diff line change
Expand Up @@ -98,15 +98,18 @@ public class BlinkSSHAgentAdd: NSObject {
let currentRunLoop = RunLoop.current

public func start(_ argc: Int32, argv: [String], session: MCPSession) -> Int32 {
do {
command = try BlinkSSHAgentAddCommand.parse(Array(argv[1...]))
do {
command = try BlinkSSHAgentAddCommand.parse(Array(argv[1...]))
} catch {
let message = BlinkSSHAgentAddCommand.message(for: error)
print(message, to: &stderr)
return -1
}

let _ = SSHDefaultAgent.instance
guard let defaultAgent = SSHDefaultAgent.instance else {
print("Default Agent is not available.", to: &stderr)
return -1
}

if command.remove {
let keyName = command.keyName ?? "id_rsa"
Expand All @@ -121,7 +124,7 @@ public class BlinkSSHAgentAdd: NSObject {
}

if command.list {
for key in SSHDefaultAgent.instance.ring {
for key in defaultAgent.ring {
let str = BKPubKey.withID(key.name)?.publicKey ?? ""
print("\(str) \(key.name)", to: &stdout)
}
Expand All @@ -137,7 +140,7 @@ public class BlinkSSHAgentAdd: NSObject {
return -1;
}

for key in SSHDefaultAgent.instance.ring {
for key in defaultAgent.ring {
if let blob = try? key.signer.publicKey.encode()[4...],
let sshkey = try? SSHKey(fromPublicBlob: blob)
{
Expand Down
7 changes: 6 additions & 1 deletion Blink/Commands/ssh/SSHConfigProvider.swift
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,12 @@ extension SSHClientConfigProvider {
}

// Link to Default Agent
agent.linkTo(agent: SSHDefaultAgent.instance)
if let defaultAgent = SSHDefaultAgent.instance {
agent.linkTo(agent: defaultAgent)
} else {
printLn("Default agent is not available.")
}

return agent
}

Expand Down
77 changes: 49 additions & 28 deletions Blink/Commands/ssh/SSHDefaultAgent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ import SSH


final class SSHDefaultAgent {
public static var instance: SSHAgent {
public static var instance: SSHAgent? {
if let agent = Self._instance {
return agent
} else {
Expand All @@ -44,6 +44,7 @@ final class SSHDefaultAgent {
}
private static var _instance: SSHAgent? = nil
private init() {}
// The pool is responsible for the location of Agents.
private static let defaultAgentFile: URL = BlinkPaths.blinkAgentSettingsURL().appendingPathComponent("default")

enum Error: Swift.Error, LocalizedError {
Expand All @@ -57,30 +58,41 @@ final class SSHDefaultAgent {
}
}

private static func load() -> SSHAgent {
let instance = SSHAgent()
Self._instance = instance
// If the Settings are not available, the agent is initialized with the default configuration.
// If there is a problem with the location, it is safe to assume that it will persist.
if let settings = try? getSettings() {
try? applySettings(settings)
} else {
try? setSettings(BKAgentSettings(prompt: .Confirm, keys: []))
// Load the (default) agent in the pool. If the Agent cannot be loaded, it will be unavailable (nil).
// If the agent doesn't exist, it will be initialized (default only).
private static func load() -> SSHAgent? {
do {
if let settings = try getSettings() {
try setAgentInstance(with: settings)
} else {
try setSettings(BKAgentSettings())
}
} catch {
return nil
}
return instance

return Self._instance
}

// Create the settings for the agent and load it in the pool.
// NOTE: For non-default agents, this would be the main initialization method.
static func setSettings(_ settings: BKAgentSettings) throws {
try BKAgentSettings.save(settings: settings, to: defaultAgentFile)
try applySettings(settings)
try setAgentInstance(with: settings)
}

private static func applySettings(_ settings: BKAgentSettings) throws {
let agent = Self.instance
agent.clear()

private static func setAgentInstance(with settings: BKAgentSettings) throws {
let bkConfig = try BKConfig()

let agent: SSHAgent
if let _instance = Self._instance {
agent = _instance
agent.clear()
} else {
agent = SSHAgent()
Self._instance = agent
}

settings.keys.forEach { key in
if let (signer, name) = bkConfig.signer(forIdentity: key) {
if let constraints = settings.constraints() {
Expand All @@ -90,18 +102,21 @@ final class SSHDefaultAgent {
}
}

static func getSettings() throws -> BKAgentSettings {
try BKAgentSettings.load(from: defaultAgentFile)
static func getSettings() throws -> BKAgentSettings? {
try BKAgentSettings.read(from: defaultAgentFile)
}

// Applying settings clears the agent first. Adding a key doesn't modify or reset previous constraints.
static func addKey(named keyName: String) throws {
let settings = try getSettings()
guard let agent = Self.instance,
let settings = try getSettings() else {
return
}

if settings.keys.contains(keyName) {
return
}

let agent = Self.instance
let bkConfig = try BKConfig()

if let (signer, name) = bkConfig.signer(forIdentity: keyName) {
Expand All @@ -124,16 +139,17 @@ final class SSHDefaultAgent {

static func removeKey(named keyName: String) throws -> Signer? {
// Remove from settings and apply
let settings = try getSettings()
guard settings.keys.contains(keyName) else {
guard let agent = Self.instance,
let settings = try getSettings(),
settings.keys.contains(keyName) else {
return nil
}

var keys = settings.keys
keys.removeAll(where: { $0 == keyName })
try BKAgentSettings.save(settings: BKAgentSettings(prompt: settings.prompt, keys: keys), to: defaultAgentFile)

return Self.instance.removeKey(keyName)
return agent.removeKey(keyName)
}
}

Expand All @@ -151,17 +167,22 @@ struct BKAgentSettings: Codable, Equatable {
let prompt: BKAgentSettingsPrompt
let keys: [String]

// init(prompt: BKAgentSettingsPrompt, keys: [String]) {
// self.prompt = prompt
// self.keys = keys
// }
init(prompt: BKAgentSettingsPrompt, keys: [String]) {
self.prompt = prompt
self.keys = keys
}

init() { self = Self(prompt: .Confirm, keys: []) }

static func save(settings: BKAgentSettings, to file: URL) throws {
let data = try JSONEncoder().encode(settings)
try data.write(to: file)
}

static func load(from file: URL) throws -> BKAgentSettings {
fileprivate static func read(from file: URL) throws -> BKAgentSettings? {
guard FileManager.default.fileExists(atPath: file.path) else {
return nil
}
let data = try Data(contentsOf: file)
return try JSONDecoder().decode(BKAgentSettings.self, from: data)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ struct DefaultAgentSettingsView: View {
if agentSettings == nil {
print("Init settings")
do {
let agentSettings = try SSHDefaultAgent.getSettings()
let agentSettings = try SSHDefaultAgent.getSettings() ?? BKAgentSettings()
self.agentSettings = agentSettings
} catch {
self.alertMessage = "Failed to get settings: \(error.localizedDescription)"
Expand Down

0 comments on commit b82e66b

Please sign in to comment.