import base64
import hashlib
import hmac
import html
import json
import os
import time
from http import cookies
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Optional
from urllib.parse import parse_qs, urlencode, urlparse


PORT = int(os.environ.get("PORT", "8100"))
SECRET = b"local-sso-python-secret"
USERS = {"admin": "123456"}
TICKET_TTL = 60
TICKETS = {}


def sign(value: str) -> str:
    return hmac.new(SECRET, value.encode("utf-8"), hashlib.sha256).hexdigest()


def encode_session(payload: dict) -> str:
    raw = base64.urlsafe_b64encode(json.dumps(payload, ensure_ascii=False).encode("utf-8")).decode("utf-8")
    return f"{raw}.{sign(raw)}"


def decode_session(token: str) -> Optional[dict]:
    if "." not in token:
        return None
    raw, signature = token.rsplit(".", 1)
    if not hmac.compare_digest(signature, sign(raw)):
        return None
    try:
        return json.loads(base64.urlsafe_b64decode(raw.encode("utf-8")))
    except Exception:
        return None


def create_ticket(username: str) -> str:
    seed = f"{username}:{time.time()}:{os.urandom(8).hex()}"
    ticket = hashlib.sha256(seed.encode("utf-8")).hexdigest()
    TICKETS[ticket] = {"username": username, "expires": time.time() + TICKET_TTL}
    return ticket


def consume_ticket(ticket: str) -> Optional[str]:
    data = TICKETS.pop(ticket, None)
    if not data or data["expires"] < time.time():
        return None
    return data["username"]


