File size: 4,308 Bytes
a5784e9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
549d337
a5784e9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
#!/usr/bin/env python3
"""Runtime bootstrap for Hugging Face Spaces."""

from __future__ import annotations

import base64
import json
import os
import re
import sys
from pathlib import Path


APP_DIR = Path(__file__).resolve().parents[2]
AUTH_DIR = APP_DIR / "auth_profiles" / "active"
KEY_FILE = APP_DIR / "auth_profiles" / "key.txt"
HF_AUTH_PATH = AUTH_DIR / "hf-auth.json"


def log(message: str) -> None:
    print(f"[hf-entrypoint] {message}", flush=True)


def _decode_auth_secret() -> str | None:
    raw_json = os.environ.get("AISTUDIO_AUTH_JSON", "").strip()
    if raw_json:
        return raw_json

    encoded_json = os.environ.get("AISTUDIO_AUTH_JSON_B64", "").strip()
    if not encoded_json:
        return None

    compact = "".join(encoded_json.split())
    try:
        return base64.b64decode(compact, validate=True).decode("utf-8")
    except Exception as exc:
        raise RuntimeError("AISTUDIO_AUTH_JSON_B64 is not valid base64 JSON") from exc


def _write_auth_from_secret() -> str | None:
    auth_secret = _decode_auth_secret()
    if not auth_secret:
        return None

    try:
        auth_data = json.loads(auth_secret)
    except json.JSONDecodeError as exc:
        raise RuntimeError("AISTUDIO auth secret is not valid JSON") from exc

    if not isinstance(auth_data, dict) or not isinstance(auth_data.get("cookies"), list):
        raise RuntimeError("AISTUDIO auth JSON must be a Playwright storage_state file")

    AUTH_DIR.mkdir(parents=True, exist_ok=True)
    HF_AUTH_PATH.write_text(
        json.dumps(auth_data, ensure_ascii=False),
        encoding="utf-8",
    )
    HF_AUTH_PATH.chmod(0o600)
    log(f"Wrote auth profile to {HF_AUTH_PATH}")
    return str(HF_AUTH_PATH)


def resolve_auth_path() -> str:
    configured_path = os.environ.get("ACTIVE_AUTH_JSON_PATH", "").strip()
    if configured_path:
        path = Path(configured_path)
        if path.exists():
            log(f"Using ACTIVE_AUTH_JSON_PATH={path}")
            return str(path)
        raise RuntimeError(f"ACTIVE_AUTH_JSON_PATH does not exist: {path}")

    generated_path = _write_auth_from_secret()
    if generated_path:
        return generated_path

    existing_profiles = sorted(AUTH_DIR.glob("*.json"))
    if existing_profiles:
        log(f"Using existing auth profile {existing_profiles[0]}")
        return str(existing_profiles[0])

    raise RuntimeError(
        "No auth profile found. Add AISTUDIO_AUTH_JSON_B64 or AISTUDIO_AUTH_JSON "
        "as a Hugging Face Space Secret."
    )


def write_api_keys() -> None:
    raw_keys = (
        os.environ.get("AISTUDIO_PROXY_API_KEYS")
        or os.environ.get("AISTUDIO_PROXY_API_KEY")
        or os.environ.get("API_KEYS")
        or os.environ.get("API_KEY")
        or ""
    )
    keys = [key.strip() for key in re.split(r"[\n,;]+", raw_keys) if key.strip()]

    KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
    if keys:
        KEY_FILE.write_text("\n".join(keys) + "\n", encoding="utf-8")
        KEY_FILE.chmod(0o600)
        log(f"Configured {len(keys)} API key(s)")
    else:
        KEY_FILE.touch(exist_ok=True)
        log("No API key secret configured; /v1 endpoints will be public")


def main() -> None:
    os.chdir(APP_DIR)
    active_auth_path = resolve_auth_path()
    write_api_keys()

    server_port = os.environ.get("PORT") or os.environ.get("SERVER_PORT") or "7860"
    stream_port = os.environ.get("STREAM_PORT", "0")

    os.environ["PORT"] = server_port
    os.environ["SERVER_PORT"] = server_port
    os.environ.setdefault("DEFAULT_FASTAPI_PORT", server_port)
    os.environ["ACTIVE_AUTH_JSON_PATH"] = active_auth_path

    command = [
        sys.executable,
        "launch_camoufox.py",
        "--headless",
        "--server-port",
        server_port,
        "--stream-port",
        stream_port,
        "--active-auth-json",
        active_auth_path,
        "--helper",
        "",
    ]

    internal_proxy = os.environ.get("INTERNAL_CAMOUFOX_PROXY")
    if internal_proxy is not None:
        command.extend(["--internal-camoufox-proxy", internal_proxy])

    log("Starting AI Studio Proxy API")
    os.execvp(command[0], command)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        log(f"Startup failed: {exc}")
        sys.exit(1)