#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from typing import Optional, TypedDict, NotRequired
from socketserver import UnixStreamServer, StreamRequestHandler, ThreadingMixIn
import os
import re


type Pointer = int
type PointerTo[T] = tuple[Pointer, T]
type Constant = tuple[int, str]
type Flags = tuple[int, list[str]]
StructTimeSpec = TypedDict('StructTimeSpec', {'tv_sec': int, 'tv_nsec': int})
StructSigAction = TypedDict('StructSigAction', {'sa_flags': Flags, 'sa_handler': NotRequired[Pointer], 'sa_sigaction': NotRequired[Pointer], 'sa_mask': list[str]})
StructSockAddr = TypedDict('StructSockAddr', {'sa_family': Constant, 'sa_data': NotRequired[bytes],
                                              'sun_path': NotRequired[bytes],
                                              'sin_addr': NotRequired[bytes], 'sin_port': NotRequired[int],
                                              'sin6_addr': NotRequired[bytes], 'sin6_port': NotRequired[int], 'sin6_scope_id': NotRequired[int]})
StructAddrInfo = TypedDict('StructAddrInfo', {})
StructMsgHdr = TypedDict('StructMsgHdr', {})


class ThreadedUnixStreamServer(ThreadingMixIn, UnixStreamServer):
    pass


