#!/usr/bin/python -tO

# transform an iptables firewall to a dot file
#
# marcus fritzsch - http://fritschy.de - 20070620
#
# this program is licensed undr the GNU GPL ver. 2 or any
# later at your option.
#
# some code was shamelessly stolen from pyroman:
# http://pyroman.alioth.debian.org/

import re, sys, os.path

# from 1.3.6. documentation
BUILTIN_TARGETS = set([ 'ACCEPT', 'DROP', 'RETURN', \
	'CLASSIFY', 'CLUSTERIP', 'CONNMARK', 'CONNSECMARK', \
	'DNAT', 'DSCP', 'ECN', 'IPV4OPTSSTRIP', 'LOG', 'MARK', 'MASQUERADE', \
	'MIRROR', 'NETMAP', 'NFQUEUE', 'NOTRACK', 'REDIRECT', 'REJECT', \
	'ROUTE', 'SAME', 'SECMARK', 'SET', 'SNAT', 'TARPIT', 'TCPMSS', 'TOS', \
	'TRACE', 'TTL', 'ULOG'])

# just in case there were no counters, disable dashed-lines
COUNTER_ERROR = False

def print_msg (out, pfx, *a):
	print >> out, "%s: %s: %s" % (os.path.basename (sys.argv[0]), \
			pfx, " ".join (map (str, a)))

def print_err (*a):
	print_msg (sys.stderr, 'ERROR', *a)

def print_warn (*a):
	print_msg (sys.stderr, 'WARNING', *a)

class Counter (object):
	def __init__ (self, p = 0, b = 0):
		self.packets, self.bytes = p, b
	def __str__ (self):
		return "[%d:%d]" % (self.packets, self.bytes)
	def __iadd__ (self, c):
		self.packets += c.packets
		self.bytes += c.bytes
		return self
	@staticmethod
	def fromstr (s):
		m = re.match (r"^\[(\d+):(\d+)\]$", s)
		if m:
			return Counter (int (m.group (1)), int (m.group (2)))
		return Counter (0, 0) # be nice

class Rule (object):
	def __init__ (self, r, t, a, n, c = None):
		self.rule, self.target, self.args, self.num, self.counter = \
				r or "", t, a or "", n, c or Counter ()
		self.rule, self.args = [ \
				''.join ([c in '<>\'"' and '\\'+c or c for c in x]) \
				for x in self.rule, self.args]
		if not len (self.rule+self.args): self.rule = "[empty]"

class Chain (object):
	def __init__ (self, n, p = ""):
		self.rules, self.policy, self.name, self.counter = \
				list (), p, n, None

class Table (object):
	def __init__ (self, n):
		self.chains, self.targets, self.name = dict (), dict (), n

def rename (n):
	"""rename a chain: replace interfering hyphens by underscores"""
	return n.replace ('-', '_')

