#!/usr/bin/env python
# Firewall
#
# Douglas Thrift
#
# $Id$

from __init__ import lock
from ccscslab.laptops.models import Laptop
import collections
import optparse
import os.path
import re
import subprocess
import tempfile

if __name__ == '__main__':
	parser = optparse.OptionParser()

	parser.add_option('-D', '--debug', action = 'store_true', dest = 'debug')
	parser.add_option('-a', '--address', dest = 'address')
	parser.add_option('-e', '--external', dest = 'external')
	parser.add_option('-i', '--internal', dest = 'internal')

	options = parser.parse_args()[0]

	if not options.address:
		parser.error('-a not specified')

	if not options.external:
		parser.error('-e not specified')

	if not options.internal:
		parser.error('-i not specified')

	lock(options.debug)

	mac_addresses = collections.OrderedDict()

	with open('/ccs/etc/infrastructure', 'rb') as infrastructure:
		for line in infrastructure:
			mac_address, name = line.rstrip().split(None, 2)

			if options.debug:
				print '%s %s:\n   inserted' % (name, mac_address)

			mac_addresses[mac_address] = name

	for laptop in Laptop.objects.order_by('person', 'mac_address'):
		person = laptop.person
		mac_address = laptop.mac_address
		fancy_name = laptop.fancy_name

		if options.debug:
			print '%s %s:' % (person, mac_address)

			if fancy_name:
				print '   %s' % fancy_name

			print '   inserted'

		mac_addresses[mac_address] = '%s (%s)' % (person, fancy_name) if fancy_name else str(person)

	assert len(mac_addresses) <= 10000

	with open('/dev/null', 'wb') as null:
		ipfw = subprocess.Popen(('ipfw', '-S', 'list', '10000-19999'), stdout = subprocess.PIPE, stderr = None if options.debug else null, close_fds = True)
		rulez = re.compile('^(1[0-9]{4}) set 1 allow tag 1 ip from any to any recv %s layer2 MAC any ([0-9A-Fa-f]{2}(?::[0-9A-Fa-f]{2}){5}) // (.*)$' % re.escape(options.internal))
		rules = set()
		mac_addressez = set()
		deletes = []

		for line in ipfw.stdout:
			match = rulez.match(line)

			if match is not None:
				rule, mac_address, name = match.groups()
				rule = int(rule)

				if mac_addresses.get(mac_address) == name:
					rules.add(rule)
					mac_addressez.add(mac_address)
				else:
					deletes.append(rule)

		ipfw.wait()

	changed = False

	for delete in deletes:
		subprocess.call(('ipfw', 'delete %u' % delete), close_fds = True)

		if options.debug:
			print '%u:\n   deleted' % delete

		changed = True

	mac_addressez = list(frozenset(mac_addresses) - mac_addressez)

	def _rule():
		return 'add %u set 1 allow tag 1 ip from any to any recv %s layer2 MAC any %s // %s' % (rule, options.internal, mac_address, name)

	for rule in xrange(10000, 20000):
		if not mac_addressez:
			break

		if rule not in rules:
			mac_address = mac_addressez.pop(0)
			name = mac_addresses[mac_address]

			subprocess.call(('ipfw', _rule()), close_fds = True)

			if options.debug:
				print '%u %s:\n   %s\n   added' % (rule, mac_address, name)

			changed = True

	if changed or not os.path.exists('/ccs/etc/firewall'):
		with tempfile.NamedTemporaryFile('wb', prefix = 'firewall.', dir = '/ccs/etc') as temp:
			print >> temp, 'add 30 set 1 allow ip from any to me recv %s' % options.internal
			print >> temp, 'add 40 set 1 forward %s ip from any to any recv %s not tagged 1' % (options.address, options.internal)
			print >> temp, 'add 50 divert natd ip4 from any to any via %s not layer2' % options.external

			for rule, (mac_address, name) in enumerate(mac_addresses.iteritems(), 10000):
				print >> temp, _rule()

			print >> temp, 'add 65000 pass ip from any to any'

			temp.delete = False

		os.rename(temp.name, '/ccs/etc/firewall')