#!/usr/bin/python

# Personal firewall script.
# Copyright (C) 2000-2008 Nicholas J. Kain <njk at aerifal dot cx>
# Licensed under GPLv2.

from __future__ import with_statement
import os, sys, commands, re, subprocess, getopt, fcntl, shutil

# stores a ruleset
class Prio:
    norm=0
    high=1
    low=2
class Rule:
    def __init__(self, udp=None, tcp=None, prio=Prio.norm):
        self.udp = udp
        self.tcp = tcp
        self.prio = prio

# maps rulesets onto names
rules = {'bnet': Rule(['4000', '6112:6119'], ['4000', '6112:6119'], Prio.high),
         'dawnofwar': Rule(['6112'], None, Prio.high),
         'dsiege': Rule(['6073'], None, Prio.high),
         'empireearth': Rule(['33334'], ['33335'], Prio.high),
         'freelancer': Rule(['2302', '2303', '2304'], None, Prio.high),
         'ftp': Rule(None, ['20'], Prio.low),
         'generals': Rule(['6601', '8086', '8088'], None, Prio.high),
         'icq': Rule(['4000', '5190'], ['4000', '5190']),
         'jediknight': Rule(['28070'], ['28070'], Prio.high),
         'kohan2': Rule(['6900', '5860'], ['6900'], Prio.high),
         'muds': Rule(['5555']),
         'qw': Rule(['27500'], None, Prio.high),
         'ron': Rule(['34987'], ['34987'], Prio.high),
         'ssam2': Rule(['25600'], ['25600'], Prio.high),
         'speakfreely': Rule(['2074:2076', '5004:5005'], None, Prio.high),
         'ut': Rule(['7777:7778'], None, Prio.high),
         'wow': Rule(None, ['3724', '6112', '6881:6999'], Prio.low)}


def GetMyIP(iface):
    "Returns the address of the specified network interface."
    a = commands.getoutput('ifconfig %s' % (iface))
    s = re.search('inet addr:(\d+\.\d+\.\d+\.\d+)', a)
    try:
        ret = s.group(1)
    except AttributeError:
        ret = None
    return ret


# note that this tmp dir should NOT be writable by anyone but root!
tmpdir  = '/lib/rc/init.d/tmp'
statedir = '/etc/firewall'

maxtries = 3

extif    = 'ppp0'
localif  = 'eth0'
extlnkif = 'eth1'
wifiif   = 'ath0'
destip   = '192.168.0.2'
extip    = GetMyIP(extif)

# the adapter that handles our trusted internal network
int_net = '192.168.0.0/24'
# our wireless access point
wifi_net = '192.168.1.0/24'
# boolean: should we allow ICMP Echo (ping) requests?
allow_ping = True
# boolean: should we allow DHCP on the external interface?
use_dhcp = False
# boolean: Do we need to clamp MSS?  Set true for PPPOE.
clamp_mss = True
# boolean: should we ratelimit ssh?
ratelimit_ssh = True
# service ports on the local machine that should be remotely accessible
tcp_services = ['smtp', 'auth', 'www']
udp_services = ['domain']
# disable ECN to these hosts
bad_ecn = ['1.2.3.4/24']
# low latency ports
tcp_services_high = ['domain', 'ntp', '6667']
udp_services_high = ['domain', 'ntp']
# bulk ports
tcp_services_low = ['ftp', 'http']
udp_services_low = []
# min cost ports
tcp_services_mincost = ['smtp']
udp_services_mincost = []
# Bogus networks: http://www.cymru.com/Bogons/
bogon_file = '/etc/bogons'
# ICMP types that are dropped by default
bad_icmp = ['5', '6', '9'] # redirect, alt host addy, router ad
# hostile hosts that we ban from communicating with us
ban_hosts = ['1.2.3.4']
# westell adsl modem diagnostic info
westell_diag = True
# uids denied access to the internal network
limited_uids = ['31337']
# netfilter modules to load
nf_modules = ['ip_nat_ftp', 'ip_nat_irc']


class mutex:
    """Mutex abstraction that uses the POSIX fcntl() locking.

       Safe to take this lock recursively because we only synchronize against
       other processes."""
    def __init__(self, name, fd=None):
        self.name = name
        self.fd = fd
    def __enter__(self):
        self.lock()
    def __exit__(self, type, value, traceback):
        self.unlock()
    def lock(self):
        if self.fd is None:
            self.fd = open(tmpdir + '/' + self.name, 'w')
            rv = fcntl.lockf(self.fd, fcntl.LOCK_EX)
    def unlock(self):
        if not self.fd is None:
            fcntl.lockf(self.fd, fcntl.LOCK_UN)
            self.fd.close()
            try:
                os.remove(tmpdir + '/' + self.name)
            except OSError:
                pass
            self.fd = None

