Two major changes in this commit:

* Moving from iptables to ipset
	* Moving from python-sh to subprocess

tools/client.py

@@ -6,9 +6,8 @@ import ipaddress
 from uuid import uuid4
 from datetime import datetime, timedelta
 
-import iptc
-
 from errors import StorageNotFound, IPTCRuleNotFound
+from helpers import run_ipset
 
 
 class Client(object):
@@ -16,11 +15,13 @@ class Client(object):
     def __init__(self, **kw):
         # Required parameters
         self.storage = kw.pop('storage')
-        self._chain = kw.pop('chain')
+        self.ipset_name = kw.pop('ipset_name')
         
         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:
@@ -37,8 +38,8 @@ class Client(object):
             )
 
         # Init iptables
-        self.table = iptc.Table(iptc.Table.MANGLE)
-        self.chain = iptc.Chain(self.table, self._chain)
+        #self.table = iptc.Table(iptc.Table.MANGLE)
+        #self.chain = iptc.Chain(self.table, self._chain)
 
         if client_data:
             self.load_client(client_data)
@@ -49,6 +50,7 @@ class Client(object):
             self.last_packets = 0
             self.last_activity = None
             self.expires = datetime.now() + timedelta(days=1)
+            self.new = True
 
 
     def load_client(self, data):
@@ -61,28 +63,6 @@ class Client(object):
         self.last_activity = data.get('last_activity')
         self.expires = data.get('expires')
 
-        # 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
-            #raise IPTCRuleNotFound(
-            #    'Could not find the iptables rule for {client_ip}'.format(
-            #        client_ip=self.ip_address
-            #    )
-            #)
-
-        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()
@@ -105,48 +85,28 @@ class Client(object):
 
 
     def remove_rule(self):
-        rule = self.find_rule(self._ip_address, self.protocol)
-        if rule:
-            self.chain.delete_rule(rule)
-
-
-    def find_rule(self, ip_address, protocol):
-        """
-        Takes an ipaddress.IPv4Interface object as ip_address argument.
-        """
-
-        if not isinstance(ip_address, ipaddress.IPv4Interface):
-            raise ValueError('Invalid argument type')
-
-        for rule in self.chain.rules:
-            src_ip = rule.src
-
-            try:
-                _ip = str(ip_address.ip)
-            except:
-                # If we can't understand the argument just return None
-                return None
-
-            if src_ip.startswith(_ip) and rule.protocol == protocol:
-                return rule
-        else:
-            return None
+        run_ipset(
+            'del',
+            '-exist',
+            self.ipset_name,
+            self.ip_address
+        )
 
 
     def commit_rule(self):
-        rule = self.find_rule(self._ip_address, self.protocol)
-        if not rule:
-            rule = iptc.Rule()
-            rule.src = self.ip_address
-            rule.protocol = self.protocol
-            rule.target = iptc.Target(rule, 'RETURN')
-            self.chain.insert_rule(rule)
+        run_ipset(
+            'add',
+            '-exist',
+            self.ipset_name,
+            self.ip_address
+        )
 
 
     @property
     def ip_address(self):
         return str(self._ip_address.ip)
 
+
     @ip_address.setter
     def ip_address(self, value):
         if isinstance(value, str):

tools/helpers.py

@@ -0,0 +1,33 @@
+
+import subprocess
+import shlex
+
+def run_ipset(command, *args, **kw):
+    use_sudo = kw.get('use_sudo', True)
+
+    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)
+    )
+
+    proc = subprocess.run(
+        shlex.split(full_command),
+        stdout=subprocess.PIPE,
+        timeout=2,
+        check=True
+    )
+
+    #proc = subprocess.Popen(
+    #    shlex.split(full_command),
+    #    stdout=subprocess.PIPE
+    #)
+
+    #(output, error) = proc.communicate(timeout=2)
+
+    return proc

tools/manage_client.py

