bta-proxy/bta_proxy/dpi.py

from asyncio.queues import Queue
import time
from asyncio import get_event_loop
from typing import TextIO
from json import dumps
from gzip import open as open_gzip

from bta_proxy.datainputstream import AsyncDataInputStream
from bta_proxy.packets import *

async def queue_writer(queue_in: Queue, queue_out: Queue, fp: TextIO):
    packet: bytes
    while (packet := await queue_in.get()) != b'':
        fp.write(dumps({
            "t": time.time(),
            "s": len(packet),
            "b": packet.hex()
        }) + "\n")
        queue_out.put_nowait(packet)

async def inspect_client(queue: Queue, addr: tuple[str, int]):
    stream_queue = Queue()
    dis = AsyncDataInputStream(stream_queue)

    last_time = time.time()

    f = open_gzip("packets-%d-%s-%d-client.txt.gz" % (int(time.time()), addr[0], addr[1]), "wt")
    get_event_loop().create_task(queue_writer(queue, stream_queue, f))

    stats: dict[int, int] = {}

    try:
        while True:
            try:
                pkt = await Packet.read_packet(dis)
            except EOFError:
                break
            now = time.time()
            delta = now - last_time
            last_time = now

            stats[pkt.packet_id] = stats.get(pkt.packet_id, 0) + 1

            match pkt.packet_id:
                case Packet10Flying.packet_id:
                    continue
                case Packet11PlayerPosition.packet_id:
                    continue
                case Packet12PlayerLook.packet_id:
                    continue
                case Packet13LookMove.packet_id:
                    continue
                case _:
                    print(f"C {delta*1000:+8.1f}ms {pkt}")
            if pkt.packet_id == Packet255KickDisconnect.packet_id:
                break
    finally:
        print("[C] Closing output file")
        for pkt_id, count in sorted(stats.items(), key=lambda kv: kv[1], reverse=True):
            print(f"[C] {Packet.REGISTRY[pkt_id].__name__}: {count}")
        f.close()


async def inspect_server(queue: Queue, addr: tuple[str, int]):
    stream_queue = Queue()
    dis = AsyncDataInputStream(stream_queue)

    last_time = time.time()

    f = open_gzip("packets-%d-%s-%d-server.txt.gz" % (int(time.time()), addr[0], addr[1]), "wt")
    get_event_loop().create_task(queue_writer(queue, stream_queue, f))

    stats: dict[int, int] = {}

    try:
        while True:
            try:
                pkt = await Packet.read_packet(dis)
            except EOFError:
                break
            now = time.time()
            delta = now - last_time
            last_time = now

            stats[pkt.packet_id] = stats.get(pkt.packet_id, 0) + 1

            match pkt.packet_id:
                case Packet53BlockChange.packet_id:
                    continue
                case Packet50PreChunk.packet_id:
                    continue
                case Packet51MapChunk.packet_id:
                    continue
                case Packet34EntityTeleport.packet_id:
                    continue
                case Packet28EntityVelocity.packet_id:
                    continue
                case Packet31RelEntityMove.packet_id:
                    continue
                case Packet32EntityLook.packet_id:
                    continue
                case Packet33RelEntityMoveLook.packet_id:
                    continue
                # case Packet73WeatherStatus.packet_id:
                #     continue
                case Packet52MultiBlockChange.packet_id:
                    continue
                case _:
                    print(f"S {delta*1000:+8.1f}ms {pkt}")
    finally:
        print("[S] Closing output file")
        for pkt_id, count in sorted(stats.items(), key=lambda kv: kv[1], reverse=True):
            print(f"[S] {Packet.REGISTRY[pkt_id].__name__}: {count}")
        f.close()