class Handler(StreamRequestHandler):
    pid: Optional[int]
    path: Optional[str]
    stack: list[tuple[int, str, tuple]]
    ret_addr: int

    def before(self) -> None: pass
    def after(self) -> None: pass
    def before_fallback(self, func_name: str, *args) -> str: pass
    def after_fallback(self, func_name: str, *args, **kwargs) -> None: pass

    def handle(self):
        first = self.rfile.readline()
        meta = {a[0]: a[1] for a in [tuple(p.decode('utf-8').split(':', 1)) for p in first.split(b' ', 3)[3].strip().split(b';')]}
        self.pid = int(meta['PID']) if 'PID' in meta else None
        self.path = meta['PATH'] if 'PATH' in meta else None
        self.stack = []
        print(f'Process with PID {self.pid} connected ({self.path})')
        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\)|NULL|null|nullptr)\s*(,|$)', 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') and len(a) > 1:
            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, pid, tid, data = msg.rstrip(b'\n').split(b' ', 3)
        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('(')]
            self.ret_addr = int(call[call.rfind(':') + 1:], 0)
            args, _ = Handler.parse_args(call[call.find('(') + 1:call.rfind(':') - 1])
            self.stack.append((self.ret_addr, func_name, args))
            try:
                func = getattr(self, f'before_{func_name}')
                if not callable(func):
                    func = None
            except AttributeError:
                func = None
            try:
                if func is None:
                    raise NotImplementedError()
                command = func(*args) or self.before_fallback(func_name, *args) or 'ok'
            except NotImplementedError:
                command = self.before_fallback(func_name, *args) or '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:].split(';')[0])
            self.ret_addr, func_name, args = self.stack.pop()
            try:
                func = getattr(self, f'after_{func_name}')
                if not callable(func):
                    func = None
            except AttributeError:
                func = None
            try:
                if func is None:
                    raise NotImplementedError()
                if ret_value is None:
                    func(*args)
                else:
                    func(*args, ret_value)
            except NotImplementedError:
                if ret_value is None:
                    self.after_fallback(func_name, *args)
                else:
                    self.after_fallback(func_name, *args, ret_value)
            print(f'[{self.pid}]  -> {ret}')

    def before_malloc(self, size: int) -> str:
        raise NotImplementedError()
    def after_malloc(self, size: int,
                     ret_value: Pointer, errno: str = None) -> None:
        raise NotImplementedError()
    def before_calloc(self, nmemb: int, size: int) -> str:
        raise NotImplementedError()
    def after_calloc(self, nmemb: int, size: int,
                     ret_value: Pointer, errno: str = None) -> None:
        raise NotImplementedError()
    def before_realloc(self, ptr: Pointer, size: int) -> str:
        raise NotImplementedError()
    def after_realloc(self, ptr: Pointer, size: int,
                      ret_value: Pointer, errno: str = None) -> None:
        raise NotImplementedError()
    def before_reallocarray(self, ptr: Pointer, nmemb: int, size: int) -> str:
        raise NotImplementedError()
    def after_reallocarray(self, ptr: Pointer, nmemb: int, size: int,
                           ret_value: Pointer, errno: str = None) -> None:
        raise NotImplementedError()
    def before_free(self, ptr: Pointer) -> str:
        raise NotImplementedError()
    def after_free(self, ptr: Pointer) -> None:
        raise NotImplementedError()
    def before_getopt(self, argc: int, argv: PointerTo[list[PointerTo[bytes]]], optstring: PointerTo[bytes]) -> str:
        raise NotImplementedError()
    def after_getopt(self, argc: int, argv: PointerTo[list[PointerTo[bytes]]], optstring: PointerTo[bytes],
                     ret_value: int) -> None:
        raise NotImplementedError()
    def before_exit(self, status: int) -> str:
        raise NotImplementedError()
    def after_exit(self, status: int) -> None:
        raise NotImplementedError()
    def before_read(self, fildes: int, buf_ptr: Pointer, nbyte: int) -> str:
        raise NotImplementedError()
    def after_read(self, fildes: int, buf_ptr: Pointer, nbyte: int,
                   ret_value: int, errno: str = None, buf: PointerTo[bytes] = None) -> None:
        raise NotImplementedError()
    def before_pread(self, fildes: int, buf_ptr: Pointer, nbyte: int, offset: int) -> str:
        raise NotImplementedError()
    def after_pread(self, fildes: int, buf_ptr: Pointer, nbyte: int, offset: int,
                    ret_value: int, errno: str = None, buf: PointerTo[bytes] = None) -> None:
        raise NotImplementedError()
    def before_write(self, fildes: int, buf: PointerTo[bytes], nbyte: int) -> str:
        raise NotImplementedError()
    def after_write(self, fildes: int, buf: PointerTo[bytes], nbyte: int,
                    ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_pwrite(self, fildes: int, buf: PointerTo[bytes], nbyte: int, offste: int) -> str:
        raise NotImplementedError()
    def after_pwrite(self, fildes: int, buf: PointerTo[bytes], nbyte: int, offste: int,
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_close(self, fildes: int) -> str:
        raise NotImplementedError()
    def after_close(self, fildes: int,
                    ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sigaction(self, sig: Constant, act: PointerTo[StructSigAction], oact_ptr: Pointer) -> str:
        raise NotImplementedError()
    def after_sigaction(self, sig: Constant, act: PointerTo[StructSigAction], oact_ptr: Pointer,
                        ret_value: int, errno: str = None, oact: StructSigAction = None) -> None:
        raise NotImplementedError()
    def before_sem_init(self, sem: Pointer, pshared: int, value: int) -> str:
        raise NotImplementedError()
    def after_sem_init(self, sem: Pointer, pshared: int, value: int,
                       ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_open(self, name: str, oflag: Flags, mode: Optional[int], value: Optional[int]) -> str:
        raise NotImplementedError()
    def after_sem_open(self, name: str, oflag: Flags, mode: Optional[int], value: Optional[int],
                       ret_value: Pointer, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_post(self, sem: Pointer) -> str:
        raise NotImplementedError()
    def after_sem_post(self, sem: Pointer,
                       ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_wait(self, sem: Pointer) -> str:
        raise NotImplementedError()
    def after_sem_wait(self, sem: Pointer,
                       ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_trywait(self, sem: Pointer) -> str:
        raise NotImplementedError()
    def after_sem_trywait(self, sem: Pointer,
                          ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_timedwait(self, sem: Pointer, abs_timeout: PointerTo[StructTimeSpec]) -> str:
        raise NotImplementedError()
    def after_sem_timedwait(self, sem: Pointer, abs_timeout: PointerTo[StructTimeSpec],
                            ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_getvalue(self, sem: Pointer, value_ptr: Pointer) -> str:
        raise NotImplementedError()
    def after_sem_getvalue(self, sem: Pointer, value_ptr: Pointer,
                           ret_value: int, errno: str = None, value: int = None) -> None:
        raise NotImplementedError()
    def before_sem_close(self, sem: Pointer) -> str:
        raise NotImplementedError()
    def after_sem_close(self, sem: Pointer,
                        ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_unlink(self, name: PointerTo[bytes]) -> str:
        raise NotImplementedError()
    def after_sem_unlink(self, name: PointerTo[bytes],
                         ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sem_destroy(self, sem: Pointer) -> str:
        raise NotImplementedError()
    def after_sem_destroy(self, sem: Pointer,
                          ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_shm_open(self, name: PointerTo[bytes], oflag: Flags, mode: int) -> str:
        raise NotImplementedError()
    def after_shm_open(self, name: PointerTo[bytes], oflag: Flags, mode: int,
                       ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_shm_unlink(self, name: PointerTo[bytes]) -> str:
        raise NotImplementedError()
    def after_shm_unlink(self, name: PointerTo[bytes],
                         ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_mmap(self, addr: Pointer, length: int, prot: Flags, flags: Flags, fildes: int, off: int) -> str:
        raise NotImplementedError()
    def after_mmap(self, addr: Pointer, length: int, prot: Flags, flags: Flags, fildes: int, off: int,
                   ret_value: Pointer, errno: str = None) -> None:
        raise NotImplementedError()
    def before_munmap(self, addr: Pointer, length: int) -> str:
        raise NotImplementedError()
    def after_munmap(self, addr: Pointer, length: int,
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_ftruncate(self, fildes: int, length: int) -> str:
        raise NotImplementedError()
    def after_ftruncate(self, fildes: int, length: int,
                        ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_fork(self) -> str:
        raise NotImplementedError()
    def after_fork(self,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_wait(self, stat_loc: Pointer) -> str:
        raise NotImplementedError()
    def after_wait(self, stat_loc: Pointer,
                   ret_value: int, errno: str = None, stat: int = None) -> None:
        raise NotImplementedError()
    def before_waitpid(self, pid: int, stat_loc: Pointer, options: Flags) -> str:
            raise NotImplementedError()
    def after_waitpid(self, pid: int, stat_loc: Pointer, options: Flags,
                      ret_value: int, errno: str = None, stat: int = None) -> None:
        raise NotImplementedError()
    def before_execl(self, pathname: PointerTo[bytes], *args: list[PointerTo[bytes]]) -> str:
        raise NotImplementedError()
    def after_execl(self, pathname: PointerTo[bytes], *args: list[PointerTo[bytes]],
                    ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_execlp(self, file: PointerTo[bytes], *args: list[PointerTo[bytes]]) -> str:
        raise NotImplementedError()
    def after_execlp(self, file: PointerTo[bytes], *args: list[PointerTo[bytes]],
                    ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_execle(self, pathname: PointerTo[bytes], *args: list[PointerTo[bytes]], envp: PointerTo[list[PointerTo[bytes]]]) -> str:
        raise NotImplementedError()
    def after_execle(self, pathname: PointerTo[bytes], *args: list[PointerTo[bytes]], envp: PointerTo[list[PointerTo[bytes]]],
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_execv(self, pathname: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]]) -> str:
        raise NotImplementedError()
    def after_execv(self, pathname: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]],
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_execvp(self, file: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]]) -> str:
        raise NotImplementedError()
    def after_execvp(self, file: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]],
                    ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_execvpe(self, file: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]], envp: PointerTo[list[PointerTo[bytes]]]) -> str:
        raise NotImplementedError()
    def after_execvpe(self, file: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]], envp: PointerTo[list[PointerTo[bytes]]],
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_execve(self, pathname: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]], envp: PointerTo[list[PointerTo[bytes]]]) -> str:
        raise NotImplementedError()
    def after_execve(self, pathname: PointerTo[bytes], argv: PointerTo[list[PointerTo[bytes]]], envp: PointerTo[list[PointerTo[bytes]]],
                      ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_fexecve(self, fd: int, argv: PointerTo[list[PointerTo[bytes]]], envp: PointerTo[list[PointerTo[bytes]]]) -> str:
        raise NotImplementedError()
    def after_fexecve(self, fd: int, argv: PointerTo[list[PointerTo[bytes]]], envp: PointerTo[list[PointerTo[bytes]]],
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_pipe(self, fildes_ptr: Pointer) -> str:
        raise NotImplementedError()
    def after_pipe(self, fildes_ptr: Pointer,
                   ret_value: int, errno: str = None, fildes: list[int] = None) -> None:
        raise NotImplementedError()
    def before_dup(self, oldfd: int) -> str:
        raise NotImplementedError()
    def after_dup(self, oldfd: int,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_dup2(self, oldfd: int, newfd: int) -> str:
        raise NotImplementedError()
    def after_dup2(self, oldfd: int, newfd: int,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_dup3(self, oldfd: int, newfd: int, flags: Flags) -> str:
        raise NotImplementedError()
    def after_dup3(self, oldfd: int, newfd: int, flags: Flags,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_socket(self, domain: Constant, socktype: Constant, protocol: int) -> str:
        raise NotImplementedError()
    def after_socket(self, domain: Constant, socktype: Constant, protocol: int,
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_bind(self, sockfd: int, address: PointerTo[StructSockAddr], address_len: int) -> str:
        raise NotImplementedError()
    def after_bind(self, sockfd: int, address: PointerTo[StructSockAddr], address_len: int,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_listen(self, sockfd: int, backlog: int) -> str:
        raise NotImplementedError()
    def after_listen(self, sockfd: int, backlog: int,
                     ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_accept(self, sockfd: int, address_ptr: Pointer, address_len_ptr: Pointer) -> str:
        raise NotImplementedError()
    def after_accept(self, sockfd: int, address_ptr: Pointer, address_len_ptr: Pointer,
                     ret_value: int, errno: str = None, address: PointerTo[StructSockAddr] = None, address_len: PointerTo[StructSockAddr] = None) -> None:
        raise NotImplementedError()
    def before_connect(self, sockfd: int, address: PointerTo[StructSockAddr], address_len: int) -> str:
        raise NotImplementedError()
    def after_connect(self, sockfd: int, address: PointerTo[StructSockAddr], address_len: int,
                      ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_getaddrinfo(self, node: PointerTo[bytes], service: PointerTo[bytes], hints: PointerTo[StructAddrInfo], res: Pointer) -> str:
        raise NotImplementedError()
    def after_getaddrinfo(self, node: PointerTo[bytes], service: PointerTo[bytes], hints: PointerTo[StructAddrInfo], res: Pointer,
                          ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_freeaddrinfo(self, res: Pointer) -> str:
        raise NotImplementedError()
    def after_freeaddrinfo(self, res: Pointer) -> None:
        raise NotImplementedError()
    def before_send(self, sockfd: int, buf: PointerTo[bytes], length: int, flags: Flags) -> str:
        raise NotImplementedError()
    def after_send(self, sockfd: int, buf: PointerTo[bytes], length: int, flags: Flags,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sendto(self, sockfd: int, buf: PointerTo[bytes], length: int, flags: Flags, dest_addr: PointerTo[StructSockAddr], addrlen: int) -> str:
        raise NotImplementedError()
    def after_sendto(self, sockfd: int, buf: PointerTo[bytes], length: int, flags: Flags, dest_addr: PointerTo[StructSockAddr], addrlen: int,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_sendmsg(self, sockfd: int, message: StructMsgHdr, flags: Flags) -> str:
        raise NotImplementedError()
    def after_sendmsg(self, sockfd: int, message: StructMsgHdr, flags: Flags,
                   ret_value: int, errno: str = None) -> None:
        raise NotImplementedError()
    def before_recv(self, sockfd: int, buf_ptr: Pointer, size: int, flags: Flags) -> str:
        raise NotImplementedError()
    def after_recv(self, sockfd: int, buf_ptr: Pointer, size: int, flags: Flags,
                   ret_value: int, errno: str = None, buf: PointerTo[bytes] = None) -> None:
        raise NotImplementedError()
    def before_recvfrom(self, sockfd: int, buf_ptr: Pointer, size: int, flags: Flags, src_addr_ptr: Pointer, addrlen_ptr: Pointer) -> str:
        raise NotImplementedError()
    def after_recvfrom(self, sockfd: int, buf_ptr: Pointer, size: int, flags: Flags, src_addr_ptr: Pointer, addrlen_ptr: Pointer,
                   ret_value: int, errno: str = None, buf: PointerTo[bytes] = None, src_addr: PointerTo[StructSockAddr] = None, addrlen: int = None) -> None:
        raise NotImplementedError()
    def before_recvmsg(self, sockfd: int, message_ptr: Pointer, flags: Flags) -> str:
        raise NotImplementedError()
    def after_recvmsg(self, sockfd: int, message_ptr: Pointer, flags: Flags,
                   ret_value: int, errno: str = None, message: PointerTo[StructMsgHdr] = None) -> None:
        raise NotImplementedError()


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