fwsniff

#!/usr/bin/env python3
# fwsniff (part of ossobv/vcutil) // wdoekes/2023 // Public Domain
#
# Sniffing live traffic for the purpose of adding a firewall.
# fwsniff has two modes:
# - traffic sniffing, in which traffic is monitored and displayed;
# - iptables generation, in which (shell) iptables rules are created
#   which add LOG to the final lines.
#
# A configuration file can be created along the lines of the
# EXAMPLE_HOSTHUBBLE_YAML example below. Pass the filename as argument
# to fwsniff for either mode.
#
# Running fwsniff is as safe as running tcpdump.
#
# BEWARE: Blindly running the generated iptables output is not safe. One
# can probably hack the input YAML to make it do unwanted stuff.
#
# NOTE: Interface support is only available for iptables-rules at the
# moment because we cannot properly detect all interfaces.
#
# TODO:
# - document timing stuff with regards to dictionary keys (tuples vs. bins)
# - move bin dictionary keys to int+str tuples
# - add something for udp?
#
import warnings
import sys

from collections import defaultdict, namedtuple
from ipaddress import IPv4Network, IPv6Network, ip_network
from os import environ, path
from struct import pack, unpack
from socket import AF_INET, inet_ntop
from time import time
from unittest import TestCase
from warnings import warn

from dpkt import NeedData, Packet, UnpackError  # python3-dpkt
from dpkt.ethernet import Ethernet  # python3-dpkt
from dpkt.ip import IP              # python3-dpkt
from dpkt.sll import SLL            # python3-dpkt
from pcapy import open_live         # python3-pcapy
from yaml import SafeLoader, load   # python3-yaml


EXAMPLE_HOSTHUBBLE_YAML = '''
aliases:
  VRF/backend/svc: 10.111.32.0/22
  k8s-masters-and-nodes:
    - k8s-masters
    - k8s-nodes
  k8s-masters:
    - master.dr.backend.example.cloud: 10.0.68.11
    - master.wp.backend.example.cloud: 10.0.87.11
    - master.zl.backend.example.cloud: 10.0.90.11
  k8s-nodes:
    - node1.dr.backend.example.cloud: 10.0.68.15
    - node1.wp.backend.example.cloud: 10.0.87.15
    - node1.zl.backend.example.cloud: 10.0.90.15
    - node2.dr.backend.example.cloud: 10.0.68.17
    - node2.wp.backend.example.cloud: 10.0.87.17
    - node2.zl.backend.example.cloud: 10.0.90.17
  other-k8s-nodes:
    - node1.dr.frontend.example.cloud: 10.0.68.129
    - node1.wp.frontend.example.cloud: 10.0.87.129
  load-balancers:
    - lb.dr.example.com: 10.1.1.68
    - lb.wp.example.com: 10.1.1.87
    - lb.zl.example.com: 10.1.1.90
  vpn:
    - vpn-internal
    - vpn-external
  vpn-external: 8.8.8.8
  vpn-internal: 10.0.87.193

services:
  # XXX: do something with interfaces?
  ssh:
    iface: [$extif]  # XXX: use me
    match: -p tcp --dport 22
    allow: [vpn]
  zabbix_agent:
    iface: [$extif]
    match: -p tcp --dport 10050
    allow: [vpn]
  master_cilium_api_in:
    match: -p tcp --dport 12379
    allow: [other-k8s-nodes]
  master_etcd_clients:
    match: -p tcp --dport 2379
    allow: [VRF/backend/svc]
  master_etcd_peers:
    match: -p tcp --dport 2380
    allow: [k8s-masters]
  master_kube_apiserver:
    match: -p tcp --dport 6443
    allow: [k8s-masters-and-nodes, load-balancers]
'''


