#!/usr/bin/env python3
"""
register_oauth_app.py
Automated OAuth app registration & Vault JWT wiring for homelab services.

This script automates the full workflow of registering a new service with
Authentik OIDC and wiring it for Vault JWT authentication.

Usage:
  python3 register_oauth_app.py \\
    --service-name my-service \\
    --namespace my-ns \\
    --redirect-uri "https://my-service.riotpiao.com/oauth2/callback" \\
    [--service-name-in-vault MY_SERVICE] \\
    [--vault-jwt-policy shell-secrets]

Requirements:
  - Authentik running in iam namespace (port-forward 7000:80 active)
  - Vault running in storage namespace
  - kubectl configured (KUBECONFIG → cluster-config/kubeconfig)
  - Environment variables set:
      AUTHENTIK_BOOTSTRAP_TOKEN       — Authentik API token
      {SERVICE_NAME}_OIDC_CLIENT_SECRET  — OAuth client secret (generated or from Vault)

Examples:

  # Simple: register a web app with default settings
  python3 register_oauth_app.py \\
    --service-name myapp \\
    --namespace apps \\
    --redirect-uri "https://myapp.riotpiao.com/callback"

  # Advanced: full service with Vault JWT auth
  python3 register_oauth_app.py \\
    --service-name myservice \\
    --namespace my-namespace \\
    --redirect-uri "https://myservice.riotpiao.com/oauth2/callback" \\
    --service-name-in-vault MYSERVICE \\
    --vault-jwt-policy shell-secrets \\
    --add-group myservice-admins \\
    --vault-jwt-bound-claims '{"groups":["myservice-admins"]}'
"""

import argparse
import json
import os
import subprocess
import sys
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from typing import Optional


BASE_AUTHENTIK_URL = "http://localhost:7000/api/v3"
VAULT_NAMESPACE = "storage"
IAM_NAMESPACE = "iam"


def _auth_headers() -> dict:
    """Return HTTP headers for Authentik API requests."""
    token = os.environ.get("AUTHENTIK_BOOTSTRAP_TOKEN", "")
    if not token:
        sys.exit(
            "ERROR: AUTHENTIK_BOOTSTRAP_TOKEN not set\n"
            "  Set it in ~/.authentik/.env or export it:\n"
            "    talos get cluster/AUTHENTIK_BOOTSTRAP_TOKEN --key AUTHENTIK_BOOTSTRAP_TOKEN"
        )
    return {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }


def _request(method: str, path: str, body: Optional[dict] = None, **params) -> dict:
    """HTTP request helper for Authentik API."""
    url = BASE_AUTHENTIK_URL + path
    if params:
        url += "?" + urllib.parse.urlencode(params)
    
    req = urllib.request.Request(
        url,
        data=json.dumps(body).encode() if body else None,
        headers=_auth_headers(),
        method=method,
    )
    
    try:
        with urllib.request.urlopen(req) as resp:
            raw = resp.read()
            return json.loads(raw) if raw else {}
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")
        print(f"ERROR {method} {path}: HTTP {exc.code} — {detail}", file=sys.stderr)
        raise


def get_pk(path: str, **filters) -> Optional[str]:
    """Fetch the PK of the first result matching filters."""
    results = _request("GET", path, **filters).get("results", [])
    return results[0]["pk"] if results else None


def get_or_create(path: str, filters: dict, data: dict) -> str:
    """Create or fetch a resource; returns PK."""
    pk = get_pk(path, **filters)
    if pk is None:
        pk = _request("POST", path, data)["pk"]
        print(f"✓ created  {path.rstrip('/')} {filters}")
    else:
        print(f"✓ exists   {path.rstrip('/')} {filters}")
    return pk


def fetch_flows_and_signing_key() -> tuple[str, str, str]:
    """Fetch required Authentik resource PKs."""
    auth_pk = get_pk(
        "/flows/instances/",
        slug="default-provider-authorization-implicit-consent"
    )
    inval_pk = get_pk(
        "/flows/instances/",
        slug="default-provider-invalidation-flow"
    )
    
    if not auth_pk or not inval_pk:
        sys.exit(
            "ERROR: required Authentik flows not found\n"
            "  Run provision_oidc.py first to initialize defaults"
        )
    
    results = _request("GET", "/crypto/certificatekeypairs/", has_key=True).get("results", [])
    signing_pk = results[0]["pk"] if results else None
    
    if not signing_pk:
        sys.exit(
            "ERROR: no signing key found in Authentik\n"
            "  Create one: Authentik UI → System → Certificates"
        )
    
    return auth_pk, inval_pk, signing_pk


