#!/usr/bin/python3

import argparse
import os
import subprocess
import sys
import traceback

from ruamel import yaml

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', default='/etc/nerve/port-forwarder.yaml')
parser.add_argument('-d', '--dry', default=False, action='store_true')
parser.add_argument('host', help="Name of the guest being acted upon.")
parser.add_argument('action', help='Action being performed, e.g. "start" or "stopped".')
parser.add_argument('sub-operation', help="Passed by libvirt, unused by this script.")
parser.add_argument('extra-argument', help="Passed by libvirt, unused by this script.")
args = parser.parse_args()

if not os.path.exists(args.config):
    sys.exit(0)

with open(args.config) as stream:
    try:
        data = yaml.safe_load(stream)
    except yaml.YAMLError as e:
        traceback.print_exc()
        print('\nError: %s: Could not parse YAML file.' % args.config)

        # NOTE: Do not exit with an error code when parsing the YAML file,
        #       this would break starting all VMs.
        sys.exit(0)

# Nothing configured for this VM, so do nothing
if args.host not in data:
    sys.exit(0)

data = data[args.host]
comment = data.get('comment', 'nerve-port-forwarder')


def run(cmd):
    if args.dry:
        print(' '.join(cmd))
    else:
        subprocess.call(cmd)


for iface, forwards in data.items():
    for forward in forwards:
        proto = forward.get('protocol', 'tcp')

        for src_port, dst_port in forward.get('ports', {}).items():
            if args.action in ['start', 'reconnect']:
                action = '-I'
            elif args.action in ['stopped', 'reconnect']:
                action = '-D'
            else:
                continue

            # /sbin/iptables -I FORWARD 1 -o virbr0 -p tcp --dport 1880
            #   -j ACCEPT -m comment --comment "from hook"
            run([
                '/sbin/iptables', action, 'FORWARD',
                '-o', forward['interface'],
                '-p', proto, '-m', proto,
                '-d', forward['address'], '--dport', str(dst_port),
                '-j', 'ACCEPT', '-m', 'comment', '--comment', comment,
            ])

            # /sbin/iptables -t nat -I PREROUTING -i br-wan
            #   -p tcp --dport $HOST_PORT -j DNAT
            #   --to $GUEST_IP:$GUEST_PORT -m comment --comment "from hook"
            run([
                '/sbin/iptables', '-t', 'nat', action, 'PREROUTING',
                '-i', iface, '-p', proto, '-m', proto, '--dport', str(src_port), '-j', 'DNAT',
                '--to', '%s:%s' % (forward['address'], dst_port),
                '-m', 'comment', '--comment', comment
            ])