class KnownRulesTest(TestCase):
    maxDiff = 8192

    @staticmethod
    def s2kv(s):
        """
        'name 1.2.3.4' -> ('name', IPv4Network('1.2.3.4'))
        """
        k, v = s.split(' ', 1)
        return k, ip_network(v)

    def test_example_yaml(self):
        known_rules = KnownRules.from_yaml(EXAMPLE_HOSTHUBBLE_YAML)
        self.assertEqual(
            known_rules.get_cidrs('k8s-masters-and-nodes'), [
                self.s2kv('master.dr.backend.example.cloud 10.0.68.11'),
                self.s2kv('master.wp.backend.example.cloud 10.0.87.11'),
                self.s2kv('master.zl.backend.example.cloud 10.0.90.11'),
                self.s2kv('node1.dr.backend.example.cloud 10.0.68.15'),
                self.s2kv('node1.wp.backend.example.cloud 10.0.87.15'),
                self.s2kv('node1.zl.backend.example.cloud 10.0.90.15'),
                self.s2kv('node2.dr.backend.example.cloud 10.0.68.17'),
                self.s2kv('node2.wp.backend.example.cloud 10.0.87.17'),
                self.s2kv('node2.zl.backend.example.cloud 10.0.90.17'),
            ]
        )

    def test_aliases_chain(self):
        known_rules = KnownRules()
        known_rules.set_aliases({
            'vpns': ['vpn-internal', 'vpn-external'],
            'vpn-internal': ['10.1.2.0', '10.1.3.0'],
            'vpn-external': ['1.2.3.4'],
        })
        self.assertEqual(
            known_rules.get_cidrs('vpns'),
            [self.s2kv('vpn-external 1.2.3.4'),
             self.s2kv('vpn-internal 10.1.2.0'),
             self.s2kv('vpn-internal 10.1.3.0')],
        )

    def test_aliases_reject_single(self):
        known_rules = KnownRules()
        known_rules.set_aliases({
            'vpn-internal': '10.1.2.0',
            'vpn-external': ['1.2.3.4', '5.5.5.5'],
        })
        self.assertEqual(
            known_rules.get_cidrs('vpn-internal'),
            [self.s2kv('vpn-internal 10.1.2.0')],
        )
        with self.assertRaises(ValueError):
            known_rules.set_aliases({
                'vpns': 'vpn-internal',
            })
        known_rules.set_aliases({
            'vpns': ['vpn-internal', 'vpn-external'],
        })
        self.assertEqual(
            known_rules.get_cidrs('vpns'),
            [self.s2kv('vpn-external 1.2.3.4'),
             self.s2kv('vpn-external 5.5.5.5'),
             self.s2kv('vpn-internal 10.1.2.0')],
        )

    def test_services_existing(self):
        known_rules = KnownRules.from_yaml(EXAMPLE_HOSTHUBBLE_YAML)
        known_service = known_rules.get_service_by_tcp_port(2379)
        self.assertEqual(known_service.name, 'master_etcd_clients')
        self.assertEqual(
            known_service.src_allowed,
            [self.s2kv('VRF/backend/svc 10.111.32.0/22')],
        )

    def test_services_non_existing(self):
        known_rules = KnownRules.from_yaml(EXAMPLE_HOSTHUBBLE_YAML)
        missing = known_rules.get_service_by_tcp_port(9999)
        self.assertEqual(missing, None)


class ServiceTest(TestCase):
    def test_service(self):
        known_service = KnownService(
            'ssh', tcp_port=22, interfaces=['$extif'], src_allowed=[
                ('vpn', ip_network('1.2.3.4'))])
        service = Service(known_service)
        self.assertEqual(service.name, 'ssh')

        for i in range(3):
            ret = service.add_count(
                b'\xff\xff\xff\xff'
                b'\x00\x16'
                b'\x01\x02\x03\x04')
        self.assertEqual(ret, (3, 'vpn'))


def ip_network_or_none(value):
    try:
        return ip_network(value)
    except ValueError:
        None


