#!/usr/bin/env python3
"""Extract credentials and tokens from a mitmproxy tnetstring flow dump.

Usage:
    python3 -I extract_tokens.py convo6_reallyfull.dump
    python3 -I extract_tokens.py convo6_reallyfull.dump --json --out tokens.json

The input is the binary tnetstring stream written by mitmproxy's flow writer
(a sequence of length-prefixed tnetstrings, each one an HTTPFlow dict).

Tokens are found in:
  * Authorization / Proxy-Authorization headers (Bearer, Basic, ...)
  * other token-like headers (X-Api-Key, X-Auth-Token, ...)
  * Cookie / Set-Cookie pairs with secret-looking names
  * JSON response/request bodies, recursively
  * form-urlencoded bodies
  * anywhere a JWT or "Bearer <token>" string occurs
"""

from __future__ import annotations

import argparse
import base64
import binascii
import gzip
import json
import re
import sys
import zlib
from datetime import datetime, timezone
from urllib.parse import parse_qsl

# --------------------------------------------------------------------------
# tnetstring decoding
# --------------------------------------------------------------------------


class TNetstringError(ValueError):
    pass


def _decode(data: bytes, pos: int):
    colon = data.find(b":", pos)
    if colon < 0:
        raise TNetstringError(f"missing length prefix at offset {pos}")
    try:
        length = int(data[pos:colon])
    except ValueError:
        raise TNetstringError(f"bad length prefix at offset {pos}") from None
    start = colon + 1
    end = start + length
    if end >= len(data):
        raise TNetstringError(f"truncated record at offset {pos}")

    payload = data[start:end]
    tag = data[end : end + 1]

    if tag == b",":
        return bytes(payload), end + 1
    if tag == b";":
        return payload.decode("utf-8", "surrogateescape"), end + 1
    if tag == b"#":
        return int(payload), end + 1
    if tag == b"^":
        return float(payload), end + 1
    if tag == b"!":
        return payload == b"true", end + 1
    if tag == b"~":
        return None, end + 1
    if tag == b"}":
        return _decode_dict(payload), end + 1
    if tag == b"]":
        return _decode_list(payload), end + 1
    raise TNetstringError(f"unknown tnetstring type {tag!r} at offset {end}")


def _decode_dict(payload: bytes) -> dict:
    out, pos = {}, 0
    while pos < len(payload):
        key, pos = _decode(payload, pos)
        value, pos = _decode(payload, pos)
        if not isinstance(key, (str, bytes)):
            raise TNetstringError("non-string dict key")
        if isinstance(key, bytes):
            key = key.decode("utf-8", "replace")
        out[key] = value
    return out


def _decode_list(payload: bytes) -> list:
    out, pos = [], 0
    while pos < len(payload):
        value, pos = _decode(payload, pos)
        out.append(value)
    return out


def iter_flows(path: str):
    """Yield (index, flow_dict) for every tnetstring record in the file."""
    with open(path, "rb") as fh:
        data = fh.read()
    pos = 0
    index = 0
    while pos < len(data):
        if data[pos : pos + 1].isspace():
            pos += 1
            continue
        flow, pos = _decode(data, pos)
        index += 1
        if isinstance(flow, dict):
            yield index, flow


# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------

SECRET_WORDS = re.compile(
    r"(?:^|_)(?:access_?token|refresh_?token|id_?token|session_?token|"
    r"auth(?:entication)?_?token|csrf_?token|xsrf_?token|token|secret|"
    r"password|passwd|pwd|api_?key|apikey|authorization|credentials?|"
    r"private_key|client_secret|session)$"
)

NON_SECRET_SUFFIXES = (
    "_type",
    "_path",
    "_url",
    "_uri",
    "_endpoint",
    "_name",
    "_label",
    "_placeholder",
    "_title",
    "_heading",
    "_message",
    "_text",
    "_length",
    "_count",
    "_enabled",
    "_timeout",
    "_expires_in",
    "_format",
    "_regex",
    "_pattern",
    "_prompt",
    "_hint",
)

# Only treated as secret on OAuth-ish endpoints, where these are real artifacts.
OAUTH_EXTRA_KEYS = {"code", "state", "authorization_code", "id_token"}

