From c7ace886dede07cab8034337f136479ef32051bc Mon Sep 17 00:00:00 2001
From: Thuban <thuban@yeuxdelibad.net>
Date: Mon, 21 Aug 2017 13:40:38 +0000
Subject: [PATCH] support logrotate

---
 vilain.py |  124 +++++++++++++++++++++++++---------------
 1 files changed, 77 insertions(+), 47 deletions(-)

diff --git a/vilain.py b/vilain.py
index a7c1b5a..d36c59f 100755
--- a/vilain.py
+++ b/vilain.py
@@ -4,12 +4,12 @@
 
 """
 Author :      thuban <thuban@yeuxdelibad.net>  
+              Vincent <vincent.delft@gmail.com>
 Licence :     MIT
 Require : python >= 3.5
 
 Description : Mimic fail2ban with pf for OpenBSD.
               Inspired from http://www.vincentdelft.be/post/post_20161106
-              with improvements of vincendelft
 
               In pf.conf, add : 
                     table <vilain_bruteforce> persist
@@ -24,44 +24,47 @@
 import configparser
 import re
 import logging
+import logging.handlers
 import subprocess
 import asyncio
 import time
 
-configfile = "/etc/vilain.conf"
-version = "0.4"
+CONFIGFILE = "/etc/vilain.conf"
+VERSION = "0.5"
 vilain_table = "vilain_bruteforce"
-logfile = "/var/log/daemon"
+LOGFILE = "/var/log/daemon"
 
 if os.geteuid() != 0:
     print("Only root can use this tool")
-    sys.exit()
+    sys.exit(1)
 
 # Configure logging
+log_handler = logging.handlers.WatchedFileHandler(LOGFILE)
+formatter = logging.Formatter(
+        '%(asctime)s %(module)s:%(funcName)s:%(message)s',
+        '%b %d %H:%M:%S')
+log_handler.setFormatter(formatter)
 logger = logging.getLogger(__name__)
-logging.basicConfig(filename=logfile,
-                    format='%(asctime)s %(module)s:%(funcName)s:%(message)s',
-                    datefmt='%H:%M:%S')
+logger.addHandler(log_handler)
 logger.setLevel(logging.INFO)
-ch = logging.StreamHandler(sys.stdout)
-logger.addHandler(ch)
 
 # functions
 def readconfig():
-    if not os.path.isfile(configfile):
+    logger.info('Read config file: {}'.format(CONFIGFILE))
+    if not os.path.isfile(CONFIGFILE):
         logging.error("Can't read config file, exiting...")
         sys.exit(1)
 
     config = configparser.ConfigParser()
-    config.read(configfile)
+    config.read(CONFIGFILE)
     return(config)
 
 def load_config():
     c = readconfig()
     d = c.defaults()
     watch_while = int(d['watch_while'])
+    VILAIN_TABLE = d['vilain_table']
     default_maxtries = int(d['maxtries'])
-    vilain_table = d['vilain_table']
     sleeptime = float(d['sleeptime'])
     ignore_ips = []
 
@@ -73,19 +76,19 @@
     c = readconfig()
     for s in c.sections():
         if c.has_option(s,'logfile'):
-            logfile = c.get(s,'logfile')
+            LOGFILE = c.get(s,'logfile')
             regex = c.get(s,'regex')
             #we take the default value of maxtries
             maxtries = c.defaults()['maxtries']
             if c.has_option(s,'maxtries'):
                 #if we have a maxtries defined in the section, we overwrite the default
                 maxtries = int(c.get(s,'maxtries'))
-            d = {'name' : s, 'logfile':logfile, 'regex':regex, 'maxtries': maxtries}
+            d = {'name' : s, 'logfile':LOGFILE, 'regex':regex, 'maxtries': maxtries}
             yield d
-
 
 class Vilain():
     def __init__(self):
+        logger.info('Start vilain version {}'.format(VERSION))
         self.loop = asyncio.get_event_loop()
         self.watch_while, self.default_maxtries, self.vilain_table, self.ignore_ips, self.sleeptime = load_config()
         self.ip_seen_at = {}
@@ -93,25 +96,28 @@
         self.bad_ip_queue = asyncio.Queue(loop=self.loop)
 
         for entry in load_sections():
-            logger.info("Start vilain for {}".format(entry['name']))
+            logger.info("Start vilain for {}".format(entry))
             asyncio.ensure_future(self.check_logs(entry['logfile'], entry['maxtries'], entry['regex'], entry['name']))
 
         asyncio.ensure_future(self.ban_ips())
+        asyncio.ensure_future(self.clean_ips())
 
     def load_bad_ips(self):
         try:
             ret = subprocess.check_output(["pfctl", "-t", self.vilain_table, "-T", "show"])
         except:
+            logger.warning("Failed to run pfctl -t {} -T show".format(self.vilain_table))
             ret = ""
         for res in ret.split():
             ip = res.strip().decode('utf-8')
-            logger.debug('Add existing banned IPs in your pf table: {}'.format(ip))
+            logger.info('Add existing banned IPs in your pf table: {}'.format(ip))
             #we assign the counter to 1, but for sure we don't know the real value 
             self.ip_seen_at[ip]={'time':time.time(),'count':1}
 
 
     def start(self):
         try:
+            logger.info('Run forever loop')
             self.loop.run_forever()
         except KeyboardInterrupt:
             self.loop.close()
@@ -128,54 +134,64 @@
             # Watch the file for changes
             stat = os.stat(logfile)
             size = stat.st_size
+            inode = stat.st_ino
             mtime = stat.st_mtime
             RE = re.compile(regex)
             while True:
                 await asyncio.sleep(self.sleeptime)
                 stat = os.stat(logfile)
-                if mtime < stat.st_mtime:
+                if size > stat.st_size and inode != stat.st_ino:
+                    logger.info("The file {} has rotated. We start from position 0".format(logfile))
+                    size = 0
+                    inode = stat.st_ino
+                if mtime < stat.st_mtime and inode == stat.st_ino:
                     logger.debug("{} has been modified".format(logfile))
                     mtime = stat.st_mtime
                     with open(logfile, "rb") as f:
-                        f.seek(size)
-                        lines = f.readlines()
-                        ul = [ u.decode() for u in lines ]
-                        line = "".join(ul).strip()
-
-                        ret = RE.match(line)
-                        logger.debug('line:{}'.format(line))
-                        if ret:
-                            bad_ip = ret.groups()[0]
-                            if bad_ip not in self.ignore_ips :
-                                logger.info('line match {} because of rule : {}'.format(bad_ip, reason))
-                                await self.bad_ip_queue.put({'ip' : bad_ip, 'reason' : reason})
-                                logger.debug('queue size: {}'.format(self.bad_ip_queue.qsize()))
-                            else:
-                                logger.info('line match {}. But IP in ignore list'.format(bad_ip))
+                        f.seek(size,0)
+                        for bline in f.readlines():
+                            line = bline.decode().strip()
+                            ret = RE.match(line)
+                            logger.debug('line:{}'.format(line))
+                            if ret:
+                                bad_ip = ret.groups()[0]
+                                if bad_ip not in self.ignore_ips :
+                                    logger.info('line match {} the {} rule'.format(bad_ip, reason))
+                                    await self.bad_ip_queue.put({'ip' : bad_ip, 'maxtries': maxtries, 'reason' : reason})
+                                    logger.debug('queue size: {}'.format(self.bad_ip_queue.qsize()))
+                                else:
+                                    logger.info('line match {}. But IP in ignore list'.format(bad_ip))
                     size = stat.st_size
 
     async def ban_ips(self):
         """
         record time when this IP has been seen in ip_seen_at = { ip:{'time':<time>,'count':<counter} }