def apply_k8s_secret(name: str, namespace: str, literals: dict) -> None:
    """Create or update a Kubernetes secret."""
    env = {**os.environ}
    manifest = subprocess.run(
        [
            "kubectl", "create", "secret", "generic", name,
            "-n", namespace, "--dry-run=client", "-o", "yaml",
            *[f"--from-literal={k}={v}" for k, v in literals.items()],
        ],
        capture_output=True, text=True, check=True, env=env,
    ).stdout
    
    subprocess.run(
        ["kubectl", "apply", "-f", "-"],
        input=manifest, text=True, check=True, env=env
    )
    print(f"✓ secret   {name} → ns/{namespace}")


def create_authentik_provider(
    service_name: str,
    client_id: str,
    client_secret: str,
    redirect_uris: list[str],
    auth_flow_pk: str,
    inval_flow_pk: str,
    signing_key_pk: str,
) -> str:
    """Create OAuth2 provider in Authentik."""
    data = {
        "name": service_name,
        "client_id": client_id,
        "client_secret": client_secret,
        "authorization_flow": auth_flow_pk,
        "invalidation_flow": inval_flow_pk,
        "redirect_uris": [{"matching_mode": "strict", "url": u} for u in redirect_uris],
        "sub_mode": "hashed_user_id",
        "include_claims_in_id_token": True,
        "property_mappings": [],  # Will be populated by caller if needed
        "grant_types": ["authorization_code", "refresh_token"],
        "signing_key": signing_key_pk,
    }
    
    pk = get_pk("/providers/oauth2/", name=service_name)
    if pk is None:
        pk = _request("POST", "/providers/oauth2/", data)["pk"]
        print(f"✓ created  provider={service_name}  client_id={client_id}")
    else:
        _request("PATCH", f"/providers/oauth2/{pk}/", data)
        print(f"✓ synced   provider={service_name}  (credentials from Vault)")
    
    return pk


def create_authentik_application(
    display_name: str,
    slug: str,
    provider_pk: str,
    launch_url: Optional[str] = None,
) -> str:
    """Create application in Authentik."""
    data = {"name": display_name, "slug": slug, "provider": provider_pk}
    if launch_url:
        data["meta_launch_url"] = launch_url
    
    return get_or_create(
        "/core/applications/",
        filters={"slug": slug},
        data=data
    )


def create_authentik_group(group_name: str) -> str:
    """Create a group in Authentik."""
    return get_or_create(
        "/core/groups/",
        filters={"name": group_name},
        data={"name": group_name, "is_superuser": False}
    )


def bind_group_to_app(app_pk: str, group_pk: str) -> None:
    """Bind a group to an application (allow-list policy)."""
    existing = _request(
        "GET", "/policies/bindings/",
        target=app_pk, group=group_pk
    ).get("results", [])
    
    if existing:
        print(f"✓ exists   binding group={group_pk} → app={app_pk}")
    else:
        _request("POST", "/policies/bindings/", {
            "target": app_pk,
            "group": group_pk,
            "enabled": True,
            "order": 0,
        })
        print(f"✓ created  binding group={group_pk} → app={app_pk}")


def create_vault_jwt_policy(
    service_name: str,
    vault_addr: str,
) -> str:
    """Create a Vault policy for the service."""
    policy_name = f"service-read-{service_name}"
    
    # Policy definition: service can read its own secrets
    policy_rules = f"""
path "secret/data/services/{service_name}/*" {{
  capabilities = ["read"]
}}
path "secret/data/cluster/*" {{
  capabilities = ["read"]
}}
"""
    
    # Write policy via vault CLI (requires auth)
    try:
        subprocess.run(
            ["vault", "policy", "write", policy_name, "-"],
            input=policy_rules,
            text=True,
            check=True,
            env={**os.environ, "VAULT_ADDR": vault_addr}
        )
        print(f"✓ created  Vault policy={policy_name}")
    except subprocess.CalledProcessError as e:
        print(f"⚠ warning  Failed to create Vault policy: {e}", file=sys.stderr)
        return ""
    
    return policy_name