PLACEHOLDER_VALUES = {
    "",
    "password",
    "passwd",
    "secret",
    "token",
    "bearer",
    "null",
    "none",
    "true",
    "false",
    "undefined",
    "***",
    "<redacted>",
}

JWT_RE = re.compile(
    r"(?<![A-Za-z0-9_-])(eyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{4,})"
)
BEARER_RE = re.compile(r"(?i)\b(Bearer\s+)([A-Za-z0-9._~+/=-]{8,})")


def camel_to_snake(name: str) -> str:
    name = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
    name = re.sub(r"[^A-Za-z0-9]+", "_", name)
    return name.strip("_").lower()


def is_secret_key(name: str) -> bool:
    if not name:
        return False
    key = camel_to_snake(name)
    if key.endswith(NON_SECRET_SUFFIXES):
        return False
    return bool(SECRET_WORDS.search(key))


def usable_value(key: str, value) -> bool:
    if not isinstance(value, str):
        return False
    value = value.strip()
    if len(value) < 3:
        return False
    if value.lower() in PLACEHOLDER_VALUES:
        return False
    if value.lower() == str(key).lower():
        return False
    return True


def to_text(value) -> str:
    if value is None:
        return ""
    if isinstance(value, bytes):
        return value.decode("utf-8", "replace")
    return str(value)


def header_items(flow: dict, part: str):
    """Yield (lower-name, value) from request or response headers."""
    msg = flow.get(part)
    if not isinstance(msg, dict):
        return
    headers = msg.get("headers")
    if not isinstance(headers, list):
        return
    for item in headers:
        if isinstance(item, (list, tuple)) and len(item) == 2:
            yield to_text(item[0]), to_text(item[1])


def headers_dict(flow: dict, part: str) -> dict:
    return {name.lower(): value for name, value in header_items(flow, part)}


def decode_body(content, headers: dict) -> bytes:
    if content is None:
        return b""
    if isinstance(content, str):
        body = content.encode("utf-8", "surrogateescape")
    else:
        body = bytes(content)
    encoding = headers.get("content-encoding", "").lower()
    try:
        if "gzip" in encoding:
            body = gzip.decompress(body)
        elif "deflate" in encoding:
            body = zlib.decompress(body)
    except (OSError, EOFError, zlib.error):
        pass
    return body


def is_jwt(value: str) -> bool:
    parts = value.split(".")
    if len(parts) != 3:
        return False
    try:
        header = json.loads(_b64url_decode(parts[0]))
    except (ValueError, binascii.Error, UnicodeDecodeError):
        return False
    return isinstance(header, dict) and "alg" in header


def _b64url_decode(part: str) -> bytes:
    part = part.strip()
    part += "=" * (-len(part) % 4)
    return base64.urlsafe_b64decode(part.encode("ascii"))


def jwt_claims(value: str):
    if not is_jwt(value):
        return None
    try:
        header = json.loads(_b64url_decode(value.split(".")[0]))
        payload = json.loads(_b64url_decode(value.split(".")[1]))
    except (ValueError, binascii.Error, UnicodeDecodeError):
        return None
    return {"header": header, "payload": payload}


def iso_time(ts) -> str:
    if not isinstance(ts, (int, float)):
        return ""
    try:
        return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat()
    except (OverflowError, OSError, ValueError):
        return ""


def flow_source(flow: dict) -> dict:
    request = flow.get("request") if isinstance(flow.get("request"), dict) else {}
    response = flow.get("response") if isinstance(flow.get("response"), dict) else {}
    scheme = to_text(request.get("scheme")) or "https"
    host = to_text(request.get("host"))
    port = request.get("port")
    path = to_text(request.get("path"))
    url = f"{scheme}://{host}"
    if port and not (
        (scheme == "https" and port == 443) or (scheme == "http" and port == 80)
    ):
        url += f":{port}"
    url += path
    ts = flow.get("timestamp_start")
    if ts is None:
        ts = request.get("timestamp_start")
    return {
        "time": iso_time(ts),
        "method": to_text(request.get("method")),
        "url": url,
        "status": response.get("status_code"),
    }


# --------------------------------------------------------------------------
# Collection
# --------------------------------------------------------------------------