ipt_mutex = mutex('ipt_mutex')
dns_mutex = mutex('dns_mutex')


def Ipt(cmd):
    "Wraps calls to iptables."
    tries = 0
    ret = 1
    while ret != 0:
        if tries > maxtries:
            print 'Iptables commmand "%s" failed.' % (cmd)
            return
        ret = subprocess.call(['iptables'] + cmd.split(' '))
        tries += 1
        if ret == 2:
            raise Usage('Invalid command line "%s" supplied to iptables.' %
                        (cmd))

def GetFileVal(fn):
    "Returns the first line of a file ignoring any trailing newline."
    if os.path.exists(fn):
        f = open(fn, 'r')
        ret = f.readline()
        ret = ret.strip()
        f.close()
        return ret
    else:
        return ''

def WriteOldIP():
    "Updates the OLDEXTIP state file to contain our current IP."
    if not extip is None:
        f = open(tmpdir + '/OLDEXTIP', 'w')
        f.write(extip)
        f.close()

def FileToList(name):
    """Converts each line of a file to a list element.  Ignores lines that
       are only newlines.  Returns the list."""
    f = open(name, 'r')
    s = f.read()
    f.close()
    s = s.split('\n')
    try:
        while True:
            s.remove('')
    except ValueError:
        pass
    return s

def RemoveFile(path, name):
    "Wrapper function that deletes a state file if it exists."
    try:
        os.remove(path + '/' + name)
    except OSError:
        pass


class iface:
    "Abstracts an interface."
    def __init__(self, name, up, down, parent=None):
        self.name = name
        self.do_up = up
        self.do_down = down
        self.parent = parent
    def up(self):
        # if a parent interface exists and is not up, bring it up
        if not self.parent is None:
            if not self.parent.isup():
                self.parent.up()
        if not self.isup():
            with ipt_mutex:
                self.do_up(self)
                f = open(tmpdir + '/' + self.name, 'w')
                f.write('up\n')
                f.close()
        else:
            print '%s firewall rules are already installed.' % (self.name)
            return
    def down(self):
        if self.isup():
            with ipt_mutex:
                self.do_down(self)
            os.remove(tmpdir + '/' + self.name)
    def isup(self):
        if os.path.exists(tmpdir + '/' + self.name):
            return True
        else:
            return False

# stack of interface objects
ifaces = []

