import re
import os
import numpy as np
from collections import Counter

f1 = open('Malware_example.txt','r') #pre-parsed malware file

convert files of words into arrays
math = open('sample_math.txt','r')
matharray = list(math.read().split())
math.close()

logic = open('sample_logic.txt','r')
logicarray = list(logic.read().split())
logic.close()

priv = open ('sample_priv.txt','r')
privarray = list(priv.read().split())
priv.close()

stack = open('sample_stack.txt','r')
stackarray = list(stack.read().split())
stack.close()

nop = open('sample_nop.txt','r')
noparray = list(nop.read().split())
nop.close()

other = open('sample_other.txt','r')
otherarray = list(other.read().split())
other.close()

memory = open('sample_memory.txt','r')
memoryarray = list(memory.read().split())
memory.close()

branch = open('sample_branch.txt','r')
brancharray = list(branch.read().split())
branch.close()

convert arrays into dictionaries
math_dict = dict()
math_dict.update(dict.fromkeys(matharray,0))

logic_dict = dict()
logic_dict.update(dict.fromkeys(logicarray,1))

priv_dict = dict()
priv_dict.update(dict.fromkeys(privarray, 2))

stack_dict = dict()
stack_dict.update(dict.fromkeys(stackarray, 3))

nop_dict = dict()
nop_dict.update(dict.fromkeys(noparray,4))

other_dict = dict()
other_dict.update(dict.fromkeys(otherarray,5))

memory_dict = dict()
memory_dict.update(dict.fromkeys(memoryarray,6))

branch_dict = dict()
branch_dict.update(dict.fromkeys(brancharray,7))

print '--------------------------------------------------------------'
print "Below is a listing of the dictionary keys:"
print math_dict.keys()
print ""
print "Below is a listing of the dictionary items"
print math_dict.items()
print ""
print '-------------------------------------------------------------'

create big dictionary of all keys
word_set = dict(math_dict.items() + logic_dict.items() + priv_dict.items() + stack_dict.items() + nop_dict.items() + other_dict.items() + memory_dict.items() + branch_dict.items())

print '--------------------------------------------------------------'
print word_set
print '--------------------------------------------------------------'

same process as above, but through os

category_dict = dict()

foldername = "samples_folder"
for txtfile in os.listdir(foldername):
	words = open(foldername+"/"+txtfile, "r").read().strip().split('\n')
	category_dict.update(dict.fromkeys(words,txtfile[:-4].split('_')[1]))

print '--------------------------------------------------------------'
print category_dict
print '--------------------------------------------------------------'

p = np.zeros((8,8), dtype=np.int)
wordfound = list()


for line in f1:
	for word in category_dict:
		for m in re.finditer(word, line):
			wordfound.append(word)

print wordfound
		
#maps to words
#statelist = map(lambda word: category_dict[word], wordfound)
#transition_counts = Counter(zip(statelist, statelist[1:]))

print '-------------------------------------------------'

#maps to numbers
statelist2 = map(lambda word: category_dict[word], wordfound)
transition_counts = Counter(zip(statelist2, statelist2[1:]))

for (x,y), c in transition_counts.iteritems():
	print "(%s,%s) = %d" % (x,y,c)
	p[x,y] = c

print '------------------------------------------------'
print ""
print "FINAL TRANSITION MATRIX"
print p

#OLD CODE
#line_number = 1
#for line in f1:
#	for m in re.finditer('esp',line):
#		print "Found", m.group(0), "line", line_number, "at", m.start(), "-", m.end()
#	for m in re.finditer('push',line):
#		print "Found", m.group(0), "line", line_number, "at", m.start(), "-", m.end()
#	for m in re.finditer('pop',line):
#		print "Found", m.group(0), "line", line_number, "at", m.start(), "-", m.end()
#	line_number += 1


f1.close()

