feat: support upstream token refresh

This commit is contained in:
liumangmang
2026-06-30 14:34:40 +08:00
parent 2d1dcb8f9f
commit c5bb63cf82
15 changed files with 895 additions and 140 deletions
+1 -1
View File
@@ -12,7 +12,7 @@ class Upstream(Base):
name: Mapped[str] = mapped_column(String(255), nullable=False)
base_url: Mapped[str] = mapped_column(String(512), nullable=False)
api_prefix: Mapped[str] = mapped_column(String(128), default="/api/v1")
# none | bearer | api_key | login_password
# none | bearer | new_api_token | nox_token | api_key | cookie | login_password
auth_type: Mapped[str] = mapped_column(String(32), default="login_password")
# JSON: {"email":"..","password":".."} or {"token":".."} etc.
auth_config_json: Mapped[str] = mapped_column(Text, default="{}")
+1 -1
View File
@@ -15,7 +15,7 @@ from app.utils.auth import get_current_user
router = APIRouter(prefix="/api/auth-capture", tags=["auth-capture"])
SENSITIVE_CANDIDATE_FIELDS = frozenset({"value", "cookie_value"})
SENSITIVE_CANDIDATE_FIELDS = frozenset({"value", "cookie_value", "refresh_token"})
class CaptureExtractResponse(BaseModel):
+34 -36
View File
@@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session
from app.database import get_db
from app.database import SessionLocal, get_db
from app.models.admin_user import AdminUser
from app.models.upstream import Upstream
from app.models.upstream_key import UpstreamGeneratedKey
@@ -25,6 +25,7 @@ from app.schemas.upstream import (
UpstreamBatchActionItem, UpstreamBatchActionSummary, UpstreamBatchActionResponse,
)
from app.services.upstream_client import UpstreamClient, UpstreamError, build_snapshot, mask_secret, _extract_key_value
from app.services.auth_config import MASK, mask_auth_config, normalize_auth_config
from app.services.snapshot_service import diff_snapshots
from app.services import scheduler as sched_svc
from app.services import webhook_service
@@ -33,20 +34,6 @@ from app.utils.auth import get_current_user
router = APIRouter(prefix="/api/upstreams", tags=["upstreams"])
MASK = "***"
SECRET_KEYS = {"password", "token", "key", "secret"}
AUTH_CONFIG_ALLOWED_KEYS = {
"none": set(),
"bearer": {"token"},
"nox_token": {"token", "user_id", "new_api_user", "provider"},
"api_key": {"key", "header"},
"cookie": {"cookie_string", "new_api_user", "user_id", "provider"},
"login_password": {"email", "password", "login_path", "username_field", "new_api_user", "user_id"},
}
def _group_id(group: dict) -> str:
for key in ("id", "group_id", "groupId"):
value = group.get(key)
@@ -80,21 +67,23 @@ def _key_response(row: UpstreamGeneratedKey, include_value: bool = False) -> Gen
)
def _mask_auth_config(auth_type: str, cfg: dict) -> dict:
masked = {}
for k, v in cfg.items():
if k.lower() in SECRET_KEYS and v:
masked[k] = MASK
else:
masked[k] = v
return masked
def _normalize_auth_config(auth_type: str, cfg: dict) -> dict:
allowed = AUTH_CONFIG_ALLOWED_KEYS.get(auth_type)
if allowed is None:
return cfg
return {k: v for k, v in cfg.items() if k in allowed}
def _persist_auth_config_update(db: Session, upstream: Upstream, updated_config: dict[str, Any]) -> None:
config_json = json.dumps(
normalize_auth_config(upstream.auth_type, updated_config),
ensure_ascii=False,
)
write_db = SessionLocal()
updated_at = datetime.now(timezone.utc)
try:
row = write_db.query(Upstream).filter(Upstream.id == upstream.id).first()
if row:
row.auth_config_json = config_json
row.updated_at = updated_at
write_db.commit()
finally:
write_db.close()
upstream.auth_config_json = config_json
upstream.updated_at = updated_at
def _extract_plaintext_key(payload: dict[str, Any] | None) -> str:
@@ -134,7 +123,7 @@ def _to_response(u: Upstream) -> UpstreamResponse:
base_url=u.base_url,
api_prefix=u.api_prefix,
auth_type=u.auth_type,
auth_config_masked=_mask_auth_config(u.auth_type, cfg),
auth_config_masked=mask_auth_config(u.auth_type, cfg),
rate_endpoint=u.rate_endpoint,
groups_endpoint=u.groups_endpoint,
enabled=u.enabled,
@@ -205,6 +194,7 @@ def list_generated_keys(uid: int, db: Session = Depends(get_db), _=Depends(get_c
auth_type=upstream.auth_type,
auth_config=auth_config,
timeout=float(upstream.timeout_seconds),
on_auth_config_update=lambda updated: _persist_auth_config_update(db, upstream, updated),
) as client:
client.login()
for prefix in website_sync._fetch_remote_managed_prefixes(db, uid):
@@ -343,8 +333,11 @@ def _is_nox_api_upstream(upstream: Upstream) -> bool:
and upstream.groups_endpoint == "/api/user/self/groups"
and (
upstream.auth_type == "nox_token"
or auth_config.get("login_path") == "/api/user/login"
or bool(auth_config.get("user_id"))
or auth_config.get("provider") == "nox-api"
or (
auth_config.get("provider") not in {"new-api", "new-api-user"}
and auth_config.get("login_path") == "/api/user/login"
)
)
)
@@ -354,7 +347,9 @@ def _is_new_api_user_upstream(upstream: Upstream) -> bool:
return (
upstream.api_prefix.strip("/") == ""
and (
bool(auth_config.get("new_api_user"))
upstream.auth_type == "new_api_token"
or auth_config.get("provider") in {"new-api", "new-api-user"}
or bool(auth_config.get("new_api_user"))
or (
upstream.groups_endpoint == "/api/user/self/groups"
and auth_config.get("login_path") != "/api/user/login"
@@ -540,6 +535,7 @@ def generate_keys_by_groups(
auth_type=u.auth_type,
auth_config=auth_config,
timeout=float(u.timeout_seconds),
on_auth_config_update=lambda updated: _persist_auth_config_update(db, u, updated),
) as client:
try:
client.login()
@@ -581,6 +577,7 @@ def _test_upstream_core(db: Session, u: Upstream) -> UpstreamBatchActionItem:
auth_type=u.auth_type,
auth_config=auth_config,
timeout=float(u.timeout_seconds),
on_auth_config_update=lambda updated: _persist_auth_config_update(db, u, updated),
) as client:
client.login()
groups = client.get_available_groups(u.groups_endpoint)
@@ -620,6 +617,7 @@ def _check_now_core(db: Session, u: Upstream) -> tuple[str, bool]:
auth_type=u.auth_type,
auth_config=auth_config,
timeout=float(u.timeout_seconds),
on_auth_config_update=lambda updated: _persist_auth_config_update(db, u, updated),
) as client:
client.login()
groups = client.get_available_groups(u.groups_endpoint)
@@ -802,7 +800,7 @@ def create_upstream(
base_url=body.base_url.rstrip("/"),
api_prefix=body.api_prefix,
auth_type=body.auth_type,
auth_config_json=json.dumps(_normalize_auth_config(body.auth_type, body.auth_config), ensure_ascii=False),
auth_config_json=json.dumps(normalize_auth_config(body.auth_type, body.auth_config), ensure_ascii=False),
rate_endpoint=body.rate_endpoint,
groups_endpoint=body.groups_endpoint,
enabled=body.enabled,
@@ -847,7 +845,7 @@ def update_upstream(
if v != MASK: # don't overwrite with mask placeholder
existing[k] = v
next_auth_type = data.get("auth_type", u.auth_type)
u.auth_config_json = json.dumps(_normalize_auth_config(next_auth_type, existing), ensure_ascii=False)
u.auth_config_json = json.dumps(normalize_auth_config(next_auth_type, existing), ensure_ascii=False)
if "base_url" in data:
data["base_url"] = data["base_url"].rstrip("/")
for k, v in data.items():
+5 -1
View File
@@ -5,6 +5,9 @@ from pydantic import BaseModel, Field
class AuthConfigBearer(BaseModel):
token: str = ""
refresh_token: str = ""
expires_in: Optional[int] = None
token_expires_at: Optional[int] = None
class AuthConfigApiKey(BaseModel):
@@ -16,13 +19,14 @@ class AuthConfigLoginPassword(BaseModel):
email: str = ""
password: str = ""
login_path: str = "/auth/login"
username_field: str = "email"
class UpstreamCreate(BaseModel):
name: str
base_url: str
api_prefix: str = "/api/v1"
auth_type: str = "login_password" # none | bearer | api_key | login_password
auth_type: str = "login_password" # none | bearer | new_api_token | nox_token | api_key | cookie | login_password
auth_config: dict[str, Any] = {}
rate_endpoint: str = "/groups/rates"
groups_endpoint: str = "/groups/available"
+43 -6
View File
@@ -9,10 +9,12 @@ from urllib.parse import urlparse
logger = logging.getLogger(__name__)
# Keys likely to contain auth tokens in storage
TOKEN_KEYS = frozenset({
"token", "access_token", "accessToken", "jwt", "auth_token", "authToken",
"refresh_token", "refreshToken", "id_token", "session_token",
ACCESS_TOKEN_KEYS = frozenset({
"token", "access_token", "accesstoken", "jwt", "auth_token", "authtoken",
"id_token", "idtoken", "session_token", "sessiontoken",
})
REFRESH_TOKEN_KEYS = frozenset({"refresh_token", "refreshtoken"})
TOKEN_KEYS = ACCESS_TOKEN_KEYS | REFRESH_TOKEN_KEYS
SECRET_KEYS = frozenset({
"secret", "api_key", "apiKey", "apikey",
})
@@ -116,6 +118,27 @@ def _curate_candidates(
extra=bundle_extra,
)
# 0.5. Sub2API browser sessions keep access token and refresh token separately.
storage_sources = [("localStorage", local_storage), ("sessionStorage", session_storage)]
access_entry = _find_storage_token_entry(storage_sources, ACCESS_TOKEN_KEYS)
refresh_entry = _find_storage_token_entry(storage_sources, REFRESH_TOKEN_KEYS)
if access_entry and refresh_entry:
access_source, access_value = access_entry
refresh_source, refresh_value = refresh_entry
_add(
candidates,
"bearer_token",
f"{access_source}+{refresh_source}",
access_value,
_preview(access_value),
"Sub2API Bearer + Refresh Token",
96,
extra={
"refresh_token": refresh_value,
"refresh_token_preview": _preview(refresh_value),
},
)
# 1. CDP-captured network headers (high confidence)
seen = set()
for h in auth_headers:
@@ -141,9 +164,10 @@ def _curate_candidates(
if not isinstance(val, str) or not val:
continue
key_lower = key.lower()
is_refresh_key = any(k in key_lower for k in REFRESH_TOKEN_KEYS)
# Explicit auth-named keys
if any(k in key_lower for k in TOKEN_KEYS):
if any(k in key_lower for k in TOKEN_KEYS) and not is_refresh_key:
preview = _preview(val)
score = 85 if "token" in key_lower and val.count(".") >= 2 else 75
_add(candidates, "bearer_token", f"{store_name}.{key}", val, preview,
@@ -153,14 +177,14 @@ def _curate_candidates(
f"{store_name}.{key}", 70)
# Looks like a JWT (xx.yy.zz format)
if val.count(".") >= 2 and 20 < len(val) < 5000:
if not is_refresh_key and val.count(".") >= 2 and 20 < len(val) < 5000:
if val not in seen:
seen.add(val)
_add(candidates, "bearer_token", f"{store_name}.{key}", val, _preview(val),
f"{store_name}.{key} (JWT)", 80)
# sk-xxx API key pattern
if val.startswith("sk-") and len(val) > 10:
if not is_refresh_key and val.startswith("sk-") and len(val) > 10:
_add(candidates, "bearer_token", f"{store_name}.{key}", val, _preview(val),
f"{store_name}.{key} (API Key)", 90)
@@ -198,6 +222,19 @@ def _find_user_header(headers: list[dict[str, str]]) -> str:
return ""
def _find_storage_token_entry(
stores: list[tuple[str, dict[str, str]]],
key_names: frozenset[str],
) -> tuple[str, str] | None:
for store_name, store in stores:
for key, value in store.items():
if key.lower() not in key_names:
continue
if isinstance(value, str) and value.strip():
return f"{store_name}.{key}", value.strip()
return None
def _find_storage_value(*stores: dict[str, str], key: str) -> str:
for store in stores:
value = store.get(key)
+34
View File
@@ -0,0 +1,34 @@
from __future__ import annotations
MASK = "***"
SECRET_KEYS = {"password", "token", "refresh_token", "access_token", "auth_token", "key", "secret"}
AUTH_CONFIG_ALLOWED_KEYS = {
"none": set(),
"bearer": {"token", "refresh_token", "expires_in", "token_expires_at"},
"new_api_token": {"token", "user_id", "new_api_user", "provider"},
"nox_token": {"token", "user_id", "new_api_user", "provider"},
"api_key": {"key", "header"},
"cookie": {"cookie_string", "new_api_user", "user_id", "provider"},
"login_password": {
"email", "password", "login_path", "username_field", "new_api_user", "user_id", "provider",
"token", "refresh_token", "expires_in", "token_expires_at",
},
}
def mask_auth_config(auth_type: str, cfg: dict) -> dict:
masked = {}
for k, v in cfg.items():
if k.lower() in SECRET_KEYS and v:
masked[k] = MASK
else:
masked[k] = v
return masked
def normalize_auth_config(auth_type: str, cfg: dict) -> dict:
allowed = AUTH_CONFIG_ALLOWED_KEYS.get(auth_type)
if allowed is None:
return cfg
return {k: v for k, v in cfg.items() if k in allowed}
+23
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import json
import logging
from datetime import datetime, timezone
from typing import Any
from apscheduler.executors.pool import ThreadPoolExecutor
from apscheduler.schedulers.background import BackgroundScheduler
@@ -12,6 +13,7 @@ from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models.upstream import Upstream
from app.models.snapshot import UpstreamRateSnapshot
from app.services.auth_config import normalize_auth_config
from app.services.upstream_client import UpstreamClient, build_snapshot
from app.services.snapshot_service import diff_snapshots, prune_snapshots
from app.services import webhook_service
@@ -27,6 +29,25 @@ def get_scheduler() -> BackgroundScheduler:
return _scheduler
def _persist_auth_config_update(upstream: Upstream, updated_config: dict[str, Any]) -> None:
config_json = json.dumps(
normalize_auth_config(upstream.auth_type, updated_config),
ensure_ascii=False,
)
updated_at = datetime.now(timezone.utc)
db = SessionLocal()
try:
row = db.query(Upstream).filter(Upstream.id == upstream.id).first()
if row:
row.auth_config_json = config_json
row.updated_at = updated_at
db.commit()
upstream.auth_config_json = row.auth_config_json
upstream.updated_at = updated_at
finally:
db.close()
def _check_upstream(upstream_id: int) -> None:
"""Full upstream check executed by scheduler (runs in thread).
@@ -55,6 +76,7 @@ def _check_upstream(upstream_id: int) -> None:
auth_type=upstream.auth_type,
auth_config=auth_config,
timeout=float(upstream.timeout_seconds),
on_auth_config_update=lambda updated: _persist_auth_config_update(upstream, updated),
) as client:
try:
client.login()
@@ -235,6 +257,7 @@ def _sync_upstream_keys(upstream_id: int, snapshot: dict[str, Any], captured_at:
auth_type=upstream.auth_type,
auth_config=auth_config,
timeout=float(upstream.timeout_seconds),
on_auth_config_update=lambda updated: _persist_auth_config_update(upstream, updated),
) as client:
client.login()
remote_key_ids = website_sync._fetch_remote_managed_key_ids(db, client, upstream_id)
+158 -18
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
import json
import time
from typing import Any, Optional
from typing import Any, Callable, Optional
from urllib.parse import urljoin
import httpx
@@ -32,6 +32,38 @@ def _find_token(value: Any) -> str:
return ""
def _find_refresh_token(value: Any) -> str:
if isinstance(value, dict):
for key in ("refresh_token", "refreshToken"):
candidate = value.get(key)
if isinstance(candidate, str) and candidate:
return candidate
for key in ("data", "result", "user", "session"):
tok = _find_refresh_token(value.get(key))
if tok:
return tok
return ""
def _find_expires_in(value: Any) -> int | None:
if isinstance(value, dict):
for key in ("expires_in", "expiresIn"):
candidate = value.get(key)
if candidate is None:
continue
try:
expires_in = int(float(candidate))
except (TypeError, ValueError):
continue
if expires_in > 0:
return expires_in
for key in ("data", "result", "user", "session"):
expires_in = _find_expires_in(value.get(key))
if expires_in is not None:
return expires_in
return None
def _clean_auth_header_value(value: Any, field_name: str) -> str:
text = str(value or "").strip()
if not text:
@@ -289,12 +321,14 @@ class UpstreamClient:
auth_type: str,
auth_config: dict[str, Any],
timeout: float = 30.0,
on_auth_config_update: Callable[[dict[str, Any]], None] | None = None,
) -> None:
self.base_url = base_url.rstrip("/")
self.api_prefix = api_prefix.strip("/")
self.auth_type = auth_type
self.auth_config = auth_config
self.timeout = timeout
self.on_auth_config_update = on_auth_config_update
self._token: str = ""
self._cookies: dict[str, str] = {}
self._new_api_user: str = ""
@@ -321,6 +355,7 @@ class UpstreamClient:
bool(self.auth_config.get("new_api_user"))
or login_path == "/api/user/login"
or self.auth_type == "cookie"
or self.auth_type == "new_api_token"
or self.auth_type == "nox_token"
)
)
@@ -349,6 +384,13 @@ class UpstreamClient:
headers["Authorization"] = f"Bearer {token}"
if user_id:
headers["Nox-Api-User"] = user_id
elif self.auth_type == "new_api_token":
token = _clean_auth_header_value(self.auth_config.get("token", ""), "New-API access token")
user_id = self._user_header_value()
if token:
headers["Authorization"] = f"Bearer {token}"
if user_id:
headers["New-Api-User"] = user_id
elif self.auth_type == "api_key":
key = _clean_auth_header_value(self.auth_config.get("key", ""), "API key")
header = self.auth_config.get("header", "Authorization")
@@ -371,28 +413,117 @@ class UpstreamClient:
headers["Nox-Api-User"] = self._new_api_user
return headers
def _is_sub2api_bearer(self) -> bool:
return self.auth_type == "bearer" and self.api_prefix == "api/v1"
def _remember_auth_tokens(
self,
token: str = "",
refresh_token: str = "",
expires_in: int | None = None,
) -> None:
changed = False
if token and self.auth_config.get("token") != token:
self.auth_config["token"] = token
changed = True
if refresh_token and self.auth_config.get("refresh_token") != refresh_token:
self.auth_config["refresh_token"] = refresh_token
changed = True
if expires_in is not None:
if self.auth_config.get("expires_in") != expires_in:
self.auth_config["expires_in"] = expires_in
changed = True
token_expires_at = int(time.time()) + expires_in
if self.auth_config.get("token_expires_at") != token_expires_at:
self.auth_config["token_expires_at"] = token_expires_at
changed = True
if changed and self.on_auth_config_update:
self.on_auth_config_update(dict(self.auth_config))
def _refresh_sub2api_bearer_token(self) -> bool:
if not self._is_sub2api_bearer():
return False
refresh_token = str(self.auth_config.get("refresh_token") or "").strip()
if not refresh_token:
return False
try:
resp = self._client.request(
"POST",
self._url("/auth/refresh"),
json={"refresh_token": refresh_token},
headers=self._headers(auth=False),
cookies=self._cookies,
)
self._cookies.update(dict(resp.cookies))
resp.raise_for_status()
payload = resp.json() if resp.content else {}
except Exception:
return False
token = _find_token(payload)
if not token:
return False
self._remember_auth_tokens(
token=token,
refresh_token=_find_refresh_token(payload) or refresh_token,
expires_in=_find_expires_in(payload),
)
return True
def _send_request(
self,
method: str,
url: str,
auth: bool = True,
allow_refresh: bool = True,
**kwargs: Any,
) -> httpx.Response:
resp = self._client.request(
method,
url,
headers=self._headers(auth),
cookies=self._cookies,
**kwargs,
)
self._cookies.update(dict(resp.cookies))
if (
getattr(resp, "status_code", None) == 401
and auth
and allow_refresh
and self._is_sub2api_bearer()
and self._refresh_sub2api_bearer_token()
):
resp = self._client.request(
method,
url,
headers=self._headers(auth),
cookies=self._cookies,
**kwargs,
)
self._cookies.update(dict(resp.cookies))
return resp
def _request(self, method: str, path: str, body: Any = None, auth: bool = True) -> Any:
if auth and self.auth_type == "cookie" and "user/self" in path and not self.auth_config.get("new_api_user"):
if auth and self.auth_type == "cookie" and "user/self" in path and not self._user_header_value():
raise UpstreamError("New-API user endpoint requires New-Api-User; re-extract the session cookie after login and save the upstream")
if auth and self.auth_type == "new_api_token" and "user/self" in path and not self._user_header_value():
raise UpstreamError("New-API endpoint requires New-Api-User; please fill user id and retry")
if auth and self.auth_type == "nox_token" and "user/self" in path and not self._user_header_value():
raise UpstreamError("Nox-API endpoint requires Nox-Api-User; please fill user id and retry")
url = self._url(path)
if body is not None:
resp = self._client.request(
resp = self._send_request(
method,
url,
json=body,
headers=self._headers(auth),
cookies=self._cookies,
auth=auth,
)
else:
resp = self._client.request(
resp = self._send_request(
method,
url,
headers=self._headers(auth),
cookies=self._cookies,
auth=auth,
)
self._cookies.update(dict(resp.cookies))
resp.raise_for_status()
ct = resp.headers.get("content-type", "")
if not resp.content:
@@ -449,14 +580,11 @@ class UpstreamClient:
return items, meta
def _request_new_api_token_list(self, path: str, params: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
resp = self._client.request(
resp = self._send_request(
"GET",
self._url(path),
params=params,
headers=self._headers(),
cookies=self._cookies,
)
self._cookies.update(dict(resp.cookies))
resp.raise_for_status()
data = resp.json()
self._ensure_api_success(data, "list New-API tokens")
@@ -585,14 +713,28 @@ class UpstreamClient:
return
email = self.auth_config.get("email", "")
password = self.auth_config.get("password", "")
login_path = self.auth_config.get("login_path", "/auth/login")
username_field = self.auth_config.get("username_field", "email")
default_login_path = "/api/user/login" if self.api_prefix == "" else "/auth/login"
login_path = self.auth_config.get("login_path") or default_login_path
default_username_field = "username" if login_path == "/api/user/login" else "email"
username_field = self.auth_config.get("username_field") or default_username_field
if not email or not password:
raise UpstreamError("login_password auth requires email and password in auth_config")
resp = self._request("POST", login_path, {username_field: email, "password": password}, auth=False)
token = _find_token(resp)
if token:
self._token = token
self._new_api_user = (
self.auth_config.get("new_api_user", "")
or self.auth_config.get("user_id", "")
or _find_user_id(resp)
)
refresh_token = _find_refresh_token(resp)
if self.api_prefix == "api/v1" or refresh_token:
self._remember_auth_tokens(
token=token,
refresh_token=refresh_token,
expires_in=_find_expires_in(resp),
)
return
if self._cookies:
self._new_api_user = self.auth_config.get("new_api_user", "") or _find_user_id(resp)
@@ -655,12 +797,10 @@ class UpstreamClient:
if status:
params["status"] = status
url = self._url(endpoint)
resp = self._client.request(
resp = self._send_request(
"GET",
url,
params=params if params else None,
headers=self._headers(),
cookies=self._cookies,
)
resp.raise_for_status()
data = resp.json()
+23
View File
@@ -3,14 +3,17 @@ from __future__ import annotations
import json
import logging
from decimal import Decimal
from datetime import datetime, timezone
from typing import Any
from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models.snapshot import UpstreamRateSnapshot
from app.models.upstream import Upstream
from app.models.upstream_key import UpstreamGeneratedKey
from app.models.website import Website, WebsiteGroupBinding, WebsiteSyncLog
from app.services.auth_config import normalize_auth_config
from app.services.website_client import Sub2ApiWebsiteClient, WebsiteError, _extract_id, calculate_target_rate
from app.services.upstream_client import UpstreamClient
from app.services import webhook_service
@@ -23,6 +26,25 @@ PRIORITY_BASE = 1
PRIORITY_STEP = 10
def _persist_upstream_auth_config(upstream: Upstream, updated_config: dict[str, Any]) -> None:
config_json = json.dumps(
normalize_auth_config(upstream.auth_type, updated_config),
ensure_ascii=False,
)
updated_at = datetime.now(timezone.utc)
db = SessionLocal()
try:
row = db.query(Upstream).filter(Upstream.id == upstream.id).first()
if row:
row.auth_config_json = config_json
row.updated_at = updated_at
db.commit()
upstream.auth_config_json = row.auth_config_json
upstream.updated_at = updated_at
finally:
db.close()
def priority_for_rate_rank(rank: int) -> int:
"""Convert a zero-based sorted rate rank to an account priority."""
return PRIORITY_BASE + rank * PRIORITY_STEP
@@ -672,6 +694,7 @@ def reconcile_upstream_keys_full(db: Session, upstream_id: int) -> bool:
auth_type=upstream.auth_type,
auth_config=auth_config,
timeout=float(upstream.timeout_seconds),
on_auth_config_update=lambda updated: _persist_upstream_auth_config(upstream, updated),
) as client:
client.login()
# 获取远端 Key 列表(支持自定义 managed_prefix
+33
View File
@@ -218,6 +218,39 @@ def test_new_api_user_propagated_to_bundle():
assert bundle.get("new_api_user") == "42"
def test_storage_access_and_refresh_token_build_sub2api_bearer_candidate():
candidates = _curate_candidates(
cookies=[],
local_storage={"auth_token": "access-token", "refresh_token": "refresh-token"},
session_storage={},
auth_headers=[],
page_url="https://sub2api.example.com/",
)
assert candidates[0]["type"] == "bearer_token"
assert candidates[0]["value"] == "access-token"
assert candidates[0]["refresh_token"] == "refresh-token"
assert candidates[0]["label"] == "Sub2API Bearer + Refresh Token"
assert not any(
c["type"] == "bearer_token" and c.get("value") == "refresh-token"
for c in candidates
)
def test_storage_access_token_without_refresh_keeps_bearer_candidate():
candidates = _curate_candidates(
cookies=[],
local_storage={"auth_token": "access-token"},
session_storage={},
auth_headers=[],
page_url="https://sub2api.example.com/",
)
assert candidates[0]["type"] == "bearer_token"
assert candidates[0]["value"] == "access-token"
assert "refresh_token" not in candidates[0]
def test_browser_import_payload_builds_cookie_bundle_with_new_api_user():
from app.services.browser_import_service import build_import_result
+46
View File
@@ -0,0 +1,46 @@
from app.services.upstream_client import UpstreamClient
def test_new_api_token_headers_include_authorization_and_new_api_user():
client = UpstreamClient(
base_url="http://newapi.local",
api_prefix="",
auth_type="new_api_token",
auth_config={"token": "access-token", "user_id": "7"},
)
headers = client._headers()
assert headers["Authorization"] == "Bearer access-token"
assert headers["New-Api-User"] == "7"
assert "Nox-Api-User" not in headers
def test_nox_token_headers_include_authorization_and_nox_api_user():
client = UpstreamClient(
base_url="http://nox.local",
api_prefix="",
auth_type="nox_token",
auth_config={"token": "nox-access-token", "user_id": "9"},
)
headers = client._headers()
assert headers["Authorization"] == "Bearer nox-access-token"
assert headers["Nox-Api-User"] == "9"
assert "New-Api-User" not in headers
def test_bearer_headers_do_not_include_new_or_nox_user_headers():
client = UpstreamClient(
base_url="http://sub2api.local",
api_prefix="/api/v1",
auth_type="bearer",
auth_config={"token": "sk-sub2api", "user_id": "11", "new_api_user": "11"},
)
headers = client._headers()
assert headers["Authorization"] == "Bearer sk-sub2api"
assert "New-Api-User" not in headers
assert "Nox-Api-User" not in headers
+173
View File
@@ -0,0 +1,173 @@
from __future__ import annotations
import httpx
import pytest
from app.services import upstream_client
from app.services.upstream_client import UpstreamClient
class FakeHttpClient:
def __init__(self, handler):
self.handler = handler
self.calls = []
def request(self, method: str, url: str, **kwargs):
request = httpx.Request(method, url)
self.calls.append({
"method": method,
"path": request.url.path,
"headers": kwargs.get("headers") or {},
"json": kwargs.get("json"),
})
return self.handler(request, kwargs)
def close(self) -> None:
pass
def _install_fake_client(monkeypatch, handler) -> FakeHttpClient:
fake = FakeHttpClient(handler)
monkeypatch.setattr(upstream_client.httpx, "Client", lambda timeout=None: fake)
return fake
def _response(request: httpx.Request, status_code: int, payload=None) -> httpx.Response:
return httpx.Response(status_code, json=payload, request=request)
def test_sub2api_bearer_refreshes_on_401_and_retries(monkeypatch):
updates: list[dict] = []
def handler(request: httpx.Request, kwargs):
if request.url.path == "/api/v1/groups/available":
auth = kwargs["headers"].get("Authorization")
if auth == "Bearer old-access":
return _response(request, 401, {"detail": "expired"})
assert auth == "Bearer new-access"
return _response(request, 200, [{"id": "g1", "name": "G1"}])
if request.url.path == "/api/v1/auth/refresh":
assert kwargs["json"] == {"refresh_token": "old-refresh"}
return _response(request, 200, {
"access_token": "new-access",
"refresh_token": "new-refresh",
"expires_in": 3600,
})
raise AssertionError(f"unexpected request: {request.method} {request.url}")
fake = _install_fake_client(monkeypatch, handler)
client = UpstreamClient(
base_url="http://sub2api.local",
api_prefix="/api/v1",
auth_type="bearer",
auth_config={"token": "old-access", "refresh_token": "old-refresh"},
on_auth_config_update=updates.append,
)
groups = client.get_available_groups("/groups/available")
assert groups == [{"id": "g1", "name": "G1"}]
assert [(c["method"], c["path"]) for c in fake.calls] == [
("GET", "/api/v1/groups/available"),
("POST", "/api/v1/auth/refresh"),
("GET", "/api/v1/groups/available"),
]
assert updates[-1]["token"] == "new-access"
assert updates[-1]["refresh_token"] == "new-refresh"
assert updates[-1]["expires_in"] == 3600
def test_sub2api_bearer_without_refresh_token_does_not_refresh(monkeypatch):
def handler(request: httpx.Request, kwargs):
assert request.url.path == "/api/v1/groups/available"
return _response(request, 401, {"detail": "expired"})
fake = _install_fake_client(monkeypatch, handler)
client = UpstreamClient(
base_url="http://sub2api.local",
api_prefix="/api/v1",
auth_type="bearer",
auth_config={"token": "old-access"},
)
with pytest.raises(httpx.HTTPStatusError):
client.get_available_groups("/groups/available")
assert [(c["method"], c["path"]) for c in fake.calls] == [
("GET", "/api/v1/groups/available"),
]
def test_non_sub2api_bearer_does_not_refresh(monkeypatch):
def handler(request: httpx.Request, kwargs):
assert request.url.path == "/groups/available"
return _response(request, 401, {"detail": "expired"})
fake = _install_fake_client(monkeypatch, handler)
client = UpstreamClient(
base_url="http://plain-bearer.local",
api_prefix="",
auth_type="bearer",
auth_config={"token": "old-access", "refresh_token": "old-refresh"},
)
with pytest.raises(httpx.HTTPStatusError):
client.get_available_groups("/groups/available")
assert [(c["method"], c["path"]) for c in fake.calls] == [
("GET", "/groups/available"),
]
def test_login_password_saves_sub2api_refresh_token(monkeypatch):
updates: list[dict] = []
def handler(request: httpx.Request, kwargs):
assert request.url.path == "/api/v1/auth/login"
assert kwargs["json"] == {"email": "user@example.com", "password": "secret"}
return _response(request, 200, {
"access_token": "login-access",
"refresh_token": "login-refresh",
"expires_in": 1800,
})
_install_fake_client(monkeypatch, handler)
client = UpstreamClient(
base_url="http://sub2api.local",
api_prefix="/api/v1",
auth_type="login_password",
auth_config={"email": "user@example.com", "password": "secret"},
on_auth_config_update=updates.append,
)
client.login()
assert client._headers()["Authorization"] == "Bearer login-access"
assert updates[-1]["token"] == "login-access"
assert updates[-1]["refresh_token"] == "login-refresh"
assert updates[-1]["expires_in"] == 1800
def test_login_password_empty_prefix_defaults_to_user_api_login(monkeypatch):
def handler(request: httpx.Request, kwargs):
assert request.url.path == "/api/user/login"
assert kwargs["json"] == {"username": "admin", "password": "secret"}
return _response(request, 200, {
"token": "nox-access",
"user": {"id": 9},
})
_install_fake_client(monkeypatch, handler)
client = UpstreamClient(
base_url="http://nox.local",
api_prefix="",
auth_type="login_password",
auth_config={"email": "admin", "password": "secret"},
)
client.login()
headers = client._headers()
assert headers["Authorization"] == "Bearer nox-access"
assert headers["Nox-Api-User"] == "9"
assert headers["New-Api-User"] == "9"