class KnownService:
    def __init__(
            self, name, src_allowed, tcp_port=None, interfaces=None):
        if not isinstance(name, str) or not name:
            raise TypeError('name should be str')
        self.name = name
        if not (
                all(isinstance(i, tuple) for i in src_allowed) and
                all(len(i) == 2 for i in src_allowed) and
                all(isinstance(i[0], str) and i[0] for i in src_allowed) and
                all(isinstance(i[1], IPv4Network) for i in src_allowed)):
            raise TypeError('src_allowed should be list of (name,net) tuples')
        self.src_allowed = src_allowed
        if not isinstance(tcp_port, (int, type(None))):
            raise TypeError('tcp_port must be int or unset')
        self.tcp_port = tcp_port
        if interfaces is not None and not (
                interfaces and isinstance(interfaces, list) and
                all(i and isinstance(i, str) for i in interfaces)):
            raise TypeError('interfaces be list of strings or unset')
        self.interfaces = interfaces


class KnownRules:
    """
    Parses a yaml like EXAMPLE_HOSTHUBBLE_YAML and exposes...

    - get_service_by_tcp_port, returning a service
    - get_cidr, returning a name+network for an IP
    """
    @classmethod
    def from_yaml_file(cls, filename):
        with open(filename) as fp:
            return cls.from_yaml(fp.read())

    @classmethod
    def from_yaml(cls, yaml):
        data = load(yaml, Loader=SafeLoader)
        if set(data.keys()) - set(['aliases', 'services']):
            raise ValueError('unexpected keys in yaml: {}'.format(data.keys()))
        obj = cls()
        obj.set_aliases(data['aliases'])
        obj.set_services(data['services'])
        return obj

    def __init__(self):
        self._aliases = {}      # alias -> [aliases]|[networks]
        self._services = {}     # service_id -> service_def
        self._tcp_ports = {}    # 1234 -> service_id

    def get_cidrs(self, alias):
        if alias not in self._aliases:
            raise ValueError(f'alias {alias!r} not found')
        values = self._aliases[alias]
        ret = []
        for value in values:
            if isinstance(value, (IPv4Network, IPv6Network)):
                ret.append((alias, value))
            else:
                ret.extend(self.get_cidrs(value))
        ret.sort()
        return ret

    def set_aliases(self, aliases):
        for key, value in aliases.items():
            if isinstance(value, str):
                # Expect a:
                # - CIDR
                # (if we accepted aliases here, we would allow aliases to
                # aliases without lists; confusing)
                self._set_alias_cidr(key, value)

            elif isinstance(value, list):
                # Expect a:
                # - list of aliases (strings)
                # - list of CIDRs (strings)
                # - list of dicts (alias: CIDR)
                types = tuple(set([type(i) for i in value]))
                if types == (str,):
                    networks = [ip_network_or_none(i) for i in value]
                    networks = [i for i in networks if i is not None]
                    if len(networks) == 0:
                        # - list of aliases
                        self._set_alias_list(key, value)
                    elif len(networks) == len(value):
                        # - list of networks
                        self._set_alias_list(key, networks)
                    else:
                        raise ValueError(
                            f'unexpected both aliases and networks in '
                            f'same list: {value!r}')

                elif types == (dict,):
                    # - list of dicts (alias: CIDR)
                    aliases = []
                    for item in value:
                        if len(item) != 1:
                            raise ValueError(
                                f'only single value dict allowed here, got '
                                f'{item!r}')
                        for alias_key, alias_value in item.items():
                            # Only takes a single CIDR
                            self._set_alias_cidr(alias_key, alias_value)
                            aliases.append(alias_key)
                    self._set_alias_list(key, aliases)

            else:
                raise ValueError(f'unexpected {key!r}: {value!r}')

    def _set_alias_cidr(self, key, value):
        if key in self._aliases:
            raise ValueError(
                f'duplicate key {key!r} with value {value!r}')
        try:
            networks = [ip_network(value)]
        except ValueError as e:
            raise ValueError(
                f'unexpected {value!r} non-network as single arg'
                ) from e
        else:
            assert isinstance(networks, list), networks
            self._aliases[key] = networks

    def _set_alias_list(self, key, list_):
        if key in self._aliases:
            raise ValueError(
                f'duplicate key {key!r} with value {list_!r}')
        assert isinstance(list_, list), list_
        self._aliases[key] = list_

    def set_services(self, services):
        for key, service in services.items():
            known_service = self._parse_service(key, service)
            assert key not in self._services, (key, self._services)
            assert known_service.tcp_port is not None, known_service
            assert known_service.tcp_port not in self._tcp_ports, (
                self._tcp_ports, known_service)
            self._services[key] = known_service
            self._tcp_ports[known_service.tcp_port] = key

    def _parse_service(self, key, service):
        if (not isinstance(service, dict) or
                (set(service.keys()) - set(['iface'])) !=
                set(['match', 'allow'])):
            raise ValueError(
                f'must has exactly match+allow[+iface] for every '
                f'service, got: {key!r}: {service!r}')

        # FIXME: right now we expect a certain format
        match_parts = service['match'].split()
        assert len(match_parts) == 4, service['match']
        assert match_parts[0] == '-p', service['match']
        assert match_parts[1] == 'tcp', service['match']
        assert match_parts[2] == '--dport', service['match']
        tcp_port = int(match_parts[3])

        # Extract IPs immediately.
        allowed = []
        assert isinstance(service['allow'], list), (key, service)
        for allow in service['allow']:
            allowed.extend(self.get_cidrs(allow))
        allowed.sort()

        interfaces = None
        if 'iface' in service:
            assert isinstance(service['iface'], list), service
            assert all(str(i) for i in service['iface']), service
            interfaces = []
            for iface in service['iface']:
                if iface.startswith('$'):
                    assert iface == '$extif', iface
                    iface = '$IFACE_extif'
                interfaces.append(iface)

        return KnownService(
            name=key,
            src_allowed=allowed,
            tcp_port=tcp_port,
            interfaces=interfaces,
        )

    def get_service_by_tcp_port(self, port):
        if port not in self._tcp_ports:
            return None
        return self._services[self._tcp_ports[port]]

    def print_iptables(self):
        print('''\
#!/bin/sh
IFACE_extif=$(ip -o route get 1.2.3.4 |
    sed -e 's/.* dev //;s/ .*//')
IP_extif=$(ip -o -4 addr show dev $IFACE_extif |
    sed -e 's/.* inet //;s@[ /].*@@')

if ! iptables -t filter -F INPUT_test 2>/dev/null; then
    iptables -t filter -N INPUT_test
    iptables -t filter -I INPUT -j INPUT_test
fi

# Do not allow just anything on localhost. Only 127/8 and the primary IP.
iptables -t filter -A INPUT_test -i lo -s 127.0.0.0/8 -d 127.0.0.0/8 -j RETURN
iptables -t filter -A INPUT_test -i lo -s $IP_extif -d $IP_extif -j RETURN
# Ignore ICMP. Ignore non-new TCP.
iptables -t filter -A INPUT_test -p icmp -j RETURN
iptables -t filter -A INPUT_test -p tcp -m tcp '!' --tcp-flags SYN,ACK SYN \\
    -j RETURN
# HACKS: For wireguard traffic (cilium_wg*)
iptables -t filter -A INPUT_test -i $IFACE_extif \\
    -p udp --sport 51871 --dport 51871 -j RETURN
# FIXME: do something with udp. unfortunately not an easy SYN to check..
# FIXME: icmpv6? tcp6? udp6?
# Ignore established (or returning!) traffic. DNS UDP responses, among others.
iptables -t filter -A INPUT_test -m state --state ESTABLISHED,RELATED -j RETURN

allow() {
    iptables -t filter -A INPUT_test \\
         $from $service -m comment --comment "$*" -j RETURN
}

''')
        # Get all possible interfaces first:
        possible_interfaces = set()
        for port in self._tcp_ports:
            known_service = self.get_service_by_tcp_port(port)
            if known_service.interfaces:
                possible_interfaces.update(known_service.interfaces)
            else:
                possible_interfaces.add('')
        possible_interfaces = list(sorted(possible_interfaces))

        # Interfaces in outer loop because this makes the iptables -nvL output
        # so much more readable when investigating.
        for iface in possible_interfaces:
            for port in self._tcp_ports:
                known_service = self.get_service_by_tcp_port(port)

                # If the known_service does not mention an interface, allow if
                # we're doing the any interface. Otherwise allow only if the
                # interfaces match.
                if iface and known_service.interfaces:
                    if iface not in known_service.interfaces:
                        continue
                elif not iface and not known_service.interfaces:
                    pass
                else:
                    continue

                for allowed_name, allowed_net in known_service.src_allowed:
                    service = f'-p tcp --dport {port}'
                    if iface:
                        service = f'-i {iface} {service}'
                    print(
                        f'service="{service}" '
                        f"from='-s {allowed_net}' "
                        f"allow '{known_service.name}: {allowed_name}'")
                print()

        print('#iptables -t filter -A INPUT_test  # just count')
        print('iptables -t filter -A INPUT_test -j LOG  # log the rest')
        print()


