diff --git a/examples/cpushark/Capture.py b/examples/cpushark/Capture.py index 457c2d8..ffc6cfb 100644 --- a/examples/cpushark/Capture.py +++ b/examples/cpushark/Capture.py @@ -1,8 +1,11 @@ import bisect from Foundation import NSAutoreleasePool, NSObject, NSThread from PyObjCTools import AppHelper +import re import struct +PROBE_CALLS = re.compile(r"^\/stalker\/probes\/(.*?)\/calls$") + class Capture(NSObject): def __new__(cls, device): return cls.alloc().initWithDevice_(device) @@ -15,7 +18,7 @@ class Capture(NSObject): self.session = None self.script = None self.modules = Modules() - self.calls = Calls(self.modules) + self.calls = Calls(self) return self def delegate(self): @@ -26,8 +29,8 @@ class Capture(NSObject): def attachToProcess_triggerPort_(self, process, triggerPort): assert self.state == CaptureState.DETACHED - self.updateState_(CaptureState.ATTACHING) - NSThread.detachNewThreadSelector_toTarget_withObject_('doAttachWithParams:', self, (process.pid, triggerPort)) + self._updateState_(CaptureState.ATTACHING) + NSThread.detachNewThreadSelector_toTarget_withObject_('_doAttachWithParams:', self, (process.pid, triggerPort)) def detach(self): assert self.state == CaptureState.ATTACHED @@ -35,14 +38,17 @@ class Capture(NSObject): script = self.script self.session = None self.script = None - self.updateState_(CaptureState.DETACHED) - NSThread.detachNewThreadSelector_toTarget_withObject_('doDetachWithParams:', self, (session, script)) + self._updateState_(CaptureState.DETACHED) + NSThread.detachNewThreadSelector_toTarget_withObject_('_doDetachWithParams:', self, (session, script)) - def updateState_(self, newState): + def _post(self, message): + NSThread.detachNewThreadSelector_toTarget_withObject_('_doPostWithParams:', self, (self.script, message)) + + def _updateState_(self, newState): self.state = newState self._delegate.captureStateDidChange() - def doAttachWithParams_(self, params): + def _doAttachWithParams_(self, params): pid, triggerPort = params pool = NSAutoreleasePool.alloc().init() session = None @@ -65,10 +71,10 @@ class Capture(NSObject): session = None script = None error = e - AppHelper.callAfter(self.attachDidCompleteWithSession_script_error_, session, script, error) + AppHelper.callAfter(self._attachDidCompleteWithSession_script_error_, session, script, error) del pool - def doDetachWithParams_(self, params): + def _doDetachWithParams_(self, params): session, script = params pool = NSAutoreleasePool.alloc().init() try: @@ -81,42 +87,52 @@ class Capture(NSObject): pass del pool - def attachDidCompleteWithSession_script_error_(self, session, script, error): + def _doPostWithParams_(self, params): + script, message = params + pool = NSAutoreleasePool.alloc().init() + try: + script.post_message(message) + except Exception, e: + print "Failed to post to script:", e + del pool + + def _attachDidCompleteWithSession_script_error_(self, session, script, error): if self.state == CaptureState.ATTACHING: self.session = session self.script = script if error is None: - self.updateState_(CaptureState.ATTACHED) + self._updateState_(CaptureState.ATTACHED) else: - self.updateState_(CaptureState.DETACHED) + self._updateState_(CaptureState.DETACHED) self.delegate.captureFailedToAttachWithError_(error) - def sessionDidDetach(self): + def _sessionDidDetach(self): if self.state == CaptureState.ATTACHING or self.state == CaptureState.ATTACHED: self.session = None - self.updateState_(CaptureState.DETACHED) + self._updateState_(CaptureState.DETACHED) - def sessionDidReceiveMessage_data_(self, message, data): + def _sessionDidReceiveMessage_data_(self, message, data): if message['type'] == 'send': stanza = message['payload'] - name = stanza['name'] fromAddress = stanza['from'] + name = stanza['name'] if fromAddress == "/process/modules" and name == '+sync': - self.modules.sync(stanza['payload']) + self.modules._sync(stanza['payload']) elif fromAddress == "/stalker/events" and name == '+add': - self.calls.add_(data) + self.calls._add_(data) elif fromAddress == "/interceptor/functions" and name == '+add': pass else: - print "Woot! Got stanza: %s from=%s" % (stanza['name'], stanza['from']) + if not self.calls._handleStanza_(stanza): + print "Woot! Got stanza: %s from=%s" % (stanza['name'], stanza['from']) else: print "Unhandled message:", message def _onSessionDetached(self): - AppHelper.callAfter(self.sessionDidDetach) + AppHelper.callAfter(self._sessionDidDetach) def _onScriptMessage(self, message, data): - AppHelper.callAfter(self.sessionDidReceiveMessage_data_, message, data) + AppHelper.callAfter(self._sessionDidReceiveMessage_data_, message, data) class CaptureState: DETACHED = 1 @@ -128,7 +144,7 @@ class Modules: self._modules = [] self._indices = [] - def sync(self, payload): + def _sync(self, payload): modules = [] for item in payload['items']: modules.append(Module(item['name'], int(item['address'], 16), item['size'])) @@ -155,15 +171,16 @@ class Module: return "(%d, %d, %s)" % (self.address, self.size, self.name) class Calls(NSObject): - def __new__(cls, modules): - return cls.alloc().initWithModules_(modules) + def __new__(cls, capture): + return cls.alloc().initWithCapture_(capture) - def initWithModules_(self, modules): + def initWithCapture_(self, capture): self = self.init() - self.modules = modules + self.capture = capture self.targetModules = [] self._targetModuleByAddress = {} self._delegate = None + self._probes = {} return self def delegate(self): @@ -172,19 +189,62 @@ class Calls(NSObject): def setDelegate_(self, delegate): self._delegate = delegate - def add_(self, data): + def addProbe_(self, func): + self.capture._post({ + 'to': "/stalker/probes", + 'name': '+add', + 'payload': { + 'address': "0x%x" % func.address + } + }) + self._probes[func.address] = func + + def removeProbe_(self, func): + self.capture._post({ + 'to': "/stalker/probes", + 'name': '+remove', + 'payload': { + 'address': "0x%x" % func.address + } + }) + self._probes.pop(func.address, None) + + def _add_(self, data): + modules = self.capture.modules for offset in range(0, len(data), 16): [t, location, target, depth] = struct.unpack("IIII", data[offset:offset + 16]) - tm = self.getTargetModuleByModule_(self.modules.lookup(target)) + tm = self.getTargetModuleByModule_(modules.lookup(target)) if tm is not None: tm.total += 1 tf = tm.getTargetFunctionByAddress_(target) tf.total += 1 self.targetModules.sort(key=lambda tm: tm.total, reverse=True) for tm in self.targetModules: - tm.functions.sort(key=lambda f: f.total, reverse=True) + tm.functions.sort(self._compareFunctions) self._delegate.callsDidChange() + def _compareFunctions(self, x, y): + if x.hasProbe == y.hasProbe: + return x.total - y.total + elif x.hasProbe: + return -1 + elif y.hasProbe: + return 1 + else: + return x.total - y.total + + def _handleStanza_(self, stanza): + m = PROBE_CALLS.match(stanza['from']) + if m is not None: + func = self._probes.get(int(m.groups()[0], 16), None) + if func is not None: + if len(func.calls) == 3: + func.calls.pop(0) + func.calls.append(FunctionCall(func, stanza['payload']['args'])) + self._delegate.callItemDidChange_(func) + return True + return False + def getTargetModuleByModule_(self, module): if module is None: return None @@ -200,6 +260,8 @@ class Calls(NSObject): return len(self.targetModules) elif isinstance(item, TargetModule): return len(item.functions) + elif isinstance(item, TargetFunction): + return len(item.calls) else: return 0 @@ -208,6 +270,8 @@ class Calls(NSObject): return False elif isinstance(item, TargetModule): return len(item.functions) > 0 + elif isinstance(item, TargetFunction): + return len(item.calls) > 0 else: return False @@ -216,20 +280,34 @@ class Calls(NSObject): return self.targetModules[index] elif isinstance(item, TargetModule): return item.functions[index] + elif isinstance(item, TargetFunction): + return item.calls[index] else: return None def outlineView_objectValueForTableColumn_byItem_(self, outlineView, tableColumn, item): + identifier = tableColumn.identifier() if isinstance(item, TargetModule): - if tableColumn.identifier() == 'name': + if identifier == 'name': return item.module.name - else: + elif identifier == 'total': return item.total - else: - if tableColumn.identifier() == 'name': + else: + return False + elif isinstance(item, TargetFunction): + if identifier == 'name': return item.name - else: + elif identifier == 'total': return item.total + else: + return item.hasProbe + else: + if identifier == 'name': + return item.summary + elif identifier == 'total': + return "" + else: + return False class TargetModule(NSObject): def __new__(cls, module): @@ -255,25 +333,76 @@ class TargetFunction(NSObject): def __new__(cls, module, offset): return cls.alloc().initWithModule_offset_(module, offset) - def initWithModule_offset_(self, module, offset): + def initWithModule_offset_(self, targetModule, offset): self = self.init() self.name = "sub_%x" % offset - self.module = module + self.module = targetModule + self.address = targetModule.module.address + offset self.offset = offset self.total = 0 + self.hasProbe = False + self.calls = [] + return self + +class FunctionCall(NSObject): + def __new__(cls, func, args): + return cls.alloc().initWithFunction_args_(func, args) + + def initWithFunction_args_(self, func, args): + self = self.init() + self.func = func + self.args = args + self.summary = "%s(%s)" % (func.name, ", ".join(args)) return self SCRIPT_TEMPLATE = """ -Stalker.trustThreshold = 2000; -Stalker.queueCapacity = 1000000; -Stalker.queueDrainInterval = 50; +var probes = Object.create(null); var initialize = function initialize() { + Stalker.trustThreshold = 2000; + Stalker.queueCapacity = 1000000; + Stalker.queueDrainInterval = 50; + sendModules(function () { interceptReadFunction('recv'); interceptReadFunction('read$UNIX2003'); interceptReadFunction('readv$UNIX2003'); }); + + recv(onStanza); +}; + +var onStanza = function onStanza(stanza) { + if (stanza.to === "/stalker/probes") { + var address = stanza.payload.address, + probeId; + switch (stanza.name) { + case '+add': + if (probes[address] === undefined) { + var probeAddress = "/stalker/probes/" + address + "/calls"; + probeId = Stalker.addCallProbe(ptr(address), function probe(args) { + var data = [ + "0x" + args[0].toString(16), + "0x" + args[1].toString(16), + "0x" + args[2].toString(16), + "0x" + args[3].toString(16) + ]; + send({ from: probeAddress, name: '+add', payload: { args: data } }); + }); + probes[address] = probeId; + } + break; + case '+remove': + probeId = probes[address]; + if (probeId !== undefined) { + Stalker.removeCallProbe(probeId); + delete probes[address]; + } + break; + } + } + + recv(onStanza); }; var sendModules = function sendModules(callback) { diff --git a/examples/cpushark/MainWindow.xib b/examples/cpushark/MainWindow.xib index c97236d..5f4fab7 100644 --- a/examples/cpushark/MainWindow.xib +++ b/examples/cpushark/MainWindow.xib @@ -135,7 +135,7 @@ DQ - + @@ -155,8 +155,8 @@ DQ - - + + @@ -168,13 +168,28 @@ DQ + + + + + + + + + + + + + + +