from flask import Flask, jsonify
import base64
import hashlib
import hmac
import json
import os
import requests

app = Flask(__name__)

TARGET = os.environ.get("TARGET", "http://proxy:8080")
FLAG = os.environ.get("FLAG", "FLAG_NOT_SET")
WEAK_SECRET = os.environ.get("WEAK_SECRET", "v1llag3!")
MIN_SECRET_LEN = int(os.environ.get("MIN_SECRET_LEN", "32"))


def b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def forge_none_token() -> str:
    header = {"alg": "none", "typ": "JWT"}
    payload = {"sub": "attacker", "role": "officer"}
    return f"{b64url(json.dumps(header, separators=(',', ':')).encode())}.{b64url(json.dumps(payload, separators=(',', ':')).encode())}."


def forge_weak_hs256_token() -> str:
    header = {"alg": "HS256", "typ": "JWT"}
    payload = {"sub": "attacker", "role": "officer"}
    head = b64url(json.dumps(header, separators=(",", ":")).encode())
    body = b64url(json.dumps(payload, separators=(",", ":")).encode())
    signing_input = f"{head}.{body}".encode()
    signature = hmac.new(WEAK_SECRET.encode(), signing_input, hashlib.sha256).digest()
    return f"{head}.{body}.{b64url(signature)}"


def login_guest_token():
    resp = requests.post(
        f"{TARGET}/api/login",
        json={"username": "guest", "password": "guest123"},
        timeout=3,
    )
    if resp.status_code != 200:
        return None
    return resp.json().get("token")


def request_profile(token: str):
    return requests.get(
        f"{TARGET}/api/profile",
        headers={"Authorization": f"Bearer {token}"},
        timeout=3,
    )


def request_control(token: str):
    return requests.get(
        f"{TARGET}/api/control",
        headers={"Authorization": f"Bearer {token}"},
        timeout=3,
    )


def request_security():
    return requests.get(
        f"{TARGET}/api/security",
        timeout=3,
    )


@app.route("/get_flag", methods=["GET"])
def get_flag():
    try:
        token = login_guest_token()
        if not token:
            return jsonify({
                "status": "blocked",
                "message": "Login is broken. Keep normal auth working.",
            })

        profile_ok = request_profile(token)
        if profile_ok.status_code != 200:
            return jsonify({
                "status": "blocked",
                "message": "Guest token is rejected. Keep valid JWT working.",
            })

        bad_sig = token[:-1] + ("A" if token[-1] != "A" else "B")
        tampered = request_profile(bad_sig)
        if tampered.status_code == 200:
            return jsonify({
                "status": "open",
                "message": "Tampered signature is accepted. JWT validation is weak.",
            })

        forged = forge_none_token()
        forged_resp = request_control(forged)
        if forged_resp.status_code == 200:
            return jsonify({
                "status": "open",
                "message": "Forged alg=none token is accepted. Fix JWT verification.",
            })

        security_resp = request_security()
        if security_resp.status_code != 200:
            return jsonify({
                "status": "blocked",
                "message": "Security policy endpoint is unavailable.",
            })

        security_data = security_resp.json()
        secret_len = int(security_data.get("jwt_secret_length", 0))
        min_len = int(security_data.get("jwt_secret_min_length", MIN_SECRET_LEN))
        if secret_len < min_len:
            return jsonify({
                "status": "open",
                "message": f"JWT secret is too short ({secret_len}). Use at least {min_len} characters.",
            })

        weak_secret_forged = forge_weak_hs256_token()
        weak_secret_resp = request_control(weak_secret_forged)
        if weak_secret_resp.status_code == 200:
            return jsonify({
                "status": "open",
                "message": "Known weak JWT secret still works. Strengthen and rotate JWT_SECRET.",
            })

        return jsonify({"status": "fixed", "flag": FLAG})
    except requests.exceptions.RequestException:
        return jsonify({
            "status": "blocked",
            "message": "Target service is unavailable.",
        })


if __name__ == "__main__":
    port = int(os.environ.get("PORT", "9090"))
    app.run(host="0.0.0.0", port=port)
