Merge pull request #65 from orlandos-nl/jo/updated-shell

Updated the SSH shell example

.gitignore

@@ -5,4 +5,5 @@
 /*.xcodeproj
 xcuserdata/
 Package.resolved
+citadel_host_key_ed25519
 .vscode/launch.json

Examples/TerminalAppServer/.gitignore

@@ -0,0 +1,8 @@
+.DS_Store
+/.build
+/Packages
+xcuserdata/
+DerivedData/
+.swiftpm/configuration/registries.json
+.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
+.netrc

Examples/TerminalAppServer/Package.swift

@@ -0,0 +1,26 @@
+// swift-tools-version: 5.10
+// The swift-tools-version declares the minimum version of Swift required to build this package.
+
+import PackageDescription
+
+let package = Package(
+    name: "TerminalAppServer",
+    platforms: [
+        .macOS(.v12),
+    ],
+    dependencies: [
+        .package(url: "https://github.com/joannis/SwiftTUI.git", branch: "jo/allow-use-with-concurrency"),
+        .package(path: "../.."),
+    ],
+    targets: [
+        // Targets are the basic building blocks of a package, defining a module or a test suite.
+        // Targets can depend on other targets in this package and products from dependencies.
+        .executableTarget(
+            name: "TerminalAppServer",
+            dependencies: [
+                .product(name: "Citadel", package: "Citadel"),
+                .product(name: "SwiftTUI", package: "SwiftTUI"),
+            ]
+        ),
+    ]
+)

Examples/TerminalAppServer/Sources/main.swift

@@ -0,0 +1,110 @@
+import Citadel
+import Crypto
+import Foundation
+import NIO
+import NIOFoundationCompat
+import NIOSSH
+import SwiftTUI
+
+@main struct ExampleSSHServer {
+    static func main() async throws {
+        let privateKey: Curve25519.Signing.PrivateKey
+        let privateKeyURL = URL(fileURLWithPath: "./citadel_host_key_ed25519")
+
+        // Read or create a private key
+        if let file = try? Data(contentsOf: privateKeyURL) {
+            // File exists, read it into a Curve25519 private key
+            privateKey = try Curve25519.Signing.PrivateKey(sshEd25519: file)
+        } else {
+            // File does not exist, create a new Curve25519 private
+            privateKey = Curve25519.Signing.PrivateKey()
+
+            // Write the private key to a file
+            try privateKey.makeSSHRepresentation().write(to: privateKeyURL, atomically: true, encoding: .utf8)
+        }
+
+        let server = try await SSHServer.host(
+            host: "localhost",
+            port: 2323,
+            hostKeys: [
+                NIOSSHPrivateKey(ed25519Key: privateKey)
+            ],
+            authenticationDelegate: LoginHandler(username: "joannis", password: "test")
+        )
+
+         server.enableShell(withDelegate: CustomAppShell())
+
+        try await server.closeFuture.get()
+    }
+}
+
+struct MyTerminalView: View {
+    var body: some View {
+        VStack {
+            Text("Hello, world!")
+                .background(.red)
+                .foregroundColor(.white)
+
+            Button("Click me") {
+                print("clicked")
+            }
+
+            Button("Don't click") {
+                print("Clicked anyways")
+            }
+        }
+        .border()
+    }
+}
+
+final class CustomAppShell: ShellDelegate {
+     @MainActor public func startShell(
+        inbound: AsyncStream<ShellClientEvent>,
+        outbound: ShellOutboundWriter,
+        context: SSHShellContext
+    ) async throws {
+        let app = Application(rootView: MyTerminalView()) { string in
+            outbound.write(ByteBuffer(string: string))
+        }
+
+        await withTaskGroup(of: Void.self) { group in
+            group.addTask { @MainActor in
+                for await message in inbound {
+                    if case .stdin(let input) = message {
+                        app.handleInput(Data(buffer: input))
+                    }
+                }
+            }
+            group.addTask { @MainActor in
+                for await windowSize in context.windowSize {
+                    app.changeWindosSize(to: Size(
+                        width: Extended(windowSize.columns),
+                        height: Extended(windowSize.rows)
+                    ))
+                }
+            }
+
+            app.draw()
+        }
+    }
+}
+
+struct LoginHandler: NIOSSHServerUserAuthenticationDelegate {
+    let username: String
+    let password: String
+
+    var supportedAuthenticationMethods: NIOSSHAvailableUserAuthenticationMethods {
+        .password
+    }
+
+    func requestReceived(
+        request: NIOSSHUserAuthenticationRequest,
+        responsePromise: EventLoopPromise<NIOSSHUserAuthenticationOutcome>
+    ) {
+        if case .password(.init(password: password)) = request.request, request.username == username {
+            return responsePromise.succeed(.success)
+        }
+
+        return responsePromise.succeed(.failure)
+    }
+}

Package.resolved

@@ -10,6 +10,15 @@
           "version": "5.3.0"
         }
       },
+      {
+        "package": "ColorizeSwift",
+        "repositoryURL": "https://github.com/mtynior/ColorizeSwift.git",
+        "state": {
+          "branch": null,
+          "revision": "2a354639173d021f4648cf1912b2b00a3a7cd83c",
+          "version": "1.6.0"
+        }
+      },
       {
         "package": "swift-atomics",
         "repositoryURL": "https://github.com/apple/swift-atomics.git",

Package.swift

@@ -21,6 +21,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
         .package(url: "https://github.com/attaswift/BigInt.git", from: "5.2.0"),
         .package(url: "https://github.com/apple/swift-crypto.git", "1.0.0" ..< "2.1.0"),
+        .package(url: "https://github.com/mtynior/ColorizeSwift.git", from: "1.5.0"),
     ],
     targets: [
         .target(name: "CCitadelBcrypt"),
@@ -33,8 +34,14 @@ let package = Package(
                 .product(name: "_CryptoExtras", package: "swift-crypto"),
                 .product(name: "BigInt", package: "BigInt"),
                 .product(name: "Logging", package: "swift-log"),
+                .productItem(name: "ColorizeSwift", package: "ColorizeSwift")
             ]
         ),
+        .executableTarget(
+            name: "CitadelServerExample",
+            dependencies: [
+                "Citadel"
+            ]),
         .testTarget(
             name: "CitadelTests",
             dependencies: [

Sources/Citadel/Exec/Client/ExecClient.swift

@@ -0,0 +1,126 @@
+import Foundation
+import NIO
+import NIOSSH
+
+final class TTYHandler: ChannelDuplexHandler {
+    typealias InboundIn = SSHChannelData
+    typealias InboundOut = ByteBuffer
+    typealias OutboundIn = ByteBuffer
+    typealias OutboundOut = SSHChannelData
+
+    let maxResponseSize: Int
+    var isIgnoringInput = false
+    var response = ByteBuffer()
+    let done: EventLoopPromise<ByteBuffer>
+
+    init(
+        maxResponseSize: Int,
+        done: EventLoopPromise<ByteBuffer>
+    ) {
+        self.maxResponseSize = maxResponseSize
+        self.done = done
+    }
+
+    func handlerAdded(context: ChannelHandlerContext) {
+        context.channel.setOption(ChannelOptions.allowRemoteHalfClosure, value: true).whenFailure { error in
+            context.fireErrorCaught(error)
+        }
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        switch event {
+        case let status as SSHChannelRequestEvent.ExitStatus:
+            if status.exitStatus != 0 {
+                done.fail(SSHClient.CommandFailed(exitCode: status.exitStatus))
+            }
+        default:
+            context.fireUserInboundEventTriggered(event)
+        }
+    }
+
+    func handlerRemoved(context: ChannelHandlerContext) {
+        done.succeed(response)
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let data = self.unwrapInboundIn(data)
+
+        guard case .byteBuffer(var bytes) = data.data, !isIgnoringInput else {
+            return
+        }
+
+        switch data.type {
+        case .channel:
+            if
+                response.readableBytes + bytes.readableBytes > maxResponseSize
+            {
+                isIgnoringInput = true
+                done.fail(CitadelError.commandOutputTooLarge)
+                return
+            }
+
+            // Channel data is forwarded on, the pipe channel will handle it.
+            response.writeBuffer(&bytes)
+            return
+        case .stdErr:
+            done.fail(TTYSTDError(message: bytes))
+        default:
+            ()
+        }
+    }
+
+    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
+        let data = self.unwrapOutboundIn(data)
+        context.write(self.wrapOutboundOut(SSHChannelData(type: .channel, data: .byteBuffer(data))), promise: promise)
+    }
+}
+
+extension SSHClient {
+    /// Executes a command on the remote server. This will return the output of the command. If the command fails, the error will be thrown. If the output is too large, the command will fail.
+    /// - Parameters:
+    ///  - command: The command to execute.
+    /// - maxResponseSize: The maximum size of the response. If the response is larger, the command will fail.
+    public func executeCommand(_ command: String, maxResponseSize: Int = .max) async throws -> ByteBuffer {
+        let promise = eventLoop.makePromise(of: ByteBuffer.self)
+
+        let channel: Channel
+
+        do {
+            channel = try await eventLoop.flatSubmit {
+                let createChannel = self.eventLoop.makePromise(of: Channel.self)
+                self.session.sshHandler.createChannel(createChannel) { channel, _ in
+                    channel.pipeline.addHandlers(
+                        TTYHandler(
+                            maxResponseSize: maxResponseSize,
+                            done: promise
+                        )
+                    )
+                }
+
+                self.eventLoop.scheduleTask(in: .seconds(15)) {
+                    createChannel.fail(CitadelError.channelCreationFailed)
+                }
+
+                return createChannel.futureResult
+            }.get()
+        } catch {
+            promise.fail(error)
+            throw error
+        }
+
+        // We need to exec a thing.
+        let execRequest = SSHChannelRequestEvent.ExecRequest(
+            command: command,
+            wantReply: true
+        )
+
+        return try await eventLoop.flatSubmit {
+            channel.triggerUserOutboundEvent(execRequest).whenFailure { [channel] error in
+                channel.close(promise: nil)
+                promise.fail(error)
+            }
+
+            return promise.futureResult
+        }.get()
+    }
+}

Sources/Citadel/SFTP/Server/SFTPServer.swift

@@ -32,6 +32,25 @@ public struct SSHContext {
     public let username: String?
 }
 
+public struct SSHShellContext {
+    public struct WindowSize {
+        public let columns: Int
+        public let rows: Int
+    }
+
+    public let session: SSHContext
+    internal let channel: Channel
+    public let windowSize: AsyncStream<WindowSize>
+
+    public var isClosed: Bool {
+        !channel.isActive
+    }
+
+    public func close(mode: CloseMode = .all) async throws {
+        try await channel.close(mode: mode)
+    }
+}
+
 /// The delegate for the SFTP subsystem. This is the interface that the SFTP subsystem uses to interact with the rest of the application. The delegate is responsible for implementing the various SFTP operations.
 public protocol SFTPDelegate {
     /// Returns the attributes for the file at the given path. This is equivalent to the `stat()` system call.
@@ -69,18 +88,18 @@ public protocol SFTPDelegate {
     func rename(oldPath: String, newPath: String, flags: UInt32, context: SSHContext) async throws -> SFTPStatusCode
 }
 
-struct SFTPServerSubsystem {
+enum SFTPServerSubsystem {
     static func setupChannelHanders(
         channel: Channel,
-        delegate: SFTPDelegate,
+        sftp: SFTPDelegate,
         logger: Logger,
         username: String?
     ) -> EventLoopFuture<Void> {
         let deserializeHandler = ByteToMessageHandler(SFTPMessageParser())
         let serializeHandler = MessageToByteHandler(SFTPMessageSerializer())
         let sftpInboundHandler = SFTPServerInboundHandler(
             logger: logger,
-            delegate: delegate,
+            delegate: sftp,
             eventLoop: channel.eventLoop,
             username: username
         )