#!/usr/bin/python 
#
# disass.py v1.1 by atlas
#
# Syntax:  ./disass.py <binary-executable>
#
# disass.py will use a few simply objdump calls to gather GOT, PLT, and Disassembly information.
# GOT addresses are tied to PLT calls and the PLT lines are tagged with function names
# Then, full disassembly is scanned for references to PLT calls, and those lines are labeled
# with the appropriate function call name.  Tested only with *nix.
#
# Yes, this may seem elementary, but I found it helpful so here it is.
# @145
#     Converted to Python from Perl in 2/2006

import os
import sys
import re

VER=1.1 


## printSyntax() does just that. ##
def printSyntax():
  print >>sys.stderr, """
disass v%0.2f Disassembling enhancer

Syntax:  %s [--txt|--html] <binary-to-disass>

"""  % (VER, sys.argv[0])
  sys.exit(1)


###################################
if (len(sys.argv) < 2 or sys.argv[1] == "-h"):
  printSyntax()

binary = sys.argv[1]

disassembly = os.popen('objdump -S '+binary).readlines()
GOT = os.popen('objdump -R '+binary).readlines()
HEADERS = os.popen('objdump -x '+binary).readlines()
SYMBOLS = os.popen('objdump -t '+binary).readlines()
SYMBOLS += os.popen('objdump -T '+binary).readlines()
LIBS = os.popen('objdump -T '+binary).readlines()
BREAK = []
CALLS = {}
JMPS = {}
CNDJMPS = {}
JMPTRGTS = {}
SUBS = {}
colors = [ "#330099","#993399","#3333FF","#009900","#FFCC33","#990000","#00CCCC","#FFFF00","#336666" ]

SRCH_PUSHEBP = "push\s+%ebp"
SRCH_MOVESPEBP = "mov\s+%esp,\s*%ebp"

### Peal off the first few lines of GOT output
if (len(GOT) > 5):
  GOT.pop(0)
  GOT.pop(0)
  GOT.pop(0)
  GOT.pop(0)
  GOT.pop(0)




########################## Stage 1: Associate GOT and PLT entries and Tag disassassembly calls to PLT with the appropriate name #############################
###  Parse GOT section from GOT
for LINE in GOT:
  LINE = LINE.strip()
  if (len(LINE) < 1):
    continue
  #
  while (LINE[0] == " " or LINE[0] == "0"):
    LINE = LINE[1:]

### Tag PLT with GOT names
  LIN = re.split("\s+",LINE)
  for each in range(len(disassembly)):
    disassembly[each] = disassembly[each].strip()
    if (re.search(LIN[0], disassembly[each]) >-1):
      disassembly[each] += "\t" + LIN[2] 

### For each PLT line found, disassembly is scanned and tagged appropriately
      if (re.search("jmp.*0x" + LIN[0], disassembly[each]) > -1):
        PLT = re.split("\s+", disassembly[each])
        PLT[0] = PLT[0][:-1]		#trim off the ":" at the end

    ### Now traverse the rest of the disassembly looking for the calls to the PLT for this entry
        for EACH in range(len(disassembly)):
          if (re.search("^.*:.*" + PLT[0],disassembly[EACH]) > -1):
            breakdown = disassembly[EACH].split(":")
            address = breakdown[0]
            re.sub("^\s*", "", address)
            BREAK.append(address)
            disassembly[EACH].strip("\n")
            disassembly[EACH] += "\t %s (brkpt: %i)\n\n" % (LIN[2],len(BREAK)-1)
          #
        #
      #
    #
  #
#

######################### Stage 2: Find all subroutines using "ret" and "push %ebp"/"mov %esp,ebp"###############################
def getSubName(address):
  # loop through SYMBOLS looking for this address
  name = None
  for i in SYMBOLS:

#    print >>sys.stderr, ("%s:%s" % (address,i))
    if (re.search(address, i ) > -1):
      name = i[45:-1]

  if (name == None):
    name = "unnamed_sub_%d" % (len(SUBS))
  return name

def tagSubs():
  for EACH in range(len(disassembly) - 1 ):
    if (re.search(SRCH_PUSHEBP, disassembly[EACH]) > -1 and re.search(SRCH_MOVESPEBP,disassembly[EACH+1]) >-1):
      subname = ""
      address = disassembly[EACH][:disassembly[EACH].index(":")]
      #chop off leading spaces...
      re.sub("^\s*", "", address)

