iptables_bpf

#!/usr/bin/env python

template = r'''
#!/bin/bash
#
# This script is ***AUTOGENERATED***
#
# This is a script for applying and removing xt_bpf iptable rule. This
# particular rule was created with:
#
#     %(bpf_cmd)s
#
# To apply the iptables BPF rule against listed destination IP's run
# this script with the flooded IP addresses as parameters:
#
#    ./%(fname)s %(sampleips)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
#
#
# For the record, here's the BPF assembly:
#
%(assembly)s
#

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`}

iptablesrule () {
    ${IPTABLES} \
        --wait \
        ${*} \
        -i ${DEFAULTINT} \
        -p udp --dport 53 \
        -m set --match-set %(ipsetname)s dst \
        -m bpf --bytecode "%(bytecode)s" \
        -m comment --comment "%(bpf_cmd_s)s" \
        -j DROP
}

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

    A=`(iptablesrule -C INPUT || echo "error") 2>/dev/null`
    if [ "${A}" != "error" ]; then
        iptablesrule -D INPUT
    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=`(iptablesrule -C INPUT || echo "error") 2>/dev/null`
    if [ "${A}" == "error" ]; then
        iptablesrule -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 rule and ipset. The iptable rule drops traffic
that matches the BPF rule, which in turn is generated from given
parameters.

See "./bpfgen --help" for more information on BPF generators.

Usage example:

  %(prog)s dns -- *.example.com

  %(prog)s dns -- --ignorecase *.example.com

  %(prog)s -6 dns -- *.example.com

  %(prog)s dns_validate -- --strict

'''.strip())
parser.add_argument('-6', '--inet6', action='store_true',
                    help='generate script for IPv6')
parser.add_argument('-c', '--comment',
                    help='Add a comment to better identify this rule in iptables -nvL'),
parser.add_argument('-i', '--ip', metavar='ip', action='append',
                    help='preset IP in the set')
parser.add_argument('-n', '--negate', action='store_true',
                    help='negate the logic')
parser.add_argument('-w', '--write', metavar='file',
                    help='name the generated script')
parser.add_argument('type', nargs=1, choices=bpftools.generator_names,
                    help='BPF generator')
parser.add_argument('parameters', nargs='*',
                    help='parameters passed to the BPF generator')

args = parser.parse_args()
if len(args.type) != 1:
    parser.print_help()
    sys.exit(-1)

inet = 4 if not args.inet6 else 6

a = []
for assembly in [False, True]:
    name, ret = bpftools.gen(args.type[0],
                                  args.parameters,
                                  assembly=assembly,
                                  l3_off=0,
                                  ipversion=inet,
                                  negate=args.negate,
                                  )
    a.append(ret)
bytecode, assembly = a

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

name = 'bpf_%s_ip%s%s%s' % (args.type[0], inet, '_' if name else '', name)

fname = args.write or name + '.sh'

if fname == '-':
    f = sys.stdout
else:
    f = open(fname, 'wb')

cmd = sys.argv[1:]

bpf_cmd_s = args.comment if args.comment else ' '.join(cmd).replace('"', "").replace("$", "")

ctx = {
    'bpf_cmd': cmd,
    'bpf_cmd_s': bpf_cmd_s,
    'bytecode': bytecode,
    'assembly': '#    ' + '\n#    '.join(assembly.split('\n')),
    'fname': fname if fname != '-' else name + '.sh',
    'ipsetname': name[:31],
    '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.write(template % ctx)
f.flush()

if f != sys.stdout:
    print "Generated file %r" % (fname,)
    os.chmod(fname, 0750)