iptables_bpf_chain

#!/usr/bin/env python

template = r'''
#!/bin/bash
#
# This script is ***AUTOGENERATED***
#
# To apply the iptables BPF rule run this script:
#
#    ./%(fname)s
#
# This script creates an ipset "%(ipsetname)s". You can manage it
# manually:
#
#     ipset add %(ipsetname)s %(sampleips)s
#
# To clean the iptables rule and ipset run:
#
#    ./%(fname)s --delete
#
#

set -o noclobber
set -o errexit
set -o nounset
set -o pipefail

: ${IPTABLES:="%(iptables)s"}
: ${IPSET:="ipset"}
: ${INPUTPLACE:="1"}
: ${DEFAULTINT:=`awk 'BEGIN {n=0} $2 == "00000000" {n=1; print $1; exit} END {if (n=0) {print "eth0"}}' /proc/net/route`}

main_match () {
    ${IPTABLES} \
        --wait \
        ${*} \
        -i ${DEFAULTINT} \
        -p udp --dport 53 \
        -m set --match-set %(ipsetname)s dst \
        -j %(chain)s
}

chain_create() {
    ${IPTABLES} --wait -N %(chain)s
%(accept_cmds)s
%(drop_cmds)s
    ${IPTABLES} --wait -A %(chain)s -j RETURN
}

chain_delete() {
    ${IPTABLES} --wait -F %(chain)s
    ${IPTABLES} --wait -X %(chain)s
}

if [ "$*" == "--delete" ]; then

    A=`(main_match -C INPUT || echo "error") 2>/dev/null`
    if [ "${A}" != "error" ]; then
        main_match -D INPUT
        chain_delete
    fi
    ${IPSET} -exist destroy %(ipsetname)s 2>/dev/null

else

    ${IPSET} -exist create %(ipsetname)s hash:net family %(ipsetfamily)s
    for IP in %(ips)s $@; do
        ${IPSET} -exist add %(ipsetname)s "$IP"
    done

    A=`(main_match -C INPUT || echo "error") 2>/dev/null`
    if [ "${A}" == "error" ]; then
        chain_create
        main_match -I INPUT ${INPUTPLACE}
    fi

fi
'''.lstrip()

import argparse
import os
import stat
import string
import sys

import bpftools


parser = argparse.ArgumentParser(
    formatter_class=argparse.RawDescriptionHelpFormatter,
    description=r'''

This program generates a bash script. The script when run will insert
(or remove) an iptable chain and ipset. The iptable chain drops
traffic that matches requests with domains given with "--accept"
options and drops domains listed with "--drop" option. Example:

  %(prog)s -a www.example.com -a ns1.example.com -d *.example.com -w example_com


'''.strip())
parser.add_argument('-6', '--inet6', action='store_true',
                    help='generate script for IPv6')
parser.add_argument('-i', '--ip', metavar='ip', action='append',
                    help='preset IP in the set')
parser.add_argument('-w', '--write', metavar='name',
                    help='name the generated script')
parser.add_argument('-a', '--accept', metavar='accept', action='append',
                    help='accept domains')
parser.add_argument('-d', '--drop', metavar='drop', action='append',
                    help='drop domains')

args = parser.parse_args()

if not args.write:
    print "set name with -w"
    sys.exit(-1)

inet = 4 if not args.inet6 else 6

fname = args.write +'_ip'+str(inet)+'.sh'

meta = []
for action, list_of_patterns in [('ACCEPT', args.accept), ('DROP', args.drop)]:
    cmds = []
    for domain in list_of_patterns:
        if domain != 'any':
            _, bytecode = bpftools.gen('dns',
                                       ['-i', domain],
                                       assembly=False,
                                       l3_off=0,
                                       ipversion=inet,
                                       )

            if int(bytecode.split(',')[0]) > 63:
                raise Exception("bytecode too long!")

            cmd = r'''
        ${IPTABLES} \
            --wait \
            -A %s \
            -m bpf --bytecode "%s" \
            -m comment --comment "%s" \
            -j %s
    ''' % (args.write.upper(), bytecode, "dns -- -i " + domain, action)
            cmds.append(cmd.strip('\n'))
        else:
            cmd = r'''
        ${IPTABLES} \
            --wait \
            -A %s \
            -m comment --comment "%s" \
            -j %s
    ''' % (args.write.upper(), "dns -- -i " + domain, action)
            cmds.append(cmd.strip('\n'))

    meta.append(cmds)

accept_cmds = meta[0]
drop_cmds = meta[1]

ctx = {
    'accept_cmds': '\n'.join(accept_cmds),
    'drop_cmds':  '\n'.join(drop_cmds),
    'fname': fname,
    'ipsetname': args.write + '_ip' + str(inet),
    'chain': args.write.upper(),
    'ips': ' '.join(repr(s) for s in (args.ip or [])),
}

if inet == 4:
    ctx.update({
            'iptables': 'iptables',
            'ipsetfamily': 'inet',
            'sampleips': '1.1.1.1/32',
            })
else:
    ctx.update({
            'iptables': 'ip6tables',
            'ipsetfamily': 'inet6',
            'sampleips': '2a00:1450:4009:803::1008/128',
            })

f = open(fname, 'wb')
f.write(template % ctx)
f.flush()
print "Generated file %r" % (fname,)
os.chmod(fname, 0750)