Merge pull request #26 from orlandos-nl/feature/jo-command-stream

Don't stop collecting on stderr chunks. Provide an AsyncThrowingStream for executeCommand data
orlandos-nl · Feb 20, 2023 · f8b7f2d · f8b7f2d
2 parents 2e6a7fa + 9e40248
commit f8b7f2d
Show file tree

Hide file tree

Showing 3 changed files with 248 additions and 56 deletions.
diff --git a/Sources/Citadel/Client.swift b/Sources/Citadel/Client.swift
@@ -1,4 +1,5 @@
 import NIO
+import Logging
 import NIOSSH
 
 public struct SSHAlgorithms {
@@ -70,10 +71,11 @@ public final class SSHClient {
     private let algorithms: SSHAlgorithms
     private let protocolOptions: Set<SSHProtocolOption>
     private var onDisconnect: (@Sendable () -> ())?
+    public let logger = Logger(label: "nl.orlandos.citadel.client")
     public var isConnected: Bool {
         session.channel.isActive
     }
-
+    
     /// The event loop that this SSH connection is running on.
     public var eventLoop: EventLoop {
         session.channel.eventLoop

diff --git a/Sources/Citadel/SSHAuthenticationMethod.swift b/Sources/Citadel/SSHAuthenticationMethod.swift
@@ -4,15 +4,24 @@ import Crypto
 
 /// Represents an authentication method.
 public struct SSHAuthenticationMethod: NIOSSHClientUserAuthenticationDelegate {
-    private let username: String
-    private let offer: NIOSSHUserAuthenticationOffer.Offer
+    private enum Implementation {
+        case custom(NIOSSHClientUserAuthenticationDelegate)
+        case user(String, offer: NIOSSHUserAuthenticationOffer.Offer)
+    }
+
+    private let implementation: Implementation
 
     internal init(
         username: String,
         offer: NIOSSHUserAuthenticationOffer.Offer
     ) {
-        self.username = username
-        self.offer = offer
+        self.implementation = .user(username, offer: offer)
+    }
+
+    internal init(
+        custom: NIOSSHClientUserAuthenticationDelegate
+    ) {
+        self.implementation = .custom(custom)
     }
 
     /// Creates a password based authentication method.
@@ -63,30 +72,39 @@ public struct SSHAuthenticationMethod: NIOSSHClientUserAuthenticationDelegate {
         return SSHAuthenticationMethod(username: username, offer: .privateKey(.init(privateKey: .init(p521Key: privateKey))))
     }
 
+    public static func custom(_ auth: NIOSSHClientUserAuthenticationDelegate) -> SSHAuthenticationMethod {
+        return SSHAuthenticationMethod(custom: auth)
+    }
+
     public func nextAuthenticationType(
         availableMethods: NIOSSHAvailableUserAuthenticationMethods,
         nextChallengePromise: EventLoopPromise<NIOSSHUserAuthenticationOffer?>
     ) {
-        switch offer {
-        case .password:
-            guard availableMethods.contains(.password) else {
-                nextChallengePromise.fail(SSHClientError.unsupportedPasswordAuthentication)
-                return
-            }
-        case .hostBased:
-            guard availableMethods.contains(.hostBased) else {
-                nextChallengePromise.fail(SSHClientError.unsupportedHostBasedAuthentication)
-                return
-            }
-        case .privateKey:
-            guard availableMethods.contains(.publicKey) else {
-                nextChallengePromise.fail(SSHClientError.unsupportedPrivateKeyAuthentication)
-                return
+        switch implementation {
+        case .user(let username, offer: let offer):
+            switch offer {
+            case .password:
+                guard availableMethods.contains(.password) else {
+                    nextChallengePromise.fail(SSHClientError.unsupportedPasswordAuthentication)
+                    return
+                }
+            case .hostBased:
+                guard availableMethods.contains(.hostBased) else {
+                    nextChallengePromise.fail(SSHClientError.unsupportedHostBasedAuthentication)
+                    return
+                }
+            case .privateKey:
+                guard availableMethods.contains(.publicKey) else {
+                    nextChallengePromise.fail(SSHClientError.unsupportedPrivateKeyAuthentication)
+                    return
+                }
+            case .none:
+                ()
             }
-        case .none:
-            ()
+
+            nextChallengePromise.succeed(NIOSSHUserAuthenticationOffer(username: username, serviceName: "", offer: offer))
+        case .custom(let implementation):
+            implementation.nextAuthenticationType(availableMethods: availableMethods, nextChallengePromise: nextChallengePromise)
         }
-
-        nextChallengePromise.succeed(NIOSSHUserAuthenticationOffer(username: username, serviceName: "", offer: offer))
     }
 }