#!/usr/bin/python

#
# ok.. this is not very robust.  but i only wrote it this evening.
# i am learning python!
#

import string
import sys
import os
import re

objdump_binary = "/usr/bin/objdump"
objdump_opt = " -d "

re_entry = re.compile("[ \t]*Entry point address:[ \t]*([0-9xa-fA-F]*)$")
re_d1 = re.compile("[^a-z0-9]*([^:]*):.*$")
re_d2 = re.compile("[^a-z0-9]*([^:]*):[\t ]*([^\t]*)[ \t]*(.*)$")
re_d3 = re.compile("[0-9a-f]{2} ")

re_cbranch = re.compile("^j[^m][^ \t]*[ \t]+0*x*([0-9a-z]*)[ \t].*$")
re_ubranch = re.compile("jmp +0*x*([0-9a-z]*)[ \t].*$")
re_call = re.compile("call +0*x*([0-9a-z]*)[ \t].*$")

class BinaryCFG:
	def __init__(self, filename, start, stop):
		self.disassembly = {}
		self.work = {}

		c = objdump_binary + objdump_opt + \
			" --start-address=" + start + \
			" --stop-address=" + stop + " " + \
			filename
		self.entry = start
		f = os.popen(c, "r")
		if not f:
			return -1
		while 1:
			line = f.readline()
			if not line:
				break
#
# FIXME: empty line, not len == 1
#
			if len(line) == 1:
				continue
			m = re_d1.match(line)
			if not m:
				print "erm> \"" + line + "\"\n"
				continue
			a = '0x' + m.group(1)
			self.disassembly[a] = line
		f.close()

	def m_bytecount(self, bytestream):
		m = re_d3.findall(bytestream)
		return len(m)

	def m_run_from(self, entry):
		k = self.disassembly.has_key(entry)
		if not k:
			return
		k = self.work.has_key(entry)
		if k:
			return
		self.work[entry] = 1
		line = self.disassembly[entry]
		m = re_d2.match(line)
		i = m.group(3)
		a = '0x' + m.group(1)

#
# my pthon is shit.. whats the switch syntax?
# here sucks anyway :)
#
		cm = re_cbranch.match(i)
		if cm:
			cn = "0x" + cm.group(1)
			print "\"" + a + "\"", "->", "\"" + cn + "\";"
			self.m_run_from(cn)
		else:
			um = re_ubranch.match(i)
			if um:
				un = "0x" + um.group(1)
				print "\"" + a + "\"", "->", "\"" + un + "\";"
				self.m_run_from(un)
				return
			else:
				lm = re_call.match(i)
				if lm:
					ln = "0x" + lm.group(1)
					print "\"" + ln + "\" [color=red];"
					print "\"" + a + "\"", "->", "\"" + ln + "\";"
					self.m_run_from(ln)

		a10 = string.atoi(a, 16)
		bs = m.group(2)
		bc = self.m_bytecount(bs)
		n10 = a10 + bc
		n = "0x%x" % (n10)
		print "\"" + a + "\"", "->", "\"" + n + "\";"
		self.m_run_from(n)

	def m_run(self):
		print "Digraph cfg {"
		print "ENTRY", "->", "\"" + self.entry + "\";"
		print "ENTRY [color=blue];"
		self.m_run_from(self.entry)
		print "}"

if len(sys.argv) != 4:
	print "usage: erm filename start stop"
	sys.exit(0)
cfg = BinaryCFG(sys.argv[1], sys.argv[2], sys.argv[3])
cfg.m_run()