310 lines
11 KiB
Python
Executable File
310 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from typing import Optional, TypedDict
|
|
from socketserver import UnixStreamServer, StreamRequestHandler, ThreadingMixIn
|
|
import argparse
|
|
import os
|
|
import re
|
|
|
|
|
|
type Pointer[T] = tuple[int, T]
|
|
type Flags = tuple[int, list[str]]
|
|
StructTimeSpec = TypedDict('StructTimeSpec', {'tv_sec': int, 'tv_nsec': int})
|
|
|
|
|
|
class ThreadedUnixStreamServer(ThreadingMixIn, UnixStreamServer):
|
|
pass
|
|
|
|
|
|
class Handler(StreamRequestHandler):
|
|
pid: int
|
|
stack: list[tuple[str, tuple]]
|
|
|
|
def before(self) -> None:
|
|
pass
|
|
|
|
def after(self) -> None:
|
|
pass
|
|
|
|
def handle(self):
|
|
first = self.rfile.readline()
|
|
self.pid = int(first.split(b':')[1])
|
|
self.stack = []
|
|
print(f'Process with PID {self.pid} connected')
|
|
self.before()
|
|
try:
|
|
while True:
|
|
msg = self.rfile.readline()
|
|
if not msg:
|
|
return
|
|
self.handle_msg(msg)
|
|
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] == '|':
|
|
m = re.match(r'^[| A-Za-z0-9_]*', argument[idx:])
|
|
flags = m.group(0)
|
|
if not flags.startswith('|') or not flags.endswith('|'):
|
|
raise ValueError()
|
|
idx += len(flags)
|
|
if idx < len(argument) and argument[idx] == ',':
|
|
idx += 1
|
|
flags = [f.strip() for f in flags[1:-1].split('|') if len(f.strip()) > 0]
|
|
return (val, flags), 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
|
|
elif argument[idx] == '{':
|
|
m = re.match(r'^[^}]*', argument[idx:])
|
|
value = m.group(0)
|
|
if not value.startswith('{') or not value.endswith('}'):
|
|
raise ValueError()
|
|
idx += len(value)
|
|
if idx < len(argument) and argument[idx] == ',':
|
|
idx += 1
|
|
entries = {}
|
|
for e in [v.strip() for v in value[1:-1].split(',') if len(e.strip()) > 0]:
|
|
k, v = e.split(':', 1)
|
|
entries[k.strip()] = int(v.strip(), 0)
|
|
return (val, entries), 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, _ = 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) or 'ok'
|
|
else:
|
|
command = 'ok'
|
|
except AttributeError:
|
|
command = 'ok'
|
|
print(f'[{self.pid}] -> {command}')
|
|
self.wfile.write(command.encode('utf-8') + b'\n')
|
|
else:
|
|
ret = data.decode('utf-8')
|
|
ret_value, _ = Handler.parse_arg(ret[7:])
|
|
func_name, args = self.stack.pop()
|
|
try:
|
|
func = getattr(self, f'after_{func_name}')
|
|
if callable(func):
|
|
if ret_value is None:
|
|
func(*args)
|
|
else:
|
|
func(*args, ret_value)
|
|
except AttributeError:
|
|
pass
|
|
print(f'[{self.pid}] -> {ret}')
|
|
|
|
def before_malloc(self, size: int) -> str: pass
|
|
def after_malloc(self, size: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_calloc(self, nmemb: int, size: int) -> str: pass
|
|
def after_calloc(self, nmemb: int, size: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_realloc(self, ptr: int, size: int) -> str: pass
|
|
def after_realloc(self, ptr: int, size: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_reallocarray(self, ptr: int, nmemb: int, size: int) -> str: pass
|
|
def after_reallocarray(self, ptr: int, nmemb: int, size: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_free(self, ptr: int) -> str: pass
|
|
def after_free(self, ptr: int) -> None: pass
|
|
def before_getopt(self, argc: int, argv: Pointer[list[Pointer[bytes]]], optstring: Pointer[bytes]) -> str: pass
|
|
def after_getopt(self, argc: int, argv: Pointer[list[Pointer[bytes]]], optstring: Pointer[bytes],
|
|
ret_value: int) -> None: pass
|
|
def before_close(self, fildes: int) -> str: pass
|
|
def after_close(self, fildes: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_sem_init(self, sem: int, pshared: int, value: int) -> str: pass
|
|
def after_sem_init(self, sem: int, pshared: int, value: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_sem_open(self, name: str, oflag: Flags, mode: Optional[int], value: Optional[int]) -> str: pass
|
|
def after_sem_open(self, name: str, oflag: Flags, mode: Optional[int], value: Optional[int],
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_sem_post(self, sem: int) -> str: pass
|
|
def after_sem_post(self, sem: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_sem_wait(self, sem: int) -> str: pass
|
|
def after_sem_wait(self, sem: int,
|
|
ret_value: int, errno: str = None) -> None: pass
|
|
def before_sem_timedwait(self, sem: int, abs_timeout: Pointer[StructTimeSpec]): pass
|
|
def after_sem_timedwait(self, sem: int, abs_timeout: Pointer[StructTimeSpec],
|
|
ret_value: int, errno: str = None): pass
|
|
|
|
|
|
class MemoryAllocationTester(Handler):
|
|
allocated: dict[int, int]
|
|
max_allocated: int
|
|
num_malloc: int
|
|
num_realloc: int
|
|
num_free: int
|
|
|
|
def before(self):
|
|
self.allocated = {}
|
|
self.max_allocated = 0
|
|
self.num_malloc = 0
|
|
self.num_realloc = 0
|
|
self.num_free = 0
|
|
|
|
def after(self):
|
|
if len(self.allocated) > 0:
|
|
print("Not free'd:")
|
|
for ptr, size in self.allocated.items():
|
|
print(f' 0x{ptr:x}: {size} bytes')
|
|
else:
|
|
print("All blocks free'd!")
|
|
print(f'Max allocated: {self.max_allocated} bytes')
|
|
|
|
def update_max_allocated(self):
|
|
total = sum(self.allocated.values())
|
|
if total > self.max_allocated:
|
|
self.max_allocated = total
|
|
|
|
def after_malloc(self, size, ret_value, errno=None) -> None:
|
|
self.num_malloc += 1
|
|
if ret_value != 0:
|
|
print(ret_value)
|
|
self.allocated[ret_value] = size
|
|
self.update_max_allocated()
|
|
|
|
def after_calloc(self, nmemb, size, ret_value, errno=None) -> None:
|
|
self.num_malloc += 1
|
|
if ret_value != 0:
|
|
self.allocated[ret_value] = nmemb * size
|
|
self.update_max_allocated()
|
|
|
|
def after_realloc(self, ptr, size, ret_value, errno=None) -> None:
|
|
self.num_realloc += 1
|
|
if ptr != 0:
|
|
if ret_value != 0:
|
|
del self.allocated[ptr]
|
|
self.allocated[ret_value] = size
|
|
self.update_max_allocated()
|
|
|
|
def after_reallocarray(self, ptr, nmemb, size, ret_value, errno=None) -> None:
|
|
self.num_realloc += 1
|
|
if ptr != 0:
|
|
if ret_value != 0:
|
|
del self.allocated[ptr]
|
|
self.allocated[ret_value] = nmemb * size
|
|
self.update_max_allocated()
|
|
|
|
def after_free(self, ptr) -> None:
|
|
self.num_free += 1
|
|
if ptr != 0:
|
|
del self.allocated[ptr]
|
|
|
|
|
|
def intercept(socket: str, handler: type[Handler]) -> None:
|
|
try:
|
|
with ThreadedUnixStreamServer(socket, handler) as server:
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
print('\nBye')
|
|
server.shutdown()
|
|
finally:
|
|
try:
|
|
os.unlink(socket)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('socket', metavar='FILE')
|
|
args = parser.parse_args()
|
|
intercept(args.socket, Handler)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|