from __future__ import print_function
import frida
import sys

def main(target_process):
    session = frida.attach(target_process)
    script = session.create_script("""
        const fastCallback = Memory.alloc(Process.pageSize);
        console.log(fastCallback);
        Memory.patchCode(fastCallback, 128, code => {
          const cw = new X86Writer(code, { pc: fastCallback });
          cw.putCmpRegI32('edi', 10);
          cw.putJccShortLabel('je', 'match', 'unlikely');

          cw.putLabel('nomatch');
          cw.putMovRegU64('rax', 42);
          cw.putJmpShortLabel('done');

          cw.putLabel('match');
          cw.putSubRegImm('rsp', 8);
          cw.putCallAddressWithAlignedArguments(slowCallback, ['edi']);
          cw.putAddRegImm('rsp', 8);

          cw.putLabel('done');
          cw.putRet();

          cw.flush();
        });
    });
""")
    print("[!] Ctrl+D on UNIX, Ctrl+Z on Windows/cmd.exe to detach from instrumented program.\n\n")
    sys.stdin.read()
    session.detach()

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Usage: %s <process name or PID>" % __file__)
        sys.exit(1)

    try:
        target_process = int(sys.argv[1])
    except ValueError:
        target_process = sys.argv[1]
    main(target_process)