"""
Processa o traffic.log gerado pelo hook.dll (D:\\RoshEmulator\\tools\\hook) -- extrai os frames
individuais (usando o mesmo framing de 4 bytes BE de tamanho + payload que ja usamos no resto do
projeto), separados por direcao (SEND/RECV) e por socket, em ordem cronologica.

Uso:
    python parse_hook_capture.py [caminho_do_traffic.log]
"""
import re
import struct
import sys
from pathlib import Path

LOG_PATH = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(r"C:\RoshHook\traffic.log")
OUT_DIR = Path(r"D:\RoshEmulator\captures\hook_parsed")
OUT_DIR.mkdir(exist_ok=True)

HEADER_RE = re.compile(r"^\[(\d{2}:\d{2}:\d{2}\.\d{3})\] (SEND|RECV) sock=0x([0-9A-Fa-f]{8}) len=(\d+)$")


def parse_log(path: Path):
    """Le o traffic.log e retorna uma lista de (timestamp, direction, sock, declared_len, bytes)."""
    entries = []
    lines = path.read_text(encoding="latin-1", errors="replace").splitlines()
    i = 0
    while i < len(lines):
        m = HEADER_RE.match(lines[i])
        if not m:
            i += 1
            continue
        ts, direction, sock, declared_len = m.group(1), m.group(2), m.group(3), int(m.group(4))
        i += 1
        hexchunks = []
        while i < len(lines) and lines[i].strip() and not HEADER_RE.match(lines[i]):
            hexchunks.append(lines[i].strip())
            i += 1
        hexstr = "".join(hexchunks).replace(" ", "")
        try:
            data = bytes.fromhex(hexstr)
        except ValueError:
            data = b""
        entries.append((ts, direction, sock, declared_len, data))
    return entries


def split_frames(stream: bytes):
    """Aplica o framing de 4 bytes BE (tamanho do payload) + payload. Retorna lista de payloads e o
    resto nao consumido (frame parcial no fim)."""
    frames = []
    off = 0
    while off + 4 <= len(stream):
        (plen,) = struct.unpack_from(">I", stream, off)
        if plen == 0 or plen > 200000:
            # provavel desalinhamento -- para aqui, guarda o resto como "sobra"
            break
        if off + 4 + plen > len(stream):
            break
        frames.append(stream[off + 4: off + 4 + plen])
        off += 4 + plen
    return frames, stream[off:]


def main():
    entries = parse_log(LOG_PATH)
    print(f"total de entradas no log: {len(entries)}")

    # agrupa por (socket, direction), concatena em ordem cronologica (a ordem do arquivo ja e'
    # cronologica, entao so precisamos manter a ordem de insercao)
    streams = {}
    for ts, direction, sock, declared_len, data in entries:
        key = (sock, direction)
        streams.setdefault(key, bytearray()).extend(data)

    summary_lines = []
    for (sock, direction), stream in streams.items():
        frames, leftover = split_frames(bytes(stream))
        summary_lines.append(f"\n=== socket 0x{sock} {direction}: {len(frames)} frames, "
                              f"{len(stream)} bytes total, {len(leftover)} bytes sobrando no fim ===")
        for idx, f in enumerate(frames):
            sig = f[:2].hex() if len(f) >= 2 else f.hex()
            summary_lines.append(f"  frame {idx}: {len(f)} bytes, sig={sig}")

        # salva os frames dessa direcao/socket num arquivo binario concatenado (com um separador
        # simples: 4 bytes de tamanho + payload, pra poder re-parsear facil depois)
        out_path = OUT_DIR / f"sock_{sock}_{direction}.bin"
        with open(out_path, "wb") as fh:
            for f in frames:
                fh.write(struct.pack(">I", len(f)))
                fh.write(f)

    summary_path = OUT_DIR / "summary.txt"
    summary_path.write_text("\n".join(summary_lines), encoding="utf-8")
    print(f"resumo salvo em {summary_path}")
    print(f"frames binarios salvos em {OUT_DIR}")


if __name__ == "__main__":
    main()
