RustProxy/tools/dev/test_compression.py

#!/usr/bin/env python3

import argparse
import gzip
import hashlib
import io
import time
import zlib

import brotli
import requests
import zstandard


TESTS = {
    "identity": "identity",
    "gzip": "gzip",
    "deflate": "deflate",
    "brotli": "br",
    "zstd": "zstd",
    "all": "zstd, br, gzip, deflate",
}


def decompress(data, content_encoding):
    encodings = [
        e.strip().lower()
        for e in content_encoding.split(",")
        if e.strip() and e.strip().lower() != "identity"
    ]

    # Content-Encoding usuwa się w odwrotnej kolejności.
    for encoding in reversed(encodings):
        if encoding == "gzip":
            data = gzip.decompress(data)

        elif encoding == "deflate":
            try:
                data = zlib.decompress(data)
            except zlib.error:
                # Niektóre serwery wysyłają raw DEFLATE.
                data = zlib.decompress(data, -zlib.MAX_WBITS)

        elif encoding == "br":
            data = brotli.decompress(data)

        elif encoding == "zstd":
            dctx = zstandard.ZstdDecompressor()
            with dctx.stream_reader(io.BytesIO(data)) as reader:
                data = reader.read()

        else:
            raise RuntimeError(f"Nieznany Content-Encoding: {encoding}")

    return data


def human_size(size):
    if size < 1024:
        return f"{size} B"
    if size < 1024 * 1024:
        return f"{size / 1024:.1f} KB"
    return f"{size / 1024 / 1024:.2f} MB"


def request_test(url, accept_encoding, verify_ssl=True, timeout=15, revalidate=False):
    headers = {
        "Accept-Encoding": accept_encoding,
        "User-Agent": "compression-tester/1.0",
    }
    if revalidate:
        headers["Cache-Control"] = "no-cache"

    start = time.perf_counter()

    r = requests.get(
        url,
        headers=headers,
        timeout=timeout,
        verify=verify_ssl,
        stream=True,
        allow_redirects=True,
    )

    r.raw.decode_content = False
    raw = r.raw.read()

    elapsed = time.perf_counter() - start

    content_encoding = r.headers.get("Content-Encoding", "identity")

    try:
        decoded = decompress(raw, content_encoding)
        error = None
    except Exception as exc:
        decoded = b""
        error = str(exc)

    return {
        "status": r.status_code,
        "requested": accept_encoding,
        "received": content_encoding,
        "raw_size": len(raw),
        "decoded_size": len(decoded),
        "time": elapsed,
        "decoded": decoded,
        "error": error,
        "headers": r.headers,
        "final_url": r.url,
    }


def main():
    parser = argparse.ArgumentParser(
        description="Tester HTTP Content-Encoding: zstd/br/gzip/deflate"
    )

    parser.add_argument("url", help="URL do przetestowania")
    parser.add_argument(
        "--insecure",
        action="store_true",
        help="Wyłącz weryfikację certyfikatu SSL",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=15,
        help="Timeout HTTP w sekundach",
    )
    parser.add_argument(
        "--revalidate",
        action="store_true",
        help="Wyslij Cache-Control: no-cache (wymusza rewalidacje i moze resetowac Age do 0)",
    )

    args = parser.parse_args()

    if args.insecure:
        requests.packages.urllib3.disable_warnings()

    print()
    print(f"URL: {args.url}")
    print("=" * 100)

    results = {}

    for name, accept_encoding in TESTS.items():
        try:
            result = request_test(
                args.url,
                accept_encoding,
                verify_ssl=not args.insecure,
                timeout=args.timeout,
                revalidate=args.revalidate,
            )
            results[name] = result
        except Exception as exc:
            results[name] = {
                "error": f"REQUEST ERROR: {exc}"
            }

    identity_body = None

    if "identity" in results:
        identity_body = results["identity"].get("decoded")

    print(
        f"{'TEST':<10}"
        f"{'HTTP':<7}"
        f"{'REQUEST':<27}"
        f"{'RESPONSE':<12}"
        f"{'RAW':<12}"
        f"{'DECODED':<12}"
        f"{'RATIO':<10}"
        f"{'TIME':<10}"
        f"{'BODY'}"
    )

    print("-" * 120)

    failed = False

    for name, result in results.items():
        if result.get("error"):
            print(f"{name:<10} ERROR: {result['error']}")
            failed = True
            continue

        raw_size = result["raw_size"]
        decoded_size = result["decoded_size"]

        if decoded_size:
            ratio = raw_size / decoded_size
            ratio_text = f"{ratio * 100:.1f}%"
        else:
            ratio_text = "-"

        if identity_body is not None and result["decoded"]:
            if result["decoded"] == identity_body:
                body_status = "OK"
            else:
                body_status = "DIFF"
        else:
            body_status = "-"

        print(
            f"{name:<10}"
            f"{result['status']:<7}"
            f"{result['requested']:<27}"
            f"{result['received']:<12}"
            f"{human_size(raw_size):<12}"
            f"{human_size(decoded_size):<12}"
            f"{ratio_text:<10}"
            f"{result['time']:.3f}s    "
            f"{body_status}"
        )

    print()
    print("=" * 100)
    print("Szczegóły nagłówków:")
    print()

    for name, result in results.items():
        if result.get("error"):
            continue

        headers = result["headers"]

        print(f"[{name}]")
        print(f"  URL finalny:       {result['final_url']}")
        print(f"  HTTP status:       {result['status']}")
        print(f"  Content-Encoding:  {headers.get('Content-Encoding', '-')}")
        print(f"  Content-Type:      {headers.get('Content-Type', '-')}")
        print(f"  Content-Length:    {headers.get('Content-Length', '-')}")
        print(f"  Vary:              {headers.get('Vary', '-')}")
        print(f"  Cache-Control:     {headers.get('Cache-Control', '-')}")
        print(f"  Age:               {headers.get('Age', '-')}")
        print(f"  X-Cache:           {headers.get('X-Cache', '-')}")
        print(f"  X-Cache-Hits:      {headers.get('X-Cache-Hits', '-')}")
        print(f"  Cache-Status:      {headers.get('Cache-Status', '-')}")
        print(f"  ETag:              {headers.get('ETag', '-')}")
        print(f"  Server:            {headers.get('Server', '-')}")
        print()

    print("SHA256 po dekompresji:")
    print()

    for name, result in results.items():
        body = result.get("decoded")

        if body:
            sha = hashlib.sha256(body).hexdigest()
            print(f"{name:<10} {sha}")

    if failed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()