def CoreUp(self):
    "Load the core firewall rules."
    if len(nf_modules) > 0:
        for i in nf_modules:
            subprocess.call(['modprobe'] + [i])

    # reset tables and delete user chains
    Ipt('-F')
    Ipt('-X')
    Ipt('-t nat -F')
    Ipt('-t nat -X')
    Ipt('-t mangle -F')
    Ipt('-t mangle -X')

    # Deny input unless explicitly accepted.
    Ipt('-P INPUT DROP')
    Ipt('-P FORWARD DROP')
    Ipt('-P OUTPUT ACCEPT')

    # Drop packets coming from obviously bogus IPs
    Ipt('-N drop-bogons')
    bogons = FileToList(bogon_file)

    for i in bogons:
        Ipt('-A drop-bogons -s %s -j DROP' % (i))

    # Drop potentially malicious ICMP
    Ipt('-N drop-icmp')
    for i in bad_icmp:
        Ipt('-A drop-icmp -p icmp --icmp-type %s -j REJECT' % (i))

    # Drop invalid connections that may be attacks
    Ipt('-N drop-invalid')
    # reject new tcp packets that aren't SYN packets
    Ipt('-A drop-invalid -p tcp ! --syn -m state --state NEW -j REJECT --reject-with tcp-reset')
    # reject syn/ack packets to prevent sequence # spoofing on other hosts
    Ipt('-A drop-invalid -p tcp --tcp-flags SYN,ACK SYN,ACK -m state --state NEW -j REJECT --reject-with tcp-reset')
    Ipt('-A drop-invalid -p tcp --tcp-flags ALL FIN,URG,PSH -m state --state NEW -j REJECT --reject-with tcp-reset')
    Ipt('-A drop-invalid -p tcp --tcp-flags ALL NONE -m state --state NEW -j REJECT --reject-with tcp-reset')
    Ipt('-A drop-invalid -p tcp --tcp-flags SYN,RST SYN,RST -m state --state NEW -j REJECT --reject-with tcp-reset')
    Ipt('-A drop-invalid -p tcp --tcp-flags SYN,FIN SYN,FIN -m state --state NEW -j REJECT --reject-with tcp-reset')
    Ipt('-A drop-invalid -m state --state INVALID -j REJECT')

    # allow a slight burst for app-marked min-delay, but throttle by size
    Ipt('-t mangle -N tosfix')
    Ipt('-t mangle -A tosfix -p tcp -m length --length 0:512 -j RETURN')
    Ipt('-t mangle -A tosfix -m limit --limit 2/s --limit-burst 10 -j RETURN')
    Ipt('-t mangle -A tosfix -j TOS --set-tos Maximize-Throughput')

    # Maximize the speed of ACKs, so long as they're not too long.
    Ipt('-t mangle -N check-acks')
    Ipt('-t mangle -A check-acks -m tos ! --tos Normal-Service -j RETURN')
    Ipt('-t mangle -A check-acks -p tcp -m length --length 0:128 -j TOS --set-tos Minimize-Delay')
    Ipt('-t mangle -A check-acks -p tcp -m length --length 128: -j TOS --set-tos Maximize-Throughput')

    # Traffic priority tagging
    Ipt('-t mangle -N mark-qos')
    for i in tcp_services_high:
        Ipt('-t mangle -A mark-qos -p tcp --dport %s -j TOS --set-tos Minimize-Delay' % (i))
    for i in udp_services_high:
        Ipt('-t mangle -A mark-qos -p udp --dport %s -j TOS --set-tos Minimize-Delay' % (i))
    for i in tcp_services_low:
        Ipt('-t mangle -A mark-qos -p tcp --dport %s -j TOS --set-tos Maximize-Throughput' % (i))
    for i in udp_services_low:
        Ipt('-t mangle -A mark-qos -p udp --dport %s -j TOS --set-tos Maximize-Throughput' % (i))
    for i in tcp_services_mincost:
        Ipt('-t mangle -A mark-qos -p tcp --dport %s -j TOS --set-tos Minimize-Cost' % (i))
    for i in udp_services_mincost:
        Ipt('-t mangle -A mark-qos -p udp --dport %s -j TOS --set-tos Minimize-Cost' % (i))

    #### FORWARD #####
    # Allow ICMP echo requests.
    if allow_ping == True:
        Ipt('-A FORWARD -p icmp --icmp-type time-exceeded -m state --state ESTABLISHED,RELATED -j ACCEPT')

    ##### INPUT ######
    # Trust state data on the Big Three Protocols
    Ipt('-A INPUT -p tcp -m state --state ESTABLISHED,RELATED -j ACCEPT')
    Ipt('-A INPUT -p udp -m state --state ESTABLISHED -j ACCEPT')
    Ipt('-A INPUT -p icmp -m state --state ESTABLISHED,RELATED -j ACCEPT')

    # Allow ICMP echo requests.
    if allow_ping == True:
        Ipt('-A INPUT -p icmp --icmp-type echo-request -j ACCEPT')

    # Allow some services to be accessed from the outside.
    for i in tcp_services:
        Ipt('-A INPUT -p tcp --dport %s -m state --state NEW -j ACCEPT' % (i))
    for i in udp_services:
        Ipt('-A INPUT -p udp --dport %s -m state --state NEW -j ACCEPT' % (i))

def CoreDown(self):
    "Unload all firewall and forwarding rules."
    SaveRules()

    # reset tables and delete user chains
    Ipt('-F')
    Ipt('-X')
    Ipt('-t nat -F')
    Ipt('-t nat -X')
    Ipt('-t mangle -F')
    Ipt('-t mangle -X')

    Ipt('-P INPUT ACCEPT')
    Ipt('-P FORWARD ACCEPT')
    Ipt('-P OUTPUT ACCEPT')

    for i in ifaces:
        if i.name != self.name:
            RemoveFile(tmpdir, i.name)

core_if = iface('core', CoreUp, CoreDown)
ifaces.append(core_if)