def make_spinner():
    """"
    Return an interator to a single character as a spinner

    Usage:

        spinner = make_spinner()
        while do_stuff():
            print(next(spinner), end='')  # '|' -> '/' -> '-' -> '\\'
    """
    cycle, cyclepos, cyclelast = r'/-\|', 0, time()

    while True:
        yield f'\r{cycle[cyclepos]}\r'

        # Less jerky cycle by only spinning if the previous iteration has been
        # seen.
        t = time()
        if (t - cyclelast) >= 0.2:
            cyclepos = (cyclepos + 1) % 4
            cyclelast = t


class NoService(namedtuple('NoService', 'service source')):
    """
    Light weight "Service" placeholder for undefined service
    """
    @classmethod
    def from_small_identifier(cls, small_identifier):
        dst_ip = inet_ntop(AF_INET, small_identifier[0:4])
        (dst_port,) = unpack('>H', small_identifier[4:6])
        src_ip = inet_ntop(AF_INET, small_identifier[6:10])
        return cls(service=f'{dst_ip}:{dst_port}', source=src_ip)


class Service:
    """
    Heavier weight "Service" with allowed networks and counts
    """
    def __init__(self, known_service):
        self.name = known_service.name
        # TODO: replace source_net with binary+mask?
        self._counts = dict((net, 0) for id_, net in known_service.src_allowed)
        self._networks_to_names = dict(
            (net, id_) for id_, net in known_service.src_allowed)

    def total_count(self):
        return sum(self._counts.values())

    def add_count(self, small_identifier):
        source_ip = ip_network(inet_ntop(AF_INET, small_identifier[6:10]))
        net = self.get_net(source_ip)
        if net:
            self._counts[net] += 1
            return self._counts[net], self._networks_to_names[net]
        return 0, ''

    def get_net(self, source_ip):
        for net in self._counts.keys():
            if source_ip.subnet_of(net):
                return net
        return None

    def __str__(self):
        return f"Service(name='{self.name}')"


