payload.py

"""Lightweight, dependency-free parsers for the application-layer payloads we
care about for threat hunting: DNS, TLS ClientHello (SNI + JA3) and HTTP
requests (Host / User-Agent).

Everything here is bounds-checked and never raises on malformed/truncated input
— it returns None or partial data instead. Captures are attacker-influenced, so
the parsers treat all input as hostile.
"""

from __future__ import annotations

import hashlib
import socket
import struct

_u16 = struct.Struct(">H").unpack_from

# ---------------------------------------------------------------- DNS --------
DNS_TYPE = {1: "A", 2: "NS", 5: "CNAME", 6: "SOA", 12: "PTR", 15: "MX",
            16: "TXT", 28: "AAAA", 33: "SRV", 65: "HTTPS"}


def _read_name(data, off, depth=0):
    """Read a (possibly compressed) DNS name. Returns (name, next_offset)."""
    labels = []
    next_off = None
    n = len(data)
    while off < n:
        ln = data[off]
        if ln == 0:
            off += 1
            break
        if (ln & 0xC0) == 0xC0:                       # compression pointer
            if off + 1 >= n or depth > 15:
                off += 2
                break
            ptr = ((ln & 0x3F) << 8) | data[off + 1]
            if next_off is None:
                next_off = off + 2
            sub, _ = _read_name(data, ptr, depth + 1)
            if sub:
                labels.append(sub)
            off = next_off
            return ".".join(labels), off
        off += 1
        labels.append(data[off:off + ln].decode("latin-1", "ignore"))
        off += ln
    return ".".join(labels), (next_off if next_off is not None else off)


def parse_dns(payload):
    """Parse a DNS message. Returns dict with qr/rcode/qname/qtype/answers."""
    if len(payload) < 12:
        return None
    try:
        flags = _u16(payload, 2)[0]
        qd = _u16(payload, 4)[0]
        an = _u16(payload, 6)[0]
    except Exception:
        return None
    qr = (flags >> 15) & 1
    rcode = flags & 0x0F
    off = 12
    qname, qtype = "", 0
    if qd >= 1:
        qname, off = _read_name(payload, off)
        if off + 4 > len(payload):
            return {"qr": qr, "rcode": rcode, "qname": qname.lower(),
                    "qtype": qtype, "answers": []}
        qtype = _u16(payload, off)[0]
        off += 4
    answers = []
    for _ in range(min(an, 30)):
        if off + 1 > len(payload):
            break
        _name, off = _read_name(payload, off)
        if off + 10 > len(payload):
            break
        atype = _u16(payload, off)[0]
        rdlen = _u16(payload, off + 8)[0]
        rdstart = off + 10
        rdata = payload[rdstart:rdstart + rdlen]
        off = rdstart + rdlen
        try:
            if atype == 1 and rdlen == 4:
                answers.append(("A", socket.inet_ntoa(rdata)))
            elif atype == 28 and rdlen == 16:
                answers.append(("AAAA", socket.inet_ntop(socket.AF_INET6, rdata)))
            elif atype == 5:
                nm, _ = _read_name(payload, rdstart)
                answers.append(("CNAME", nm.lower()))
            elif atype == 16 and rdlen >= 1:
                answers.append(("TXT", rdata[1:].decode("latin-1", "ignore")[:160]))
        except Exception:
            continue
    return {"qr": qr, "rcode": rcode, "qname": qname.lower(),
            "qtype": qtype, "answers": answers}


# ---------------------------------------------------------------- TLS --------
def parse_tls_client_hello(payload):
    """Extract SNI + JA3 from a TLS ClientHello. Returns dict or None."""
    if len(payload) < 6 or payload[0] != 0x16:   # not a handshake record
        return None
    try:
        p = 5
        if payload[p] != 0x01:                   # not ClientHello
            return None
        p += 4                                    # handshake type + length
        client_version = _u16(payload, p)[0]
        p += 2 + 32                               # version + random
        sid_len = payload[p]
        p += 1 + sid_len
        cs_len = _u16(payload, p)[0]
        p += 2
        ciphers = [_u16(payload, p + i)[0] for i in range(0, cs_len, 2)]
        p += cs_len
        comp_len = payload[p]
        p += 1 + comp_len
        if p + 2 > len(payload):
            ext_total = 0
        else:
            ext_total = _u16(payload, p)[0]
        p += 2
        end = min(len(payload), p + ext_total)
        sni, exts, curves, ecpf = None, [], [], []
        while p + 4 <= end:
            etype = _u16(payload, p)[0]
            elen = _u16(payload, p + 2)[0]
            p += 4
            edata = payload[p:p + elen]
            p += elen
            exts.append(etype)
            if etype == 0 and len(edata) >= 5:               # server_name
                nlen = _u16(edata, 3)[0]
                sni = edata[5:5 + nlen].decode("latin-1", "ignore")
            elif etype == 10 and len(edata) >= 2:            # supported_groups
                gl = _u16(edata, 0)[0]
                curves = [_u16(edata, 2 + i)[0] for i in range(0, gl, 2)]
            elif etype == 11 and len(edata) >= 1:            # ec_point_formats
                fl = edata[0]
                ecpf = list(edata[1:1 + fl])
    except Exception:
        return None
    ja3 = "%d,%s,%s,%s,%s" % (
        client_version,
        "-".join(map(str, ciphers)),
        "-".join(map(str, exts)),
        "-".join(map(str, curves)),
        "-".join(map(str, ecpf)),
    )
    return {"sni": sni.lower() if sni else None, "ja3": ja3,
            "ja3_hash": hashlib.md5(ja3.encode()).hexdigest()}


# --------------------------------------------------------------- HTTP --------
_METHODS = (b"GET ", b"POST ", b"HEAD ", b"PUT ", b"DELETE ", b"OPTIONS ",
            b"PATCH ", b"CONNECT ", b"TRACE ")


def parse_http_request(payload):
    """Extract method/path/Host/User-Agent from an HTTP/1.x request."""
    if not payload.startswith(_METHODS):
        return None
    text = payload[:4096].decode("latin-1", "ignore")
    lines = text.split("\r\n")
    if not lines or " " not in lines[0]:
        return None
    parts = lines[0].split(" ")
    method, path = parts[0], (parts[1] if len(parts) > 1 else "")
    host = ua = None
    for ln in lines[1:]:
        low = ln.lower()
        if low.startswith("host:"):
            host = ln[5:].strip()
        elif low.startswith("user-agent:"):
            ua = ln[11:].strip()
    return {"method": method, "path": path[:256],
            "host": host.lower() if host else None, "user_agent": ua}