Two major changes in this commit:

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

plugins.cfg

@@ -7,6 +7,8 @@ enabled = True
 debug = True
 mandatory = True
 
+# ***iptables plugin DEPRECATED in favor of ipset***
+#
 # Runs an iptables command to add a rule, commonly used for captive portal
 # firewalls.
 [iptables]
@@ -21,3 +23,13 @@ arping = -f -c 1 -w 30 -I eth0 {ip_address}
 # This is a command to run to create iptables rules. Two arguments are
 # passed and replace these two placeholders.
 iptables_cmd = /usr/local/sbin/cp_iptables.sh "{ip_address}" "{mac_address}"
+
+
+[ipset]
+mandatory = False
+enabled = False
+debug = True
+
+# Simply prepend sudo here if you won't run rq worker as root
+ipset_add_cmd = ipset -exist add authenticated-clients {client_ip}
+ipset_name = authenticated-clients

plugins/ipset.py

@@ -0,0 +1,113 @@
+# Add client IP into ipset
+
+import re
+import socket
+import subprocess
+import shlex
+from io import BytesIO, StringIO
+from logging import getLogger, DEBUG, WARN, INFO
+from datetime import datetime
+
+try:
+    from configparser import RawConfigParser
+except ImportError:
+    from ConfigParser import RawConfigParser
+
+# By default run ipset through sudo, so the worker process must run with
+# a user that is allowed to execute '/sbin/ipset add' command with NOPASSWD.
+#from sh import sudo, Command, ErrorReturnCode
+
+from portal import logHandler, logFormatter
+
+
+def run(arg):
+    # Some info from the plugin dispatcher.
+    environ = arg['environ']
+    plugin_config = arg['config']
+
+    config = RawConfigParser(defaults=plugin_config)
+    config.add_section('ipset')
+    config._sections['ipset'] = plugin_config
+
+    # Setup plugin logging
+    l = getLogger('plugin_ipset')
+    l.addHandler(logHandler)
+    if config.getboolean('ipset', 'debug'):
+        l.setLevel(DEBUG)
+        l.debug('debug logging enabled')
+
+    # Get client IP from webapp, try HTTP_X_FORWARDED_FOR and fallback on
+    # REMOTE_ADDR.
+    client_ip = environ.get(
+        'HTTP_X_FORWARDED_FOR',
+        environ.get('REMOTE_ADDR')
+    )
+    error_msg = None
+    plugin_failed = True
+    start_time = datetime.now()
+
+    # Verify client IP
+    try:
+        socket.inet_aton(client_ip)
+    except socket.error:
+        l.error('Client IP-address is invalid')
+        return {
+            'error': str(e),
+            'failed': True
+        }
+
+    ipset_name = config.get('ipset', 'ipset_name')
+    use_sudo = config.getboolean('ipset', 'use_sudo')
+    #output, error = StringIO(), StringIO()
+
+    if client_ip:
+        ipset_cmd = config.get(
+            'ipset',
+            'ipset_add_cmd'
+        ).format(
+            client_ip=client_ip
+        )
+
+        proc = subprocess.Popen(
+            shlex.split(ipset_cmd)
+        )
+
+        try:
+            (output, error) = proc.communicate(timeout=2)
+
+        except Exception as e:
+            error_msg = str(e)
+            l.warn('{cmd}: failed call: {error}'.format(
+                cmd=ipset_cmd,
+                error=str(e)
+            ))
+            raise
+
+        end_time = datetime.now()
+
+        if proc.returncode == 0:
+            l.info('Added ip:"{ip}" to set:"{set}" in "{duration}": {stdout}'.format(
+                ip=client_ip,
+                set=ipset_name,
+                duration=end_time-start_time,
+                stdout=output
+            ))
+            plugin_failed = False
+        else:
+            l.warn('{cmd}: failed[{ret}] in "{duration}": {stderr}'.format(
+                cmd=ipset_cmd,
+                ret=proc.returncode,
+                duration=end_time-start_time,
+                stderr=error
+            ))
+            raise Exception('Plugin failed')
+
+    else:
+        l.info('No client IP, no action taken')
+        error_msg = 'No client IP'
+        raise Exception('Plugin failed')
+
+    return {
+        'error': error_msg,
+        'failed': plugin_failed
+    }

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.table = iptc.Table(iptc.Table.MANGLE)
-        self.chain = iptc.Chain(self.table, self._chain)
+        #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)
+        run_ipset(
-        if rule:
+            'del',
-            self.chain.delete_rule(rule)
+            '-exist',
-
+            self.ipset_name,
-
+            self.ip_address
-    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
 
 
     def commit_rule(self):
-        rule = self.find_rule(self._ip_address, self.protocol)
+        run_ipset(
-        if not rule:
+            'add',
-            rule = iptc.Rule()
+            '-exist',
-            rule.src = self.ip_address
+            self.ipset_name,
-            rule.protocol = self.protocol
+            self.ip_address
-            rule.target = iptc.Target(rule, 'RETURN')
+        )
-            self.chain.insert_rule(rule)
 
 
     @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((
+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 '
+    Handle clients in the captive portal. Default mode of operation is to
-    'client. And last mode is to --delete the client completely.'
+    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',
+    '--refresh',
-    required=True,
+    action='store_true',
-    choices=['tcp', 'udp'],
+    help='Refresh client ipset data first'
-    help='Protocol for client'
 )
 
 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,
+        ip_address=src_ip,
-        protocol=args.protocol,
+        ipset_name=config.get('ipset', 'set_name')
-        chain=config.get('iptables', 'chain')
     )
-except errors.StorageNotFound:
-    print('Client not found')
-    exit(1)
 
-if args.disable:
+    if args.delete:
-    enabled = False
+        # This both deletes the ipset entry AND the client entry from DB. Normally
-else:
+        # excessive and disable is better.
-    enabled = True
+        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