tunnel.swift

import Dispatch
import Foundation

import NIO
import NIOSSH


let LocalHost = "127.0.0.1"
let LocalPort = 10080

let ServerHost = "192.168.0.7"
let ServerPort = 22
let ServerUser = "stefan"

let RemoteHost = "127.0.0.1"
let RemotePort = 80






final class GlueHandler {
    private var partner: GlueHandler?

    private var context: ChannelHandlerContext?

    private var pendingRead: Bool = false

    private init() {}
}

extension GlueHandler {
    static func matchedPair() -> (GlueHandler, GlueHandler) {
        let first = GlueHandler()
        let second = GlueHandler()

        first.partner = second
        second.partner = first

        return (first, second)
    }
}

extension GlueHandler {
    private func partnerWrite(_ data: NIOAny) {
        self.context?.write(data, promise: nil)
    }

    private func partnerFlush() {
        self.context?.flush()
    }

    private func partnerWriteEOF() {
        self.context?.close(mode: .output, promise: nil)
    }

    private func partnerCloseFull() {
        self.context?.close(promise: nil)
    }

    private func partnerBecameWritable() {
        if self.pendingRead {
            self.pendingRead = false
            self.context?.read()
        }
    }

    private var partnerWritable: Bool {
        self.context?.channel.isWritable ?? false
    }
}

extension GlueHandler: ChannelDuplexHandler {
    typealias InboundIn = NIOAny
    typealias OutboundIn = NIOAny
    typealias OutboundOut = NIOAny

    func handlerAdded(context: ChannelHandlerContext) {
        self.context = context

        // It's possible our partner asked if we were writable, before, and we couldn't answer.
        // Consider updating it.
        if context.channel.isWritable {
            self.partner?.partnerBecameWritable()
        }
    }

    func handlerRemoved(context: ChannelHandlerContext) {
        self.context = nil
        self.partner = nil
    }

    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        self.partner?.partnerWrite(data)
    }

    func channelReadComplete(context: ChannelHandlerContext) {
        self.partner?.partnerFlush()
    }

    func channelInactive(context: ChannelHandlerContext) {
        self.partner?.partnerCloseFull()
    }

    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
        if let event = event as? ChannelEvent, case .inputClosed = event {
            // We have read EOF.
            self.partner?.partnerWriteEOF()
        }
    }

    func errorCaught(context: ChannelHandlerContext, error: Error) {
        self.partner?.partnerCloseFull()
    }

    func channelWritabilityChanged(context: ChannelHandlerContext) {
        if context.channel.isWritable {
            self.partner?.partnerBecameWritable()
        }
    }

    func read(context: ChannelHandlerContext) {
        if let partner = self.partner, partner.partnerWritable {
            context.read()
        } else {
            self.pendingRead = true
        }
    }
}






enum SSHClientError: Swift.Error {
    case passwordAuthenticationNotSupported
    case commandExecFailed
    case invalidChannelType
    case invalidData
}




/// A client user auth delegate that provides an interactive prompt for password-based user auth.
final class InteractivePasswordPromptDelegate: NIOSSHClientUserAuthenticationDelegate {
    private let queue: DispatchQueue

    private var username: String?

    private var password: String?

    init(username: String?, password: String?) {
        self.queue = DispatchQueue(label: "io.swiftnio.ssh.InteractivePasswordPromptDelegate")
        self.username = username
        self.password = password
    }

    func nextAuthenticationType(availableMethods: NIOSSHAvailableUserAuthenticationMethods, nextChallengePromise: EventLoopPromise<NIOSSHUserAuthenticationOffer?>) {
        guard availableMethods.contains(.password) else {
            print("Error: password auth not supported")
            nextChallengePromise.fail(SSHClientError.passwordAuthenticationNotSupported)
            return
        }

        self.queue.async {
            if self.username == nil {
                print("Username: ", terminator: "")
                self.username = readLine() ?? ""
            }

            if self.password == nil {
                #if os(Windows)
                print("Password: ", terminator: "")
                self.password = readLine() ?? ""
                #else
                self.password = String(cString: getpass("Password: "))
                #endif
            }

            nextChallengePromise.succeed(NIOSSHUserAuthenticationOffer(username: self.username!, serviceName: "", offer: .password(.init(password: self.password!))))
        }
    }
}


final class AcceptAllHostKeysDelegate: NIOSSHClientServerAuthenticationDelegate {
    func validateHostKey(hostKey: NIOSSHPublicKey, validationCompletePromise: EventLoopPromise<Void>) {
        // Do not replicate this in your own code: validate host keys! This is a
        // choice made for expedience, not for any other reason.
        validationCompletePromise.succeed(())
    }
}




