"""
Get details about packets blocked by ufw (uses journald)
"""

import datetime
import re
from systemd import journal

from logparse import config
from logparse.formatting import *
from logparse.load_parsers import Parser
from logparse.util import resolve

PROTOCOLS = ["TCP", "UDP", "UDP-Lite", "ICMP", "ICMPv6", "AH", "SCTP", "MH"]

class Packet():
    """
    Class to hold variables for each packet. Also parses incoming log messages
    on object initialisation.
    """

    def __init__(self, msg):
        """
        Determine fields in log message. If any of the fields are missing the
        log is considered malformed and is discarded. Also the protocol can be
        specified as either an integer or a string - see `man ufw`.
        """
        try:
            self.inif, self.outif, self.mac, self.src, self.dst, self.len, \
                    self.proto, self.spt, self.dpt = \
                    re.search(r"IN=(?P<inif>\w*).*OUT=(?P<outif>\w*).*"
                        "MAC=(?P<mac>\S*).*SRC=(?P<src>\S*).*DST=(?P<dst>\S*)"
                        ".*LEN=(?P<length>\d*).*PROTO=(?P<proto>\S*)"
                        "(?:\sSPT=(?P<spt>\d*))?(?:\sDPT=(?P<dpt>\d*))?", msg
                    ).groupdict().values()
            if self.proto and self.proto.isdigit():
                self.proto = PROTOCOLS[int(self.proto)-1]
        except Exception as e:
            logger.warning("Malformed packet log: {0}. Error message: {1}"
                    .format(msg, str(e)))
            return None

class UfwJournald(Parser):

    def __init__(self):
        super().__init__()
        self.name = "ufw"
        self.info = "Get details about packets blocked by ufw"

    def parse_log(self):

        logger.debug("Starting ufw section")
        section = Section("ufw")

        # Find applicable log entries

        j = journal.Reader()
        j.this_machine()
        j.add_match(_TRANSPORT='kernel')
        j.add_match(PRIORITY=4)
        j.seek_realtime(section.period.startdate)
        
        logger.debug("Searching for messages")

        blocked_packets = [Packet(entry["MESSAGE"]) for entry in j
                if "MESSAGE" in entry and "UFW BLOCK" in entry["MESSAGE"]]

        # Parse messages

        logger.debug("Parsing messages")

        inbound_interfaces = []
        outbound_interfaces = []
        n_inbound = n_outbond = 0
        src_ips = []
        dst_ips = []
        src_ports = []
        dst_ports = []
        protocols = {'UDP': 0, 'TCP': 0}
        src_macs = []

        for pkt in blocked_packets:
            if pkt.inif:
                inbound_interfaces.append(pkt.inif)
            elif pkt.outif:
                outbound_interfaces.append(pkt.outif)
            if pkt.src: src_ips.append(resolve(pkt.src,
                config.prefs.get("ufw", "ufw-resolve-domains")))
            if pkt.dst: dst_ips.append(resolve(pkt.dst,
                config.prefs.get("ufw", "ufw-resolve-domains")))
            if pkt.spt: src_ports.append(pkt.spt)
            if pkt.dpt: dst_ports.append(pkt.dpt)
            if pkt.proto: protocols[pkt.proto] += 1

        # Format data objects

        section.append_data(Data(subtitle="{} blocked ({} UDP, {} TCP)".format(
                plural("packet", len(blocked_packets)),
                protocols['UDP'], protocols['TCP'])))

        src_port_data = Data(items=src_ports)
        src_port_data.orderbyfreq()
        src_port_data.subtitle = plural("source port", len(src_port_data.items))
        src_port_data.truncl(config.prefs.getint("logparse", "maxlist"))
        section.append_data(src_port_data)

        dst_port_data= Data(items=dst_ports)
        dst_port_data.orderbyfreq()
        dst_port_data.subtitle = plural("destination port",
                len(dst_port_data.items))
        dst_port_data.truncl(config.prefs.getint("logparse", "maxlist"))
        section.append_data(dst_port_data)

        src_ips_data= Data(items=src_ips)
        src_ips_data.orderbyfreq()
        src_ips_data.subtitle = plural("source IP", len(src_ips_data.items))
        src_ips_data.truncl(config.prefs.getint("logparse", "maxlist"))
        section.append_data(src_ips_data)

        dst_ips_data= Data(items=dst_ips)
        dst_ips_data.orderbyfreq()
        dst_ips_data.subtitle = plural("destination IP",
                len(dst_ips_data.items))
        dst_ips_data.truncl(config.prefs.getint("logparse", "maxlist"))
        section.append_data(dst_ips_data)

        logger.info("Finished ufw section")
        return section

    def check_dependencies(self):
        """
        Basic dependency check to determine if there are any logs to parse
        """

        ufw_cmdline = "ufw --version"
        if self._check_dependency_command(ufw_cmdline)[0] != 0:
            return (False, ["ufw"])
        else:
            return (True, None)