class IPFrame(Packet):
    __hdr__ = tuple()

    def unpack(self, buf):
        super().unpack(buf)
        try:
            self.data = IP(self.data)
            self.ip = self.data
        except (KeyError, UnpackError):
            pass


class HostHubble:
    SNAPLEN = 16 + 60 + 20      # ethernet/sll + IP(max) + TCP(min..)
    PROMISC = False             # no need for promiscuous mode
    BUFFER_TIMEOUT_MS = 200     # MUST be non-zero
    TCP_SYN_ONLY = '(tcp[tcpflags]&tcp-syn)!=0 and (tcp[tcpflags]&tcp-ack)=0'
    SLL_TYPES_TO_US = (
        # 3=promisc, 4=from_us
        0,  # specifically sent to us by somebody else
        1,  # broadcast by somebody else
        2,  # multicast, but not broadcast, by somebody else
    )
    ARPHRD_ETHER = 1
    ARPHRD_LOOPBACK = 772
    ARPHRD_NONE = 0xFFFE
    LOOPBACK_MAC = b'\x00\x00\x00\x00\x00\x00'

    def __init__(self, known_rules, listen_interface='any'):
        # Args
        self._known_rules = known_rules
        self._listen_interface = listen_interface

        # Store the counts
        self._services = {}
        self._noservices = defaultdict(int)

    def main(self):
        cap = open_live(
            self._listen_interface, self.SNAPLEN, self.PROMISC,
            self.BUFFER_TIMEOUT_MS)

        # Here we filter so we only get TCP SYNs. We may want to expand on this
        # later.
        cap.setfilter(self.TCP_SYN_ONLY)

        if self._listen_interface == 'any':
            decode_frame = self._decode_sll  # 16 bytes..
            minlen = 16 + 20 + 20
        elif self._listen_interface.startswith('cilium_wg'):
            decode_frame = self._decode_ip   # 0 bytes
            minlen = 14 + 20 + 20
        else:
            decode_frame = self._decode_ethernet  # 14 bytes..
            minlen = 20 + 20

        spinner = (make_spinner() if sys.stderr.isatty() else (lambda: ''))

        while True:
            # Show spinning wheel
            print(next(spinner), end='', file=sys.stderr)

            header, packet = cap.next()
            if len(packet) < minlen:
                warn(f'skipping small packet {packet!r}')
                continue

            # #print('packet', repr(packet))
            try:
                outer_frame = decode_frame(packet)
                if not outer_frame:
                    continue

                ip = outer_frame.ip
                tcp = ip.tcp
                small_identifier = self.compress(ip, tcp)
            except (AttributeError, NeedData):
                print('packet', repr(packet))
                print('outer_frame', repr(outer_frame))
                raise

            if not self.add_known_service(small_identifier):
                self.add_unknown_service(small_identifier)

    def _decode_ethernet(self, packet):
        outer_frame = Ethernet(packet)
        # We could improve this by finding our network interfaces and selecting
        # only those with dst_ether us and broadcast (and multicast).
        # Do we need 'ip -brief -oneline link show' here?
        return outer_frame

    def _decode_ip(self, packet):
        # Here we could only filter if we knew which IPs are ours.
        return IPFrame(packet)

    def _decode_sll(self, packet):
        # Here we filter everything that is:
        # - not to us
        # - is sent on the loopback interface (with 127.0.0.0/8)
        outer_frame = SLL(packet)
        if outer_frame.type not in self.SLL_TYPES_TO_US:
            warn(f'skipping not-to-us packet type {outer_frame.type}')
            return None

        if outer_frame.hrd == self.ARPHRD_LOOPBACK:
            assert outer_frame.hlen == 6, repr(outer_frame)
            assert outer_frame.hdr[0:6] == self.LOOPBACK_MAC, (
                repr(outer_frame))  # source from localhost device
            if outer_frame.ip.src[0] == outer_frame.ip.src[0] == 127:
                warn('skipping loopback packet with 127.0.0.0/8')
                return None
        elif outer_frame.hrd == self.ARPHRD_NONE:
            print('ARPHRD_NONE:', repr(outer_frame), repr(outer_frame.ip))
            assert outer_frame.hlen == 0, repr(outer_frame)
        elif outer_frame.hrd == self.ARPHRD_ETHER:
            pass
        else:
            assert False, (
                f'ARPHRD_{outer_frame.hrd}', outer_frame.hlen,
                ':'.join(f'{i:02x}' for i in outer_frame.hdr),
                repr(outer_frame))

        return outer_frame

    def add_known_service(self, small_identifier):
        service = self.get_known_service(small_identifier)
        if not service:
            return False  # not found by dst port

        count, name = service.add_count(small_identifier)
        if not count:
            return False  # not found by src ip

        if count == 1:
            print(f'* expected {service.name!r} <- {name!r}')

        return True

    def add_unknown_service(self, small_identifier):
        if small_identifier not in self._noservices:
            nosvc = NoService.from_small_identifier(small_identifier)
            print(f'* UNEXPECTED {nosvc.service!r} <- {nosvc.source!r}')
        self._noservices[small_identifier] += 1

    def get_known_service(self, small_identifier):
        smaller_identifier = small_identifier[0:6]  # dstip+dport
        try:
            service = self._services[smaller_identifier]
        except KeyError:
            (dport,) = unpack('>H', small_identifier[4:6])
            known_service = self._known_rules.get_service_by_tcp_port(dport)
            if not known_service:
                return None

            service = self._services[smaller_identifier] = (
                Service(known_service))
        return service

    def compress(self, ip, tcp):
        return ip.dst + pack('>H', tcp.dport) + ip.src

    def report(self):
        print('')
        print('# reporting for duty!')
        print()
        print('## known services:')
        for smaller_identifier, service in sorted(self._services.items()):
            count = service.total_count()
            print(f'{count:4d}  {service.name}:')
            for net, count in sorted(service._counts.items()):
                net_name = service._networks_to_names[net]
                print(f'{count:4d}    {net_name}')
        print()
        print('## unknown services:')
        for small_identifier, count in sorted(self._noservices.items()):
            noservice = NoService.from_small_identifier(small_identifier)
            print(f'{count:4d}  {noservice}')
        print()