def create_vault_jwt_role(
    service_name: str,
    vault_addr: str,
    policies: list[str],
    bound_claims: Optional[dict] = None,
) -> None:
    """Create a Vault JWT auth role for the service."""
    bound_claims_json = json.dumps(bound_claims) if bound_claims else "{}"
    
    role_config = {
        "role_type": "jwt",
        "bound_audiences": ["vault"],
        "user_claim": "sub",
        "bound_claims": bound_claims,
        "token_policies": policies,
        "token_ttl": "4h",
        "token_max_ttl": "8h",
    }
    
    try:
        subprocess.run(
            ["vault", "write", f"auth/jwt/role/{service_name}", "-"],
            input=json.dumps(role_config),
            text=True,
            check=True,
            env={**os.environ, "VAULT_ADDR": vault_addr}
        )
        print(f"✓ created  Vault JWT role={service_name}  policies={policies}")
    except subprocess.CalledProcessError as e:
        print(f"⚠ warning  Failed to create Vault JWT role: {e}", file=sys.stderr)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Register a new OAuth service with Authentik + Vault",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    
    parser.add_argument(
        "--service-name",
        required=True,
        help="Service name (used as Authentik provider/app slug)"
    )
    parser.add_argument(
        "--namespace",
        required=True,
        help="Kubernetes namespace where service runs"
    )
    parser.add_argument(
        "--redirect-uri",
        required=True,
        help="OAuth2 redirect URI (e.g., https://myapp.riotpiao.com/callback)"
    )
    parser.add_argument(
        "--service-name-in-vault",
        default=None,
        help="Service name for Vault (defaults to uppercase service-name)"
    )
    parser.add_argument(
        "--vault-jwt-policy",
        default="shell-secrets",
        help="Vault policy to attach to JWT role (comma-separated for multiple)"
    )
    parser.add_argument(
        "--add-group",
        default=None,
        help="Create and bind a group to the app (e.g., 'myservice-admins')"
    )
    parser.add_argument(
        "--vault-jwt-bound-claims",
        default=None,
        help='JSON string of claims to bind JWT role (e.g., \'{"groups":["myservice-admins"]}\')'
    )
    parser.add_argument(
        "--no-vault",
        action="store_true",
        help="Skip Vault JWT role creation (just register with Authentik)"
    )
    parser.add_argument(
        "--vault-addr",
        default=None,
        help="Vault address (defaults to localhost:8200 via port-forward)"
    )
    
    args = parser.parse_args()
    
    service_name = args.service_name.lower()
    namespace = args.namespace
    redirect_uri = args.redirect_uri
    vault_service_name = (args.service_name_in_vault or service_name).upper()
    vault_addr = args.vault_addr or "http://127.0.0.1:8200"
    vault_policies = [p.strip() for p in args.vault_jwt_policy.split(",")]
    
    # Fetch or generate client secret from environment / Vault
    client_secret_env = f"{vault_service_name}_OIDC_CLIENT_SECRET"
    client_secret = os.environ.get(client_secret_env, "")
    
    if not client_secret:
        sys.exit(
            f"ERROR: {client_secret_env} not set\n"
            f"  Generate and store in Vault:\n"
            f"    talos put cluster/{client_secret_env} "
            f"{client_secret_env}=$(openssl rand -hex 32)"
        )
    
    print(f"\n{'='*70}")
    print(f"OAuth App Registration: {service_name}")
    print(f"{'='*70}\n")
    
    print(f"📋 Configuration:")
    print(f"  Service Name:     {service_name}")
    print(f"  Namespace:        {namespace}")
    print(f"  Redirect URI:     {redirect_uri}")
    print(f"  Vault Service:    {vault_service_name}")
    print(f"  Vault Policies:   {', '.join(vault_policies)}")
    if args.add_group:
        print(f"  Group:            {args.add_group}")
    print()
    
    # Fetch Authentik resources
    print("🔍 Fetching Authentik configuration...")
    auth_flow_pk, inval_flow_pk, signing_key_pk = fetch_flows_and_signing_key()
    print(f"✓ flows & signing key found\n")
    
    # Create Authentik provider
    print("🔐 Creating Authentik OAuth2 Provider...")
    provider_pk = create_authentik_provider(
        service_name=service_name,
        client_id=service_name,
        client_secret=client_secret,
        redirect_uris=[redirect_uri],
        auth_flow_pk=auth_flow_pk,
        inval_flow_pk=inval_flow_pk,
        signing_key_pk=signing_key_pk,
    )
    print()
    
    # Create Authentik application
    print("📱 Creating Authentik Application...")
    app_pk = create_authentik_application(
        display_name=service_name.replace("-", " ").title(),
        slug=service_name,
        provider_pk=provider_pk,
        launch_url=f"https://{service_name}.riotpiao.com"
    )
    print()
    
    # Create and bind group (if requested)
    if args.add_group:
        print(f"👥 Creating Group: {args.add_group}")
        group_pk = create_authentik_group(args.add_group)
        bind_group_to_app(app_pk, group_pk)
        print()
    
    # Create K8s secret
    print("🔑 Creating Kubernetes Secret...")
    apply_k8s_secret(
        f"{service_name}-oidc",
        namespace,
        {
            "client_id": service_name,
            "client_secret": client_secret,
            "issuer_url": f"http://authentik-server.iam.svc.cluster.local/application/o/{service_name}/",
            "redirect_uri": redirect_uri,
        }
    )
    print()
    
    # Create Vault resources (if not disabled)
    if not args.no_vault:
        print("🔐 Configuring Vault JWT Authentication...")
        
        # Parse bound claims if provided
        bound_claims = None
        if args.vault_jwt_bound_claims:
            try:
                bound_claims = json.loads(args.vault_jwt_bound_claims)
            except json.JSONDecodeError as e:
                print(f"⚠ warning  Invalid JSON for --vault-jwt-bound-claims: {e}", file=sys.stderr)
        
        # Create Vault policy
        policy_name = create_vault_jwt_policy(service_name, vault_addr)
        
        # Create Vault JWT role
        if policy_name:
            all_policies = list(vault_policies) + [policy_name]
            create_vault_jwt_role(
                service_name,
                vault_addr,
                all_policies,
                bound_claims
            )
        print()
    
    # Summary
    print(f"{'='*70}")
    print(f"✅ Registration Complete!")
    print(f"{'='*70}\n")
    
    print("📝 Next Steps:\n")
    
    print(f"1. Configure your service with these environment variables:")
    print(f"   export OIDC_ISSUER_URL='http://authentik-server.iam.svc.cluster.local/application/o/{service_name}/'")
    print(f"   export OIDC_CLIENT_ID='{service_name}'")
    print(f"   export OIDC_CLIENT_SECRET='$({vault_service_name}_OIDC_CLIENT_SECRET)'")
    print(f"   export OIDC_REDIRECT_URI='{redirect_uri}'")
    print()
    
    print(f"2. Mount the Kubernetes secret in your Helm values:")
    print(f"   env:")
    print(f"     - name: OIDC_CLIENT_SECRET")
    print(f"       valueFrom:")
    print(f"         secretKeyRef:")
    print(f"           name: {service_name}-oidc")
    print(f"           key: client_secret")
    print()
    
    if not args.no_vault:
        print(f"3. Test Vault JWT authentication:")
        print(f"   kubectl port-forward -n {VAULT_NAMESPACE} svc/vault 8200:8200 &")
        print(f"   vault login -method=oidc role=homelab")
        print(f"   vault read auth/jwt/role/{service_name}")
        print()
    
    print(f"4. Verify OAuth flow:")
    print(f"   kubectl port-forward -n {IAM_NAMESPACE} svc/authentik-server 7000:80 &")
    print(f"   # Open Authentik UI: http://localhost:7000/if/admin/")
    print(f"   # Check: Applications → {service_name}")
    print()


if __name__ == "__main__":
    main()