Improve cpushark example to also cover call probes

2026-06-08 14:16:17 +00:00 · 2013-10-18 00:04:59 +02:00
parent 5ac817a742
commit d5873eb448
3 changed files with 202 additions and 44 deletions
@@ -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) {
@@ -135,7 +135,7 @@ DQ
                            <rect key="frame" x="1" y="17" width="498" height="172"/>
                            <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
                            <subviews>
-                                <outlineView verticalHuggingPriority="750" allowsExpansionToolTips="YES" columnAutoresizingStyle="lastColumnOnly" multipleSelection="NO" autosaveColumns="NO" headerView="poO-79-hdg" indentationPerLevel="16" outlineTableColumn="98s-hc-Ws2" id="ouC-rl-JvY">
+                                <outlineView verticalHuggingPriority="750" allowsExpansionToolTips="YES" columnAutoresizingStyle="lastColumnOnly" multipleSelection="NO" autosaveColumns="NO" headerView="poO-79-hdg" indentationPerLevel="16" autoresizesOutlineColumn="YES" outlineTableColumn="98s-hc-Ws2" id="ouC-rl-JvY">
                                    <rect key="frame" x="0.0" y="0.0" width="498" height="172"/>
                                    <autoresizingMask key="autoresizingMask"/>
                                    <size key="intercellSpacing" width="3" height="2"/>
@@ -155,8 +155,8 @@ DQ
                                            </textFieldCell>
                                            <tableColumnResizingMask key="resizingMask" resizeWithTable="YES" userResizable="YES"/>
                                        </tableColumn>
-                                        <tableColumn identifier="total" width="391" minWidth="40" maxWidth="1000" id="JdC-cM-gLh">
-                                            <tableHeaderCell key="headerCell" lineBreakMode="truncatingTail" borderStyle="border" alignment="left">
+                                        <tableColumn identifier="total" width="100" minWidth="40" maxWidth="1000" id="JdC-cM-gLh">
+                                            <tableHeaderCell key="headerCell" lineBreakMode="truncatingTail" borderStyle="border" alignment="left" title="Total calls">
                                                <font key="font" metaFont="smallSystem"/>
                                                <color key="textColor" name="headerTextColor" catalog="System" colorSpace="catalog"/>
                                                <color key="backgroundColor" white="0.33333298560000002" alpha="1" colorSpace="calibratedWhite"/>
@@ -168,13 +168,28 @@ DQ
                                            </textFieldCell>
                                            <tableColumnResizingMask key="resizingMask" resizeWithTable="YES" userResizable="YES"/>
                                        </tableColumn>
+                                        <tableColumn identifier="action" width="64" minWidth="10" maxWidth="3.4028234663852886e+38" id="8sf-9V-Hq9">
+                                            <tableHeaderCell key="headerCell" lineBreakMode="truncatingTail" borderStyle="border" alignment="left" title="Has probe">
+                                                <font key="font" metaFont="smallSystem"/>
+                                                <color key="textColor" name="headerTextColor" catalog="System" colorSpace="catalog"/>
+                                                <color key="backgroundColor" name="headerColor" catalog="System" colorSpace="catalog"/>
+                                            </tableHeaderCell>
+                                            <buttonCell key="dataCell" type="check" bezelStyle="regularSquare" imagePosition="left" alignment="left" inset="2" id="FcW-HG-7vy">
+                                                <behavior key="behavior" changeContents="YES" doesNotDimImage="YES" lightByContents="YES"/>
+                                                <font key="font" metaFont="system"/>
+                                                <connections>
+                                                    <action selector="toggleTracing:" target="-2" id="7dt-Hn-GuQ"/>
+                                                </connections>
+                                            </buttonCell>
+                                            <tableColumnResizingMask key="resizingMask" resizeWithTable="YES" userResizable="YES"/>
+                                        </tableColumn>
                                    </tableColumns>
                                </outlineView>
                            </subviews>
                            <color key="backgroundColor" name="controlBackgroundColor" catalog="System" colorSpace="catalog"/>
                        </clipView>
                        <scroller key="horizontalScroller" hidden="YES" verticalHuggingPriority="750" id="YYd-vO-1l4">
-                            <rect key="frame" x="1" y="1" width="0.0" height="16"/>
+                            <rect key="frame" x="1" y="173" width="498" height="16"/>
                            <autoresizingMask key="autoresizingMask"/>
                        </scroller>
                        <scroller key="verticalScroller" hidden="YES" verticalHuggingPriority="750" doubleValue="1" id="vbI-E6-0lE">
@@ -1,5 +1,5 @@
 import frida
-from Capture import Capture, CaptureState
+from Capture import Capture, CaptureState, TargetFunction
 from Cocoa import NSRunCriticalAlertPanel, NSUserDefaults, NSWindowController, objc
 from ProcessList import ProcessList

@@ -80,6 +80,17 @@ class MainWindowController(NSWindowController):
        self.capture.detach()

    @objc.IBAction
+    def toggleTracing_(self, sender):
+        item = sender.itemAtRow_(sender.selectedRow())
+        if isinstance(item, TargetFunction):
+            func = item
+            if func.hasProbe:
+                self.capture.calls.removeProbe_(func)
+            else:
+                self.capture.calls.addProbe_(func)
+            func.hasProbe = not func.hasProbe
+            self.callTableView.reloadItem_(func)
+
    def updateAttachForm_(self, sender):
        isDetached = self.capture.state == CaptureState.DETACHED
        hasProcess = self.selectedProcess() is not None
@@ -110,3 +121,6 @@ class MainWindowController(NSWindowController):
    def callsDidChange(self):
        self.callTableView.reloadData()

+    def callItemDidChange_(self, item):
+        self.callTableView.reloadItem_reloadChildren_(item, True)
+