def formatwarning(message, category, filename, lineno, line=None):
    """
    Override default Warning layout, from:

        /PATH/TO/dutree.py:326: UserWarning:
            [Errno 2] No such file or directory: '/0.d/05.d'
          warnings.warn(str(e))

    To:

        dutree.py:330: UserWarning:
            [Errno 2] No such file or directory: '/0.d/05.d'
    """
    return '{basename}:{lineno}: {category}: {message}\n'.format(
        basename=path.basename(filename), lineno=lineno,
        category=category.__name__, message=message)
warnings.formatwarning = formatwarning  # noqa


if __name__ == '__main__':
    if sys.argv[1:2] in (['-h'], ['--help']):
        argv0 = sys.argv[0]
        print(f'Usage: {argv0} [YAMLFILE] [--make-iptables]')
        print(f'or: RUNTESTS=1 {sys.argv[0]}')
        print('')
        print('Testing using iptables:')
        print(f'# {argv0} YAMLFILE --make-iptables > iptables && sh iptables')
        print('# tail /var/log/kern.log')
        print('')
        print('Testing using fwsniff packet capture:')
        print(f'# {argv0} [YAMLFILE]')
        print('')
        exit(0)

    if environ.get('RUNTESTS', '') not in ('', '0', 'n'):
        from unittest import main
        main()
        raise RuntimeError('unreachable code')

    if path.exists(' '.join(sys.argv[1:2])):
        known_rules = KnownRules.from_yaml_file(' '.join(sys.argv[1:2]))
        sys.argv.pop(1)
    else:
        known_rules = KnownRules()

    if sys.argv[1:2] == ['--make-iptables']:
        known_rules.print_iptables()
        exit(0)

    if sys.argv[1:]:
        raise RuntimeError('unknown args')

    hosthubble = HostHubble(known_rules)
    try:
        hosthubble.main()
    except KeyboardInterrupt:
        hosthubble.report()