final class PortForwardingServer {
    private var serverChannel: Channel?
    private let serverLoop: EventLoop
    private let group: EventLoopGroup
    private let bindHost: String
    private let bindPort: Int
    private let forwardingChannelConstructor: (Channel) -> EventLoopFuture<Void>

    init(group: EventLoopGroup, bindHost: String, bindPort: Int, _ forwardingChannelConstructor: @escaping (Channel) -> EventLoopFuture<Void>) {
        self.serverLoop = group.next()
        self.group = group
        self.forwardingChannelConstructor = forwardingChannelConstructor
        self.bindHost = bindHost
        self.bindPort = bindPort
    }

    func run() -> EventLoopFuture<Void> {
        ServerBootstrap(group: self.serverLoop, childGroup: self.group)
            .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
            .childChannelInitializer(self.forwardingChannelConstructor)
            .bind(host: self.bindHost, port: self.bindPort)
            .flatMap {
                self.serverChannel = $0
                return $0.closeFuture
            }
    }

    func close() -> EventLoopFuture<Void> {
        self.serverLoop.flatSubmit {
            guard let server = self.serverChannel else {
                // The server wasn't created yet, so we can just shut down straight away and let
                // the OS clean us up.
                return self.serverLoop.makeSucceededFuture(())
            }

            return server.close()
        }
    }
}


/// A simple handler that wraps data into SSHChannelData for forwarding.
final class SSHWrapperHandler: ChannelDuplexHandler {
    typealias InboundIn = SSHChannelData
    typealias InboundOut = ByteBuffer
    typealias OutboundIn = ByteBuffer
    typealias OutboundOut = SSHChannelData

    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
        let data = self.unwrapInboundIn(data)

        guard case .channel = data.type, case .byteBuffer(let buffer) = data.data else {
            context.fireErrorCaught(SSHClientError.invalidData)
            return
        }

        context.fireChannelRead(self.wrapInboundOut(buffer))
    }

    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
        let data = self.unwrapOutboundIn(data)
        let wrapped = SSHChannelData(type: .channel, data: .byteBuffer(data))
        context.write(self.wrapOutboundOut(wrapped), promise: promise)
    }
}




final class ErrorHandler: ChannelInboundHandler {
    typealias InboundIn = Any

    func errorCaught(context: ChannelHandlerContext, error: Error) {
        print("Error in pipeline: \(error)")
        context.close(promise: nil)
    }
}




// Start a server

let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
defer {
    try! group.syncShutdownGracefully()
}

let bootstrap = ClientBootstrap(group: group)
    .channelInitializer { channel in
        channel.pipeline.addHandlers([NIOSSHHandler(role: .client(.init(userAuthDelegate: InteractivePasswordPromptDelegate(username: ServerUser, password: ProcessInfo.processInfo.environment["PASSWORD"]), serverAuthDelegate: AcceptAllHostKeysDelegate())), allocator: channel.allocator, inboundChildChannelInitializer: nil), ErrorHandler()])
    }
    .channelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
    .channelOption(ChannelOptions.socket(SocketOptionLevel(IPPROTO_TCP), TCP_NODELAY), value: 1)

let channel = try bootstrap.connect(host: ServerHost, port: ServerPort).wait()

let server = PortForwardingServer(group: group, bindHost: LocalHost, bindPort: LocalPort) { inboundChannel in
    channel.pipeline.handler(type: NIOSSHHandler.self).flatMap { sshHandler in
        let promise = inboundChannel.eventLoop.makePromise(of: Channel.self)

        let directTCPIP = SSHChannelType.DirectTCPIP(targetHost: RemoteHost, targetPort: RemotePort, originatorAddress: inboundChannel.remoteAddress!)
        sshHandler.createChannel(promise, channelType: .directTCPIP(directTCPIP)) { childChannel, channelType in
            guard case .directTCPIP = channelType else {
                return channel.eventLoop.makeFailedFuture(SSHClientError.invalidChannelType)
            }

            let (ours, theirs) = GlueHandler.matchedPair()
            return childChannel.pipeline.addHandlers([SSHWrapperHandler(), ours, ErrorHandler()]).flatMap {
                inboundChannel.pipeline.addHandlers([theirs, ErrorHandler()])
            }
        }
        return promise.futureResult.map { _ in }
    }
}


// Run the server until done

try! server.run().wait()