@@ -7,8 +7,10 @@ from argparse import ArgumentParser, FileType, ArgumentTypeError
 from pprint import pprint as pp
 from configparser import RawConfigParser
 from datetime import datetime, timedelta
+from io import BytesIO
 
 import errors
+from helpers import run_ipset
 from storage import StoragePostgres
 from client import Client
 
@@ -32,11 +34,20 @@ def valid_datetime_type(arg_datetime_str):
         raise ArgumentTypeError(msg) 
 
 
-parser = ArgumentParser((
-    'Handle clients in the captive portal. Default mode of operation is to'
-    ' create new clients and enable them. Other mode is to --disable the '
-    'client. And last mode is to --delete the client completely.'
-))
+parser = ArgumentParser(
+    '''
+    Handle clients in the captive portal. Default mode of operation is to
+    create new clients and enable them.
+    '''
+)
+
+parser.add_argument(
+    '-v', '--verbose',
+    action='count',
+    default=False,
+    dest='verbose',
+    help='Verbose output, use more v\'s to increase verbosity'
+)
 
 parser.add_argument(
     '--expires',
@@ -60,10 +71,9 @@ parser.add_argument(
 )
 
 parser.add_argument(
-    '--protocol',
-    required=True,
-    choices=['tcp', 'udp'],
-    help='Protocol for client'
+    '--refresh',
+    action='store_true',
+    help='Refresh client ipset data first'
 )
 
 parser.add_argument(
@@ -75,6 +85,7 @@ parser.add_argument(
 
 parser.add_argument(
     'src_ip',
+    nargs='*',
     help='Client source IP to add'
 )
 
@@ -84,27 +95,75 @@ config = RawConfigParser()
 config.readfp(args.config)
 
 sr = StoragePostgres(config=config)
-try:
+
+if args.refresh:
+    # Sync clients and packet counters from ipset into storage.
+    proc = run_ipset(
+        'list',
+        config.get('ipset', 'set_name'),
+        '-output',
+        'save'
+    )
+
+    for _line in proc.stdout.splitlines():
+        # Convert from bytestring first
+        line = _line.decode('utf-8')
+
+        if not line.startswith('add'):
+            continue
+
+        (
+            cmd,
+            set_name,
+            client_ip,
+            packets_str,
+            packets_val,
+            bytes_str,
+            bytes_val
+        ) = line.split()
+
+        try:
+            client = Client(
+                storage=sr,
+                ip_address=client_ip,
+                ipset_name=config.get('ipset', 'set_name')
+            )
+        except Exception as e:
+            if args.verbose:
+                print('Failed to init client:{ip}: {error}'.format(
+                    ip=client_ip,
+                    error=str(e)
+                ))
+            continue
+
+        if client.new:
+            client.commit()
+
+        if int(packets_val) != client.last_packets:
+            client.last_activity = datetime.now()
+            client.last_packets = packets_val
+            client.commit()
+
+for src_ip in args.src_ip:
+    # Get client by IP or create a new one.
     client = Client(
         storage=sr,
-        ip_address=args.src_ip,
-        protocol=args.protocol,
-        chain=config.get('iptables', 'chain')
+        ip_address=src_ip,
+        ipset_name=config.get('ipset', 'set_name')
     )
-except errors.StorageNotFound:
-    print('Client not found')
-    exit(1)
 
-if args.disable:
-    enabled = False
-else:
-    enabled = True
+    if args.delete:
+        # This both deletes the ipset entry AND the client entry from DB. Normally
+        # excessive and disable is better.
+        client.delete()
+        exit(0)
+
+    if args.disable:
+        client.enabled = False
+    else:
+        client.enabled = True
 
-if args.delete:
-    # This both deletes the iptables rule AND the client entry from DB.
-    client.delete()
-else:
     if args.expires:
         client.expires = args.expires
-    client.enabled = enabled
+
     client.commit()

tools/requirements.txt

@@ -0,0 +1 @@
+psycopg2