def get_iptables (dump_file = None):
	global COUNTER_ERROR
	from popen2 import popen3

	lines = []

	if dump_file:
		infile = file (dump_file)
		lines = [x.rstrip ().lstrip () for x in infile.readlines ()]
		infile.close ()
	else:
		o, i, e = popen3 ("iptables-save -c")
		lines = [x.rstrip ().lstrip () for x in o.readlines ()]
		[x.close () for x in (o, e, i)]

	tables = dict ()
	cur_tbl = None

	match_line = re.compile (\
			r"^(?:\[(\d+):(\d+)\] )?(?:-A ([^ ]+))(?: (.*))?(?: -j ([^ ]+)(?: (.*))?)$")
	PACKETS, BYTES, CHAIN, RULE, TARGET, T_ARGS = range(1,7)

	ruleCounter, line_ctr = 1, 0
	for line in lines:

		line_ctr += 1

		if line [0] == '#': # ignore comments
			continue

		if line [0] == '*': # new table
			s = rename (line[1:])
			cur_tbl = Table (s)
			continue

		if line [0] == ':': # new chain for current table
			s = line[1:].split ()
			assert len (s) in (2, 3)
			policy, counter = s[1], None
			if policy == '-': policy = None
			else: counter = Counter.fromstr (s[2])
			chain = rename (s[0])
			cur_tbl.chains [chain] = Chain (chain, policy)
			cur_tbl.chains [chain].counter = counter
			continue

		if line == 'COMMIT': # end current table
			tables [cur_tbl.name] = cur_tbl
			cur_tbl = None
			continue

		assert cur_tbl

		# m.group (N): line, packets, bytes, chain, rule, target, target-args
		m = match_line.match (line)
		if m:
			ctr = m.group (PACKETS), m.group (BYTES)
			if not COUNTER_ERROR and (not ctr[0] or not ctr[1]):
				print_warn ("counters not available!")
				COUNTER_ERROR = True
			ctr = (ctr[0] or '0', ctr[1] or '0')
			counter = Counter (int (ctr[0]), int (ctr[1]))
			rule = Rule (m.group (RULE), rename(m.group (TARGET)), \
					m.group (T_ARGS), ruleCounter, counter)
			ruleCounter += 1
			cur_tbl.chains [rename (m.group(CHAIN))].rules.append (rule)
			cur_tbl.targets [rename (m.group(TARGET))] = \
					cur_tbl.targets.get(rename (m.group(TARGET)), 0) + 1
		else:
			print_err ("line %d could not be parsed: `%s'" % (line_ctr, line))
	# end for line in lines

	return tables

def dot_output (table, out, min_line_width, max_line_width):
	if not len (table.targets):
		print_warn ("table `%s' is empty, doing nothing" % table.name)
		return

	def print_dot (*a):
		out.write (" ".join (a))

	def get_policy (chain):
		return chain.policy and chain.policy or 'RETURN'

	def record_prologue (chain):
		print_dot ('  %s [label="%s ' % (chain.name, chain.name))

	def record_epilogue (chain):
		print_dot ('| ')
		pol = get_policy (chain)
		print_dot ("<policy_%s> %s " % (pol, pol))
		edges ['%s:policy_%s' % (chain.name, pol)] = (pol, chain.counter)
		print_dot ('"];\n')

	class MangleName (object):
		def __init__ (self):
			self.__ctr, self.__lut = 1, dict ()
		def get (self, chain, rule):
			name = chain+rule.rule+rule.args+str(rule.num)
			if not self.__lut.get (name):
				self.__lut [name] = self.__ctr
				self.__ctr += 1
			return 'name_' + str (self.__lut [name])
	nameMangler = MangleName ()

	print_dot ("digraph %s {\n" % table.name)
	print_dot ("  rankdir=LR;\n")
	print_dot ("  edge [splines=true];\n")
	print_dot ("  node [shape=record];\n\n")

	edges = {} # needed at the end to draw all edges
	targets = {} # needed to identify not yet drawn targets
	displayed = {}

	# precompute counters, i.e. add all rules counters to their target
	counters, max_bytes = {}, 0
	for chain in table.chains.itervalues ():
		for rule in chain.rules:
			if not counters.get (rule.target): counters [rule.target] = Counter ()
			counters [rule.target] += rule.counter

	for chain in table.chains.itervalues ():
		targets [get_policy (chain)] = targets.get (get_policy (chain), 0) + 1

		# fix user defined chains with the precomputed counter
		if not chain.counter:
			if counters.get (chain.name):
				chain.counter = counters [chain.name]
			else:
				chain.counter = Counter () # last escape, i.e. unreferenced chains
		max_bytes = max (chain.counter.bytes, max_bytes)

		if not len (chain.rules):
			continue

		record_prologue (chain)
		displayed [chain.name] = 1
		for rule in chain.rules:
			max_bytes = max (rule.counter.bytes, max_bytes)
			print_dot ('| ')
			print_dot ('<%s> %s ' % (nameMangler.get (chain.name, rule), \
					rule.rule+(rule.args and " "+rule.args or "")))
			edges ['%s:%s' % (chain.name, nameMangler.get (chain.name, rule))] = \
					(rule.target, rule.counter)
			targets [rule.target] = targets.get (rule.target, 0) + 1
		record_epilogue (chain)

	# explicitly print nodes not yet printed, including ellipsed
	# builtin targets
	for t in targets.iterkeys ():
		if t in BUILTIN_TARGETS:
			print_dot ('  %s [label=%s shape=ellipse];\n' % (t, t))
		elif not displayed.get (t):
			record_prologue (table.chains [t])
			record_epilogue (table.chains [t])

	# one more newline...
	print_dot ('\n')

	import math
	log_base = math.log (max_bytes) / (max_line_width - min_line_width)
	for start, end_and_ctr in edges.iteritems ():
		end, ctr = end_and_ctr
		style = ''
		if ctr.bytes == 0 and not COUNTER_ERROR:
			style += "dashed, setlinewidth(%f)" % min_line_width
		else:
			style += "solid, setlinewidth(%f)" % (min_line_width + \
					math.log (float (ctr.bytes+1)) / log_base)
		print_dot ('  %s -> %s [style="%s"];\n' % (start, end, style))

	print_dot ('}\n')

