#
#   logactio  -  simple framework for doing configured action on certain
#                log file events
#
#   Copyright 2012 - 2015   David Sommerseth <dazo@eurephia.org>
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation, version 2
#   of the License.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
#   For the avoidance of doubt the "preferred form" of this code is one which
#   is in an open unpatent encumbered format. Where cryptographic key signing
#   forms part of the process of creating an executable the information
#   including keys needed to generate an equivalently functional executable
#   are deemed to be part of the source code.
#

import sys, os, subprocess, tempfile, re, time
import LogActio.Message, LogActio.ReporterQueue


class IPTipset(LogActio.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

        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 == "true" or create == "1") and True or False
            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 == "true" or counters == "1") and True or False

        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 and logger or 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
        LogActio.ReporterQueue.ReporterQueue.__init__(self,
                                                      "IPTipset",
                                                      "IPTables IPset processor",
                                                      self.__processqueue)


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


    def __parse_cmd_log(self, cmd, logfp):
        logfp.seek(0)
        for line in logfp:
            self.__log(3, "[IPTipset] %s: %s" % (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, "[IPTipset] Executing: %s" % " ".join(args))
        cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
        res = cmd.wait()
        self.__parse_cmd_log("ipset:%s" % mode, tmplog)

        # Clean up
        tmplog.close()
        del tmplog

        return res


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


    def __load_ipset_state(self):
        try:
            f = open(self.__ipset_save, "r")
            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 IOError as e:
            # Ignore "No such file or directory", as the file may not exist
            if e.errno != 2:
                raise e


    def __save_ipset_state(self):
        args = ["ipset", "save", self.__ipsetname]
        f = open(self.__ipset_save, "w")
        self.__log(4, "[IPTipset]: Saving state - Executing %s" % " ".join(args))
        subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=f)
        f.close()


    def __parse_already_registered(self):
        args = ["ipset", "save", self.__ipsetname]
        tmplog = tempfile.SpooledTemporaryFile(mode="w+b")
        self.__log(4, "[IPTipset] Executing: %s" % " ".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("^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()
        del tmplog
        return retlist


    def __prepare_iptables(self):
        for chain in self.__iptchains:
            # Prepare iptables command line
            args = False
            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, "[IPTipset] Executing: %s" % " ".join(args))
            cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
            cmd.wait()
            self.__parse_cmd_log("iptables:%s" % chain, tmplog)
            tmplog.close()
            del tmplog


    def __cleanup_iptables(self):
        for chain in self.__iptchains:
            # Prepare iptables command line
            args = False
            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, "[IPTipset] Executing: %s" % " ".join(args))
            cmd = subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=tmplog, stderr=tmplog)
            cmd.wait()
            self.__parse_cmd_log("iptables:%s" % chain, tmplog)
            tmplog.close()
            del tmplog

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


    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() == LogActio.Message.MSG_SHUTDOWN ):
                # Prepare for shutdown
                done = True

            elif( msg.MessageType() == LogActio.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, "[IPTipset] IP address %s was removed from ipset '%s'.  Will re-add it." % (m["ipaddress"], self.__ipsetname))
                        registered.remove(m["ipaddress"])
                        raise ValueError

                except ValueError:
                    self.__log(2, "[IPTipset] {Rule %s} Adding IP address %s to ipset '%s' based on entry in log file '%s' with the threshold %i after %i hits" %
                               (m["rulename"], m["ipaddress"], self.__ipsetname, m["logfile"], m["threshold"], m["count"]))
                    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)