"""
vuln_app.py - Deliberately vulnerable Flask application for Chapter 3 labs.

AUTHORISATION GUARD: Run this ONLY on your own machine or in the provided
Docker container. Do NOT expose port 5000 to a public network.
Scanning or attacking systems you do not own is unlawful; see Chapter 9.

Vulnerabilities present (intentional, for lab use):
  - SQL injection in /login (string concatenation, CWE-89)
  - OS command injection in /ping (shell=True, CWE-78)
  - Information leakage in /search (raw exception returned, CWE-209)
  - Fail-open error handling in /admin-check (except: pass, CWE-636)
  - JWT alg:none and RS256->HS256 confusion in /jwt-verify
  - Missing Pydantic strict mode in /order (lax type coercion)
"""

import sqlite3
import subprocess
import os
import logging
import uuid
from pathlib import Path

from flask import Flask, request, jsonify, g
import jwt
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from pydantic import BaseModel, field_validator, ValidationError
from pydantic import conint
from typing import Literal

app = Flask(__name__)
app.config["DEBUG"] = False

# ── Logging (server-side only) ───────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)

DB_PATH = "/tmp/ch03_lab.db"

# ── RSA key pair for JWT labs ─────────────────────────────────────────────────
_rsa_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
RSA_PRIVATE_PEM = _rsa_key.private_bytes(
    serialization.Encoding.PEM,
    serialization.PrivateFormat.TraditionalOpenSSL,
    serialization.NoEncryption(),
)
RSA_PUBLIC_PEM = _rsa_key.public_key().public_bytes(
    serialization.Encoding.PEM,
    serialization.PublicFormat.SubjectPublicKeyInfo,
)
HMAC_SECRET = "lab-hmac-secret-do-not-use-in-production"


# ── Database initialisation ───────────────────────────────────────────────────
def init_db():
    con = sqlite3.connect(DB_PATH)
    cur = con.cursor()
    cur.execute(
        "CREATE TABLE IF NOT EXISTS users"
        " (id INTEGER PRIMARY KEY, name TEXT, password TEXT, role TEXT)"
    )
    cur.execute(
        "CREATE TABLE IF NOT EXISTS notes"
        " (id INTEGER PRIMARY KEY, owner TEXT, content TEXT)"
    )
    # Seed data
    users = [
        ("alice", "password123", "user"),
        ("bob", "hunter2", "user"),
        ("admin", "supersecret", "admin"),
    ]
    for name, pw, role in users:
        cur.execute(
            "INSERT OR IGNORE INTO users (name, password, role) VALUES (?, ?, ?)",
            (name, pw, role),
        )
    notes = [
        ("alice", "Alice's private note"),
        ("admin", "Admin config: DB_HOST=db.internal"),
    ]
    for owner, content in notes:
        cur.execute(
            "INSERT OR IGNORE INTO notes (owner, content) VALUES (?, ?)",
            (owner, content),
        )
    con.commit()
    con.close()


def get_db():
    if "db" not in g:
        g.db = sqlite3.connect(DB_PATH)
        g.db.row_factory = sqlite3.Row
    return g.db


@app.teardown_appcontext
def close_db(_exc=None):
    db = g.pop("db", None)
    if db is not None:
        db.close()


# ════════════════════════════════════════════════════════════════════════
# Lab 1 / Lab 2 - VULNERABLE login (SQL injection)
# ════════════════════════════════════════════════════════════════════════

@app.route("/login", methods=["POST"])
def login_vuln():
    """VULNERABLE: string-concatenated SQL query (CWE-89). Lab 2 target."""
    user = request.form.get("user", "")
    password = request.form.get("password", "")
    db = get_db()
    # INTENTIONAL FLAW: never build queries this way
    query = f"SELECT * FROM users WHERE name = '{user}' AND password = '{password}'"
    logger.info("login_vuln query: %s", query)
    row = db.execute(query).fetchone()
    if row:
        return jsonify({"status": "ok", "role": row["role"]})
    return jsonify({"status": "fail"}), 401