def ExtUp(self):
    "Brings up the internet exposed interface."

    #### FORWARD #####
    Ipt('-N fi_ext') # must insert at head
    Ipt('-N fo_ext')

    # Clamp MSS on external interfaces if required.
    if clamp_mss == True:
        Ipt('-A fo_ext -p tcp --tcp-flags SYN,RST SYN -m tcpmss --mss 1452:65535 -j TCPMSS --clamp-mss-to-pmtu')

    # Drop invalid packets ASAP.
    Ipt('-A fi_ext -j drop-invalid')

    # Deny bogus networks as input from external interfaces.
    Ipt('-A fi_ext -j drop-bogons')

    # Detect and drop potentially malicious ICMP
    Ipt('-A fi_ext -j drop-icmp')

    ##### INPUT ######

    Ipt('-N i_ext') # must insert at head

    # Drop invalid packets ASAP.
    Ipt('-A i_ext -j drop-invalid')

    # Deny bogus networks as input from external interfaces.
    Ipt('-A i_ext -j drop-bogons')

    # We don't like these hosts.
    for i in ban_hosts:
        Ipt('-A i_ext -s %s -j REJECT' % (i))

    # Detect and drop potentially malicious ICMP
    Ipt('-A i_ext -j drop-icmp')

    # pass in DHCP, keeping state
    if use_dhcp == True:
        Ipt('-A i_ext -p udp -s $DHCPSIP --dport bootpc -m state --state NEW -j ACCEPT')

    # rate limit ssh connections if necessary
    if ratelimit_ssh == True:
        Ipt('-A i_ext -p tcp --dport ssh -m state --state ESTABLISHED,RELATED -j ACCEPT')
        Ipt('-A i_ext -p tcp --dport ssh -m state --state NEW -m limit --limit 3/min --limit-burst 3 -j ACCEPT')

    ##### OUTPUT #####
    Ipt('-t mangle -N mo_ext')
    
    # TTL normalization
    default_ttl = GetFileVal('/proc/sys/net/ipv4/ip_default_ttl')
    Ipt('-t mangle -A mo_ext -s %s -j TTL --ttl-set %s' % (int_net, default_ttl))
    Ipt('-t mangle -A mo_ext -s %s -j TTL --ttl-set %s' % (wifi_net, default_ttl))

    ### POSTROUTING ###
    Ipt('-t mangle -N mp_ext')
    Ipt('-t nat -N np_ext')

    # IP Masquerading
    Ipt('-t nat -A np_ext -s %s -j MASQUERADE' % (int_net))
    Ipt('-t nat -A np_ext -s %s -j MASQUERADE' % (wifi_net))

    # Some broken places don't like ECN; we can disable ECN for them.
    for i in bad_ecn:
        Ipt('-t mangle -A mp_ext -p tcp -d %s -j ECN --ecn-tcp-remove' % (i))

    # Sanitize packets that are marked low cost by apps
    Ipt('-t mangle -A mp_ext -p tcp -m tos --tos Minimize-Delay -j tosfix')

    # Prioritize ACKs, so long as they are not too long
    Ipt('-t mangle -A mp_ext -p tcp -m tcp --tcp-flags SYN,RST,ACK ACK -j check-acks')

    # Prioritize TCP SYN and RST
    Ipt('-t mangle -A mp_ext -p tcp -m tcp --tcp-flags ! SYN,RST,ACK ACK -j TOS --set-tos Minimize-Delay')

    # Deprioritize long-lived, high bandwidth connections
    Ipt('-t mangle -A mp_ext -p tcp -m connbytes --connbytes 500000: --connbytes-dir both --connbytes-mode bytes -j TOS --set-tos Minimize-Cost')

    # Initialize custom priority chain
    Ipt('-t mangle -A mp_ext -j mark-qos')

    # map TOS bit onto classes
    Ipt('-t mangle -A mp_ext -m tos --tos Minimize-Delay -j CLASSIFY --set-class 1:10')
    Ipt('-t mangle -A mp_ext -m tos --tos Maximize-Throughput -j CLASSIFY --set-class 1:20')
    Ipt('-t mangle -A mp_ext -m tos --tos Minimize-Cost -j CLASSIFY --set-class 1:30')
    Ipt('-t mangle -A mp_ext -m tos --tos Maximize-Reliability -j CLASSIFY --set-class 1:30')

    Ipt('-I FORWARD -i %s -j fi_ext' % (self.name))
    Ipt('-I FORWARD -o %s -j fo_ext' % (self.name))
    Ipt('-I INPUT -i %s -j i_ext' % (self.name))
    Ipt('-t mangle -A OUTPUT -o %s -j mo_ext' % (self.name))
    Ipt('-t nat -A POSTROUTING -o %s -j np_ext' % (self.name))
    Ipt('-t mangle -A POSTROUTING -o %s -j mp_ext' % (self.name))

    RestoreRules()

