#
#   logactio  -  simple framework for doing configured action on certain
#                log file events
#
# SPDX-FileCopyrightText: Copyright 2012-  David Sommerseth <dazo@eurephia.org>
# SPDX-License-Identifier: GPL-2.0-only
#

import json
import subprocess
import sys
import tempfile
from LogActio import Message, ReporterQueue


class NftablesSet(ReporterQueue.ReporterQueue):
    """LogActio reporter module which adds IP addresses to a pre-existing nftables set

    Example configuration to be used in /etc/logactio.cfg

        [Reporter:nftblock]
        module: NftablesSet
        nft-family: inet
        nft-table: filter
        nft-set: blocklist
        nft-save: /var/lib/nftables/logactio-nftset.save

        nft-family is optional and defaults to 'inet'.
        nft-table and nft-set are mandatory.

        The nftables set must be pre-created before starting logactio.  This
        module will not create or destroy the set.

        If nft-save is set, the JSON output of 'nft -j list set' is written to
        the given file on every update.  Restoring the set from this file at
        startup is left to an external tool (e.g. a systemd service or init
        script).

        The example above will result in these commands being executed on each
        new event:

            # nft add element inet filter blocklist { <ip> }
    """

    def __init__(self, config, logger=None):
        self.__nft_save = False
        self.__idx = 0

        if "nft-table" not in config:
            raise Exception("NftablesSet is missing the nft-table setting")
        if "nft-set" not in config:
            raise Exception("NftablesSet is missing the nft-set setting")

        self.__family = config.get("nft-family", "inet")
        self.__table = config["nft-table"]
        self.__setname = config["nft-set"]

        if "nft-save" in config:
            self.__nft_save = config["nft-save"]

        self.__log = logger if logger is not None else self.__logfnc

        super().__init__("NftablesSet", "nftables set processor", self.__processqueue)


    def __logfnc(self, lvl, msg):
        print(msg)
        sys.stdout.flush()


    def __parse_cmd_log(self, cmd, logfp):
        logfp.seek(0)
        for line in logfp:
            self.__log(3, f"[NftablesSet] {cmd}: {line.decode('utf-8')}")


    def __call_nft(self, args):
        tmplog = tempfile.SpooledTemporaryFile(mode="w+b")
        self.__log(4, f"[NftablesSet] Executing: {' '.join(args)}")
        cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
        res = cmd.wait()
        self.__parse_cmd_log("nft", tmplog)
        tmplog.close()
        return res


    def __add_element(self, ipaddress):
        return self.__call_nft(["nft", "add", "element",
                                self.__family, self.__table, self.__setname,
                                "{", ipaddress, "}"])


    def __nft_list_json(self):
        args = ["nft", "-j", "list", "set", self.__family, self.__table, self.__setname]
        tmplog = tempfile.SpooledTemporaryFile(mode="w+b")
        self.__log(4, f"[NftablesSet] Executing: {' '.join(args)}")
        cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
        cmd.wait()
        tmplog.seek(0)
        output = tmplog.read().decode("utf-8")
        tmplog.close()
        return output


    def __parse_already_registered(self):
        output = self.__nft_list_json()
        try:
            data = json.loads(output)
            for obj in data.get("nftables", []):
                s = obj.get("set")
                if s and s.get("name") == self.__setname:
                    return [e["elem"] if isinstance(e, dict) else e
                            for e in s.get("elem", [])]
        except (json.JSONDecodeError, KeyError):
            pass
        return []


    def __save_state(self):
        args = ["nft", "list", "set", self.__family, self.__table, self.__setname]
        self.__log(4, f"[NftablesSet] Executing: {' '.join(args)}")
        self.__log(4, f"[NftablesSet] Saving state to {self.__nft_save}")
        with open(self.__nft_save, "w") as f:
            subprocess.run(args, stdin=subprocess.DEVNULL, stdout=f, check=False)


    def __processqueue(self):
        self.__log(1, "[NftablesSet] Ready.")
        registered = self.__parse_already_registered()

        done = False
        while not done:
            msg = self._QueueGet()

            if msg.MessageType() == Message.MSG_SHUTDOWN:
                done = True

            elif msg.MessageType() == Message.MSG_SEND:
                m = msg.Message()

                if m["ipaddress"] not in registered:
                    self.__idx += 1

                    self.__log(2, f"[NftablesSet] {{{self.__idx}}} {{Rule {m['rulename']}}} "
                                  f"Logfile: {m['logfile']}, Threshold: {m['threshold']}, "
                                  f"Count: {m['count']} -- Adding IP address {m['ipaddress']} "
                                  f"to nftables set "
                                  f"'{self.__family} {self.__table} {self.__setname}'")

                    self.__add_element(m["ipaddress"])
                    registered.append(m["ipaddress"])

                    if self.__nft_save:
                        self.__save_state()

        self.__log(3, "[NftablesSet] Module shut down")


    def ProcessEvent(self, logfile, rulename, msg, count, threshold):
        msg = {"rulename": rulename, "threshold": threshold,
               "ipaddress": msg[0], "logfile": logfile,
               "count": count}
        self._QueueMsg(0, msg)


def InitReporter(config, logger=None):
    return NftablesSet(config, logger)
