intercept.py: Properly parse function arguments
This commit is contained in:
2
proj/server/.gitignore
vendored
Normal file
2
proj/server/.gitignore
vendored
Normal file
@@ -0,0 +1,2 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
@@ -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)':
|
||||
size = int(args[0], 0)
|
||||
self.allocated[int(ret_value, 0)] = size
|
||||
if ret_value != 0:
|
||||
print(ret_value)
|
||||
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)':
|
||||
size = int(args[0], 0) * int(args[1], 0)
|
||||
self.allocated[int(ret_value, 0)] = size
|
||||
if ret_value != 0:
|
||||
self.allocated[ret_value] = nmemb * 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)':
|
||||
new_size = int(args[1], 0)
|
||||
if ret_value != '(nil)':
|
||||
del self.allocated[int(args[0], 0)]
|
||||
self.allocated[int(ret_value, 0)] = new_size
|
||||
if ptr != 0:
|
||||
if ret_value != 0:
|
||||
del self.allocated[ptr]
|
||||
self.allocated[ret_value] = 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)':
|
||||
del self.allocated[int(args[0], 0)]
|
||||
if ptr != 0:
|
||||
del self.allocated[ptr]
|
||||
|
||||
|
||||
def intercept(socket: str, handler: type[Handler]) -> None:
|
||||
|
||||
@@ -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:
|
||||
if value != '(nil)':
|
||||
self.allocated[int(value, 0)] = int(args[0], 0)
|
||||
def after_malloc(self, size: int, ret_value: int) -> None:
|
||||
if ret_value != 0:
|
||||
self.allocated[ret_value] = size
|
||||
|
||||
def before_free(self, args: list[str]) -> str:
|
||||
if args[0] != '(nil)':
|
||||
del self.allocated[int(args[0], 0)]
|
||||
def before_free(self, ptr: int) -> str:
|
||||
if ptr != 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__':
|
||||
|
||||
Reference in New Issue
Block a user