def ExtDown(self):
    "Brings down the internet exposed interface."
    # This is ugly, but we have to handle the case where we have outdated
    # rules from an old IP still loaded.
    CheckIP()
    SaveRules()
    ClearRules()

    Ipt('-D FORWARD -i %s -j fi_ext' % (self.name))
    Ipt('-D FORWARD -o %s -j fo_ext' % (self.name))
    Ipt('-D INPUT -i %s -j i_ext' % (self.name))
    Ipt('-t mangle -D OUTPUT -o %s -j mo_ext' % (self.name))
    Ipt('-t nat -D POSTROUTING -o %s -j np_ext' % (self.name))
    Ipt('-t mangle -D POSTROUTING -o %s -j mp_ext' % (self.name))

    Ipt('-F fi_ext')
    Ipt('-F fo_ext')
    Ipt('-F i_ext')
    Ipt('-t mangle -F mo_ext')
    Ipt('-t nat -F np_ext')
    Ipt('-t mangle -F mp_ext')
    Ipt('-X fi_ext')
    Ipt('-X fo_ext')
    Ipt('-X i_ext')
    Ipt('-t mangle -X mo_ext')
    Ipt('-t nat -X np_ext')
    Ipt('-t mangle -X mp_ext')

ext_if = iface(extif, ExtUp, ExtDown, core_if)
ifaces.append(ext_if)


def LocalUp(self):
    "Brings up the ethernet interface that hosts the wired LAN."

    #### FORWARD #####
    Ipt('-N fi_local')
    Ipt('-N fo_local')

    # We trust things that don't come from outside.
    Ipt('-A fi_local -j ACCEPT')

    # LAN NAT rules
    Ipt('-A fi_local -s %s -m state --state NEW,ESTABLISHED -j ACCEPT' % (int_net))
    Ipt('-A fo_local -i ! %s -d %s -m state --state ESTABLISHED -j ACCEPT' % (self.name, int_net))

    ##### INPUT ######
    Ipt('-N i_local')

    # We trust things that don't come from outside.
    Ipt('-A i_local -j ACCEPT')

    ##### OUTPUT #####
    Ipt('-N o_local')

    # Ban remote accounts from accessing the internal network.
    for i in limited_uids:
        Ipt('-A o_local -m owner --uid-owner %s -j REJECT --reject-with icmp-port-unreachable' % (i))

    ##################
    Ipt('-A FORWARD -i %s -j fi_local' % (self.name))
    Ipt('-A FORWARD -o %s -j fo_local' % (self.name))
    Ipt('-A INPUT -i %s -j i_local' % (self.name))
    Ipt('-A OUTPUT -o %s -j o_local' % (self.name))

def LocalDown(self):
    "Brings down the ethernet interface that hosts the wired LAN."

    Ipt('-D FORWARD -i %s -j fi_local' % (self.name))
    Ipt('-D FORWARD -o %s -j fo_local' % (self.name))
    Ipt('-D INPUT -i %s -j i_local' % (self.name))
    Ipt('-D OUTPUT -o %s -j o_local' % (self.name))

    Ipt('-F fi_local')
    Ipt('-F fo_local')
    Ipt('-F i_local')
    Ipt('-F o_local')
    Ipt('-X fi_local')
    Ipt('-X fo_local')
    Ipt('-X i_local')
    Ipt('-X o_local')

ifaces.append(iface(localif, LocalUp, LocalDown, core_if))


def ExtlinkUp(self):
    "Brings up the ethernet interface that hosts the PPPOE link."

    Ipt('-N i_extlnk')

    # allow diagnostic multicast packets from the westell modem
    if westell_diag == True:
        Ipt('-A i_extlnk -p udp -d 224.73.193.62/32 -s 192.168.1.1/32 --sport 1875 --dport 1875 -j ACCEPT')

    Ipt('-A INPUT -i %s -j i_extlnk' % (self.name))

def ExtlinkDown(self):
    "Brings down the ethernet interface that hosts the PPPOE link."

    Ipt('-D INPUT -i %s -j i_extlnk' % (self.name))

    Ipt('-F i_extlnk')
    Ipt('-X i_extlnk')

ifaces.append(iface(extlnkif, ExtlinkUp, ExtlinkDown, core_if))


