#!/usr/bin/env python3
"""
Superset Auth Diagnostics Script — Extended v3

Проверяет все доступные способы аутентификации в Superset
ПОЛНЫМ ЦИКЛОМ: логин → сессия → загрузка дашборда → верификация контента.

Использование:
  python3 superset_auth_diag.py --url https://superset.example.com --username admin

Опции:
  --insecure, -k    отключить проверку SSL
  --dashboard-id    конкретный дашборд для теста

Требует: pip install httpx
"""

import argparse
import json
import re
import sys
from getpass import getpass
from urllib.parse import urlparse

try:
    import httpx
except ImportError:
    print("Установите httpx: pip install httpx")
    sys.exit(1)


# ── helpers ──────────────────────────────────────────────────

def banner(msg):
    print(f"\n{'='*60}")
    print(f"  {msg}")
    print(f"{'='*60}")

def ok(msg):
    print(f"  ✅ {msg}")

def fail(msg):
    print(f"  ❌ {msg}")

def warn(msg):
    print(f"  ⚠️  {msg}")

def section(msg):
    print(f"\n  --- {msg} ---")

def http_client(verify: bool = True) -> httpx.Client:
    import warnings
    if not verify:
        warnings.filterwarnings("ignore", message=".*verify=False.*", category=UserWarning)
    return httpx.Client(verify=verify, timeout=30, follow_redirects=False)

def extract_session_cookie(set_cookie_header: str) -> dict | None:
    """Извлекает session=... из Set-Cookie заголовка."""
    if not set_cookie_header:
        return None
    for part in set_cookie_header.split(";"):
        part = part.strip()
        if part.startswith("session="):
            value = part.split("=", 1)[1]
            return {"name": "session", "value": value}
    return None


# ── 1. Reachability ──────────────────────────────────────────

def check_reachability(base_url: str, verify: bool) -> bool:
    banner("1. Доступность Superset")
    with http_client(verify) as c:
        try:
            r = c.get(f"{base_url}/health/")
            ok(f"GET /health/ → HTTP {r.status_code}")
            return True
        except httpx.RequestError:
            pass
        try:
            r = c.get(base_url)
            ok(f"GET {base_url} → HTTP {r.status_code}")
            return True
        except httpx.RequestError as e:
            fail(f"Недоступен: {e}")
            return False


# ── 2. JWT Login ─────────────────────────────────────────────

def check_jwt_login(base_url: str, username: str, password: str, verify: bool) -> str | None:
    banner("2. JWT API Login")
    with http_client(verify) as c:
        for provider in ("db", "ldap"):
            try:
                r = c.post(
                    f"{base_url}/api/v1/security/login",
                    json={"username": username, "password": password, "provider": provider, "refresh": True},
                )
                data = r.json() if r.text else {}
                if r.status_code == 200 and data.get("access_token"):
                    ok(f"provider='{provider}' → JWT OK")
                    if data.get("refresh_token"):
                        ok(f"  Refresh token OK")
                    return data["access_token"]
                fail(f"provider='{provider}' → HTTP {r.status_code}: {str(data)[:200]}")
            except Exception as e:
                fail(f"provider='{provider}' → ошибка: {e}")
    return None


# ── 3. JWT endpoints ─────────────────────────────────────────

def check_jwt_endpoints(base_url: str, token: str, verify: bool):
    banner("3. JWT — доступ к API endpoints")
    headers = {"Authorization": f"Bearer {token}"}
    with http_client(verify) as c:
        for path in ("/api/v1/me/", "/api/v1/dashboard/", "/api/v1/security/csrf_token/"):
            try:
                r = c.get(f"{base_url}{path}", headers=headers)
                ok(f"GET {path} → HTTP {r.status_code}") if r.status_code in (200, 201) \
                    else fail(f"GET {path} → HTTP {r.status_code}")
            except Exception as e:
                fail(f"GET {path} → ошибка: {e}")


# ── 3b. CSRF token + session cookie ─────────────────────────

