#!/usr/bin/python3
import frida
import sys
import json

#Message from js code. Contains stack dump
def on_message(message,data):
#	print(message)	

	if(message['payload']=='log'):
		print("log:")		
	
	#Stack dump. Count zeros in stack, less is bettter.
	if(message['payload']=='dataz'):
		c=0;
		for x in data:
			if(x==0):
				c=c+1
	#		sys.stdout.write(hex(x))
	#		sys.stdout.write(',')
	#		c=c+1
	#		if(c%30==0):
	#			print("")
		print(c,sys.argv[1])
		with open("test.txt", "a") as myfile:
			myfile.write("%s %s\n\r" % (c,sys.argv[1]))
		
def on_detached():
#	print("on_detached")
	sys.exit(0)	
	

pid = frida.spawn(['./dist'])
session = frida.attach(pid)

script = session.create_script("""

//At begin of execution, NOP out scanf and zero stack memory.
Interceptor.attach(ptr("0x040059F"), {
    onEnter: function(args) {
		var start=ptr("0x4005D9")
		var maxPatchSize = 5;
		Memory.patchCode(start, maxPatchSize, function (code) {
		  var cw = new X86Writer(code, { pc: start });
		  cw.putNop();
		  cw.putNop();
		  cw.putNop();
		  cw.putNop();
		  cw.putNop();
		  cw.flush();
		});	

		//This is slow
		for (i = -4 ; i > -0x25350; i-=4) {
			Memory.writeU32(ptr(this.context['rbp']+i).add(i), 0);
		} 				
    }
});

//Instead of scanf, write flag to memory here.
Interceptor.attach(ptr("0x004005C9"), {
    onEnter: function(args) {
		Memory.writeByteArray(ptr(this.context['rbp']).add(-96), %s);		
    }
});

//Execution done. Dump stack to python code.
Interceptor.attach(ptr("0x5C2AC0"), {	
    onEnter: function(args) {
        //send(JSON.stringify(this.context['pc']));
		send('dataz',Memory.readByteArray(ptr(this.context['rbp']-0x25350), 0x25350))	

    }
});


""" % str(list(map(ord,sys.argv[1]))))
script.on('message', on_message)
session.on('detached', on_detached)
script.load()
frida.resume(pid)
sys.stdin.read()