def LoopUp(self):
    "Brings up the local loop interface."

    Ipt('-N fi_lo')

    # We trust things that don't come from outside.
    Ipt('-A fi_lo -j ACCEPT')

    Ipt('-N i_lo')

    # We trust things that don't come from outside.
    Ipt('-A i_lo -j ACCEPT')

    Ipt('-A FORWARD -i %s -j fi_lo' % (self.name))
    Ipt('-A INPUT -i %s -j i_lo' % (self.name))

def LoopDown(self):
    "Brings down the local loop interface."

    Ipt('-D FORWARD -i %s -j fi_lo' % (self.name))
    Ipt('-D INPUT -i %s -j i_lo' % (self.name))

    Ipt('-F fi_lo')
    Ipt('-F i_lo')
    Ipt('-X fi_lo')
    Ipt('-X i_lo')

ifaces.append(iface('lo', LoopUp, LoopDown, core_if))


def WifiUp(self):
    "Brings up the wifi interface."

    Ipt('-N fi_wifi')
    Ipt('-N fo_wifi')

    # We trust things that don't come from outside.
    Ipt('-A fi_wifi -j ACCEPT')

    # VPN Tunnel rules
    Ipt('-A fi_wifi -s %s -m state --state NEW,ESTABLISHED -j ACCEPT' % (wifi_net))
    Ipt('-A fo_wifi -i ! %s -d %s -m state --state ESTABLISHED -j ACCEPT' % (self.name, wifi_net))

    Ipt('-N i_wifi')

    # We trust things that don't come from outside.
    Ipt('-A i_wifi -j ACCEPT')

    Ipt('-N o_wifi')

    # Ban remote accounts from accessing the internal network.
    for i in limited_uids:
        Ipt('-A o_wifi -m owner --uid-owner %s -j REJECT --reject-with icmp-port-unreachable' % (i))

    Ipt('-A FORWARD -i %s -j fi_wifi' % (self.name))
    Ipt('-A FORWARD -o %s -j fo_wifi' % (self.name))
    Ipt('-A INPUT -i %s -j i_wifi' % (self.name))
    Ipt('-A OUTPUT -o %s -j o_wifi' % (self.name))

def WifiDown(self):
    "Brings down the wifi interface."

    Ipt('-D FORWARD -i %s -j fi_wifi' % (self.name))
    Ipt('-D FORWARD -o %s -j fo_wifi' % (self.name))
    Ipt('-D INPUT -i %s -j i_wifi' % (self.name))
    Ipt('-D OUTPUT -o %s -j o_wifi' % (self.name))

    Ipt('-F fi_wifi')
    Ipt('-F fo_wifi')
    Ipt('-F i_wifi')
    Ipt('-F o_wifi')
    Ipt('-X fi_wifi')
    Ipt('-X fo_wifi')
    Ipt('-X i_wifi')
    Ipt('-X o_wifi')

ifaces.append(iface(wifiif, WifiUp, WifiDown, core_if))


def IfUp(name):
    "Brings up the listed interface."
    found = False
    for i in ifaces:
        if i.name == name:
            found = True
            if i.isup() == False:
                i.up()
            else:
                print '%s is already up.' % (name)
    if found == False:
        print 'No rule exists for "%s".' % (name)

def IfDown(name):
    "Brings down the listed interface."
    found = False
    for i in ifaces:
        if i.name == name:
            found = True
            if i.isup() == True:
                i.down()
            else:
                print '%s is already down.' % (name)
    if found == False:
        print 'No rules exist for "%s".' % (name)
    

def CheckIP():
    """If we've not yet updated our firewall state to reflect a changed
    external IP address, clear and remake the external interface rules,
    then update and restart the authoritative DNS server."""
    global extip

    if extip is None:
        return False

    oldip = GetFileVal(tmpdir + '/OLDEXTIP')

    # order of operations here is very important
    if oldip != extip:

        WriteOldIP()
        
        if ext_if.isup():
            t = extip
            extip = oldip
            SaveRules()
            ClearRules()
            extip = t
            RestoreRules()
        
        UpdateDNS()
    return True

