function_compare.py

"""
Remove text comparison code from here.
"""

import sys
import os
import argparse

list_functions= {}
def get_functions(ida_input_dir, filename):
    """
    Read strings from files generated by an IDAPython script and store them in a list for further processing.
    """
    with open(ida_input_dir+'/'+filename,'rU') as f:
        t1= ''
        t2= ''
        t1= [line for line in f.readlines()]
        for function in t1:
            t2= function.split('$$$')
            list_functions[filename+'---'+t2[0]]= t2[1][:-1]

    return list_functions

def get_exe_to_reverse():
    """
    Read command line arguments.
    """
    parser= argparse.ArgumentParser()
    parser.add_argument('file', nargs=1, help='Enter the name of the file you want to reverse engineer here')
    parser.add_argument('-m', action="store_true", help='Use this if you want to compare only by mnemonics')
    args = parser.parse_args()

    return args.file[0], args.m

def function_compare_by_mnem(exe_to_reverse, list_functions):
    """
    Compare each function from the binary you want to reverse with all the functions from all other files
    """

    """
    Separate the functions of the binary and all the libraries so its easy to compare content.
    """
    libraries= {}

    for key in list_functions.keys():
        if not key.startswith(exe_to_reverse+'---'):
            libraries[key]= list_functions[key]
            del list_functions[key]

    # Compare just the mnemonics of functions. The danger here is false positives. Lets see though.
    matched_functions= {}
    all_matches      = {}
    
    for key1, value1 in list_functions.items():
        key1= key1.split('---')[1]
        flag= 0

        for key2, value2 in libraries.items():
            key2= key2.split('---')[1]
            
            value1= value1.rstrip()
            value2= value2.rstrip()
            t1= value1.split('<')
            t2= value2.split('<')
            
            if t1[0] == t2[0]:
                flag= 1
                if not matched_functions.has_key(key1+'<'+t1[1]):
                    matched_functions[key1+'<'+t1[1]]= key2
                elif matched_functions.has_key(key1+'<'+t1[1]):
                    matched_functions[key1+'<'+t1[1]]+= ','+key2
                #if all_matches.has_key(key1+'<'+t1[1]):
                #    all_matches[key1+'<'+t1[1]]+= ','+key2
                #elif not all_matches.has_key(key1+'<'+t1[1]):
                #    all_matches[key1+'<'+t1[1]]= key2
                #break
            else:
                continue

        #if all_matches.has_key(key1+'<'+t1[1]):
        #    all_matches[key1+'<'+t1[1]] += '\n'*2
        #elif not all_matches.has_key(key1+'<'+t1[1]):
        #    all_matches[key1+'<'+t1[1]]  = 'No matches found'+'\n'*2

    #return matched_functions, all_matches
    return matched_functions   

#def write_results(output_file, matched_functions, all_matches, non_matching_instructions):
def write_results(output_file, matched_functions, non_matching_instructions):
    """
    Write results to a file that should then be loaded into IDA
    """
    #all_matches_output_file= "all_matches.txt"
    
    with open(output_file,'a') as f:
        for key,val in matched_functions.items():
            f.write(key+"\t"+str(val))
            f.write("\n")

    #with open(all_matches_output_file,'a') as f:
    #    for key,val in all_matches.items():
    #        f.write(key+"\t"+str(val))
    #        f.write("\n")
            
def check_file_existence(ida_input_dir, exe_to_reverse):
    fullpath= ida_input_dir+'/'+exe_to_reverse
    if os.path.exists(fullpath):
        return 1
    else:
        return 0

if __name__ == "__main__":
    # Identify which the actual binary is and check if you want to compare by mnemonics
    exe_to_reverse, compare_by_mnem= get_exe_to_reverse()

    # Check if the disassembly of the binary you want to reverse has actually been saved
    ida_input_dir= 'input'
    does_file_exist= check_file_existence(ida_input_dir, exe_to_reverse)

    if not does_file_exist:
        print "File %s does not exist in the %s directory. Did you run save_disasm.py correctly?" %(exe_to_reverse, ida_input_dir)
        sys.exit(0)
    else:
        print "Found input for this function :). Will compare now. Make sure you have the output for all the files you want to compare this", \
              "against also in the input directory.\n"

    t2={}
    list_functions= {}

    #Read all the function names and bodies returned after parsing the IDB files of the binary and all the libraries
    if compare_by_mnem:
        for filename in os.listdir(ida_input_dir):
            if filename.endswith('_mnem.txt'):
                t2= get_functions(ida_input_dir, filename)
                list_functions= dict(t2.items()+list_functions.items())
            else:
                continue
    else:
        print "\nMake sure you add the -m switch while running this program\n"
        sys.exit(0)

    # This is what does the actual comparison
    if exe_to_reverse.endswith('_mnem.txt'):
        #matched_functions, all_matches = function_compare_by_mnem(exe_to_reverse, list_functions)
        matched_functions = function_compare_by_mnem(exe_to_reverse, list_functions)
    else:
        print "Ignoring file ",exe_to_reverse,"\n"
        print "Ensure you pass the filename as an argument correctly. The format is <binaryname>_mnem.txt"
        sys.exit(0)

    # Write IDA results to file
    output_file= 'input_to_rename_function.txt'
    #write_results(output_file, matched_functions, all_matches, 'NA')
    write_results(output_file, matched_functions, 'NA')