refactored client lib, so it can be used by the new portalclient module

portalclientlib/client.py

@@ -0,0 +1,122 @@
+"""
+Handles "clients" in IPtables for captive portal.
+"""
+
+import ipaddress
+from uuid import uuid4
+from datetime import datetime, timedelta
+
+from portalclientlib import errors
+from portalclientlib.helpers import run_ipset
+
+
+class Client(object):
+
+    def __init__(self, **kw):
+        # Required parameters
+        self.storage = kw.pop('storage')
+        self.ipset_name = kw.pop('ipset_name')
+        self.use_sudo = kw.pop('use_sudo', False)
+        
+        self.ip_address = kw.pop('ip_address', '127.0.0.1')
+        self.protocol = kw.pop('protocol', 'tcp')
+
+        self.new = False
+
+        # First try to get an existing client by ID
+        self.client_id = kw.pop('client_id', None)
+        if self.client_id:
+            client_data = self.storage.get_client_by_id(self.client_id)
+
+            # If ID is specified then we raise exception if client isn't
+            # found.
+            if client_data is None:
+                raise errors.StorageNotFound('Client not found')
+        else:
+            client_data = self.storage.get_client(
+                self._ip_address,
+                self.protocol
+            )
+
+        # Init iptables
+        #self.table = iptc.Table(iptc.Table.MANGLE)
+        #self.chain = iptc.Chain(self.table, self._chain)
+
+        if client_data:
+            self.load_client(client_data)
+        else:
+            self.client_id = str(uuid4())
+            self.created = datetime.now()
+            self.enabled = False
+            self.last_packets = 0
+            self.last_activity = None
+            self.expires = datetime.now() + timedelta(days=1)
+            self.new = True
+
+
+    def load_client(self, data):
+        self.client_id = data.get('client_id')
+        self.created = data.get('created')
+        self.ip_address = data.get('ip_address')
+        self.protocol = data.get('protocol')
+        self.enabled = data.get('enabled')
+        self.last_packets = data.get('last_packets')
+        self.last_activity = data.get('last_activity')
+        self.expires = data.get('expires')
+
+
+    def commit(self):
+        self.commit_client()
+
+        if self.enabled:
+            self.commit_rule()
+        else:
+            self.remove_rule()
+
+
+    def commit_client(self):
+        self.storage.write_client(
+            self
+        )
+
+
+    def delete(self):
+        self.remove_rule()
+        self.storage.remove_client(self)
+
+
+    def remove_rule(self):
+        run_ipset(
+            'del',
+            '-exist',
+            self.ipset_name,
+            self.ip_address,
+            use_sudo=self.use_sudo
+        )
+
+
+    def commit_rule(self):
+        run_ipset(
+            'add',
+            '-exist',
+            self.ipset_name,
+            self.ip_address,
+            use_sudo=self.use_sudo
+        )
+
+
+    @property
+    def ip_address(self):
+        return str(self._ip_address.ip)
+
+
+    @ip_address.setter
+    def ip_address(self, value):
+        if isinstance(value, str):
+            self._ip_address = ipaddress.IPv4Interface(value)
+        elif isinstance(value, ipaddress.IPv4Interface):
+            self._ip_address = value
+        else:
+            raise ValueError('Cannot set invalid value')
+
+

portalclientlib/errors.py

@@ -0,0 +1,8 @@
+class StorageNotFound(Exception):
+    pass
+
+class IPTCRuleExists(Exception):
+    pass
+
+class IPTCRuleNotFound(Exception):
+    pass

portalclientlib/helpers.py

@@ -0,0 +1,25 @@
+
+import subprocess
+import shlex
+
+def run_ipset(command, *args, **kw):
+    use_sudo = kw.get('use_sudo', True)
+    timeout = kw.get('timeout', 2)
+
+    if use_sudo:
+        ipset_cmd = 'sudo ipset'
+    else:
+        ipset_cmd = 'ipset'
+
+    full_command = '{ipset} {command} {args}'.format(
+        ipset=ipset_cmd,
+        command=command,
+        args=' '.join(args)
+    )
+
+    output = subprocess.check_output(
+        shlex.split(full_command),
+        timeout=timeout
+    )
+
+    return output

portalclientlib/storage.py

@@ -0,0 +1,121 @@
+"""
+Database storage backends for client.py.
+"""
+
+import json
+from datetime import datetime
+
+import psycopg2
+from psycopg2.extras import DictCursor, register_ipaddress, Inet
+from redis import Redis
+
+from portalclientlib.client import Client
+
+
+class StoragePostgres(object):
+    """
+    This requires python 3 for inet data type.
+    """
+
+    def __init__(self, **kw):
+        config = kw.pop('config')
+
+        self.conn = psycopg2.connect(
+            host=config.get('postgres', 'hostname'),
+            user=config.get('postgres', 'username'),
+            password=config.get('postgres', 'password'),
+            dbname=config.get('postgres', 'database'),
+            port=config.getint('postgres', 'port'),
+            sslmode='disable',
+            cursor_factory=DictCursor
+        )
+        self.cur = self.conn.cursor()
+        register_ipaddress()
+    
+
+    def client_ids(self):
+        self.cur.execute(
+            'select client_id from authenticated_clients'
+        )
+        return self.cur.fetchall()
+
+
+    def get_client_by_id(self, client_id):
+        self.cur.execute(
+            'select * from authenticated_clients where client_id=%s',
+            (client_id,)
+        )
+        return self.cur.fetchone()
+
+
+    def get_client(self, ip_address, protocol):
+        """
+        Expects an ipaddress.IPv4Interface as ip_address argument.
+        """
+
+        self.cur.execute(
+            'select * from authenticated_clients where ip_address=%s and protocol=%s',
+            (ip_address, protocol, )
+        )
+        return self.cur.fetchone()
+
+
+    def write_client(self, client):
+        query = (
+            'insert into authenticated_clients (client_id, created, ip_address, protocol, '
+            'enabled, last_packets, last_activity, expires) values '
+            '(%s, %s, %s, %s, %s, %s, %s, %s) on conflict (client_id, '
+            'ip_address, protocol) do update set (enabled, last_packets, '
+            'last_activity, expires) = (EXCLUDED.enabled, EXCLUDED.last_packets, '
+            'EXCLUDED.last_activity, EXCLUDED.expires)'
+        )
+        self.cur.execute(
+            query,
+            (
+                client.client_id,
+                client.created,
+                client.ip_address,
+                client.protocol,
+                client.enabled,
+                client.last_packets,
+                client.last_activity,
+                client.expires
+            )
+        )
+        self.conn.commit()
+
+
+    def remove_client(self, client):
+        query = 'delete from authenticated_clients where client_id=%s'
+        self.cur.execute(query, (client.client_id,))
+        self.conn.commit()
+
+
+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):
+    """
+    Note: Abandoned this storage backend for Postgres.
+    """
+
+    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):
+        raise NotImplemented