class TokenCollector:
    def __init__(self):
        # value -> entry; deduplicates identical tokens seen many times
        self.entries: dict[str, dict] = {}

    def add(self, value, kind: str, key: str, source: dict, where: str):
        if not usable_value(key, value):
            return
        value = value.strip()
        entry = self.entries.get(value)
        if entry is None:
            entry = {
                "value": value,
                "kinds": [],
                "keys": [],
                "occurrences": [],
                "jwt": jwt_claims(value),
            }
            self.entries[value] = entry
        if kind not in entry["kinds"]:
            entry["kinds"].append(kind)
        if key and key not in entry["keys"]:
            entry["keys"].append(key)
        occurrence = dict(source, part=where)
        signature = tuple(sorted(occurrence.items()))
        if all(tuple(sorted(o.items())) != signature for o in entry["occurrences"]):
            entry["occurrences"].append(occurrence)
        return entry

    def add_jwts(self, text: str, source: dict, where: str):
        """Pick up JWTs and Bearer values embedded in arbitrary text."""
        for match in JWT_RE.finditer(text):
            self.add(match.group(1), "jwt", "embedded-jwt", source, where)
        for match in BEARER_RE.finditer(text):
            self.add(match.group(2), "bearer", "embedded-bearer", source, where)


def scan_headers(collector: TokenCollector, flow: dict, part: str, source: dict):
    for name, value in header_items(flow, part):
        lname = name.lower()
        if lname in ("authorization", "proxy-authorization"):
            scheme, _, credentials = value.partition(" ")
            scheme = scheme.strip().lower()
            credentials = credentials.strip()
            if scheme == "basic" and credentials:
                try:
                    decoded = (
                        base64.b64decode(credentials + "=" * (-len(credentials) % 4))
                        .decode("utf-8", "replace")
                    )
                except (ValueError, binascii.Error):
                    decoded = ""
                if decoded:
                    collector.add(
                        decoded,
                        "authorization.basic",
                        name,
                        source,
                        f"{part} header {name} (decoded)",
                    )
                    collector.add_jwts(
                        decoded, source, f"{part} header {name} (decoded)"
                    )
                continue
            if credentials:
                collector.add(
                    credentials,
                    f"authorization.{scheme or 'unknown'}",
                    name,
                    source,
                    f"{part} header {name}",
                )
                collector.add_jwts(
                    credentials, source, f"{part} header {name}"
                )
            continue

        if lname in ("cookie", "set-cookie"):
            for chunk in value.split(";"):
                cname, sep, cvalue = chunk.partition("=")
                if sep and is_secret_key(cname.strip()):
                    collector.add(
                        cvalue.strip().strip('"'),
                        "cookie",
                        cname.strip(),
                        source,
                        f"{part} header {name}",
                    )
            collector.add_jwts(value, source, f"{part} header {name}")
            continue

        if is_secret_key(lname) and value:
            collector.add(
                value, "header", name, source, f"{part} header {name}"
            )
            collector.add_jwts(value, source, f"{part} header {name}")


def walk_json(
    collector: TokenCollector,
    obj,
    path: str,
    source: dict,
    where: str,
    oauth_context: bool = False,
):
    if isinstance(obj, dict):
        for key, value in obj.items():
            child = f"{path}.{key}" if path else str(key)
            extra = oauth_context and str(key).lower() in OAUTH_EXTRA_KEYS
            if isinstance(value, str) and (is_secret_key(str(key)) or extra):
                collector.add(value, "json", child, source, f"{where} JSON")
            walk_json(collector, value, child, source, where, oauth_context)
    elif isinstance(obj, list):
        for index, value in enumerate(obj):
            walk_json(
                collector, value, f"{path}[{index}]", source, where, oauth_context
            )
    elif isinstance(obj, str):
        collector.add_jwts(obj, source, f"{where} text")


