import logging
import select
import socketserver
from typing import Optional

from paramiko.transport import Transport


class ForwardServer(socketserver.ThreadingTCPServer):
    daemon_threads = True
    allow_reuse_address = True


class Handler(socketserver.BaseRequestHandler):
    ssh_transport: Optional[Transport] = None
    chain_host: Optional[str] = None
    chain_port: Optional[int] = None
    logger = logging.Logger(__name__)

    def handle(self):
        try:
            chan = self.ssh_transport.open_channel(
                "direct-tcpip",
                (self.chain_host, self.chain_port),
                self.request.getpeername(),
            )
        except Exception as e:
            self.logger.warning(
                'Incoming request to {host}:{port} failed: {error}'.format(
                    host=self.chain_host, port=self.chain_port, error=repr(e)
                )
            )
            return

        if chan is None:
            self.logger.warning(
                'Incoming request to {host}:{port} was rejected by the SSH server'.format(
                    host=self.chain_host, port=self.chain_port
                )
            )
            return

        self.logger.info(
            'Connected! Tunnel open {} -> {} -> {}'.format(
                self.request.getpeername(),
                chan.getpeername(),
                (self.chain_host, self.chain_port),
            )
        )

        while True:
            r, w, x = select.select([self.request, chan], [], [])
            if self.request in r:
                data = self.request.recv(1024)
                if len(data) == 0:
                    break
                chan.send(data)
            if chan in r:
                data = chan.recv(1024)
                if len(data) == 0:
                    break
                self.request.send(data)

        peer = self.request.getpeername()
        chan.close()
        self.request.close()
        self.logger.info('Tunnel closed from {}'.format(peer))


def forward_tunnel(
    local_port, remote_host, remote_port, transport, bind_addr=''
) -> ForwardServer:
    class SubHandler(Handler):
        ssh_transport = transport
        chain_host = remote_host
        chain_port = remote_port

    return ForwardServer((bind_addr, local_port), SubHandler)


# vim:sw=4:ts=4:et:
