proxy.py

#!/usr/bin/env python3

import sys
import os
import getpass
import socket
import select
import threading
import logging
import yaml
from urllib.parse import urlparse

import paramiko
import click

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(os.path.basename(__file__))
logging.getLogger('paramiko').setLevel(logging.WARNING)

nginx_confdir = '/etc/nginx/forward-agent.d'

LOGLEVEL = {
    'info': logging.INFO,
    'warn': logging.WARN,
    'debug': logging.DEBUG
}

default = {
    'user': getpass.getuser(),
    'key': os.path.expanduser('~/.ssh/id_rsa'),
    'proxy': "ssh://proxy.demo.com:22",
    'service': "http://localhost:8080",
    'config': os.path.expanduser('~/.forward.conf'),
    'verbosity': 'info'
}

nginx_template = '''
server {
    server_name %(server_names)s;
    listen 80;

    proxy_set_header Host $http_host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
    proxy_set_header X-Forwarded-Proto $scheme;

    location / {
        proxy_pass %(scheme)s://127.0.0.1:%(port)d;
    }
}
'''

def handler(chan, host, port):
    sock = socket.socket()
    try:
        sock.connect((host, port))
    except Exception as e:
        log.error(f'Forwarding request to {host}:{port} failed: {e}')
        return

    log.debug(f'Tunnel {chan.origin_addr} -> {chan.getpeername()} -> {host}:{port} established.')

    while True:
        r, w, x = select.select([sock, chan], [], [])
        if sock in r:
            data = sock.recv(1024)
            if len(data) == 0:
                break
            chan.send(data)
        if chan in r:
            data = chan.recv(1024)
            if len(data) == 0:
                break
            sock.send(data)
    chan.close()
    sock.close()
    log.debug(f'Tunnel closed from {chan.origin_addr}')

def write_nginx_config(app, names, service, client, port):
    names.insert(0, f'{app}{port}.demo.com')

    remote_file = os.path.join(nginx_confdir, f'{port}.conf')
    sftp_client = client.open_sftp()
    sftp_client.open(remote_file, 'w').write(
        nginx_template % {
            'server_names': ', '.join(names),
            'scheme': service.scheme,
            'port': port
        }
    )
    sftp_client.close()
    restart_nginx(client)
    for n in names:
        log.info(f"{service.scheme}://{n} is active.")

def remove_nginx_config(client, port):
    log.info('Cleaning up proxy configuration...')
    remote_file = os.path.join(nginx_confdir, f'{port}.conf')
    sftp_client = client.open_sftp()
    sftp_client.remove(remote_file)
    sftp_client.close()
    restart_nginx(client)

def restart_nginx(client):
    client.exec_command('/usr/sbin/nginx -s reload')

def reverse_tunnel(app, names, service, client):
    try:
        transport = client.get_transport()
        internal_port = transport.request_port_forward('', 0)
        write_nginx_config(app, names, service, client, internal_port)
        log.debug(f"Tunnel successfully established on internal port {internal_port}.")

        while True:
            chan = transport.accept(1000)
            if chan is None:
                continue
            thread = threading.Thread(
                target = handler,
                args   = (chan, service.hostname, service.port)
            )
            thread.setDaemon(True)
            thread.start()

    except KeyboardInterrupt:
        log.info('Port forwarding shutting down...')

    finally:
        remove_nginx_config(client, internal_port)


try:
    yaml.load(default['config'])
except FileNotFoundError:
    log.warn(f"{default['config']}: file not found. Using defaults.")


@click.command()
@click.option('--app', '-a',
    type    = str,
    default = default['user'],
    help    = f"Application name to be prepended to proxy URI [{default['user']}]"
)
@click.option('--name', '-n',
    type     = list,
    multiple = True,
    help     = "Alternative DNS names for this service [empty list]"
)
@click.option('--user', '-u',
    type    = str,
    default = default['user'],
    help    = f"SSH username [{default['user']}]"
)
@click.option('--key', '-k',
    type    = str,
    default = default['key'],
    help    = f"SSH private key file to use [{default['key']}]"
)
@click.option('--proxy', '-p',
    type    = str,
    default = default['proxy'],
    help    = f"Proxy server [{default['proxy']}]"
)
@click.option('--service', '-s',
    type    = str,
    default = default['service'],
    help    = f"URI of service being proxied [{default['service']}]"
)
@click.option('--verbosity', '-v',
    type    = click.Choice(LOGLEVEL),
    default = default['verbosity'],
    help    = f"Verbosity [{default['verbosity']}]"
)
def main(app, name, user, key, proxy, service, verbosity):

    log.setLevel(LOGLEVEL[verbosity])

    proxy = urlparse(proxy)
    service = urlparse(service)

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

    log.debug(f'Connecting to proxy host {proxy.hostname}:{proxy.port} ...')
    try:
        client.connect(proxy.hostname, proxy.port, username=user, key_filename=key)
    except Exception as e:
        log.error(f'Failed to connect to {proxy.hostname}:{proxy.port}: {e}')
        raise SystemError

    log.debug(f'Setting up forwarding to {service.hostname} ...')
    reverse_tunnel(app, list(name), service, client)


if __name__ == '__main__':
    main()