#!/usr/bin/env python3
#
# Open port to specified env or default env
#
import argparse
import re
import subprocess

import psutil

###########
# CONFIGS #
###########
default_user = 'mg'
default_env = 'box01'

aliases = {
    '1C': {'port': '6000'},
    'ADMIN': {'port': '6001'},
    'SEARCH': {'port': '6002'},
    'API': {'port': '6003'},
    'CRM': {'port': '6004'},
    'CRMBACKEND': {'port': '6004'},
    'CRM_BACKEND': {'port': '6004'},
    'IMAGES': {'port': '6005'},
    'CRMFRONT': {'port': '7000'},
    'CRMFRONTEND': {'port': '7000'},
    'CRM_FRONTEND': {'port': '7000'},
    'FRONTEND': {'port': '7000'},
    'HRD': {'port': '7000'},
    'REDIS': {'port': '6379'},
    'MYSQL': {'port': '3306'},
    'MONGO': {'port': '27017'},
    'MONGODB': {'port': '27017'},
}

reverse = {
    '6000': '1c',
    '6001': 'Admin',
    '6002': 'Search',
    '6003': 'Api',
    '6004': 'Crm backend',
    '6005': 'Images',
    '7000': 'Crm frontend',
    '7001': 'Frontend',
    '7002': 'Hrd',
    '6379': 'Redis',
    '3306': 'MySQL',
    '27017': 'MongoDB',
}

# Globals
ssh_pattern = "ssh -J jumper@bastion.mygig.tech:32323 %s@%s.ru-central1.internal"
search_pattern = ".*-(?P<dest>[L|R]) (?P<lport>\d{2,5}):(?P<rhost>.+):(?P<rport>\d{2,5}) %s@(?P<env>%s)"
forward_pattern = "ssh -J jumper@bastion.mygig.tech:32323 -nfNT -L %s:%s:%s %s@%s.ru-central1.internal"
reverse_pattern = "ssh -J jumper@bastion.mygig.tech:32323 -nfNT -R %s:%s:%s %s@%s.ru-central1.internal"


class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
    pass


class Tunnel:
    pass


def tunnel_print(t):
    if t.dest == 'L' and t.rhost == 'localhost':
        dest = '--FORWARD-->'
    elif t.dest == 'L' and t.rhost != 'localhost':
        dest = '---PROXY--->'
    elif t.dest != 'L' and t.rhost == 'localhost':
        dest = '\033[1;31m<----BACK---\033[1;m'
    else:
        dest = '\033[1;31m<--RPROXY---\033[1;m'
    print('\033[1;32m%6d\033[1;m  \033[1;31m%-15s\033[1;m  %5s  %s  %5s  %s' % (t.pid, t.name, t.lport, dest, t.rport, t.env))


def tunnels_print(tunnel_list):
    if not tunnel_list:
        return

    print('\033[1m   PID  NAME             LPORT  DESTINATION   RPORT  ENV\033[0m')
    for t in tunnel_list:
        tunnel_print(t)


def get_tunnels(uname, env, ports_map):
    tunnels = []
    for p in psutil.process_iter(attrs=['pid', 'name', 'cmdline']):
        cmdline = ' '.join(p.info['cmdline'])
        matches = re.match(search_pattern % (uname, env), cmdline)
        if not matches or (len(ports_map) and not any(d['port'] == matches.group('rport') for d in ports_map)):
            continue

        tunnel = Tunnel()
        tunnel.pid = p.info['pid']
        tunnel.env = matches.group('env')
        tunnel.name = reverse.get(matches.group('rport'), 'UNKNOWN')
        tunnel.dest = matches.group('dest')
        tunnel.lport = matches.group('lport')
        tunnel.rhost = matches.group('rhost')
        tunnel.rport = matches.group('rport')

        tunnels.append(tunnel)

    return tunnels


def kill_tunnels(tunnels):
    if not tunnels:
        return

    for t in tunnels:
        print('Terminating pid %s (%s on port %s)...' % (t.pid, t.name, t.lport), end='')
        proc = psutil.Process(t.pid)
        proc.terminate()
        print('[\033[1;32mOK\033[1;m]')


def forward_tunnels(uname, env, ports_map):
    if not ports_map:
        return

    kill_tunnels(get_tunnels(uname, env, ports_map))
    for p in ports_map:
        print('Making forward tunnel for %s on port %s...' % (reverse.get(p['port'], 'UNKNOWN'), p['port']), end='')
        psutil.Popen([forward_pattern % (p['port'], p.get('host', 'localhost'), p['port'], uname, env)], shell=True)
        print('[\033[1;32mOK\033[1;m]')


def back_tunnels(uname, env, ports_map):
    if not ports_map:
        return

    kill_tunnels(get_tunnels(uname, env, ports_map))
    for p in ports_map:
        print('Making backdoor tunnel for %s on port %s...' % (reverse.get(p['port'], 'UNKNOWN'), p['port']), end='')
        psutil.Popen([reverse_pattern % (p['port'], p.get('host', 'localhost'), p['port'], uname, env)], shell=True)
        print('[\033[1;32mOK\033[1;m]')


def ssh_connect(uname, env):
    subprocess.call(ssh_pattern % (uname, env), shell=True)


# MAIN
description = '-= Tunnel maker =-'
epilog = 'Examples:\n' \
         '  tunnel                                      - show all active tunnels\n' \
         '  tunnel -t                                   - open ssh connection to specified env\n' \
         '  tunnel -p redis                             - show active tunnels for redis\n' \
         '  tunnel -e box01 -l -p redis 6000            - forward ports for redis and 6000 port to sandbox-01\n' \
         '  tunnel -r -p 6003                           - backdoor port for api service to default env\n' \
         '  tunnel -k                                   - kill all tunnels\n' \
         '  tunnel -kp api                              - kill tunnels for api service'

parser = argparse.ArgumentParser(description=description, epilog=epilog, formatter_class=CustomFormatter)

parser.add_argument('-e', '--env', action='store', default=None, dest='env',
                    help='specify environment for tunnels')
parser.add_argument('-u', '--user', action='store', default=None, dest='user',
                    help='specify username for tunnels')
parser.add_argument('-s', '--show', action='store_true', default=False, dest='show',
                    help='show list of active tunnels')
parser.add_argument('-t', '--tty', action='store_true', default=False, dest='tty',
                    help='open ssh connection to specified env')
parser.add_argument('-k', '--kill', action='store_true', default=False, dest='kill',
                    help='terminate active tunnels')
parser.add_argument('-r', '--remote', action='store_true', default=False, dest='remote',
                    help='backdoor remote port to local')
parser.add_argument('-l', '--local', action='store_true', default=False, dest='local',
                    help='map local port to remote')
parser.add_argument('-p', '--port', action='append', default=[], dest='ports', nargs='+',
                    help='ports list or known applications aliases')

results = parser.parse_args()

tty = results.tty
kill = results.kill
show = results.show
user = results.user or default_user
envir = results.env or default_env
local = results.local
ports = [port for group in results.ports for port in group]
remote = results.remote
port_map = [aliases.get(port.upper(), {'port': port}) for port in ports]

if tty:
    ssh_connect(user, envir)
    exit(0)

if show:
    tunnels_print(get_tunnels(user, envir, port_map))
    exit(0)

if kill:
    kill_tunnels(get_tunnels(user, envir, port_map))
    exit(0)

if local:
    forward_tunnels(user, envir, port_map)
    exit(0)

if remote:
    back_tunnels(user, envir, port_map)
    exit(0)

tunnels_print(get_tunnels(user, envir, port_map))
