import frida
import sys

def on_message(message, data):
    if message['type'] == 'send':
        print(message['payload'])
    elif message['type'] == 'error':
        print(message['stack'])
    else:
     print(message)

pid = frida.spawn("C:\Temp\stage0.exe")
session = frida.attach(pid)

script = """
var pNtAllocateVirtualMemory = Module.findExportByName("ntdll.dll", 'NtAllocateVirtualMemory')

Interceptor.attach(pNtAllocateVirtualMemory, {
    onEnter: function (args) {
        this.ProcessHandle = args[0];
        this.BaseAddress = args[1];
        this.ZeroBits = args[2];
        this.RegionSize = args[3];
        this.AllocationType = args[4];
        this.Protect = args[5];
    },
    onLeave: function (args) {
        if(!(this.ProcessHandle == 0xffffffff || this.ProcessHandle == 0xffffffffffffffff)){
            send("[-] I saw you call NtAllocateVirtualMemory");
            send("Process Handle: " + this.ProcessHandle);
            send("BaseAddress: " + this.BaseAddress);
            send("ZeroBits: " + this.ZeroBits);
            send("RegionSize: " + this.RegionSize);
            send("AllocationType: " + this.AllocationType);
            send("Protect: " + this.Protect);
        }
    }
});

var pNtWriteVirtualMemory = Module.findExportByName("ntdll.dll", 'NtWriteVirtualMemory')

Interceptor.attach(pNtWriteVirtualMemory, {
    onEnter: function (args) {
        this.Handle = args[0];
        this.BaseAddress = args[1];
        this.Buffer = args[2];
        this.NumberOfBytesToWrite = args[3];
        this.NumberOfBytesWritten = args[4];
    },
    onLeave: function (args) {
        if(!(this.Handle == 0xffffffff)){
            send("[-] I saw you call NtWriteVirtualMemory");
            send("Handle: " + this.Handle);
            send("BaseAddress: " + this.BaseAddress);
            send("Buffer: " + this.Buffer);
            send("NumberOfBytesToWrite: " + this.NumberOfBytesToWrite);
            send("NumberOfBytesWritten: " + this.NumberOfBytesWritten);
        }
    }
});
"""

script = session.create_script(script)
frida.resume(pid)
script.on('message', on_message)
script.load()

try:
    while True:
        pass
except KeyboardInterrupt:
    session.detach()
    sys.exit(0)