@app.route("/login-safe", methods=["POST"])
def login_safe():
    """FIXED version: parameterised query. Use this to verify the fix in Lab 2."""
    user = request.form.get("user", "")
    password = request.form.get("password", "")
    db = get_db()
    row = db.execute(
        "SELECT * FROM users WHERE name = ? AND password = ?", (user, password)
    ).fetchone()
    if row:
        return jsonify({"status": "ok", "role": row["role"]})
    return jsonify({"status": "fail"}), 401


# ════════════════════════════════════════════════════════════════════════
# Lab 2 - VULNERABLE /ping (OS command injection)
# ════════════════════════════════════════════════════════════════════════

@app.route("/ping")
def ping_vuln():
    """VULNERABLE: shell=True with user-controlled string (CWE-78). Lab 2 target."""
    host = request.args.get("host", "127.0.0.1")
    # INTENTIONAL FLAW: shell=True lets host carry shell metacharacters
    result = subprocess.run(
        f"ping -c 1 {host}",
        shell=True,
        capture_output=True,
        text=True,
        timeout=5,
    )
    return jsonify({"stdout": result.stdout, "stderr": result.stderr})


@app.route("/ping-safe")
def ping_safe():
    """FIXED version: argv list, shell=False. Use to verify the fix in Lab 2."""
    host = request.args.get("host", "127.0.0.1")
    result = subprocess.run(
        ["ping", "-c", "1", host],
        shell=False,
        capture_output=True,
        text=True,
        timeout=5,
    )
    return jsonify({"stdout": result.stdout, "stderr": result.stderr})


# ════════════════════════════════════════════════════════════════════════
# Lab 3 - Pydantic strict mode (NoSQL-style operator injection via lax coercion)
# ════════════════════════════════════════════════════════════════════════

class OrderLax(BaseModel):
    """Lax mode - coerces types silently (intentionally unsafe for lab)."""
    quantity: conint(gt=0)
    status: Literal["pending", "paid", "shipped"]


class OrderStrict(BaseModel):
    """Strict mode - rejects type mismatches, including object-where-str-expected."""
    model_config = {"strict": True}
    quantity: conint(gt=0)
    status: Literal["pending", "paid", "shipped"]

    @field_validator("status")
    @classmethod
    def normalise(cls, v):
        return v.lower()


@app.route("/order-lax", methods=["POST"])
def order_lax():
    """VULNERABLE: lax Pydantic - accepts object payloads silently (Lab 3)."""
    try:
        order = OrderLax(**request.get_json(force=True))
        return jsonify({"accepted": True, "order": order.model_dump()})
    except ValidationError as exc:
        return jsonify({"error": exc.errors()}), 422


@app.route("/order-strict", methods=["POST"])
def order_strict():
    """FIXED: strict Pydantic - rejects {\"$ne\":\"x\"} because it is not a str (Lab 3)."""
    try:
        order = OrderStrict(**request.get_json(force=True))
        return jsonify({"accepted": True, "order": order.model_dump()})
    except ValidationError as exc:
        cid = uuid.uuid4().hex[:8]
        logger.warning("Validation rejected (cid=%s): %s", cid, exc.errors())
        return jsonify({"error": "Invalid request", "ref": cid}), 422


# ════════════════════════════════════════════════════════════════════════
# Lab 4 - JWT: alg:none and algorithm-confusion
# ════════════════════════════════════════════════════════════════════════

@app.route("/jwt-issue")
def jwt_issue():
    """Issue a legitimate RS256-signed JWT with a role claim."""
    import time
    payload = {
        "sub": request.args.get("user", "alice"),
        "role": request.args.get("role", "user"),
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,
        "iss": "ch03-lab",
        "aud": "ch03-lab",
    }
    token = jwt.encode(payload, RSA_PRIVATE_PEM, algorithm="RS256")
    return jsonify({"token": token})


@app.route("/jwt-verify-naive", methods=["POST"])
def jwt_verify_naive():
    """VULNERABLE: no algorithms= arg - accepts alg:none and HS256 confusion (Lab 4)."""
    token = (request.get_json(force=True) or {}).get("token", "")
    try:
        # INTENTIONAL FLAW: no algorithms list - lets the token choose its own alg
        claims = jwt.decode(token, RSA_PUBLIC_PEM, options={"verify_signature": False})
        return jsonify({"status": "accepted", "claims": claims})
    except Exception as exc:
        return jsonify({"status": "rejected", "reason": str(exc)}), 401


