started writing better tools

portal.py

@@ -1,4 +1,6 @@
-# Captiveportal web application using Bottle.py
+# Captiveportal
+# This is the web API the portal website speaks to and requests new clients
+# to be added by running jobs in rq.
 
 import json
 from pprint import pprint as pp

tools/README.md

@@ -0,0 +1,5 @@
+# Client tools
+
+These tools deal with clients in IPtables backed by a PostgreSQL database.
+
+A class is defined in client.py that handles creating and looking up clients.

tools/add_client.py

@@ -2,31 +2,43 @@
 # Python helper tool to add IPtables rule using the iptc library. This must
 # of course run as root for iptc to work.
 
-from argparse import ArgumentParser
+from argparse import ArgumentParser, FileType
 from pprint import pprint as pp
+from configparser import RawConfigParser
 
-import iptc
+from storage import StorageRedis
+from client import Client
 
 parser = ArgumentParser()
 
-parser.add_argument('--chain', required=True)
-parser.add_argument('--protocol', required=True)
-parser.add_argument('--src-ip', required=True)
+parser.add_argument(
+    '--protocol',
+    required=True,
+    choices=['tcp', 'udp'],
+    help='Protocol for client'
+)
+
+parser.add_argument(
+    '--config',
+    type=FileType('r'),
+    required=True,
+    help='Configuration file'
+)
+
+parser.add_argument(
+    'src_ip',
+    help='Client source IP to add'
+)
 
 args = parser.parse_args()
 
-table = iptc.Table(iptc.Table.MANGLE)
-chain = iptc.Chain(table, args.chain)
+config = RawConfigParser()
+config.readfp(args.config)
 
-# Check if rule exists
-for rule in chain.rules:
-    src_ip = rule.src
-    if src_ip.startswith(args.src_ip) and rule.protocol == args.protocol:
-        print('Rule exists')
-        break
-else:
-    rule = iptc.Rule()
-    rule.src = args.src_ip
-    rule.protocol = args.protocol
-    rule.target = iptc.Target(rule, 'RETURN')
-    chain.insert_rule(rule)
+sr = StorageRedis(config=config)
+client = Client(
+    storage=sr,
+    client_id=args.src_ip,
+    protocol=args.protocol,
+    chain=config.get('iptables', 'chain')
+)

tools/client.py

@@ -0,0 +1,69 @@
+"""
+Handles "clients" in IPtables for captive portal.
+"""
+
+from datetime import datetime
+
+#import iptc
+
+
+class Client(object):
+
+    def __init__(self, **kw):
+        self.storage = kw.pop('storage')
+        self.client_id = kw.pop('client_id')
+        self.protocol = kw.pop('protocol')
+        self.chain = kw.pop('chain')
+
+        # Default values for client data
+        self.data = {
+            'client_id': self.client_id,
+            'protocol': self.protocol,
+            'created': datetime.now(),
+            'bytes': 0,
+            'packets': 0,
+            'last_activity': None
+        }
+
+        self.client_exists = False
+
+        # Attempt to fetch client from storage
+        client_data = self.storage.get_client(self.client_id)
+        if client_data:
+            self.data = client_data
+            self.exists = True
+
+
+    def commit(self):
+        self.commit_client()
+        #self.commit_rule()
+
+
+    def commit_client(self):
+        if self.exists:
+            self.storage.update_client(
+                self.client_id,
+                **self.data
+            )
+        else:
+            self.storage.add_client(
+                self.client_id,
+                **self.data
+            )
+
+    def commit_rule(self):
+        table = iptc.Table(iptc.Table.MANGLE)
+        chain = iptc.Chain(table, self.chain)
+
+        # Check if rule exists
+        for rule in chain.rules:
+            src_ip = rule.src
+            if src_ip.startswith(self.client_id) and rule.protocol == self.protocol:
+                print('Rule exists')
+                break
+        else:
+            rule = iptc.Rule()
+            rule.src = self.client_id
+            rule.protocol = self.protocol
+            rule.target = iptc.Target(rule, 'RETURN')
+            chain.insert_rule(rule)

tools/storage.cfg

@@ -0,0 +1,10 @@
+[pgsql]
+hostname=localhost
+username=captiveportal
+password=secret.
+database=captiveportal
+
+[redis]
+hostname=localhost
+port=6379
+db=0

tools/storage.py

@@ -0,0 +1,35 @@
+"""
+Database storage backends for client.py.
+"""
+
+import json
+from datetime import datetime
+
+from redis import Redis
+
+
+class DateTimeEncoder(json.JSONEncoder):
+    """
+    json.JSONEncoder sub-class that converts all datetime objects to
+    epoch timestamp integer values.
+    """
+    def default(self, o):
+        if isinstance(o, datetime):
+            return int(o.strftime('%s'))
+        return json.JSONEncoder.default(self, o)
+
+
+class StorageRedis(object):
+
+    def __init__(self, **kw):
+        config = kw.pop('config')
+
+        self.r = Redis(
+            host=config.get('redis', 'hostname'),
+            port=config.getint('redis', 'port'),
+            db=config.getint('redis', 'db')
+        )
+
+    
+    def add_client(self, client_id, **kw):
+        pass