diff --git a/src/frida/discoverer.py b/src/frida/discoverer.py index ad6e766..25976ba 100644 --- a/src/frida/discoverer.py +++ b/src/frida/discoverer.py @@ -1,10 +1,13 @@ +from frida.core import ModuleFunction + class Discoverer(object): - def __init__(self): + def __init__(self, reactor): + self._reactor = reactor self._script = None - def start(self, process): + def start(self, process, ui): def on_message(message, data): - print message, data + self._reactor.schedule(lambda: self._process_message(message, data, process, ui)) source = self._create_discover_script() self._script = process.session.create_script(source) self._script.on("message", on_message) @@ -20,69 +23,196 @@ class Discoverer(object): def _create_discover_script(self): return """ -var pending = []; -var active = []; -Process.enumerateThreads({ - onMatch: function (thread) { - pending.push(thread); - }, - onComplete: function () { - var currentThreadId = Process.getCurrentThreadId(); - pending = pending.filter(function (thread) { - return thread.id !== currentThreadId; - }); - var processNext = function processNext() { - active.forEach(function (thread) { - send("unfollow(" + thread.id + ")"); - Stalker.unfollow(thread.id); +var Sampler = function Sampler() { + var total = 0; + var pending = []; + var active = []; + var samples = {}; + Process.enumerateThreads({ + onMatch: function (thread) { + pending.push(thread); + }, + onComplete: function () { + var currentThreadId = Process.getCurrentThreadId(); + pending = pending.filter(function (thread) { + return thread.id !== currentThreadId; }); - active = pending.splice(0, 4); - active.forEach(function (thread) { - send("follow(" + thread.id + ")"); - Stalker.follow(thread.id, { - events: { call: true }, - onCallSummary: function (summary) { - send("summary from " + thread.id); - } + total = pending.length; + var processNext = function processNext() { + active.forEach(function (thread) { + Stalker.unfollow(thread.id); }); - }); - if (active.length > 0) { - setTimeout(processNext, 2000); - setTimeout(Stalker.garbageCollect, 2100); - } - }; - setTimeout(processNext, 0); - } -}); + active = pending.splice(0, 4); + if (active.length > 0) { + var begin = total - pending.length - active.length; + send({ + from: "/sampler", + name: '+progress', + payload: { + begin: begin, + end: begin + active.length - 1, + total: total + } + }); + } else { + for (var address in samples) { + if (samples.hasOwnProperty(address)) { + var counts = samples[address].counts; + var sum = 0; + for (var i = 0; i !== counts.length; i++) { + sum += counts[i]; + } + var callsPerSecond = Math.round(sum / (counts.length * 0.25)); + samples[address] = callsPerSecond; + } + } + send({ + from: "/sampler", + name: '+result', + payload: { + samples: samples + } + }); + samples = null; + } + active.forEach(function (thread) { + Stalker.follow(thread.id, { + events: { call: true }, + onCallSummary: function (summary) { + if (samples === null) { + return; + } + + for (var address in summary) { + if (summary.hasOwnProperty(address)) { + var sample = samples[address]; + if (sample === undefined) { + sample = { counts: [] }; + samples[address] = sample; + } + sample.counts.push(summary[address]); + } + } + } + }); + }); + if (active.length > 0) { + setTimeout(processNext, 2000); + setTimeout(Stalker.garbageCollect, 2100); + } + }; + setTimeout(processNext, 0); + } + }); +}; +sampler = new Sampler(); """ + def _process_message(self, message, data, process, ui): + if message['type'] == 'send': + stanza = message['payload'] + name = stanza['name'] + payload = stanza['payload'] + if stanza['from'] == "/sampler": + if name == '+progress': + ui.on_sample_progress(payload['begin'], payload['end'], payload['total']) + elif name == '+result': + module_functions = {} + dynamic_functions = [] + for address, rate in payload['samples'].iteritems(): + address = int(address, 16) + function = process.ensure_function(address) + if isinstance(function, ModuleFunction): + functions = module_functions.get(function.module, []) + if len(functions) == 0: + module_functions[function.module] = functions + functions.append(function) + else: + dynamic_functions.append(function) + ui.on_sample_result(module_functions, dynamic_functions) + else: + print message, data + else: + print message, data + else: + print message, data + + +class UI(object): + def on_sample_progress(self, begin, end, total): + pass + + def on_sample_result(self, module_functions, dynamic_functions): + pass + def main(): + import colorama + from colorama import Fore, Back, Style import frida + from frida.core import Reactor from optparse import OptionParser import sys + colorama.init(autoreset=True) + usage = "usage: %prog [options] process-name-or-id" parser = OptionParser(usage=usage) (options, args) = parser.parse_args() if len(args) != 1: parser.error("process name or id must be specified") + reactor = Reactor(raw_input) + + class ConsoleUI(UI): + def __init__(self): + self._status_updated = False + + def update_status(self, message): + if self._status_updated: + cursor_position = "\033[A" + else: + cursor_position = "" + print(cursor_position + Style.BRIGHT + message) + self._status_updated = True + + def on_sample_progress(self, begin, end, total): + self.update_status("Sampling %d threads: %d through %d..." % (total, begin, end)) + + def on_sample_result(self, module_functions, dynamic_functions): + print "Module functions:" + for module, functions in module_functions.iteritems(): + print "\t%s" % module.name + for function in functions: + print "\t\t%s" % function + + if len(dynamic_functions) > 0: + print + + print "Dynamic functions:" + for function in dynamic_functions: + print "\t%s" % function + + reactor.stop() + + ui = ConsoleUI() + try: target = int(args[0]) except: target = args[0] try: + ui.update_status("Attaching...") process = frida.attach(target) except Exception, e: - print >> sys.stderr, "Failed to attach: %s" % e + ui.update_status("Failed to attach: %s" % e) sys.exit(1) - d = Discoverer() - d.start(process) - print "Discovery started" - print "Press ENTER to stop" - raw_input() + ui.update_status("Injecting script...") + + d = Discoverer(reactor) + d.start(process, ui) + reactor.run() print "Stopping..." d.stop() process.detach() diff --git a/src/setup.py b/src/setup.py index 60543ba..5be8e16 100755 --- a/src/setup.py +++ b/src/setup.py @@ -40,5 +40,8 @@ setup( ext_modules=[Extension('_frida', [])], cmdclass={ 'build_ext': FridaPrebuiltExt - } + }, + install_requires=[ + "colorama >= 0.2.7" + ] )