b77a8cadb2
Some checks failed
Upstream Sync / Sync latest commits from upstream repo (push) Has been cancelled
240 lines
8.5 KiB
Python
240 lines
8.5 KiB
Python
"""
|
||
YouTube OAuth 2.0 + Live Streaming API：授权后拉取串流密钥并写入 config/youtube.ini。
|
||
需在 Google Cloud Console 为本 OAuth 客户端添加「已授权的重定向 URI」：
|
||
http://127.0.0.1:8001/youtube/oauth/callback
|
||
（若本地服务端口不是 8001，请同时改 redirect_uri 与 uvicorn 端口。）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Optional
|
||
|
||
from google.auth.transport.requests import Request
|
||
from google.oauth2.credentials import Credentials
|
||
from google_auth_oauthlib.flow import Flow
|
||
from googleapiclient.discovery import build
|
||
from googleapiclient.errors import HttpError
|
||
|
||
BASE_DIR = Path(__file__).resolve().parent
|
||
CONFIG_DIR = BASE_DIR / "config"
|
||
GOOGLE_OAUTH_JSON = CONFIG_DIR / "google_oauth.json"
|
||
TOKEN_JSON = CONFIG_DIR / "youtube_oauth_token.json"
|
||
YOUTUBE_INI = CONFIG_DIR / "youtube.ini"
|
||
|
||
# 须与 Google 在换 token 时返回的 scope 集合一致，否则 oauthlib 会报
|
||
# “Scope has changed from … to …”。include_granted_scopes 会合并历史授权，也易触发该错误。
|
||
SCOPES = [
|
||
"https://www.googleapis.com/auth/youtube",
|
||
"https://www.googleapis.com/auth/youtube.force-ssl",
|
||
"https://www.googleapis.com/auth/youtubepartner",
|
||
]
|
||
|
||
# state -> (Flow, created_ts)；回调必须用同一 Flow 实例换 token
|
||
_oauth_pending: dict[str, tuple[Flow, float]] = {}
|
||
STATE_TTL = 600.0
|
||
|
||
|
||
def _cleanup_pending() -> None:
|
||
now = time.time()
|
||
dead = [s for s, (_, t) in _oauth_pending.items() if now - t > STATE_TTL]
|
||
for s in dead:
|
||
_oauth_pending.pop(s, None)
|
||
|
||
|
||
def load_client_config_dict() -> tuple[dict[str, Any], str]:
|
||
if not GOOGLE_OAUTH_JSON.exists():
|
||
raise FileNotFoundError(
|
||
"缺少 config/google_oauth.json。请复制 config/google_oauth.example.json 并填入客户端 ID/Secret。"
|
||
)
|
||
data = json.loads(GOOGLE_OAUTH_JSON.read_text(encoding="utf-8"))
|
||
redirect_uri = data.get("redirect_uri") or "http://127.0.0.1:8001/youtube/oauth/callback"
|
||
if "web" in data:
|
||
cfg = {"web": data["web"]}
|
||
elif "installed" in data:
|
||
inst = data["installed"]
|
||
cfg = {
|
||
"web": {
|
||
"client_id": inst["client_id"],
|
||
"client_secret": inst["client_secret"],
|
||
"auth_uri": inst.get("auth_uri", "https://accounts.google.com/o/oauth2/auth"),
|
||
"token_uri": inst.get("token_uri", "https://oauth2.googleapis.com/token"),
|
||
"redirect_uris": [redirect_uri],
|
||
}
|
||
}
|
||
else:
|
||
raise ValueError("google_oauth.json 需包含 web 或 installed 字段")
|
||
return cfg, redirect_uri
|
||
|
||
|
||
def build_flow() -> tuple[Flow, str]:
|
||
cfg, redirect_uri = load_client_config_dict()
|
||
flow = Flow.from_client_config(cfg, scopes=SCOPES, redirect_uri=redirect_uri)
|
||
return flow, redirect_uri
|
||
|
||
|
||
def oauth_authorization_url() -> tuple[str, str]:
|
||
import secrets
|
||
|
||
_cleanup_pending()
|
||
flow, _ = build_flow()
|
||
state = secrets.token_urlsafe(32)
|
||
result = flow.authorization_url(
|
||
access_type="offline",
|
||
state=state,
|
||
prompt="consent",
|
||
)
|
||
authorization_url = result[0] if isinstance(result, (list, tuple)) else result
|
||
_oauth_pending[state] = (flow, time.time())
|
||
return authorization_url, state
|
||
|
||
|
||
def oauth_exchange_code(code: str, state: str) -> None:
|
||
_cleanup_pending()
|
||
if state not in _oauth_pending:
|
||
raise ValueError("授权会话无效或已过期，请在本软件中重新点击「Google 授权」。")
|
||
flow, t0 = _oauth_pending.pop(state)
|
||
if time.time() - t0 > STATE_TTL:
|
||
raise ValueError("授权已超时，请重新点击「Google 授权」。")
|
||
# oauthlib 要求令牌里的 scope 与客户端声明完全一致；Google 常返回多枚 scope 或与声明略有差异
|
||
_prev_relax = os.environ.get("OAUTHLIB_RELAX_TOKEN_SCOPE")
|
||
os.environ["OAUTHLIB_RELAX_TOKEN_SCOPE"] = "1"
|
||
try:
|
||
flow.fetch_token(code=code)
|
||
finally:
|
||
if _prev_relax is None:
|
||
os.environ.pop("OAUTHLIB_RELAX_TOKEN_SCOPE", None)
|
||
else:
|
||
os.environ["OAUTHLIB_RELAX_TOKEN_SCOPE"] = _prev_relax
|
||
creds = flow.credentials
|
||
CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
||
TOKEN_JSON.write_text(creds.to_json(), encoding="utf-8")
|
||
|
||
|
||
def load_credentials() -> Optional[Credentials]:
|
||
if not TOKEN_JSON.exists():
|
||
return None
|
||
return Credentials.from_authorized_user_file(str(TOKEN_JSON), SCOPES)
|
||
|
||
|
||
def ensure_valid_credentials() -> Credentials:
|
||
creds = load_credentials()
|
||
if not creds:
|
||
raise RuntimeError("尚未完成 Google 授权，请先点击「Google 授权」。")
|
||
if creds.valid:
|
||
return creds
|
||
if creds.expired and creds.refresh_token:
|
||
creds.refresh(Request())
|
||
TOKEN_JSON.write_text(creds.to_json(), encoding="utf-8")
|
||
return creds
|
||
raise RuntimeError("授权已失效且无 refresh_token，请重新授权（勾选同意离线访问）。")
|
||
|
||
|
||
def oauth_status() -> dict[str, Any]:
|
||
creds = load_credentials()
|
||
if not creds:
|
||
return {"connected": False, "has_refresh": False}
|
||
return {
|
||
"connected": True,
|
||
"has_refresh": bool(creds.refresh_token),
|
||
"expired": bool(creds.expired),
|
||
}
|
||
|
||
|
||
def _pick_ingestion(items: list[dict[str, Any]]) -> Optional[dict[str, Any]]:
|
||
for it in items:
|
||
cdn = it.get("cdn") or {}
|
||
ing = cdn.get("ingestionInfo")
|
||
if ing and ing.get("streamName"):
|
||
return ing
|
||
return None
|
||
|
||
|
||
def fetch_stream_key_and_write_ini() -> dict[str, Any]:
|
||
"""
|
||
列出当前账号的直播流；若无则创建一路，将 streamName 写入 youtube.ini 的 key=。
|
||
"""
|
||
creds = ensure_valid_credentials()
|
||
youtube = build("youtube", "v3", credentials=creds, cache_discovery=False)
|
||
try:
|
||
resp = youtube.liveStreams().list(part="snippet,cdn,status", mine=True, maxResults=10).execute()
|
||
except HttpError as e:
|
||
raise RuntimeError(f"YouTube API 错误: {e}") from e
|
||
|
||
items = resp.get("items", [])
|
||
ing = _pick_ingestion(items)
|
||
|
||
if not ing:
|
||
body = {
|
||
"snippet": {"title": "直播录制助手"},
|
||
"cdn": {
|
||
"ingestionType": "rtmp",
|
||
"frameRate": "variable",
|
||
"resolution": "variable",
|
||
},
|
||
}
|
||
try:
|
||
created = youtube.liveStreams().insert(part="snippet,cdn", body=body).execute()
|
||
except HttpError as e:
|
||
raise RuntimeError(f"创建直播流失败: {e}") from e
|
||
ing = (created.get("cdn") or {}).get("ingestionInfo")
|
||
if not ing or not ing.get("streamName"):
|
||
raise RuntimeError("已创建直播流但未返回串流密钥，请到 YouTube 工作室检查。")
|
||
|
||
stream_name = str(ing["streamName"])
|
||
write_youtube_key_to_ini(stream_name)
|
||
|
||
rtmp = ing.get("ingestionAddress") or "rtmp://a.rtmp.youtube.com/live2"
|
||
rtmps = ing.get("rtmpsAddress") or "rtmps://a.rtmp.youtube.com/live2"
|
||
return {
|
||
"ok": True,
|
||
"stream_name": stream_name,
|
||
"ingestion_address": rtmp,
|
||
"rtmps_address": rtmps,
|
||
"message": "已写入 config/youtube.ini 的 key= 项。",
|
||
}
|
||
|
||
|
||
def write_youtube_key_to_ini(stream_key: str) -> None:
|
||
CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
||
if not YOUTUBE_INI.exists():
|
||
YOUTUBE_INI.write_text(
|
||
"[youtube]\n"
|
||
"# 由 Google 授权自动写入，也可手动修改\n"
|
||
f"key = {stream_key}\n",
|
||
encoding="utf-8-sig",
|
||
)
|
||
return
|
||
|
||
raw = YOUTUBE_INI.read_text(encoding="utf-8-sig")
|
||
lines = raw.splitlines()
|
||
out: list[str] = []
|
||
replaced = False
|
||
for line in lines:
|
||
stripped = line.lstrip()
|
||
if stripped.startswith("#") or stripped.startswith(";"):
|
||
out.append(line)
|
||
continue
|
||
if "=" in line:
|
||
k = line.split("=", 1)[0].strip().lower()
|
||
if k == "key":
|
||
indent = line[: len(line) - len(line.lstrip())]
|
||
out.append(f"{indent}key = {stream_key}")
|
||
replaced = True
|
||
continue
|
||
out.append(line)
|
||
|
||
if not replaced:
|
||
inserted = False
|
||
final: list[str] = []
|
||
for line in out:
|
||
final.append(line)
|
||
if line.strip().lower() == "[youtube]" and not inserted:
|
||
final.append(f"key = {stream_key}")
|
||
inserted = True
|
||
out = final if inserted else out + ["", f"key = {stream_key}"]
|
||
|
||
YOUTUBE_INI.write_text("\n".join(out) + ("\n" if out and not out[-1].endswith("\n") else ""), encoding="utf-8-sig")
|