def scan_body(collector: TokenCollector, flow: dict, part: str, source: dict):
    msg = flow.get(part)
    if not isinstance(msg, dict):
        return
    headers = headers_dict(flow, part)
    body = decode_body(msg.get("content"), headers)
    if not body:
        return
    text = body.decode("utf-8", "replace")

    content_type = headers.get("content-type", "").lower()
    oauth_context = any(
        marker in source.get("url", "").lower()
        for marker in ("oauth", "authorize", "/token", "/auth", "login")
    )
    parsed_json = False
    if "json" in content_type or text.lstrip().startswith(("{", "[")):
        try:
            walk_json(collector, json.loads(text), "", source, part, oauth_context)
            parsed_json = True
        except ValueError:
            parsed_json = False

    if not parsed_json:
        # form-urlencoded and other key=value bodies
        if "application/x-www-form-urlencoded" in content_type:
            try:
                for key, value in parse_qsl(text, keep_blank_values=False):
                    if is_secret_key(key):
                        collector.add(
                            value, "form", key, source, f"{part} form body"
                        )
            except ValueError:
                pass
        collector.add_jwts(text, source, f"{part} body")
        for key, value in re.findall(
            r"(?i)((?:access|refresh|id|session|auth)[_-]?token)\s*=\s*([^&\s\"']+)",
            text,
        ):
            collector.add(value, "body", key, source, f"{part} body")


def extract(path: str) -> tuple[TokenCollector, int]:
    collector = TokenCollector()
    flow_count = 0
    for _, flow in iter_flows(path):
        flow_count += 1
        source = flow_source(flow)
        scan_headers(collector, flow, "request", source)
        scan_headers(collector, flow, "response", source)
        scan_body(collector, flow, "request", source)
        scan_body(collector, flow, "response", source)
    return collector, flow_count


# --------------------------------------------------------------------------
# Output
# --------------------------------------------------------------------------


def shorten(value: str, limit: int) -> str:
    if limit and len(value) > limit:
        return value[: limit // 2] + f"...[{len(value)} chars]..." + value[-limit // 4 :]
    return value


def print_human(collector: TokenCollector, flow_count: int, truncate: int):
    entries = list(collector.entries.values())
    print(f"Scanned {flow_count} flow(s); {len(entries)} unique token value(s).\n")
    for number, entry in enumerate(entries, 1):
        kinds = ", ".join(entry["kinds"])
        print(f"[{number}] {kinds}")
        print(f"    value: {shorten(entry['value'], truncate)}")
        print(f"    keys : {', '.join(entry['keys']) or '-'}")
        if entry["jwt"]:
            payload = entry["jwt"]["payload"]
            interesting = {
                k: payload[k]
                for k in (
                    "exp",
                    "iat",
                    "nbf",
                    "username",
                    "upn",
                    "preferred_username",
                    "name",
                    "aud",
                    "iss",
                    "scp",
                    "roles",
                    "appid",
                    "tid",
                )
                if k in payload
            }
            if "exp" in interesting:
                interesting["exp_utc"] = iso_time(interesting["exp"])
            print(f"    jwt  : {json.dumps(interesting, ensure_ascii=False)}")
        first = entry["occurrences"][0] if entry["occurrences"] else {}
        print(
            f"    seen : {len(entry['occurrences'])} occurrence(s); first: "
            f"{first.get('method', '')} {first.get('url', '')} "
            f"({first.get('part', '')}, {first.get('time', '')})"
        )
        print()


def build_report(collector: TokenCollector, flow_count: int, path: str) -> dict:
    return {
        "file": path,
        "flows": flow_count,
        "token_count": len(collector.entries),
        "tokens": list(collector.entries.values()),
    }


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(
        description="Extract tokens/credentials from a mitmproxy tnetstring dump."
    )
    parser.add_argument("dump", help="mitmproxy flow dump file")
    parser.add_argument(
        "--json", action="store_true", help="print the report as JSON"
    )
    parser.add_argument(
        "--out", metavar="FILE", help="also write the JSON report to FILE"
    )
    parser.add_argument(
        "--truncate",
        type=int,
        default=0,
        metavar="N",
        help="truncate long values in human output to about N characters",
    )
    args = parser.parse_args(argv)

    try:
        collector, flow_count = extract(args.dump)
    except (OSError, TNetstringError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    report = build_report(collector, flow_count, args.dump)
    if args.out:
        with open(args.out, "w", encoding="utf-8") as fh:
            json.dump(report, fh, indent=2, ensure_ascii=False)
        print(f"wrote {args.out}", file=sys.stderr)

    if args.json:
        json.dump(report, sys.stdout, indent=2, ensure_ascii=False)
        sys.stdout.write("\n")
    else:
        print_human(collector, flow_count, args.truncate)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())