@app.route("/jwt-verify-safe", methods=["POST"])
def jwt_verify_safe():
    """FIXED: explicit algorithm allow-list + iss/aud validation (Lab 4)."""
    token = (request.get_json(force=True) or {}).get("token", "")
    try:
        claims = jwt.decode(
            token,
            RSA_PUBLIC_PEM,
            algorithms=["RS256"],
            issuer="ch03-lab",
            audience="ch03-lab",
        )
        return jsonify({"status": "accepted", "claims": claims})
    except jwt.InvalidTokenError as exc:
        return jsonify({"status": "rejected", "reason": str(exc)}), 401


@app.route("/jwt-pubkey")
def jwt_pubkey():
    """Expose the RSA public key so students can attempt HS256 confusion (Lab 4)."""
    return RSA_PUBLIC_PEM, 200, {"Content-Type": "text/plain"}


# ════════════════════════════════════════════════════════════════════════
# Lab 5 - Error handling: CWE-209 (info leak) and CWE-636 (fail-open)
# ════════════════════════════════════════════════════════════════════════

@app.route("/search-leaky")
def search_leaky():
    """VULNERABLE: returns raw exception to client (CWE-209). Lab 5 target."""
    term = request.args.get("q", "")
    db = get_db()
    try:
        # INTENTIONAL FLAW: raw exception text sent to client
        rows = db.execute(
            f"SELECT content FROM notes WHERE owner = '{term}'"
        ).fetchall()
        return jsonify({"results": [r["content"] for r in rows]})
    except Exception as exc:
        # CWE-209: leaks schema, paths, query internals
        return str(exc), 500


@app.route("/search-safe")
def search_safe():
    """FIXED: generic error + correlation ID; full detail server-side only (Lab 5)."""
    term = request.args.get("q", "")
    db = get_db()
    try:
        rows = db.execute(
            "SELECT content FROM notes WHERE owner = ?", (term,)
        ).fetchall()
        return jsonify({"results": [r["content"] for r in rows]})
    except Exception:
        cid = uuid.uuid4().hex[:8]
        logger.exception("search_safe failed cid=%s", cid)
        return jsonify({"error": "An error occurred", "ref": cid}), 500


@app.route("/admin-check-failopen")
def admin_check_failopen():
    """VULNERABLE: except: pass swallows auth check, grants access on error (CWE-636)."""
    token = request.args.get("token", "")
    authorised = False
    try:
        claims = jwt.decode(
            token, RSA_PUBLIC_PEM, algorithms=["RS256"],
            issuer="ch03-lab", audience="ch03-lab"
        )
        authorised = claims.get("role") == "admin"
        if not authorised:
            return jsonify({"access": "denied"}), 403
    except Exception:
        # INTENTIONAL FLAW: swallowing the exception means authorised stays False
        # but the route falls through and the response below is returned
        pass
    # Because pass leaves authorised=False, this demonstrates how fail-open looks:
    # replace 'pass' with 'authorised = True' to show an even worse variant.
    return jsonify({"access": "granted (fail-open demo)" if not authorised else "granted"})


@app.route("/admin-check-safe")
def admin_check_safe():
    """FIXED: fail closed - exception denies access and logs with correlation ID (Lab 5)."""
    token = request.args.get("token", "")
    try:
        claims = jwt.decode(
            token, RSA_PUBLIC_PEM, algorithms=["RS256"],
            issuer="ch03-lab", audience="ch03-lab"
        )
        if claims.get("role") != "admin":
            return jsonify({"access": "denied"}), 403
        return jsonify({"access": "granted"})
    except Exception:
        cid = uuid.uuid4().hex[:8]
        logger.exception("admin-check auth error cid=%s", cid)
        return jsonify({"access": "denied", "ref": cid}), 403


# ════════════════════════════════════════════════════════════════════════
# Health check
# ════════════════════════════════════════════════════════════════════════

@app.route("/healthz")
def healthz():
    return jsonify({"status": "ok", "lab": "ch03"})


if __name__ == "__main__":
    init_db()
    app.run(host="0.0.0.0", port=5000, debug=False)
