#!/usr/bin/env python3

"""Libvirt port-forwarding hook.

Libvirt hook for setting up iptables port-forwarding rules when using NAT-ed
networking.
"""
__author__ = "Sascha Peilicke <saschpe@gmx.de>"
__version__ = "0.2.0"

import os
import json
import re
import subprocess
import sys

if "check_output" not in dir(subprocess):
    def f(*popenargs, **kwargs):
        if 'stdout' in kwargs:
            raise ValueError('stdout argument not allowed, it will be overridden.')
        process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
        output, unused_err = process.communicate()
        retcode = process.poll()
        if retcode:
            cmd = kwargs.get("args")
            if cmd is None:
                cmd = popenargs[0]
            raise subprocess.CalledProcessError(retcode, cmd)
        return output
    subprocess.check_output = f

CONFIG_PATH = os.getenv('CONFIG_PATH') or os.path.dirname(
    os.path.abspath(__file__))
CONFIG_FILENAME = os.getenv('CONFIG_FILENAME') or os.path.join(
    CONFIG_PATH, "hooks.json")
CONFIG_SCHEMA_FILENAME = os.getenv('CONFIG_SCHEMA_FILENAME') or os.path.join(
    CONFIG_PATH, "hooks.schema.json")
IPTABLES_BINARY = os.getenv('IPTABLES_BINARY') or subprocess.check_output([
    "which", "iptables"]).strip()


# Allow comments in json, copied from https://github.com/getify/JSON.minify
def json_minify(string, strip_space=True):
    tokenizer = re.compile('"|(/\*)|(\*/)|(//)|\n|\r')
    end_slashes_re = re.compile(r'(\\)*$')

    in_string = False
    in_multi = False
    in_single = False

    new_str = []
    index = 0

    for match in re.finditer(tokenizer, string):

        if not (in_multi or in_single):
            tmp = string[index:match.start()]
            if not in_string and strip_space:
                # replace white space as defined in standard
                tmp = re.sub('[ \t\n\r]+', '', tmp)
            new_str.append(tmp)
        elif not strip_space:
            # Replace comments with white space so that the JSON parser reports
            # the correct column numbers on parsing errors.
            new_str.append(' ' * (match.start() - index))

        index = match.end()
        val = match.group()

        if val == '"' and not (in_multi or in_single):
            escaped = end_slashes_re.search(string, 0, match.start())

            # start of string or unescaped quote character to end string
            if not in_string or (escaped is None or len(escaped.group()) % 2 == 0):  # noqa
                in_string = not in_string
            index -= 1  # include " character in next catch
        elif not (in_string or in_multi or in_single):
            if val == '/*':
                in_multi = True
            elif val == '//':
                in_single = True
        elif val == '*/' and in_multi and not (in_string or in_single):
            in_multi = False
            if not strip_space:
                new_str.append(' ' * len(val))
        elif val in '\r\n' and not (in_multi or in_string) and in_single:
            in_single = False
        elif not ((in_multi or in_single) or (val in ' \r\n\t' and strip_space)):  # noqa
            new_str.append(val)

        if not strip_space:
            if val in '\r\n':
                new_str.append(val)
            elif in_multi or in_single:
                new_str.append(' ' * len(val))

    new_str.append(string[index:])
    return ''.join(new_str)


def host_ip():
    """Returns the default route interface IP (if any).

    In other words, the public IP used to access the virtualization host. It
    is used as default public IP for guest forwarding rules should they not
    specify a different public IP to forward from.
    """
    if not hasattr(host_ip, "_host_ip"):
        cmd = "ip route | grep default | cut -d' ' -f5 | head -n1"
        default_route_interface = subprocess.check_output(
            cmd, shell=True).decode().strip()
        cmd = "ip addr show {0} | grep -E 'inet .*{0}' | cut -d' ' -f6 | cut -d'/' -f1".format(
            default_route_interface)
        host_ip._host_ip = subprocess.check_output(
            cmd, shell=True).decode().strip()
    return host_ip._host_ip.split('\n')[0]


