#!/usr/bin/env python

#
# Future: More meta info, references, bb links, func args, var bt
#

import sys, os, time, string, re, stat

sys.path.append("C:\\code\\python\\paimei")

import pida

def symlink(source, target):
    if windows:
        try:
            wsh = win32com.client.Dispatch("WScript.Shell")
        except:
            print "[!] Couldnt instantiate WScript.Shell"
            return False
            
        link = wsh.CreateShortcut(source + ".lnk")
        link.TargetPath = re.sub('/', '\\\\', target) # Windows forces me to do this gay shit
        link.Save()
        
        return True
    else:
        try:
            os.symlink(target, source)
        except:
            print "[!] Problem creating symlink %s -> %s" % (source, target)
            return False
        
        return True
    
    return False

if len(sys.argv) < 3:
    print "%s <pida file> <destination dir>" % sys.argv[0]
    sys.exit(-1)

pidafilename = sys.argv[1]
path = sys.argv[2]

if sys.platform == "win32":
    import win32com.client

    windows = True
else:
    windows = False

print "[*] Checking if we have %s" % pidafilename
if os.access(pidafilename, os.F_OK) != True:
    print "[!] Problem accessing %s" % pidafilename
    sys.exit(-1)
else:
    print "[*] Loading %s pida module" % pidafilename
    module = pida.load(pidafilename)

if os.access(path, os.F_OK) != True:
    print "[!] Problem accessing %s" % path
    sys.exit(-1)
else:
    print "[*] Changing to %s" % path
    try:
        os.chdir(path)
    except:
        print "[!] Problem with chdir"

#
# Iterate functions making dirs, then bb in dir dumping bb info by addr
#

symlinks = {}

for function in module.nodes.values():
    funcname = hex(function.ea_start)
    
    try:
        os.mkdir(funcname)
    except:
        print "[!] Problem making %s" % (funcname)
        sys.exit(-1)
    
    try:
        os.chdir(funcname)
    except:
        print "[!] Problem chdir to %s" % (funcname)
        
    for bb in function.nodes.values():
        bbname = hex(bb.ea_start)
        
        try:
            os.mkdir(bbname)
        except:
            print "[!] Problem making %s" % (bbname)    
        
        try:
            os.chdir(bbname)
        except:
            print "[!] Problem chdir to %s" % (bbname)
        
        for instruction in bb.sorted_instructions():
            if instruction.mnem == "call" and re.compile("sub_.*").match(instruction.op1):
                symlinks["%s/%s/%s/%s_0x%s" % (path, funcname, bbname, hex(instruction.ea), instruction.op1.rsplit("_")[-1])] = "%s/0x%s" % (path, instruction.op1.rsplit("_")[-1])
                continue
            
            insname = hex(instruction.ea)
            
            fh = open(insname, "w")
            fh.write("%s %s %s %s" % (instruction.mnem, instruction.op1, instruction.op2, instruction.op3))
            fh.close()
            
            os.chmod(insname, stat.S_IRWXU|stat.S_IRWXG|stat.S_IRWXO)
        
        try:
            os.chdir("..")
        except:
            print "[!] Problem chdir to parent directory"
            sys.exit(-1)
    
    try:        
        os.chdir(path)
    except:
        print "[!] Problem chdir to %s" % (path)

#
# Do symlinks last so we ensure the dirs are there
#
for key in symlinks.keys():
    print "[*] Creating SymLink %s -> %s" % (key, symlinks[key])
    symlink(key, symlinks[key])