#
#   logactio  -  simple framework for doing configured action on certain
#                log file events
#
#   Copyright 2013   David Sommerseth <dazo@users.sourceforge.net>
#
#   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; either 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307 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
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:QPID]
        module: IPTipset
        ipset-name: BlockList
        ipset-create: True
        ipset-hashsize: 2048
        ipset-timeout: 3600
        ipset-counters: True
        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 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

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

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

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

        # 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()

        # 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(2, "[IPTipset] %s: %s" % (cmd, line))


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

        nullfp = os.open("/dev/null", os.O_RDWR)
        tmplog = tempfile.SpooledTemporaryFile(mode="rw+b")
        self.__log(4, "[IPTipset] Executing: %s" % " ".join(args))
        cmd = subprocess.Popen(args, stdin=nullfp, stdout=tmplog, stderr=tmplog)
        cmd.wait()
        self.__parse_cmd_log("ipset:%s" % mode, tmplog)

        # Clean up
        tmplog.close()
        del tmplog
        os.close(nullfp);


    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 __parse_already_registered(self):
        args = ["ipset", "save", self.__ipsetname]
        nullfp = os.open("/dev/null", os.O_RDWR)
        tmplog = tempfile.SpooledTemporaryFile(mode="rw+b")
        self.__log(4, "[IPTipset] Executing: %s" % " ".join(args))
        cmd = subprocess.Popen(args, stdin=nullfp, 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())
            if m:
                rgm = m.groups()
                if rgm[0] == self.__ipsetname:
                    retlist.append(rgm[1])
        tmplog.close()
        del tmplog
        os.close(nullfp)
        del nullfp
        return retlist


    def __prepare_iptables(self):
        nullfp = os.open("/dev/null", os.O_RDWR)

        for chain in self.__iptchains:
            # Prepare iptables command line
            args = False
            if self.__iptchaininserts and self.__iptchaininserts.has_key(chain):
                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="rw+b")
            self.__log(4, "[IPTipset] Executing: %s" % " ".join(args))
            cmd = subprocess.Popen(args, stdin=nullfp, stdout=tmplog, stderr=tmplog)
            cmd.wait()
            self.__parse_cmd_log("iptables:%s" % chain, tmplog)
            tmplog.close()
            del tmplog

        # Clean up
        os.close(nullfp)
        del nullfp


    def __cleanup_iptables(self):
        nullfp = os.open("/dev/null", os.O_RDWR)

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

        # Clean up
        os.close(nullfp)
        del nullfp


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


        # self.__cleanup_iptables()  # Not working - not getting a match ... iptables bug?

        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, "logfile": logfile,
               "count": count}

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


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