def UpdateDNS():
    "Updates the data files of the auth DNS server to reflect our external IP."
    execdir = '/usr/local/sbin/'
    confdir = '/etc/nsd'
    dnsdir  = confdir + '/domains'

    if extip is None:
        return False

    with dns_mutex:
        s = re.search(r'(\d+)\.(\d+)\.(\d+)\.(\d+)', extip)
        revip = '%s.%s.%s.%s' % (s.group(4), s.group(3), s.group(2), s.group(1))

        def update(fn):
            fi = open(fn + '.src', 'r')
            fo = open(fn, 'w')
            fc = fi.read()
            fi.close()
            fc = re.sub('_DARKSWORDIP_', extip, fc)
            fc = re.sub('_DARKSWORDIPR_', revip, fc)
            fo.write(fc)
            fo.close()

        oum = os.umask(022)
        update(confdir + '/nsd.conf')
        update(dnsdir + '/kain.us.zone')
        subprocess.call([execdir + 'nsdc', 'rebuild'])
        subprocess.call([execdir + 'nsdc', 'restart'])
        os.umask(oum)
    return True


def ModifyRuleset(name, key):
    "Either adds or removes forwarding rulesets using iptables."

    if extip is None:
        return False

    if key == 'add':
        k = 'A'
    elif key == 'del':
        k = 'D'
    else:
        raise Exception

    tos = None
    if rules[name].prio == Prio.low:
        tos = 'Minimize-Delay'
    if rules[name].prio == Prio.high:
        tos = 'Minimize-Cost'

    for j in rules[name].udp:
        Ipt('-%s FORWARD -i %s -p UDP -d %s --dport %s -m state --state NEW,ESTABLISHED -j ACCEPT' % (k, extif, destip, j))
        Ipt('-t nat -%s PREROUTING -i %s -p UDP -d %s --dport %s -j DNAT --to-destination %s' % (k, extip, extip, j, destip))
    for j in rules[name].tcp:
        Ipt('-%s FORWARD -i %s -p TCP -d %s --dport %s -m state --state NEW,ESTABLISHED -j ACCEPT' % (k, extif, destip, j))
        Ipt('-t nat -%s PREROUTING -i %s -p TCP -d %s --dport %s -j DNAT --to-destination %s' % (k, extif, extip, j, destip))

    if tos != None:
        for j in rules[name].udp:
            Ipt('-t mangle -%s mark-qos -p udp --sport %s -j TOS --set-tos %s' % (k, j, tos))
        for j in rules[name].tcp:
            Ipt('-t mangle -%s mark-qos -p tcp --sport %s -j TOS --set-tos %s' % (k, j, tos))
    return True

def AddRuleset(name):
    "Add the forwarding ruleset to the kernel firewall rules."
    undofile = statedir + '/UNDO'
    # Check and see if the ruleset is already loaded; if so, then stop.
    if os.path.exists(undofile):
        f = open(undofile, 'r')
        a = f.read()
        f.close()
        s = re.search('\b%s\b' % name, a)
        if s != None:
            raise Usage('Nothing need be done: "%s" already exists.' % name)

    # The rules don't already exist, so add them.
    f = open(undofile, 'a')
    f.write(name + '\n')
    f.close()

    ModifyRuleset(name, 'add')

def DelRuleset(name):
    "Delete the forwarding ruleset from the kernel firewall rules."
    undofile = statedir + '/UNDO'
    # Remove entry from undofile. If undofile is blank, remove it.
    if os.path.exists(undofile):
        f = open(undofile, 'r')
        a = f.read()
        f.close()

        # If no entry exists in the undo file, then the rules don't exist.
        s = re.search(r'\b%s\b' % name, a)
        if s is None:
            raise Usage('Rule "%s" does not exist.  Nothing to delete.' % name)

        # Remove the entry.  Delete the file if it is empty.
        a = re.sub('%s' % (name) + r'\n', '', a)
        if a != '':
            f = open(undofile, 'w')
            f.write(a)
            f.close()
        else:
            RemoveFile(statedir, 'UNDO')
        
    ModifyRuleset(name, 'del')

def SaveRules():
    "Copies UNDO to REDO."
    undofile = statedir + '/UNDO'
    redofile = statedir + '/REDO'
    if os.path.exists(undofile):
        shutil.copyfile(undofile, redofile)
        RemoveFile(statedir, 'UNDO')

def ClearRules():
    "Unloads all rules based on UNDO state."
    undofile = statedir + '/UNDO'
    if os.path.exists(undofile):
        for i in FileToList(undofile):
            DelRuleset(i)

def RestoreRules():
    "Loads each rule listed in REDO.  Deletes REDO."
    redofile = statedir + '/REDO'
    undofile = statedir + '/UNDO'
    # if machine crashed, we'll have UNDO but not REDO; handle it properly
    if os.path.exists(redofile):
        pass
    elif os.path.exists(undofile):
        SaveRules()
    else:
        return
    for i in FileToList(redofile):
        AddRuleset(i)
        RemoveFile(statedir, 'REDO')


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg

