Unifi Security Gateway iptables and ipsets helper scripts

ipsets.py

#!/usr/bin/env python

import sys
import subprocess
import json

def range_string(start, end):
    if start != end:
        return '%s-%s' % (start, end)
    else:
        return str(start)

ipsets = {}
in_set = False
in_members = False

set_name = ""

in_range = False
port_start = 0
port_end = 0


p = subprocess.Popen(['sudo', 'ipset', 'list'],stdout=subprocess.PIPE)

while True:
    line = p.stdout.readline()
    if line != b'':
        line = line.strip()
        if line == '':
            if in_members and in_range:
                ipsets[set_name]['Members'].append(range_string(port_start, port_end))
            in_set = False
            in_members = False
            in_range = False
            continue
        if not in_set:
            k,v = line.split(":", 1)
            if k == 'Name':
                in_set = True
                set_name = v.strip()
                ipsets[set_name] = {}
        else:
            if not in_members:
                k,v = line.split(":", 1)
                if k == 'Members':
                    in_members = True
                    ipsets[set_name]['Members'] = []
                else:
                    ipsets[set_name][k.strip()] = v.strip()
            else:
                v = line.strip()
                try:
                    v = int(v)
                    if not in_range:
                        in_range = True
                        port_start = port_end = v
                    elif v == (port_end + 1):
                        port_end = v
                    else:
                        ipsets[set_name]['Members'].append(range_string(port_start, port_end))
                        port_start = port_end = v
                except ValueError:
                    if in_range:
                        in_range = False
                        ipsets[set_name]['Members'].append(range_string(port_start, port_end))
                    ipsets[set_name]['Members'].append(v)
    else:
        if in_members and in_range:
            ipsets[set_name]['Members'].append(range_string(port_start, port_end))
        break

p.wait()

p = subprocess.Popen(['mca-ctrl', '-t', 'dump-cfg'], stdout=subprocess.PIPE)
unifi_config = json.load(p.stdout)
p.wait()
unifi_groups = {}
unifi_groups.update(unifi_config['firewall']['group']['address-group'])
unifi_groups.update(unifi_config['firewall']['group']['ipv6-address-group'])
unifi_groups.update(unifi_config['firewall']['group']['port-group'])

unifi_groups = { k: v['description'][11:] for (k,v) in unifi_groups.items() if v.has_key('description') and v['description'][:11] == 'customized-' }

name_width   = max(30, 3 + len(max(ipsets, key=len)))
ugrp_width   = max(30, 3 + len(max(unifi_groups.values(), key=len)))
type_width   = max(15, 3 + len(max([ipsets[s]['Type'] for s in ipsets], key=len)))
member_width = max(30, 3 + len(max((x for y in [ipsets[s]['Members'] for s in ipsets] for x in y), key=len)))


def print_set_line(n, u, t, m):
    print '{name:<{name_w}}{ugrp:<{ugrp_w}}{type:<{type_w}}{members:<{member_w:}}'.format(
        name=n,    name_w=name_width,
        ugrp=u,    ugrp_w=ugrp_width,
        type=t,    type_w=type_width,
        members=m, member_w=member_width
        )

print_set_line("Name", "Unifi Group", "Type", "Members")
print '_' * (name_width + ugrp_width + type_width + member_width)
for name in sorted(ipsets):
    ugrp    = unifi_groups[name] if name in unifi_groups else ''
    type    = ipsets[name]['Type']
    members = ipsets[name]['Members']
    print_set_line(name, ugrp, type, (members[0] if members else ''))
    if len(members) > 1:
        for member in members[1:]:
            print_set_line('', '', '', member)

iptables.py

#!/usr/bin/env python

import sys
import subprocess
import json

p = subprocess.Popen(['mca-ctrl', '-t', 'dump-cfg'], stdout=subprocess.PIPE)
unifi_config = json.load(p.stdout)
p.wait()
unifi_groups = {}
unifi_groups.update(unifi_config['firewall']['group']['address-group'])
unifi_groups.update(unifi_config['firewall']['group']['ipv6-address-group'])
unifi_groups.update(unifi_config['firewall']['group']['port-group'])

sed_cmd = ['sed']
for (k,v) in unifi_groups.items():
    if v.has_key('description') and v['description'][:11] == 'customized-':
        sed_cmd.append('-e')
        sed_cmd.append('s/{}/{}/'.format(k, v['description'][11:]))

ipt_cmd = ['sudo', 'iptables'] + sys.argv[1:]

p_iptables = subprocess.Popen(ipt_cmd,                          stdout=subprocess.PIPE)
p_sed      = subprocess.Popen(sed_cmd, stdin=p_iptables.stdout, stdout=subprocess.PIPE)

(sed_stdout, sed_stderr) = p_sed.communicate()

print(sed_stdout)