#! /usr/bin/env python

#############################################################################
##                                                                         ##
## egggraph.py --- output a dotfile from a shellcode                       ##
##              see http://www.secdev.org/                                 ##
##              for more informations                                      ##
##                                                                         ##
## Copyright (C) 2004  Philippe Biondi <phil@secdev.org>                   ##
##                                                                         ##
## This program is free software; you can redistribute it and/or modify it ##
## under the terms of the GNU General Public License as published by the   ##
## Free Software Foundation; either version 2, or (at your option) any     ##
## later version.                                                          ##
##                                                                         ##
## This program is distributed in the hope that it will be useful, but     ##
## WITHOUT ANY WARRANTY; without even the implied warranty of              ##
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU       ##
## General Public License for more details.                                ##
##                                                                         ##
#############################################################################



import getopt,sys,os,warnings,re

warnings.filterwarnings("ignore","tempnam",RuntimeWarning, __name__)

# $Id$

VERSION="$Version$"

def usage():
    print "Syntax: egggraph [-i file]"
    
def error(x,err=1,printusage=0):
    printinfo(0,"ERROR: %s" % x)
    if printusage:
        usage()
    raise SystemExit(err)


infile = sys.stdin
verb = 0
arch = "i386"

try:
    opts=getopt.getopt(sys.argv[1:], "hVi:a:")
    for opt,optarg in opts[0]:
        if opt == "-h":
            usage()
        if opt == "-V":
            printinfo(0, "egggraph v%s\nCopyright (C) 2004  Philippe Biondi <phil@secdev.org>" % VERSION)
            sys.exit(0)
            raise SystemExit
        elif opt == "-v":
            try:
                verb = int(optarg)
            except ValueError,msg:
                raise getopt.GetoptError(str(msg),None)
        elif opt == "-i":
            try:
                infile = open(optarg)
            except OSError,msg:
                raise getopt.GetoptError(str(msg),None)
        elif opt == "-a":
            arch = optarg
except getopt.GetoptError,msg:
    error(msg, printusage=1)


eggfile = os.tempnam("/tmp/","egg")
f=open(eggfile,"w+")
f.write(infile.read())
f.close()

asmf = os.popen('objdump -b binary -m %s -D "%s"' % (arch,eggfile))
a = asmf.read()
status = asmf.close()
if status is not None:
    if status & 0xff:
        print "killed by signal %i" % (status & 0xff)
    else:
        print "exited with status %i" % (status >> 8)


def getjmp_i386(cmd,args):
    if cmd.startswith("j"):
        return int(args,0)
    if cmd.startswith("call"):
        return int(args,0)
    return None
    



re_ins = re.compile("^\s+([0-9a-f]+):\s+([a-z0-9][ 0-9a-z]*[a-z0-9])\s+([^\s][^\s]*)\s*(.*)$")

brk = [0]
brk2 = []
code = []
cline = 0
brkadr = {}
jmp = []
adr2line = {}
for l in a.splitlines():
    m = re_ins.match(l)
    if m:
#        print l
#        print m.groups()
        
        adr,opcode,cmd,arg = m.groups()
        code.append("%4s: %-6s %-30s" % (adr,cmd,arg))
        adr = int(adr,16)
        adr2line[adr] = cline
        cline += 1
        j = getjmp_i386(cmd,arg)
        if j is not None:
            if cline not in brk:
                brk.append(cline)
            if j not in brk2:
                brk2.append(j)
            brkadr[cline]=j
#            print >>sys.stderr,"JUMP to %x"%j
#    else:
#        print "##",l

for i in brk2:
    try:
        i = adr2line[i]
    except:
        pass
    if i not in brk:
        brk.append(i)

brk.sort()
last=0
jmp = map(lambda x,y: (x,y),brk[:-1],brk[1:])

for i,j in jmp[:]:
    try:
        jmp.append((i,adr2line[brkadr[j]]))
    except:
        pass

#for b in range(
#for s,d in brkadr:
#    try:
#        brk.append(adr2line[d])
#    except KeyError:
#        pass
#for 
#    try:
#        jmp.append((last,adr2line[d]))
#    except KeyError:
#        pass
#    last = s
#    
#
last = 0;
print 'digraph egggraph {\n\tnode [ shape = box,fontname = "fixed" ];'
for i in range(len(code)):
    if i in brk:
        if i:
            print '"];'
        print '\tblock%i [label = "' % i,
#        jmp.append((last,i))
        last = i
    print "%s\\n" % code[i],

print '"];'

for i in jmp:
    print "block%i -> block%i;" % i

    
        
print "}"