-
-        check old ip in ip_seen_at : remove older than watch_while
+        and ban with pf
         """
-        logger.info('ban_ips sarted with sleeptime={}'.format(self.sleeptime))
+        logger.info('ban_ips started')
         while True:
-            await asyncio.sleep(self.sleeptime)
             ip_item = await self.bad_ip_queue.get()
             logger.debug('ban_ips awake')
             ip = ip_item['ip']
             reason = ip_item['reason']
             maxtries = ip_item['maxtries']
-            self.ip_seen_at.setdefault(ip,{'time':time.time(),'count':0})
+            self.ip_seen_at.setdefault(ip, {'time':time.time(),'count':0})
             self.ip_seen_at[ip]['count'] += 1
             n_ip = self.ip_seen_at[ip]['count']
             logger.info("{} detected, reason {}, count: {}, maxtries: {}".format(ip, reason, n_ip, maxtries))
             if n_ip >= maxtries:
                 ret = subprocess.call(["pfctl", "-t", self.vilain_table, "-T", "add", ip])
                 logger.info("Blacklisting {}, return code:{}".format(ip, ret))
-                self.ip_seen_at.pop(ip)
+            #for debugging, this line allow us to see if the script run until here
+            logger.debug('ban_ips end:{}'.format(self.ip_seen_at))
+
+    async def clean_ips(self):
+        """
+        check old ip in ip_seen_at : remove older than watch_while
+        """
+        logger.info('clean_ips started with sleeptime={}'.format(self.sleeptime))
+        while True:
+            await asyncio.sleep(self.watch_while)
             to_remove = []
             for recorded_ip, data in self.ip_seen_at.items():
                 if time.time() - data['time'] >= self.watch_while:
@@ -185,12 +201,7 @@
             for ip in to_remove:
                 self.ip_seen_at.pop(ip)
             #for debugging, this line allow us to see if the script run until here
-            logger.debug('ban_ips end:{}'.format(self.ip_seen_at))
-
-
-
-
-
+            logger.debug('clean_ips end:{}'.format(self.ip_seen_at))
 
 
 
@@ -202,7 +213,26 @@
     return 0
 
 if __name__ == '__main__':
-	main()
+    import argparse
+    parser = argparse.ArgumentParser(description="Vilain mimic fail2ban with pf for OpenBSD")
+    parser.add_argument('--debug','-d', action="store_true", help="run in debug mode")
+    parser.add_argument('--conf','-c', nargs="?", help="location of the config file")
+    parser.add_argument('--version','-v', action="store_true", help="Show the version and exit")
+    args = parser.parse_args()
+    if args.debug:
+        print("run in debug")
+        logger.setLevel(logging.DEBUG)
+        ch = logging.StreamHandler(sys.stdout)
+        logger.addHandler(ch)
+    if args.conf:
+        CONFIGFILE = args.conf
+    if args.version:
+        print("Version: ", VERSION)
+        sys.exit(0)
+    main()
+
+
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
 
 
 # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4

--
Gitblit v1.9.3