class PythonSsoHandler(BaseHTTPRequestHandler):
    def send_html(self, body: str, status: int = 200, headers: Optional[dict] = None) -> None:
        content = body.encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(content)))
        if headers:
            for key, value in headers.items():
                self.send_header(key, value)
        self.end_headers()
        self.wfile.write(content)

    def redirect(self, target: str, headers: Optional[dict] = None) -> None:
        self.send_response(302)
        self.send_header("Location", target)
        if headers:
            for key, value in headers.items():
                self.send_header(key, value)
        self.end_headers()

    def read_cookies(self) -> cookies.SimpleCookie:
        jar = cookies.SimpleCookie()
        jar.load(self.headers.get("Cookie", ""))
        return jar

    def get_auth_user(self) -> Optional[str]:
        jar = self.read_cookies()
        token = jar.get("sso_session")
        if not token:
            return None
        payload = decode_session(token.value)
        return payload.get("username") if payload else None

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/":
            self.render_home()
        elif parsed.path == "/sso/login":
            self.render_login(parsed.query)
        elif parsed.path == "/app":
            self.render_app(parsed.query)
        elif parsed.path == "/app/callback":
            self.handle_callback(parsed.query)
        elif parsed.path == "/logout":
            self.handle_logout()
        else:
            self.send_html("not found", 404)

    def do_POST(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/sso/login":
            self.handle_login()
        else:
            self.send_html("not found", 404)

    def render_home(self) -> None:
        self.send_html(page_layout("Python 单点登录示例", """
          <section class="hero">
            <h1>Python SSO 单点登录 Demo</h1>
            <p>同一个服务模拟认证中心和业务应用：未登录访问业务应用会跳到认证中心，登录成功后用一次性 ticket 回跳，业务应用消费 ticket 后写入自己的会话。</p>
            <a class="button" href="/app">进入业务应用</a>
          </section>
        """ + docs_html("Python")))

    def render_login(self, query: str) -> None:
        params = parse_qs(query)
        redirect_uri = params.get("redirect_uri", ["/app/callback"])[0]
        service = params.get("service", ["demo-app"])[0]
        user = self.get_auth_user()
        if user:
            ticket = create_ticket(user)
            self.redirect(f"{redirect_uri}?{urlencode({'ticket': ticket})}")
            return

        self.send_html(page_layout("认证中心登录", f"""
          <section class="panel narrow">
            <h1>认证中心</h1>
            <p>账号：admin，密码：123456。登录后认证中心写入 <code>sso_session</code>，并签发一次性 ticket 给业务应用。</p>
            <form method="post" action="/sso/login">
              <input type="hidden" name="redirect_uri" value="{html.escape(redirect_uri)}">
              <input type="hidden" name="service" value="{html.escape(service)}">
              <label>用户名<input name="username" value="admin"></label>
              <label>密码<input name="password" type="password" value="123456"></label>
              <button type="submit">登录并回跳</button>
            </form>
          </section>
        """))

    def handle_login(self) -> None:
        length = int(self.headers.get("Content-Length", "0"))
        payload = parse_qs(self.rfile.read(length).decode("utf-8"))
        username = payload.get("username", [""])[0]
        password = payload.get("password", [""])[0]
        redirect_uri = payload.get("redirect_uri", ["/app/callback"])[0]
        if USERS.get(username) != password:
            self.send_html(page_layout("登录失败", "<section class='panel'><h1>登录失败</h1><p>用户名或密码错误。</p><a href='/app'>重新登录</a></section>"), 401)
            return

        ticket = create_ticket(username)
        session = encode_session({"username": username, "iat": int(time.time())})
        self.redirect(f"{redirect_uri}?{urlencode({'ticket': ticket})}", {
            "Set-Cookie": f"sso_session={session}; HttpOnly; Path=/; SameSite=Lax"
        })

    def render_app(self, query: str) -> None:
        jar = self.read_cookies()
        token = jar.get("app_session")
        payload = decode_session(token.value) if token else None
        if not payload:
            target = "/sso/login?" + urlencode({"service": "demo-app", "redirect_uri": "/app/callback"})
            self.redirect(target)
            return

        username = payload["username"]
        self.send_html(page_layout("业务应用", f"""
          <section class="hero">
            <h1>业务应用已登录</h1>
            <p>当前用户：<strong>{html.escape(username)}</strong>。业务应用没有保存密码，只保存通过 ticket 换来的本地会话。</p>
            <a class="button" href="/logout">退出登录</a>
          </section>
        """ + docs_html("Python")))

    def handle_callback(self, query: str) -> None:
        ticket = parse_qs(query).get("ticket", [""])[0]
        username = consume_ticket(ticket)
        if not username:
            self.send_html(page_layout("ticket 无效", "<section class='panel'><h1>ticket 无效或已过期</h1><a href='/app'>重新登录</a></section>"), 401)
            return

        session = encode_session({"username": username, "iat": int(time.time())})
        self.redirect("/app", {"Set-Cookie": f"app_session={session}; HttpOnly; Path=/; SameSite=Lax"})

    def handle_logout(self) -> None:
        self.redirect("/", {
            "Set-Cookie": "app_session=; Max-Age=0; Path=/; HttpOnly; SameSite=Lax, sso_session=; Max-Age=0; Path=/; HttpOnly; SameSite=Lax"
        })


def docs_html(language: str) -> str:
    return f"""
      <section class="grid">
        <article class="panel"><h2>实现了哪些功能</h2><ul><li>认证中心登录</li><li>业务应用登录拦截</li><li>一次性 ticket 回跳</li><li>业务应用本地会话</li><li>单点退出</li></ul></article>
        <article class="panel"><h2>怎么实现</h2><p>访问业务应用时先检查 app_session。没有会话就跳转到认证中心。认证中心登录成功后签发短期 ticket，业务应用在 callback 中消费 ticket 并生成自己的 session。</p></article>
        <article class="panel"><h2>主要代码段</h2><pre><code>ticket = create_ticket(username)
redirect('/app/callback?ticket=' + ticket)

username = consume_ticket(ticket)
set_cookie('app_session', sign(username))</code></pre></article>
        <article class="panel"><h2>整体逻辑</h2><ol><li>用户访问 /app</li><li>业务应用发现未登录，跳到 /sso/login</li><li>认证中心校验账号密码</li><li>认证中心签发一次性 ticket 并回跳</li><li>业务应用消费 ticket，写入 app_session</li></ol></article>
        <article class="panel"><h2>注意事项</h2><ul><li>ticket 必须短期有效且只能使用一次</li><li>回跳地址必须做白名单校验</li><li>Cookie 应设置 HttpOnly、Secure、SameSite</li><li>真实项目应使用 OAuth2/OIDC、CAS 或 SAML</li><li>分布式环境需要集中存储 ticket 和 session</li></ul></article>
      </section>
    """


def page_layout(title: str, content: str) -> str:
    return f"""
<!DOCTYPE html>
<html lang="zh-CN">
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>{html.escape(title)}</title>
  <style>
    * {{ box-sizing: border-box; margin: 0; padding: 0; }}
    body {{ min-height: 100vh; background: #f5f2ea; color: #1e252b; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", "Microsoft YaHei", sans-serif; line-height: 1.7; }}
    a {{ color: #155f8f; }}
    .page {{ width: min(1120px, calc(100% - 32px)); margin: 0 auto; padding: 34px 0 48px; }}
    .hero, .panel {{ border: 1px solid #d7cdbc; border-radius: 8px; background: #fffdf8; padding: 24px; }}
    .hero {{ margin-bottom: 18px; }}
    h1 {{ font-size: clamp(30px, 5vw, 52px); line-height: 1.1; letter-spacing: 0; margin-bottom: 12px; }}
    h2 {{ font-size: 20px; margin-bottom: 8px; }}
    p, li {{ color: #53606b; }}
    .button, button {{ display: inline-block; margin-top: 16px; border: 0; border-radius: 6px; background: #155f8f; color: white; padding: 10px 16px; font: inherit; font-weight: 800; text-decoration: none; cursor: pointer; }}
    .grid {{ display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; }}
    .narrow {{ max-width: 520px; margin: 60px auto; }}
    label {{ display: block; margin-top: 12px; color: #394650; }}
    input {{ width: 100%; margin-top: 4px; border: 1px solid #c9d1d8; border-radius: 6px; padding: 10px 12px; font: inherit; }}
    ul, ol {{ padding-left: 20px; }}
    pre {{ overflow-x: auto; border-radius: 6px; background: #111827; color: #d7e7ff; padding: 14px; font: 13px/1.55 Consolas, Monaco, monospace; }}
    code {{ font-family: Consolas, Monaco, monospace; }}
    @media (max-width: 760px) {{ .grid {{ grid-template-columns: 1fr; }} }}
  </style>
</head>
<body><main class="page">{content}</main></body>
</html>
    """


if __name__ == "__main__":
    print(f"Python SSO demo running at http://127.0.0.1:{PORT}")
    ThreadingHTTPServer(("127.0.0.1", PORT), PythonSsoHandler).serve_forever()