def usage ():
	print """Usage: %s [-wsch] [ip-tables]
  -W num    set max line width to num (default: 5)
  -w num    set zero-traffic line width to num (default: 1)
  -s str    set the root_name suffix to str (default: none)
  -c        causes the dot file to be printed to stdout (default: off)
	-i file   read an iptables-save dump from file
  -f        force overwriting of existing files
  -h        show this help message

example:
  %s -s _test
  will save every table to a table_test.dot file

By default only the filter table will be dotted.
The table selection will be set to exactly the ones
on the command line if given.""" % (sys.argv[0], sys.argv[0])
	sys.exit (1)

if __name__ == '__main__':
	from getopt import gnu_getopt as getopt
	import os.path as path

	line_width0, line_width1 = 1, 5
	to_stdout = force = False
	root_suffix = ""
	in_file = None

	opts = getopt (sys.argv [1:], 'w:W:s:i:fch')
	for opt, arg in opts[0]:
		if opt == '-w': line_width0 = float (arg)
		elif opt == '-W': line_width1 = float (arg)
		elif opt == '-c': to_stdout = True
		elif opt == '-s': root_suffix = arg
		elif opt == '-i': in_file = arg
		elif opt == '-f': force = True
		elif opt == '-h': usage ()
	
	def check_sanity (check, message, exit=True):
		if check:
			print_err (message)
			if exit: sys.exit (1)
	
	check_sanity (line_width0 <= 0, "min line width is invalid")
	check_sanity (line_width1 <= 0, "max line width is invalid")
	check_sanity (line_width0 >= line_width1, \
			"min line width > max line width")
	check_sanity (to_stdout and len (opts[1]) > 1, \
			"writing more than 1 table to stdout!", False)

	# check file name and look for the force
	def get_file_name (f):
		f = path.abspath (f)
		if path.exists (f):
			if not force:
				print_err ("output file %s exists use -f to force execution"%f)
				return
		return f

	iptables = {}
	try:
		iptables = get_iptables (in_file)
	except IOError:
		print_err (sys.exc_info ()[1])
		sys.exit (1)

	for tbl in len (opts [1]) and opts [1] or ['filter']:
		try:
			if to_stdout:
				dot_output (iptables [tbl], sys.stdout, line_width0, line_width1)
			else:
				outfile = get_file_name ("%s%s.dot" % (tbl, root_suffix))
				if not outfile: continue
				dot_output (iptables [tbl], file (outfile, 'w'), \
						line_width0, line_width1)
		except OSError:
			print_err (sys.exc_info ()[1])
		# want to see other errors here