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):