#      print >>sys.stderr,(address)

      subname = getSubName(address)
      SUBS.setdefault(address, subname)
      disassembly[EACH] = "-\n-\n------------==================== subroutine: %s =======================------------\n%s"%(subname,disassembly[EACH])

      BREAK.append("0x%s" % address)


tagSubs()

######################### Stage 2: Find locals and parm calls within each sub ###############################
def tagSubVars():
  sortedsubs = SUBS.keys()
  sortedsubs.sort()
  subiter = iter(sortedsubs)
  VARS = []

  sub = "BEGINNING"
  try: 
    sub = subiter.next()
  except StopIteration:
    pass

  head = -1
  for EACH in range(len(disassembly) ):
    line = disassembly[EACH]
    print >>sys.stderr,("Searching for %s : %s"%(sub,line))
    if (re.search("%s:" % (sub), line) != None):
      print >>sys.stderr,("NewSub for Vars")
      if (head > 0):
        VARS.sort()
        for var in VARS:
          disassembly[head] += "\n	%s ___________________________________" %(var)
      try:
        sub = subiter.next()
      except:
        pass
      head = EACH
      VARS = []
    
    for segment in line.split(','):
      varidx = -1
      varstart = -1
      varend = -1
      try:
        try:
          varidx = segment.index('(%ebp)')
        except:
          varidx = segment.find('[ebp')
      
        if (varidx > -1):
          if (segment[varidx] == '('):
            varstart = segment.index('0x')
            varend = segment.index('%ebp)')+4
          elif (segment[varidx] == '['):
            varstart = varidx
            varend = segment.index(']')
          var = segment[varstart:varend+1]
          try:
            VARS.index(var)
          except:
            VARS.append(var)
            
          print >>sys.stderr,(" Found:  %s:%s"%(line[varstart:varend+1],line))
      except:
        print >>sys.stderr,("tagSubVars:  Something broke at %d" % (EACH))


tagSubVars()

######################### Stage 2: Find all arbitrary memory calls and label.  Use PEFILE and PyElf to determine Strings ###############################
def tagMemAccess():
  pass

######################### Stage 2: Find all non-PLT calls and tag them with name of subroutine they are calling ###############################
def tagCalls():
  pass

######################### Stage 3: Find all jmp, je, jz, jl, jg, jge, jle, jnz, jns calls #############################
def tagJMPs():
  pass

  ### Check whether inside current sub
  ##### If same sub and a jump backward, mark that section as a loop
  ##### ASCII art the jmps?   HTML?
    # For each jmp block, check disassembly[line][somechar] for [/\|].  
    ## If found, check [somechar+1] until an empty char is found
    # Insert a "|" for each line in between and a /- or \- as appropriate  (HTML and colors?)

    # Store "destination" in array
    ## Add a blank line after each destination for easier reading







def printTXT(outfile):
  print >>outfile,("DISASSEMBLY:")
  for i in  disassembly:
    print >>outfile,(i.strip())
  print >>outfile,("\n\nGOT:")
  for i in GOT:
    print >>outfile,(i.strip()) 
  print >>outfile,("\n\nHEADERS:")
  for i in HEADERS:
    print >>outfile,(i.strip())
  print >>outfile,("\n\nSYMBOLS:")
  for i in SYMBOLS:
    print >>outfile,(i.strip())
  print >>outfile,("\n\nBreakpoints for each \"call\":\n");
  for brk in BREAK:
    print >>outfile,(" break *0x" + brk)
  #

  print >>outfile,("""\nDISPLAY SETTINGS/Basic
   display/i $pc
   display/x $edx
   display/x $ecx
   display/x $ebx
   display/x $eax
   display/32wx $ebp-92
   display/32xw $esp 
""")





def printHTML(outfile):
  print >>outfile,("DISASSEMBLY:")
  for i in  disassembly:
    print(i.strip())
  print >>outfile,("\n\nGOT:")
  for i in GOT:
    print >>outfile,(i.strip()) 
  print >>outfile,("\n\nHEADERS:")
  for i in HEADERS:
    print >>outfile,(i.strip())
  print >>outfile,("\n\nSYMBOLS:")
  for i in SYMBOLS:
    print >>outfile,(i.strip())
  print >>outfile,("\n\nBreakpoints for each \"call\":\n");
  for brk in BREAK:
    print >>outfile,(" break *0x" + brk)
  #

  print >>outfile,("""\nDISPLAY SETTINGS/Basic
   display/i $pc
   display/x $edx
   display/x $ecx
   display/x $ebx
   display/x $eax
   display/32wx $ebp-92
   display/32xw $esp 
""")




printTXT(sys.stdout)
