#
#   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 os
import re
import subprocess
import sys
import tempfile
import time
from LogActio import Message, ReporterQueue


class IPTipset(ReporterQueue.ReporterQueue):
    """LogActio reporter modules which adds IP addresses to an iptalbes IP set chain

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

        [Reporter:ipsetblock]
        module: IPTipset
        ipset-name: BlockList
        ipset-create: True
        ipset-hashsize: 2048
        ipset-timeout: 3600
        ipset-counters: True
        ipset-save: /var/lib/ipset/logactio-ipset.save
        iptables-chains: INPUT,OUTPUT,FORWARD
        iptables-insert-points: INPUT:2,FORWARD:4
        iptables-jump: DROP

        ipset-name contains the reference string used when calling 'ipset'.  This
        field is mandatory.

        If ipset-create is True, true or 1, it will attempt to create this ipset set
        when starting up.  In this case, the ipset-hashsize will be used, if set.  See
        the ipset(8) man page for more informaion.  The ipset set created will be
        of the 'hash:ip' type.  The ipset-create is optional and defaults to False.
        If the ipset set does not exist and ipset-create is not enabled, it will fail
        when testing the ipset set at startup.

        The ipset-timeout parameter adds the a default timeout for all added entries,
        see the ipset(8) man page for more information about this feature.

        The ipset-counters parameter adds a counter field for each added entry.  This
        must be True, true or 1 to be considered set.  Otherwise it is disabled.  See
        the ipset(8) man page for more information about this feature too.

        If ipset-save is set with a file name, ipset will preserve the ipset state when
        shutting down.  It will also reload the state upon start-up.

        If iptables-chains is set, it will insert an iptables rule which checks the
        ipset set.  If iptables-insert-point is set (optinal), the rule will be inserted
        at the given point, otherwise it will be inserted at the top.  The iptables-jump
        is mandatory with iptables-chains and adds the jump destiation when a match
        against the ipset set is found.

        The example config above will result in these commands being executed when
        starting:

             # ipset --exist create BlockList hash:ipset hashsize 2048 timeout 3600
             # iptables -I INPUT 2 -m set --match-set BlockList src -j DROP
             # iptables -I OUTPUT -m set --match-set BlockList src -j DROP
             # iptables -I FORWARD 4 -m set --match-set BlockList src -j DROP
    """

    def __init__(self, config, logger=None):

        # Configuration parsing
        self.__create = False
        self.__hashsize = "1024"
        self.__timeout = 0
        self.__counters = False
        self.__iptchains = False
        self.__iptchainsjump = False
        self.__iptchaininserts = False
        self.__ipset_save = False
        self.__idx = 0

        if "ipset-name" not in config:
            raise Exception("IPTipset is missing in ipset name")
        else:
            self.__ipsetname = config["ipset-name"]

        if "ipset-create" in config:
            create = config["ipset-create"].lower()
            self.__create = create in ("true", "1")
            if self.__create and "ipset-hashsize" in config:
                self.__hashsize = str(config["ipset-hashsize"])
            if self.__create and "ipset-timeout" in config:
                self.__timeout = str(config["ipset-timeout"])
            if self.__create and "ipset-counters" in config:
                counters = config["ipset-counters"].lower()
                self.__counters = counters in ("true", "1")

        if "iptables-chains" in config:
            self.__iptchains = [v.strip() for v in config["iptables-chains"].split(",")]
            if "iptables-jump" not in config:
                raise Exception("IPTipset needs the iptables-jump variable when iptables-chains is set")
            self.__iptchainsjump = config["iptables-jump"]
            if "iptables-insert-points" in config:
                self.__iptchaininserts = {}
                for inspoint in config["iptables-insert-points"].split(","):
                    (chain, point) = inspoint.split(":")
                    self.__iptchaininserts[chain.strip()] = str(point)

        if "ipset-save" in config:
            self.__ipset_save = config["ipset-save"]

        # Prepare this object, ipset and iptables
        self.__log = logger if logger is not None else self.__logfnc
        if self.__create:
            self.__prepare_ipset()
        if self.__iptchains:
            self.__prepare_iptables()
        if self.__ipset_save:
            self.__load_ipset_state()

        # Register this module as a reporter module
        super().__init__("IPTipset", "IPTables IPset 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"[IPTipset] {cmd}: {line.decode('utf-8')}")


    def __call_ipset(self, mode, args=None):
        if mode == "create":
            args = ["ipset", "--exist", "create",  self.__ipsetname, "hash:ip"] + args
        else:
            if args is None:
                args = ["ipset", mode, self.__ipsetname]
            elif isinstance(args, list):
                args = ["ipset", mode, self.__ipsetname] + args
            else:
                args = ["ipset", mode, self.__ipsetname, args]

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

        # Clean up
        tmplog.close()

        return res


    def __prepare_ipset(self):
        params = []
        params += ["hashsize", self.__hashsize] if self.__hashsize else []
        params += ["timeout", self.__timeout] if self.__timeout else []
        params += ["counters"] if self.__counters else []
        self.__call_ipset("create", params)


    def __load_ipset_state(self):
        try:
            with open(self.__ipset_save, "r") as f:
                for line in f:
                    s = line.split()
                    # Only care about the add lines, we've already created the ipset list
                    if s[0] != 'add':
                        continue
                    self.__call_ipset(s[0], s[2:])
        except OSError as e:
            # Ignore "No such file or directory", as the file may not exist
            if e.errno != 2:
                raise


    def __save_ipset_state(self):
        args = ["ipset", "save", self.__ipsetname]
        self.__log(4, f"[IPTipset]: Saving state - Executing {' '.join(args)}")
        with open(self.__ipset_save, "w") as f:
            subprocess.run(args, stdin=subprocess.DEVNULL, stdout=f, check=False)


    def __parse_already_registered(self):
        args = ["ipset", "save", self.__ipsetname]
        tmplog = tempfile.SpooledTemporaryFile(mode="w+b")
        self.__log(4, f"[IPTipset] Executing: {' '.join(args)}")
        cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
        cmd.wait()

        # Process all "add" lines which matches our ipset set name
        tmplog.seek(0)
        rg = re.compile(r"^add (.*) \b((?:[0-9]{1,3}\.){3}[0-9]{1,3})\b")
        retlist = []
        for line in tmplog:
            m = rg.match(line.strip().decode('utf-8'))
            if m:
                rgm = m.groups()
                if rgm[0] == self.__ipsetname:
                    retlist.append(rgm[1])
        tmplog.close()
        return retlist


    def __prepare_iptables(self):
        for chain in self.__iptchains:
            # Prepare iptables command line
            if self.__iptchaininserts and chain in self.__iptchaininserts:
                args = ["iptables", "-I", chain, self.__iptchaininserts[chain],
                        "-m", "set", "--match-set", self.__ipsetname,
                        "-j", self.__iptchainsjump]
            else:
                args = ["iptables", "-I", chain,
                        "-m", "set", "--match-set", self.__ipsetname, "src",
                        "-j", self.__iptchainsjump]

            # Call iptables and wait for it to complete and log the output
            tmplog = tempfile.SpooledTemporaryFile(mode="w+b")
            self.__log(4, f"[IPTipset] Executing: {' '.join(args)}")
            cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
            cmd.wait()
            self.__parse_cmd_log(f"iptables:{chain}", tmplog)
            tmplog.close()


    def __cleanup_iptables(self):
        for chain in self.__iptchains:
            # Prepare iptables command line
            if self.__iptchaininserts and chain in self.__iptchaininserts:
                args = ["iptables", "-D", chain, self.__iptchaininserts[chain],
                        "-m", "set", "--match-set", self.__ipsetname,
                        "-j", self.__iptchainsjump]
            else:
                args = ["iptables", "-D", chain,
                        "-m", "set", "--match-set", self.__ipsetname, "src",
                        "-j", self.__iptchainsjump]

            # Call iptables and wait for it to complete and log the output
            tmplog = tempfile.SpooledTemporaryFile(mode="w+b")
            self.__log(4, f"[IPTipset] Executing: {' '.join(args)}")
            cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
            cmd.wait()
            self.__parse_cmd_log(f"iptables:{chain}", tmplog)
            tmplog.close()

        # Allow iptables to complete its job before indicating we're done
        time.sleep(5)


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

        # Process the internal message queue
        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()

                try:
                    registered.index(m["ipaddress"])

                    # Check if this IP address is still in ipset, if not register it again
                    if self.__call_ipset("test", m["ipaddress"]) == 1:
                        self.__log(4, f"[IPTipset] IP address {m['ipaddress']} was removed from "
                                      f"ipset '{self.__ipsetname}'.  Will re-add it.")
                        registered.remove(m["ipaddress"])
                        raise ValueError

                except ValueError:
                    self.__idx += 1

                    self.__log(2, f"[IPTipset] {{{self.__idx}}} {{Rule {m['rulename']}}} "
                                  f"Logfile: {m['logfile']}, Threshold: {m['threshold']}, "
                                  f"Count: {m['count']} -- Adding IP address {m['ipaddress']} "
                                  f"to ipset '{self.__ipsetname}'")

                    self.__call_ipset("add", m["ipaddress"])
                    registered.append(m["ipaddress"])


        if self.__iptchains:
            self.__cleanup_iptables()
        if self.__ipset_save:
            self.__save_ipset_state()
        if self.__iptchains and self.__ipset_save:
            self.__call_ipset("destroy")
        self.__log(3, "[IPTipset] Module shut down")


    def ProcessEvent(self, logfile, rulename, msg, count, threshold):
        # FIXME:  Ensure the IP address is infact an IP address (regex check)

        # Format the report message
        msg = {"rulename": rulename, "threshold": threshold,
               "ipaddress": msg[0], "logfile": logfile,
               "count": count}

        # Queue the message for sending
        self._QueueMsg(0, msg)


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