From 61e950dd429f68fcee3563d97367d2d5dcccbd4c Mon Sep 17 00:00:00 2001
From: Stefan Midjich <swehack@gmail.com>
Date: Fri, 29 Sep 2017 18:37:30 +0200
Subject: [PATCH] testing new purge script

---
 tools/client.py        | 20 ++++++++++++++++++
 tools/errors.py        |  3 +++
 tools/purge_clients.py | 47 ++++++++++++++++++++++++++++++++++++++++++
 tools/storage.py       |  7 +++++++
 4 files changed, 77 insertions(+)
 create mode 100644 tools/purge_clients.py

diff --git a/tools/client.py b/tools/client.py
index 6121d01..044434a 100644
--- a/tools/client.py
+++ b/tools/client.py
@@ -62,6 +62,26 @@ class Client(object):
         self.last_packets = data.get('last_packets')
         self.last_activity = data.get('last_activity')
 
+        # Try and find a rule for this client and with that rule also packet
+        # count. Don't rely on it existing though.
+        rule = None
+        try:
+            rule = self.find_rule(self.ip_address, self.protocol)
+        except Exception as e:
+            # TODO: This should raise an exception and be handled further up
+            # the stack by logging the error.
+            #raise IPTCRuleNotFound('Could not find the iptables rule for {client_ip}'.format(
+            # client_ip=self.ip_address
+            #))
+            return None
+
+        if rule:
+            (packet_count, byte_count) = rule.get_counters()
+
+            if self.last_packets < packet_count:
+                self.last_activity = datetime.now()
+                self.last_packets = packet_count
+
 
     def commit(self):
         self.commit_client()
diff --git a/tools/errors.py b/tools/errors.py
index e92eece..2609dd7 100644
--- a/tools/errors.py
+++ b/tools/errors.py
@@ -3,3 +3,6 @@ class StorageNotFound(Exception):
 
 class IPTCRuleExists(Exception):
     pass
+
+class IPTCRuleNotFound(Exception):
+    pass
\ No newline at end of file
diff --git a/tools/purge_clients.py b/tools/purge_clients.py
new file mode 100644
index 0000000..b56e70d
--- /dev/null
+++ b/tools/purge_clients.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+# Python helper tool to purge expired clients from DB and iptables. Requires
+# root privileges for iptc to work.
+
+from sys import exit
+from argparse import ArgumentParser, FileType
+from pprint import pprint as pp
+from configparser import RawConfigParser
+from datetime import datetime, timedelta
+
+import errors
+from storage import StoragePostgres
+from client import Client
+
+
+parser = ArgumentParser((
+    'Purge expired clients by disabling them.'
+))
+
+parser.add_argument(
+    '--config',
+    type=FileType('r'),
+    required=True,
+    help='Configuration file'
+)
+
+args = parser.parse_args()
+
+config = RawConfigParser()
+config.readfp(args.config)
+
+sr = StoragePostgres(config=config)
+
+for client_id in sr.client_ids():
+    client = Client(
+        storage=sr,
+        chain=config.get('iptables', 'chain'),
+        client_id=client_id
+    )
+
+    if datetime.now() > client.expires:
+        client.enabled = False
+        client.commit()
+    else:
+        # Simply commit whatever was loaded during Client.__init__(), like
+        # up-to-date packet count stats for example.
+        client.commit()
\ No newline at end of file
diff --git a/tools/storage.py b/tools/storage.py
index 0da4e53..bef4d9d 100644
--- a/tools/storage.py
+++ b/tools/storage.py
@@ -31,6 +31,13 @@ class StoragePostgres(object):
         )
         self.cur = self.conn.cursor()
         register_ipaddress()
+    
+
+    def client_ids(self):
+        self.cur.execute(
+            'select client_id from client'
+        )
+        return self.cur
 
 
     def get_client_by_id(self, client_id):