bta-proxy/tools/packetreader.py

# x-run: PYTHONPATH=.. python packetreader.py ../packets-127.0.0.1-54356-server.txt.gz
from asyncio.queues import Queue

import asyncio
from bta_proxy.datainputstream import AsyncDataInputStream
from bta_proxy.packets import Packet
from sys import argv
from gzip import open as open_gzip
from json import loads
import logging

loggers = [
    logging.getLogger(name)
    for name in logging.root.manager.loggerDict
    if name.startswith("bta_proxy")
]

class CustomFormatter(logging.Formatter):
    grey = "\x1b[38;20m"
    yellow = "\x1b[93;20m"
    red = "\x1b[91;20m"
    bold_red = "\x1b[91;1m"
    reset = "\x1b[0m"
    fmt = "(%(filename)s:%(lineno)d) %(name)s - %(message)s"

    FORMATS = {
        logging.DEBUG: "\x1b[92m" + fmt + reset,
        logging.INFO: "\x1b[94m" + fmt + reset,
        logging.WARNING: yellow + fmt + reset,
        logging.ERROR: red + fmt + reset,
        logging.CRITICAL: bold_red + fmt + reset
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)

streamhandler = logging.StreamHandler()
streamhandler.setLevel(logging.DEBUG)

streamhandler.setFormatter(CustomFormatter())

for logger in loggers:
    logger.setLevel(logging.DEBUG)


logger = logging.getLogger("packetreader")
logger.setLevel(logging.DEBUG)
logging.getLogger().addHandler(streamhandler)



async def amain(stream: AsyncDataInputStream):
    while True:
        try:
            pkt = await Packet.read_packet(stream)
            logger.info(f"we just got a package {pkt!r}")
        except EOFError:
            logger.warning("EOFError")
            break
        except Exception as e:
            logger.error(e)
            raise e
    logger.info("exiting")


def main(filename: str):
    queue = Queue()
    with open_gzip(filename, "rt") as fp:
        for line in fp:
            data = loads(line.strip())
            queue.put_nowait(bytes.fromhex(data["b"]))
    queue.put_nowait(None)
    stream = AsyncDataInputStream(queue)
    loop = asyncio.get_event_loop()
    loop.run_until_complete(amain(stream))

    while True:
        pass


if __name__ == "__main__":
    main(argv[1])