#
#   ufw_journald.py
#   
#   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():

    def __init__(self, msg):
        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_journald"
        self.info = "Get details about packets blocked by ufw"

    def parse_log(self):

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

        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"]]

        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
                
        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