def config(validate=True):
    """Returns the hook configuration.

    Assumes that the file /etc/libvirt/hooks/qemu.json exists and contains
    JSON-formatted configuration data. Optionally tries to validate the
    configuration if the 'jsonschema' module is available.

    Args:
        validate: Use JSON schema validation
    """
    if not hasattr(config, "_conf"):
        with open(CONFIG_FILENAME, "r") as f:
            config._conf = json.loads(json_minify(f.read()))
        if validate:
            # Try schema validation but avoid hard 'jsonschema' requirement:
            try:
                import jsonschema
                with open(CONFIG_SCHEMA_FILENAME, "r") as f:
                    config._schema = json.load(f)
                jsonschema.validate(config._conf,
                                    config._schema,
                                    format_checker=jsonschema.FormatChecker())
            except ImportError:
                pass
    return config._conf


def create_chain(table, name):
    """ Creates the named chain. """
    subprocess.call([IPTABLES_BINARY, "-t", table, "-N", name])


def delete_chain(table, name):
    """ Flushes and deletes the named chain. """
    subprocess.call([IPTABLES_BINARY, "-t", table, "-F", name])
    subprocess.call([IPTABLES_BINARY, "-t", table, "-X", name])


def populate_chains(dnat_chain, snat_chain, fwd_chain, public_ip, private_ip, domain, source_ip=None):
    """ Fills the two custom chains with the port mappings. """
    port_map = domain["port_map"]
    for protocol in port_map:
        for ports in port_map[protocol]:
            # a single integer 80 is equivalent to [80, 80]
            public_port, private_port = ports if isinstance(ports, list) else [
                ports, ports]
            dest = "{0}:{1}".format(private_ip, str(private_port))
            subprocess.call([IPTABLES_BINARY, "-t", "nat", "-A", dnat_chain, "-p", protocol,
                             "-d", public_ip, "--dport", str(public_port), "-j", "DNAT", "--to", dest] +
                            (["-s", source_ip] if source_ip else []))
            subprocess.call([IPTABLES_BINARY, "-t", "nat", "-A", snat_chain, "-p", protocol,
                             "-s", private_ip, "--dport", str(private_port), "-j", "SNAT", "--to-source", public_ip])
            subprocess.call([IPTABLES_BINARY, "-t", "nat", "-A", snat_chain, "-p", protocol,
                             "-s", private_ip, "-d", private_ip, "--dport", str(public_port), "-j", "MASQUERADE"])
            interface = ["-o", domain["interface"]
                         ] if "interface" in domain else []
            subprocess.call([IPTABLES_BINARY, "-t", "filter", "-A", fwd_chain, "-p", protocol,
                             "-d", private_ip, "--dport", str(private_port), "-j", "ACCEPT"] + interface)

    # Iterate over all port ranges
    if "port_range" in domain:
        for port_range in domain["port_range"]:
            ports_range = (str(port_range["init_port"]) + ":" +
                           str(port_range["init_port"] + port_range["ports_num"] - 1))
            dest = "{0}:{1}".format(
                private_ip, ports_range.replace(":", "-", 1))
            subprocess.call([IPTABLES_BINARY, "-t", "nat", "-A", dnat_chain, "-p", port_range["protocol"],
                             "-d", public_ip, "--dport", ports_range, "-j", "DNAT", "--to", dest])
            subprocess.call([IPTABLES_BINARY, "-t", "nat", "-A", snat_chain, "-p", port_range["protocol"],
                             "-s", private_ip, "--dport", ports_range, "-j", "SNAT", "--to-source", public_ip])
            subprocess.call([IPTABLES_BINARY, "-t", "nat", "-A", snat_chain, "-p", port_range["protocol"],
                             "-s", private_ip, "-d", private_ip, "--dport", ports_range, "-j", "MASQUERADE"])
            interface = ["-o", domain["interface"]
                         ] if "interface" in domain else []
            subprocess.call([IPTABLES_BINARY, "-t", "filter", "-A", fwd_chain, "-p", port_range["protocol"],
                             "-d", private_ip, "--dport", ports_range, "-j", "ACCEPT"] + interface)


