import os
import sys
import re

walk_dir = sys.argv[1]  # folder with MASTIFF generated files for each binary

print('walk_dir = ' + walk_dir)

# If your current working directory may change during script execution, it's recommended to
# immediately convert program arguments to an absolute path. Then the variable root below will
# be an absolute path as well. Example:
walk_dir = os.path.abspath(walk_dir)
print('walk_dir (absolute) = ' + os.path.abspath(walk_dir))


cnt = 0

for root, subdirs, files in os.walk(walk_dir):
    list_file_path = os.path.join(root, '')

    filename = "peinfo-quick.txt"
    # filename = "peinfo-full.txt" #90/109

    cnt += 1
    # 61,688 created out of 63,628
    if filename in files:  # 58,856 found out of 63,628
        file_path = os.path.join(root, filename)
        vx_hash = os.path.basename(os.path.normpath(root))

        with open(file_path) as f:
            # lines = re.findall(r"\b.+dll.*\t.*\b", f.read().lower()) #dll, call, address
            lines = re.findall(r"\b.+dll.*\t.*\t\b", f.read().lower())

            if lines:
                # print "==================================> DLLS LINES found"
                # pprint.pprint(lines)

                # print "=====================================================> SYS CALLS"
                for line in lines:
                    lineStr = line.split("\t")
                    # gen call list for each sample
                    with open("/home/tarun/Desktop/DLL/SysCalls/" + vx_hash + ".txt", 'a') as w:
                        # if lineStr[1] not in w.read().lower():
                        # print lineStr[1]
                        w.write(lineStr[1]+'\n')
                # print cnt #print number of unique calls per exe
            else:
                #  write empty file if calls not found
                # vx_hash = os.path.basename(os.path.normpath(root))
                print "NO CALL INFO FOUND for file: ", vx_hash

                with open("/home/tarun/Desktop/DLL/SysCalls/" + vx_hash + ".txt", 'a') as w:
                    w.write('\n')
    else:
        #  write empty file if PE info not found
        # Remember to del extra files (about 15) generated for folder structure before creating matrix
        # (eg: MASTIFF_final_dataset.txt, Adware.txt,Backdoor.txt)
        vx_hash = os.path.basename(os.path.normpath(root))
        print "NO PE INFO FOUND =====================================================>", vx_hash

        with open("/home/tarun/Desktop/DLL/SysCalls/" + vx_hash + ".txt", 'a') as w:
            w.write('\n')

print cnt  #print number of files created
