feat: support upstream token refresh
This commit is contained in:
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user