intercept.py: Properly parse function arguments

proj/server/.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

proj/server/src/intercept.py

@@ -4,6 +4,7 @@
 from socketserver import UnixStreamServer, StreamRequestHandler, ThreadingMixIn
 import argparse
 import os
+import re
 
 
 class ThreadedUnixStreamServer(ThreadingMixIn, UnixStreamServer):
@@ -12,7 +13,7 @@ class ThreadedUnixStreamServer(ThreadingMixIn, UnixStreamServer):
 
 class Handler(StreamRequestHandler):
     pid: int
-    stack: list[tuple[str, list[str]]]
+    stack: list[tuple[str, tuple]]
 
     def before(self) -> None:
         pass
@@ -35,18 +36,104 @@ class Handler(StreamRequestHandler):
         finally:
             self.after()
 
+    @staticmethod
+    def parse_str(argument: str) -> tuple[str or bytes, int]:
+        if not ((len(argument) >= 2 and argument[0] == '"') or (len(argument) >= 3 and argument[0] == 'b' and argument[1] == '"')):
+            raise ValueError()
+        idx = 1
+        esc, fin = False, False
+        data = b'' if argument[0] == 'b' else ''
+        tmp = None
+        for ch in argument[1:]:
+            idx += 1
+            if fin:
+                if ch in (' ', '\t'):
+                    continue
+                elif ch in (',', ']'):
+                    idx -= 1
+                    break
+            elif tmp:
+                tmp += ch
+                if len(tmp) == 2:
+                    data += bytes([int(tmp, 16)]) if argument[0] == 'b' else chr(int(tmp, 16))
+                    tmp = None
+            elif esc:
+                if ch in ('\\', '"'):
+                    data += ch.encode('ascii') if argument[0] == 'b' else ch
+                elif ch == 'x':
+                   tmp = ''
+                   esc = False
+                else:
+                    raise ValueError()
+            elif ch == '"':
+                fin = True
+            elif ch == '\\':
+                esc = True
+            else:
+                data += ch.encode('utf-8') if argument[0] == 'b' else ch
+        if not fin:
+            raise ValueError()
+        return data, idx
+
+    @staticmethod
+    def parse_arg(argument: str) -> tuple[any, int]:
+        if argument == '':
+            return None, 0
+        m = re.match(r'\s*\(nil\)', argument)
+        if m:
+            return 0, len(m.group(0))
+        m = re.match(r'^\s*(.*?)([,:]|$)', argument)
+        a, e = m.group(1), m.group(2)
+        idx = len(m.group(0))
+        if a.startswith('0x'):
+            val = int(a[2:], 16)
+        elif a.startswith('0'):
+            val = int(a[1:], 8)
+        else:
+            val = int(a, 10)
+        if e in (',', ''):
+            return val, idx
+        if argument[idx] == '[':
+            idx += 1
+            l, i = Handler.parse_args(argument[idx:])
+            idx += i
+            if idx < len(argument) and argument[idx] == ',':
+                idx += 1
+            return (val, list(l)), idx
+        elif argument[idx] == '"':
+            s, i = Handler.parse_str(argument[idx:])
+            idx += i
+            if idx < len(argument) and argument[idx] == ',':
+                idx += 1
+            return (val, s), idx
+        else:
+            raise ValueError()
+
+    @staticmethod
+    def parse_args(arguments: str) -> tuple[tuple, int]:
+        args = []
+        idx = 0
+        while idx < len(arguments):
+            if arguments[idx] == ']':
+                idx += 1
+                break
+            val, i = Handler.parse_arg(arguments[idx:])
+            args.append(val)
+            idx += i
+        return tuple(args), idx
+
     def handle_msg(self, msg: bytes):
         timestamp, data = msg.rstrip(b'\n').split(b' ', 1)
         if not data.startswith(b'return ') and not data == b'return':
             call = data.decode('utf-8')
             print(f'[{self.pid}] {call}')
             func_name = call[:call.find('(')]
-            args = call[call.find('(') + 1:-1].split(', ')
+            args, _ = Handler.parse_args(call[call.find('(') + 1:-1])
             self.stack.append((func_name, args))
             try:
                 func = getattr(self, f'before_{func_name}')
                 if callable(func):
