#!/usr/bin/env python

# Copyright (C) 2003-2007  Robey Pointer <robeypointer@gmail.com>
#
# This file is part of paramiko.
#
# Paramiko is free software; you can redistribute it and/or modify it under the
# terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 2.1 of the License, or (at your option)
# any later version.
#
# Paramiko is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Paramiko; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA.


import sys
import select
import socketserver as SocketServer
import paramiko
from shipyard.models import SshRemote

SSH_PORT = 22
DEFAULT_PORT = 4000

g_verbose = True


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


class Handler(SocketServer.BaseRequestHandler):
    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:
            verbose(
                "Incoming request to %s:%d failed: %s"
                % (self.chain_host, self.chain_port, repr(e))
            )
            return
        if chan is None:
            verbose(
                "Incoming request to %s:%d was rejected by the SSH server."
                % (self.chain_host, self.chain_port)
            )
            return

        verbose(
            "Connected!  Tunnel open %r -> %r -> %r"
            % (
                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)

        peername = self.request.getpeername()
        chan.close()
        self.request.close()
        verbose("Tunnel closed from %r" % (peername,))


def forward_tunnel(local_port, remote_host, remote_port, transport):
    # this is a little convoluted, but lets me configure things for the Handler
    # object.  (SocketServer doesn't give Handlers any way to access the outer
    # server normally.)
    class SubHander(Handler):
        chain_host = remote_host
        chain_port = remote_port
        ssh_transport = transport

    ForwardServer(("", local_port), SubHander).serve_forever()


def verbose(s):
    if g_verbose:
        print(s)


def open_tunnel(
    ssh: SshRemote,
    local_port: int = DEFAULT_PORT,
    remote_port: int = DEFAULT_PORT,
    verb: bool = False,
):
    host = ssh.url.host if ssh.url.host else ""
    port = int(ssh.url.port) if ssh.url.port else SSH_PORT
    user = ssh.url.user if ssh.url.user else ""
    global g_verbose
    g_verbose = verb

    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.set_missing_host_key_policy(paramiko.WarningPolicy())

    print(f"Connecting to {host}:{port} ...")

    try:
        client.connect(
            host,
            port,
            username=user,
            look_for_keys=True,
            password=None,
        )
    except Exception as e:
        print("*** Failed to connect to %s:%s: %r" % (host, port, e))
        sys.exit(1)

    print("Forwarding local port %d to %s:%d ..." % (local_port, host, remote_port))

    try:
        forward_tunnel(local_port, "localhost", remote_port, client.get_transport())
    except KeyboardInterrupt:
        print(" Port forwarding stopped.")
        sys.exit(0)