def insert_chains(action, dnat_chain, snat_chain, fwd_chain, public_ip, private_ip):
    """ inserts (action='-I') or removes (action='-D') the custom chains."""
    subprocess.call([IPTABLES_BINARY, "-t", "nat", action,
                     "OUTPUT", "-d", public_ip, "-j", dnat_chain])
    subprocess.call([IPTABLES_BINARY, "-t", "nat", action,
                     "PREROUTING", "-d", public_ip, "-j", dnat_chain])
    # TODO: Find solution for connections from different private_ip to public_ip
    # maybe use private_ip_net as source instead
    # WORKAROUND: remove `"-s", private_ip, `
    subprocess.call([IPTABLES_BINARY, "-t", "nat", action, "POSTROUTING",
                     "-s", private_ip, "-d", private_ip, "-j", snat_chain])
    subprocess.call([IPTABLES_BINARY, "-t", "filter", action,
                     "FORWARD", "-d", private_ip, "-j", fwd_chain])

# the snat_chain doesn't work unless we turn off filtering bridged packets
# https://wiki.libvirt.org/page/Net.bridge.bridge-nf-call_and_sysctl.conf


def disable_bridge_filtering():
    proc_file = '/proc/sys/net/bridge/bridge-nf-call-iptables'
    if(os.path.isfile(proc_file)):
        with open(proc_file, 'w') as brnf:
            brnf.write('0\n')


def start_forwarding(dnat_chain, snat_chain, fwd_chain, public_ip, private_ip, domain, source_ip=None):
    """ sets up iptables port-forwarding rules based on the port_map. """
    disable_bridge_filtering()
    create_chain("nat", dnat_chain)
    create_chain("nat", snat_chain)
    create_chain("filter", fwd_chain)
    populate_chains(dnat_chain, snat_chain, fwd_chain,
                    public_ip, private_ip, domain, source_ip)

    insert_chains("-I", dnat_chain, snat_chain,
                  fwd_chain, public_ip, private_ip)


def stop_forwarding(dnat_chain, snat_chain, fwd_chain, public_ip, private_ip):
    """ tears down the iptables port-forwarding rules. """
    insert_chains("-D", dnat_chain, snat_chain,
                  fwd_chain, public_ip, private_ip)
    delete_chain("nat", dnat_chain)
    delete_chain("nat", snat_chain)
    delete_chain("filter", fwd_chain)


def substitute_domain_name(domain_name, index_str):
    # handle the 28 char limit of iptables
    # we need to take the possible 5 chars of "DNAT-" and
    # each char of the index string plus the separator into account
    return domain_name[0:28 - (5 + len(index_str) + 1)] + "-" + index_str


def handle_domain(action, domain, vir_domain):
    dnat_chain = "DNAT-{0}".format(vir_domain)
    snat_chain = "SNAT-{0}".format(vir_domain)
    fwd_chain = "FWD-{0}".format(vir_domain)
    public_ip = domain.get("public_ip", host_ip())
    private_ip = domain["private_ip"]
    source_ip = domain.get("source_ip")

    if action in ["stopped", "reconnect"]:
        stop_forwarding(dnat_chain, snat_chain,
                fwd_chain, public_ip, private_ip)
    if action in ["start", "reconnect"]:
        start_forwarding(dnat_chain, snat_chain, fwd_chain,
                public_ip, private_ip, domain, source_ip)


if __name__ == "__main__":
    vir_domain, action = sys.argv[1:3]
    domain = config().get(vir_domain)
    if domain is None:
        sys.exit(0)

    if not isinstance(domain, list):
        handle_domain(action, domain, vir_domain[0:28 - 5])
    else:
        i = 0
        for actual_domain in domain:
            handle_domain(action, actual_domain, substitute_domain_name(vir_domain, str(i)))
            i += 1