def check_csrf_session(base_url: str, token: str, verify: bool) -> tuple[str | None, dict | None]:
    """
    GET /api/v1/security/csrf_token/ с JWT.
    Возвращает (csrf_token, session_cookie_dict).
    """
    section("3b. CSRF token + Set-Cookie session")
    csrf_token = None
    csrf_session_cookie = None
    with http_client(verify) as c:
        try:
            r = c.get(f"{base_url}/api/v1/security/csrf_token/",
                      headers={"Authorization": f"Bearer {token}"})
            data = r.json() if r.text else {}
            if r.status_code == 200:
                csrf_token = (data.get("result") or data.get("csrf_token") or "").strip('"')
                if csrf_token:
                    ok(f"CSRF token: {csrf_token[:30]}...")
                # Извлекаем session cookie из Set-Cookie
                set_cookie = r.headers.get("set-cookie", "")
                csrf_session = extract_session_cookie(set_cookie)
                if csrf_session:
                    ok(f"Session cookie с CSRF: {csrf_session['value'][:30]}...")
                    csrf_session_cookie = csrf_session
                else:
                    warn("Set-Cookie с session не найден в ответе CSRF endpoint")
            else:
                fail(f"GET csrf_token/ → HTTP {r.status_code}")
        except Exception as e:
            warn(f"GET csrf_token/ ошибка: {e}")

    # Также проверим, есть ли session cookie от GET / (anonymous)
    section("3c. GET / (anonymous) — Set-Cookie")
    with http_client(verify) as c:
        try:
            r = c.get(base_url, follow_redirects=False)
            sc = r.headers.get("set-cookie", "")
            if "session" in sc:
                sess = extract_session_cookie(sc)
                ok(f"GET / → HTTP {r.status_code}, session cookie есть (anonymous)")
                if sess:
                    print(f"         session: {sess['value'][:40]}...")
            else:
                warn(f"GET / → HTTP {r.status_code}, session cookie НЕТ")
        except Exception as e:
            warn(f"GET / ошибка: {e}")

    return csrf_token, csrf_session_cookie


# ── 4. Guest Token ───────────────────────────────────────────

def check_guest_token(base_url: str, token: str,
                      csrf_token: str | None, csrf_session: dict | None,
                      verify: bool, dashboard_ids: list[str] | None) -> list[str] | None:
    banner("4. Guest Token")
    bearer = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

    with http_client(verify) as c:
        # Получаем список дашбордов если не задан
        if not dashboard_ids:
            section("Получаем список дашбордов")
            try:
                r = c.get(f"{base_url}/api/v1/dashboard/", headers=bearer)
                data = r.json() if r.text else {}
                if r.status_code == 200:
                    dashboards = data.get("result", [])
                    ok(f"Найдено дашбордов: {len(dashboards)}")
                    for d in dashboards[:3]:
                        print(f"         id={d['id']}  {d.get('dashboard_title','')}")
                    dashboard_ids = [str(d["id"]) for d in dashboards[:1]]
                else:
                    warn(f"HTTP {r.status_code}, используем dashboard_id=1")
                    dashboard_ids = ["1"]
            except Exception as e:
                warn(f"Ошибка: {e}, используем dashboard_id=1")
                dashboard_ids = ["1"]

        for did in dashboard_ids:
            payload = {
                "user": {"username": "admin"},
                "resources": [{"type": "dashboard", "id": did}],
                "rls": [],
            }

            # Способ 1: только JWT (без CSRF)
            try:
                r = c.post(f"{base_url}/api/v1/security/guest_token/", headers=bearer, json=payload)
                if r.status_code == 200:
                    ok(f"dashboard #{did}: guest token (без CSRF) 🎉")
                    continue
                warn(f"dashboard #{did}: HTTP {r.status_code} (без CSRF)")
            except Exception as e:
                warn(f"dashboard #{did}: ошибка (без CSRF): {e}")

            # Способ 2: JWT + X-CSRFToken header
            if csrf_token:
                h = {**bearer, "X-CSRFToken": csrf_token, "Referer": f"{base_url}/"}
                try:
                    r = c.post(f"{base_url}/api/v1/security/guest_token/", headers=h, json=payload)
                    if r.status_code == 200:
                        ok(f"dashboard #{did}: guest token (CSRF header) 🎉")
                        continue
                    warn(f"dashboard #{did}: HTTP {r.status_code} (CSRF header): {r.text[:100]}")
                except Exception as e:
                    warn(f"dashboard #{did}: ошибка (CSRF header): {e}")

            # Способ 3: JWT + X-CSRFToken + session cookie (из CSRF endpoint)
            if csrf_token and csrf_session:
                h = {**bearer, "X-CSRFToken": csrf_token, "Referer": f"{base_url}/"}
                try:
                    print(f"\n         >>> DEBUG: sending guest_token request...")
                    print(f"         >>> X-CSRFToken: {csrf_token[:20]}...")
                    print(f"         >>> Cookie: session={csrf_session['value'][:20]}...")
                    r = c.post(f"{base_url}/api/v1/security/guest_token/",
                               headers=h, json=payload,
                               cookies={"session": csrf_session["value"]})
                    print(f"         >>> Response: HTTP {r.status_code}")
                    print(f"         >>> Set-Cookie: {r.headers.get('set-cookie', 'NONE')[:100]}")
                    if r.status_code == 200:
                        ok(f"dashboard #{did}: guest token (JWT+CSRF+sess) 🎉")
                        continue
                    else:
                        print(f"         >>> Body: {r.text[:200]}")
                        warn(f"dashboard #{did}: HTTP {r.status_code} (JWT+CSRF+sess)")
                except Exception as e:
                    warn(f"dashboard #{did}: ошибка (JWT+CSRF+sess): {e}")
            else:
                warn(f"dashboard #{did}: пропускаем JWT+CSRF+sess — нет csrf_token или csrf_session")

    return dashboard_ids