-                    command = func(args)
+                    command = func(*args)
                 else:
                     command = 'ok'
             except AttributeError:
@@ -55,12 +142,15 @@ class Handler(StreamRequestHandler):
             self.wfile.write(command.encode('utf-8') + b'\n')
         else:
             ret = data.decode('utf-8')
-            ret_value = ret[7:]
+            ret_value, _ = Handler.parse_arg(ret[7:])
             func_name, args = self.stack.pop()
             try:
                 func = getattr(self, f'after_{func_name}')
                 if callable(func):
-                    func(args, ret_value)
+                    if ret_value is None:
+                        func(*args)
+                    else:
+                        func(*args, ret_value)
             except AttributeError:
                 pass
             print(f'[{self.pid}]  -> {ret}')
@@ -94,33 +184,31 @@ class MemoryAllocationTester(Handler):
         if total > self.max_allocated:
             self.max_allocated = total
 
-    def after_malloc(self, args: list[str], ret_value: str) -> None:
+    def after_malloc(self, size: int, ret_value: int) -> None:
         self.num_malloc += 1
-        if ret_value != '(nil)':
+        if ret_value != 0:
-            size = int(args[0], 0)
+            print(ret_value)
-            self.allocated[int(ret_value, 0)] = size
+            self.allocated[ret_value] = size
             self.update_max_allocated()
 
-    def after_calloc(self, args: list[str], ret_value: str) -> None:
+    def after_calloc(self, nmemb: int, size: int, ret_value: int) -> None:
         self.num_malloc += 1
-        if ret_value != '(nil)':
+        if ret_value != 0:
-            size = int(args[0], 0) * int(args[1], 0)
+            self.allocated[ret_value] = nmemb * size
-            self.allocated[int(ret_value, 0)] = size
             self.update_max_allocated()
 
-    def after_realloc(self, args: list[str], ret_value: str) -> None:
+    def after_realloc(self, ptr: int, size: int, ret_value: int) -> None:
         self.num_realloc += 1
-        if args[0] != '(nil)':
+        if ptr != 0:
-            new_size = int(args[1], 0)
+            if ret_value != 0:
-            if ret_value != '(nil)':
+                del self.allocated[ptr]
-                del self.allocated[int(args[0], 0)]
+                self.allocated[ret_value] = size
-                self.allocated[int(ret_value, 0)] = new_size
                 self.update_max_allocated()
 
-    def after_free(self, args: list[str], ret_value: str) -> None:
+    def after_free(self, ptr: int) -> None:
         self.num_free += 1
-        if args[0] != '(nil)':
+        if ptr != 0:
-            del self.allocated[int(args[0], 0)]
+            del self.allocated[ptr]
 
 
 def intercept(socket: str, handler: type[Handler]) -> None:

proj/server/src/server.py

@@ -6,7 +6,7 @@ import argparse
 import intercept
 
 
-class Handler(intercept.Handler):
+class TestHandler(intercept.Handler):
     allocated: dict[int, int]
 
     def before(self):
@@ -20,21 +20,26 @@ class Handler(intercept.Handler):
         else:
             print("All blocks free'd!")
 
-    def after_malloc(self, args: list[str], value: str) -> None:
+    def after_malloc(self, size: int, ret_value: int) -> None:
-        if value != '(nil)':
+        if ret_value != 0:
-            self.allocated[int(value, 0)] = int(args[0], 0)
+            self.allocated[ret_value] = size
 
-    def before_free(self, args: list[str]) -> str:
+    def before_free(self, ptr: int) -> str:
-        if args[0] != '(nil)':
+        if ptr != 0:
-            del self.allocated[int(args[0], 0)]
+            del self.allocated[ptr]
         return 'ok'
 
 
+class GetoptHandler(intercept.Handler):
+    def after_getopt(self, argc: int, argv: tuple[int, list[tuple[int, bytes]]], optstring: tuple[int, bytes], ret_value: int) -> None:
+        print(argc, [v[1] for v in argv[1]], optstring[1], ret_value)
+
+
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument('socket', metavar='FILE')
     args = parser.parse_args()
-    intercept.intercept(args.socket, intercept.MallocTester)
+    intercept.intercept(args.socket, intercept.MemoryAllocationTester)
 
 
 if __name__ == '__main__':