def main(argv=None):
    """Firewall script.  Commands:

    add <set1> <set2> ... <setN> : loads nat ruleset(s) to the firewall
    del <set1> <set2> ... <setN> : unloads nat ruleset(s) from the firewall
    clear : clears all nat rules and memorizes them for later restore
    restore : restores previous cleared nat rules
    list : show a list of known nat rules
    data|details : show actual nat rule data
    dnsupdate : perform an update of DNS records
    checkip : check to see if the external ip has change and update if needed
    up all|<if1> <if2> ... <ifN> : loads firewall rules for interfaces
    down all|<if1> <if2> ... <ifN> : unloads firewall rules for interfaces"""
    if argv is None:
        argv = sys.argv
    try:
        try:
            opts, args = getopt.getopt(sys.argv[1:], 'h', ['help', ])
        except getopt.error, msg:
            raise Usage(msg)

        # Handle switches.
        for o, a in opts:
            if o in ('-h', '--help'):
                print main.__doc__
                return 1

        if len(args) == 0:
            print main.__doc__
            return 1

        # Print detailed ruleset data if requested.
        if args[0] == 'data' or args[0] == 'details':
            for n, r in rules.iteritems():
                print '%s:' % (n)
                print '\tUDP:  %s' % (r.udp)
                print '\tTCP:  %s' % (r.tcp)
                if r.prio == Prio.norm:
                    print '\tPrio: Normal'
                elif r.prio == Prio.high:
                    print '\tPrio: High'
                else:
                    print '\tPrio: Low'
            return 0

        # Print list of rulesets if requested.
        if args[0] == 'list':
            print rules.keys()
            return 0

        # Add a ruleset to the firewall if requested.
        elif args[0] == 'add':
            if len(args) == 1:
                raise Usage('No ruleset specified!  Use "list" to see choices.')
            for i in args[1:]:
                if not rules.has_key(i):
                    raise Usage('Rule "%s" does not exist.' % (i))
            with ipt_mutex:
                CheckIP()
                for i in args[1:]:
                    AddRuleset(i)
            return 0

        # Delete a ruleset from the firewall if requested.
        elif args[0] == 'del' or args[0] == 'delete' or args[0] == 'remove':
            if len(args) == 1:
                raise Usage('No ruleset specified!  Use "list" to see choices.')
            for i in args[1:]:
                if not rules.has_key(i):
                    raise Usage('Rule "%s" does not exist.' % (i))
            with ipt_mutex:
                CheckIP()
                for i in args[1:]:
                    DelRuleset(i)
            return 0

        # Clear all rules.
        elif args[0] == 'clear':
            with ipt_mutex:
                CheckIP()
                SaveRules()
                ClearRules()
            return 0

        # Restore all rules.
        elif args[0] == 'restore':
            with ipt_mutex:
                CheckIP()
                RestoreRules()
            return 0

        # Update DNS rules to reflect new IP.
        elif args[0] == 'dnsupdate':
            with ipt_mutex:
                CheckIP()
            UpdateDNS()
            return 0

        # Check to see if we need to update our DNS rules and external IF FW.
        elif args[0] == 'checkip':
            with ipt_mutex:
                CheckIP()
            return 0

        elif args[0] == 'up':
            with ipt_mutex:
                CheckIP()
                if len(args) == 1:
                    for i in ifaces:
                        i.up()
                elif args[1] == 'all':
                    for i in ifaces:
                        i.up()
                else:
                    for i in args[1:]:
                        IfUp(i)
            return 0

        elif args[0] == 'down':
            with ipt_mutex:
                CheckIP()
                if len(args) == 1:
                    IfDown('core')
                elif args[1] == 'all':
                    IfDown('core')
                else:
                    for i in args[1:]:
                        IfDown(i)
            return 0

        # No valid command, so print an error and exit.
        else:
            print >>sys.stderr, 'Invalid command "%s" specified.' % (args[0])
            return 1
    except Usage, err:
        print >>sys.stderr, err.msg
        print >>sys.stderr, 'for help use --help'
        return 1


    for arg in args[1:]:
        print arg

# This little helper returns an error code matching that of main(), but
# only if we're not running inside of another python toplevel.
if __name__ == "__main__":
    sys.exit(main())