# ── 5. Form Login ────────────────────────────────────────────

def check_form_login(base_url: str, username: str, password: str, verify: bool) -> str | None:
    """POST /login/, возвращает сырой Set-Cookie или None."""
    banner("5. Form Login (POST /login/)")
    with http_client(verify) as c:
        # GET /login/ — что показывает
        try:
            r_get = c.get(f"{base_url}/login/", follow_redirects=False)
            warn(f"GET /login/ → HTTP {r_get.status_code}")
            if r_get.status_code in (301, 302):
                loc = r_get.headers.get("location", "")
                warn(f"  Редирект на: {loc.split('?')[0]}")
                if "adfs" in loc.lower():
                    warn("  ADFS — db auth отключён на странице, POST может не работать")
        except Exception as e:
            warn(f"GET /login/ ошибка: {e}")

        # POST с credentials
        try:
            r = c.post(f"{base_url}/login/",
                       data={"username": username, "password": password},
                       follow_redirects=False)
            set_cookie = r.headers.get("set-cookie", "")
            session_cookie = extract_session_cookie(set_cookie)

            if session_cookie:
                ok(f"POST /login/ → HTTP {r.status_code}, session cookie получена")
                print(f"         session: {session_cookie['value'][:40]}...")
                return set_cookie

            if r.status_code in (301, 302):
                loc = r.headers.get("location", "")
                if "/login" not in loc:
                    ok(f"POST /login/ → HTTP {r.status_code} → {loc.split('?')[0]} — login успешен")
                    if set_cookie:
                        return set_cookie
                fail(f"POST /login/ → редирект на login")
            else:
                fail(f"POST /login/ → HTTP {r.status_code}, session не найдена")
        except Exception as e:
            fail(f"POST /login/ ошибка: {e}")
    return None


# ── 6. Session Cookie → Dashboard ────────────────────────────

def check_session_to_dashboard(base_url: str, set_cookie_header: str | None,
                               verify: bool, dashboard_ids: list[str] | None):
    banner("6. Session Cookie → Dashboard (ПОЛНЫЙ ЦИКЛ)")

    if not set_cookie_header:
        fail("Нет session cookie — проверка невозможна")
        return

    session_cookie = extract_session_cookie(set_cookie_header)
    if not session_cookie:
        fail("Не удалось извлечь session cookie")
        return

    with http_client(verify) as c:
        # 6a. Главная страница
        section("6a. Главная страница с session cookie")
        try:
            r = c.get(base_url, cookies={"session": session_cookie["value"]}, follow_redirects=False)
            if r.status_code == 200:
                ok(f"GET / → HTTP 200, SPA загружен")
            elif r.status_code in (301, 302):
                loc = r.headers.get("location", "")
                if "login" in loc.lower():
                    fail(f"GET / → редирект на login — session cookie НЕ РАБОТАЕТ")
                else:
                    ok(f"GET / → HTTP {r.status_code} → {loc.split('?')[0]} (редирект, не login)")
            else:
                warn(f"GET / → HTTP {r.status_code}")
        except Exception as e:
            fail(f"GET / ошибка: {e}")

        # 6b. Dashboard
        section("6b. Dashboard с session cookie")
        if not dashboard_ids:
            dashboard_ids = ["1"]
        for did in dashboard_ids[:2]:
            dash_url = f"{base_url}/superset/dashboard/{did}/?standalone=true"
            try:
                r = c.get(dash_url, cookies={"session": session_cookie["value"]}, follow_redirects=False)
                if r.status_code == 200:
                    ok(f"dashboard #{did} → HTTP 200 — ДАШБОРД ЗАГРУЖЕН 🎉")
                    has_spa = "spa.html" in r.text or "bootstrap" in r.text
                    ok("  SPA/bootstrap_data найден") if has_spa else warn("  SPA/bootstrap_data НЕ найден")
                    print(f"         Размер HTML: {len(r.text)} bytes")
                elif r.status_code in (301, 302):
                    loc = r.headers.get("location", "")
                    fail(f"dashboard #{did} → редирект на {loc.split('?')[0]} — АУТЕНТИФИКАЦИЯ НЕ РАБОТАЕТ ❌")
                elif r.status_code == 404:
                    fail(f"dashboard #{did} → HTTP 404 — дашборд не найден")
                else:
                    warn(f"dashboard #{did} → HTTP {r.status_code}")
            except Exception as e:
                fail(f"dashboard #{did} → ошибка: {e}")


