#!/usr/bin/env python3
"""
Mock Airflow AWS Auth Manager - Demonstrates Host Header Injection vulnerability
"""

from flask import Flask, request, redirect, make_response
from onelogin.saml2.auth import OneLogin_Saml2_Auth
from onelogin.saml2.idp_metadata_parser import OneLogin_Saml2_IdPMetadataParser
import sys

app = Flask(__name__)

SAML_METADATA_URL = sys.argv[1] if len(sys.argv) > 1 else ""
PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 8080

def _prepare_request_VULNERABLE(req):
    """
    VULNERABLE: Uses Host header directly (like real Airflow)
    """
    host = req.headers.get("Host", req.host)
    
    # Parse host and port from Host header
    if ":" in host:
        hostname, port = host.rsplit(":", 1)
    else:
        hostname = host
        port = "443" if req.scheme == "https" else "80"
    
    return {
        "https": "on" if req.scheme == "https" else "off",
        "http_host": hostname,
        "server_port": port,
        "script_name": req.path,
        "get_data": req.args.to_dict(),
        "post_data": req.form.to_dict(),
    }

def _init_saml_auth(req):
    request_data = _prepare_request_VULNERABLE(req)
    
    print(f"[DEBUG] http_host={request_data['http_host']}, server_port={request_data['server_port']}")
    
    settings = {
        "debug": True,
        "sp": {
            "entityId": "aws-auth-manager-saml-client",
            "assertionConsumerService": {
                "url": f"http://{request_data['http_host']}:{request_data['server_port']}/login_callback",
                "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
            },
        },
    }
    idp_data = OneLogin_Saml2_IdPMetadataParser.parse_remote(SAML_METADATA_URL)
    merged = OneLogin_Saml2_IdPMetadataParser.merge_settings(idp_data, settings)
    return OneLogin_Saml2_Auth(request_data, merged)

@app.route("/")
def index():
    return f"""
    <h1>Mock Airflow - Port {PORT}</h1>
    <p><a href="/login">Login with SAML</a></p>
    <p>SAML Metadata: {SAML_METADATA_URL}</p>
    <hr>
    <p><b>Host Header received:</b> {request.headers.get('Host')}</p>
    """

@app.route("/login")
def login():
    print(f"[LOGIN] Host header: {request.headers.get('Host')}")
    saml_auth = _init_saml_auth(request)
    redirect_url = saml_auth.login()
    return redirect(redirect_url)

@app.route("/login_callback", methods=["POST"])
def login_callback():
    host = request.headers.get("Host")
    print(f"\n{'='*60}")
    print(f"[CALLBACK] Host header: {host}")
    print(f"[CALLBACK] SAMLResponse received")
    print(f"{'='*60}")
    
    saml_auth = _init_saml_auth(request)
    
    try:
        saml_auth.process_response()
        
        if not saml_auth.is_authenticated():
            error_reason = saml_auth.get_last_error_reason()
            print(f"[CALLBACK] Auth FAILED: {error_reason}")
            return f"Authentication failed: {error_reason}", 401
        
        nameid = saml_auth.get_nameid()
        print(f"[CALLBACK] Auth SUCCESS! User: {nameid}")
        
        response = make_response(f"<h1>SUCCESS!</h1><p>User: {nameid}</p><p>Server Port: {PORT}</p>")
        response.set_cookie("session", f"authenticated_{nameid}", httponly=True)
        return response
        
    except Exception as e:
        print(f"[CALLBACK] Exception: {e}")
        return f"SAML Error: {e}", 500

if __name__ == "__main__":
    if not SAML_METADATA_URL:
        print("Usage: python mock_airflow.py <SAML_METADATA_URL> [PORT]")
        sys.exit(1)
    
    print(f"\n{'='*60}")
    print(f"  Mock Airflow - VULNERABLE to Host Header Injection")
    print(f"  Port: {PORT}")
    print(f"  Metadata: {SAML_METADATA_URL}")
    print(f"{'='*60}\n")
    
    app.run(host="0.0.0.0", port=PORT, debug=True)