#!/usr/bin/env python3

"""Simple script to forward TCP connections.  The impetuous for this
script is to allow me to connect from a private network to a different
private network running on the same host.  This is needed for
connecting to Vagrant dev VMs (which have their SSH port forwarding
bound to 127.0.0.1) from an Intel NUC attached to the Vagrant host but
on a private network.

Author: Paul Serice
License: Public Domain.

"""

import argparse
import datetime as dt
import os
import socket
import sys
import threading

# Version.
__version__ = "1.0.0"

def nowstr():
    """Convert the date and time right now to a string having the format
    the following format:

        %Y-%m-%d %H:%M:%S

    :param date: date and time
    :type date: datetime.datetime

    :return: date and time as a string
    :rtype: str

    """

    return dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def copy(src, dst):
    """Copy bytes from the source socket to the destination socket.  Note
    that no bytes are copied from the destination to the source to
    avoid deadlock by trying to read and write in series.  If you need
    bidirectional copying, call this function from two different
    threads.

    :param src: source socket
    :type src: socket.socket

    :param dst: destination socket
    :type dst: socket.socket

    """

    try:

        while True:

            # Read the next block of bytes.
            buf = src.recv(16384)

            # Check for EOF.
            if not buf:
                # Done.
                break

            # Write the next block of bytes.
            dst.send(buf)

    except Exception as e:

        print("%s: Error: %s" % (nowstr(), e), file=sys.stderr, flush=True)

    finally:

        try:

            # Notify the destination that we are finished writing.
            dst.shutdown(socket.SHUT_WR)

        except Exception:

            # Quash the exception.  It just means dst is already
            # closed presumably because of an exception that
            # dst.send() raised.
            pass


def handle_connection(peername1, client1, client2):
    """Handle the connection between the first client (which made the
    initial connection to the local service) and the second client
    (which is a connection we made to the remote service).  All this
    functions does is copy data between the two sockets.

    :param peername1: peer name of the first client (used for error reporting)
    :type peername1: str

    :param client1: first client
    :type client1: socket.socket

    :param client2: second client
    :type client2: socket.socket

    """

    try:

        # Copy date from client1 to client2.
        t1 = threading.Thread(target=copy, args=(client1, client2))
        t1.start()

        # Copy date from client2 to client1.
        t2 = threading.Thread(target=copy, args=(client2, client1))
        t2.start()

        # Wait for both sockets to shut down their write ends.
        t1.join()
        t2.join()

    except Exception as e:

        print("%s: Error: %s" % (nowstr(), e), file=sys.stderr, flush=True)

    finally:

        # Close the sockets.
        print("%s:  Closing connection: %s" %
              (nowstr(), peername1), flush=True)

        try:
            client1.close()
        except:
            pass

        try:
            client2.close()
        except:
            pass


def forward_ports(local_host, local_port, remote_host, remote_port):

    """Forward ports from the local host and TCP port to the remote host
    and TCP port.

    :param local_host: local host
    :type local_host: str

    :param local_port: local port
    :type local_port: int

    :param remote_host: remote host
    :type remote_host: str

    :param remote_port: remote port
    :type remote_port: int

    """

    try:

        # Create the server socket.
        server = socket.socket()
        server.bind((local_host, local_port))
        server.listen(16)

        while True:

            # Accept incoming connections.
            client1 = server.accept()[0]
            try:
                # Get the peername while the client is still likely to
                # be connected.
                peername1 = client1.getpeername()
                print("%s: Accepted connection: %s" %
                      (nowstr(), peername1), flush=True)
            except Exception:
                # Client already disconnected.  There is nothing to do.
                print("%s: Accepted connection but client has already " \
                      "disconnected.", (nowstr(),), flush=True)
                continue

            try:
                # Open a connection to the remote server.
                client2 = socket.socket()
                client2.connect((remote_host, remote_port))
            except Exception as e:
                # Print error.
                print("%s: Error: Unable to connect to (%s, %s): %s" %
                      (nowstr(), remote_host, remote_port, e),
                      file=sys.stderr, flush=True)
                # Close client1.
                print("%s:  Closing connection: %s" %
                      (nowstr(), peername1), flush=True)
                client1.close()
                # Wait for the next connection.
                continue

            # Handle the connection in its own thread.
            t = threading.Thread(target=handle_connection,
                                 args=(peername1, client1, client2),
                                 daemon=True)
            t.start()

    except Exception as e:

        print("%s: Error: %s" % (nowstr(), e), file=sys.stderr, flush=True)


def main():

    # Create a new ArgrumentParser instance that shows the default
    # values for options that have them as part of the help message.
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Simple port-forwarding application.")

    # -l (local address)
    parser.add_argument("-l",
                        required=True,
                        dest="local_address",
                        metavar="HOST:PORT",
                        help="local address in the form <host>:<port>")

    # -r (remote address)
    parser.add_argument("-r",
                        required=True,
                        dest="remote_address",
                        metavar="HOST:PORT",
                        help="remote address in the form <host>:<port>")

    # -v (version)
    parser.add_argument("-v",
                        action="version",
                        version="%s %s" %
                                (os.path.basename(sys.argv[0]),
                                 __version__),
                        help="version infomation")

    # Parse command-line arguments.
    args = parser.parse_args()

    # Split the local address into host and port parts.
    (local_host, local_port) = args.local_address.split(":")
    local_port = int(local_port)

    # Split the remote address into host and port parts.
    (remote_host, remote_port) = args.remote_address.split(":")
    remote_port = int(remote_port)

    # Start forwarding ports.
    forward_ports(local_host, local_port, remote_host, remote_port)


if __name__ == "__main__":
    main()