# ── 7. JWT → UI dashboard ────────────────────────────────────

def check_jwt_ui(base_url: str, token: str, verify: bool, dashboard_ids: list[str] | None):
    banner("7. JWT Bearer → Dashboard UI (без session cookie)")
    if not dashboard_ids:
        dashboard_ids = ["1"]
    with http_client(verify) as c:
        for did in dashboard_ids[:1]:
            try:
                r = c.get(f"{base_url}/superset/dashboard/{did}/?standalone=true",
                          headers={"Authorization": f"Bearer {token}"},
                          follow_redirects=False)
                if r.status_code == 200:
                    ok(f"dashboard #{did} → HTTP 200 — JWT РАБОТАЕТ ДЛЯ UI!")
                elif r.status_code in (301, 302):
                    loc = r.headers.get("location", "")
                    fail(f"dashboard #{did} → редирект на {loc.split('?')[0]} — JWT НЕ РАБОТАЕТ")
                else:
                    warn(f"dashboard #{did} → HTTP {r.status_code}")
            except Exception as e:
                warn(f"dashboard #{did} → ошибка: {e}")


# ── MAIN ──────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Superset Auth Diagnostics (Extended v3)")
    parser.add_argument("--url", required=True, help="Superset base URL")
    parser.add_argument("--username", default="admin", help="Username")
    parser.add_argument("--password", help="Password (будет запрошен)")
    parser.add_argument("--dashboard-id", help="ID дашборда для теста")
    parser.add_argument("--insecure", "-k", action="store_true", help="Отключить проверку SSL")

    args = parser.parse_args()
    base_url = args.url.rstrip("/")
    username = args.username
    password = args.password or getpass(f"Пароль для {username}: ")
    verify = not args.insecure
    dashboard_ids = [args.dashboard_id] if args.dashboard_id else None

    if not verify:
        import urllib3
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    print(f"\n🔍 Superset Auth Diagnostics (Extended v3)")
    print(f"   URL:      {base_url}")
    print(f"   Username: {username}")
    print(f"   SSL:      {'ON' if verify else 'OFF (--insecure)'}")

    # 1. Reachability
    if not check_reachability(base_url, verify):
        print("\n❌ Superset недоступен.")
        sys.exit(1)

    # 2. JWT Login
    jwt_token = check_jwt_login(base_url, username, password, verify)

    csrf_token = None
    csrf_session_cookie = None

    if jwt_token:
        # 3. JWT endpoints
        check_jwt_endpoints(base_url, jwt_token, verify)

        # 3b. CSRF token + session cookie
        csrf_token, csrf_session_cookie = check_csrf_session(base_url, jwt_token, verify)

        # 4. Guest Token
        dashboard_ids = check_guest_token(base_url, jwt_token, csrf_token, csrf_session_cookie,
                                          verify, dashboard_ids)

        # 7. JWT → UI
        check_jwt_ui(base_url, jwt_token, verify, dashboard_ids)

    # 5. Form Login
    session_cookie = check_form_login(base_url, username, password, verify)

    # 6. Session → Dashboard
    check_session_to_dashboard(base_url, session_cookie, verify, dashboard_ids)

    # ── ИТОГ ──
    banner("РЕЗЮМЕ")
    jwt_ok = jwt_token is not None
    form_ok = session_cookie is not None

    ok("JWT API Login: РАБОТАЕТ") if jwt_ok else fail("JWT API Login: НЕ РАБОТАЕТ")
    ok("Form Login: РАБОТАЕТ") if form_ok else fail("Form Login: НЕ РАБОТАЕТ")

    print()
    if jwt_ok and form_ok:
        print("🏆  ОБА ПУТИ РАБОТАЮТ")
        print()
        print("   Реализация ScreenshotService:")
        print("   1. POST /api/v1/security/login → JWT")
        print("   2. POST /login/ через request_context → session cookie")
        print("   3. context.add_cookies([session_cookie])")
        print("   4. page.goto('/superset/dashboard/{id}/?standalone=true')")
        print("   5. Скриншот")
    elif jwt_ok:
        print("⚠️  Только JWT — нужен EMBEDDED_SUPERSET")
    else:
        print("❌  НИ ОДИН МЕТОД НЕ РАБОТАЕТ")


if __name__ == "__main__":
    main()
