401 lines
15 KiB
Python
401 lines
15 KiB
Python
|
# ============================================================
|
|||
|
|
# Pi Agent 长期只读凭证（moduleId: auth-agent-tokens, 可重生 ✅）
|
|||
|
|
# 方向 A · P1：/api/agent/* 机器凭证。
|
|||
|
|
# - HMAC(base64url body + secret) 与 licenses.py/providers.py 同风格；
|
|||
|
|
# - token 本身只携带 tokenId/kind/scope/时间；principal/tenant 从登记文件解析；
|
|||
|
|
# - APS_AGENT_TOKEN_SECRET 缺失时签发与鉴权显式失败，绝不回退匿名/本地管理员。
|
|||
|
|
# ============================================================
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import hashlib
|
|||
|
|
import hmac
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
import uuid
|
|||
|
|
from base64 import urlsafe_b64decode as _urlsafe_b64decode
|
|||
|
|
from base64 import urlsafe_b64encode as _urlsafe_b64encode
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from server.auth.context import IdentityContext
|
|||
|
|
from server.auth.providers import AuthError
|
|||
|
|
|
|||
|
|
AGENT_TOKEN_VERSION = 1
|
|||
|
|
AGENT_TOKEN_KIND = "agent.read"
|
|||
|
|
AGENT_TOKEN_SCOPE_READ = "read"
|
|||
|
|
AGENT_TOKEN_FILE_NAME = "agent-tokens.json"
|
|||
|
|
AGENT_TOKEN_MAX_TTL_DAYS = 3650
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _b64encode(raw: bytes) -> str:
|
|||
|
|
return _urlsafe_b64encode(raw).decode("ascii").rstrip("=")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _b64decode_text(raw: str) -> bytes:
|
|||
|
|
return _urlsafe_b64decode(raw + "=" * (-len(raw) % 4))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _configured_secret() -> bytes:
|
|||
|
|
value = (os.environ.get("APS_AGENT_TOKEN_SECRET") or "").strip()
|
|||
|
|
if not value:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_NOT_CONFIGURED",
|
|||
|
|
"Pi Agent 凭证服务未配置：请设置 APS_AGENT_TOKEN_SECRET",
|
|||
|
|
503,
|
|||
|
|
)
|
|||
|
|
return value.encode("utf-8")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _default_path() -> Path:
|
|||
|
|
try:
|
|||
|
|
from server.aps_home import data_dir
|
|||
|
|
|
|||
|
|
return data_dir() / AGENT_TOKEN_FILE_NAME
|
|||
|
|
except (OSError, RuntimeError, ImportError):
|
|||
|
|
return (
|
|||
|
|
Path(os.environ.get("APS_DATA_DIR", "server/data")) / AGENT_TOKEN_FILE_NAME
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _configured_path() -> Path:
|
|||
|
|
raw = (os.environ.get("APS_AGENT_TOKEN_PATH") or "").strip()
|
|||
|
|
path = Path(raw).expanduser() if raw else _default_path()
|
|||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
return path.resolve()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _issuer_roles(issuer: IdentityContext | None) -> set[str]:
|
|||
|
|
if issuer is None:
|
|||
|
|
return {"system", "admin"}
|
|||
|
|
roles = {str(role).lower() for role in (issuer.roles or ())}
|
|||
|
|
if issuer.auth_kind in {"disabled", "local-admin"}:
|
|||
|
|
roles.update({"system", "admin"})
|
|||
|
|
return roles
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _can_manage(issuer: IdentityContext | None, record: dict[str, Any]) -> bool:
|
|||
|
|
if issuer is None:
|
|||
|
|
return True
|
|||
|
|
if str(record.get("tenantUuid") or "") != issuer.tenant_uuid:
|
|||
|
|
return False
|
|||
|
|
roles = _issuer_roles(issuer)
|
|||
|
|
return bool(roles & {"system", "admin"})
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AgentTokenService:
|
|||
|
|
"""File-backed agent read-token registry with HMAC-signed tokens."""
|
|||
|
|
|
|||
|
|
def __init__(self, secret: bytes | None = None, path: Path | None = None) -> None:
|
|||
|
|
self.secret = secret
|
|||
|
|
self.path = (path or _configured_path()).resolve()
|
|||
|
|
self._lock = threading.RLock()
|
|||
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
def _require_secret(self) -> bytes:
|
|||
|
|
if not self.secret:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_NOT_CONFIGURED",
|
|||
|
|
"Pi Agent 凭证服务未配置：请设置 APS_AGENT_TOKEN_SECRET",
|
|||
|
|
503,
|
|||
|
|
)
|
|||
|
|
return self.secret
|
|||
|
|
|
|||
|
|
def _load(self) -> dict[str, Any]:
|
|||
|
|
try:
|
|||
|
|
if not self.path.is_file():
|
|||
|
|
return {"version": AGENT_TOKEN_VERSION, "tokens": []}
|
|||
|
|
with self.path.open("r", encoding="utf-8") as handle:
|
|||
|
|
document = json.load(handle)
|
|||
|
|
if not isinstance(document, dict) or not isinstance(
|
|||
|
|
document.get("tokens"), list
|
|||
|
|
):
|
|||
|
|
raise TypeError("registry schema")
|
|||
|
|
return document
|
|||
|
|
except AuthError:
|
|||
|
|
raise
|
|||
|
|
except Exception as exc:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_STORE_CORRUPT",
|
|||
|
|
f"Pi Agent 凭证登记文件不可读：{self.path}",
|
|||
|
|
503,
|
|||
|
|
) from exc
|
|||
|
|
|
|||
|
|
def _save(self, document: dict[str, Any]) -> None:
|
|||
|
|
import tempfile
|
|||
|
|
|
|||
|
|
with self._lock:
|
|||
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
fd = None
|
|||
|
|
tmp = None
|
|||
|
|
try:
|
|||
|
|
fd, tmp_name = tempfile.mkstemp(
|
|||
|
|
dir=str(self.path.parent), prefix=".agent-tokens-", suffix=".tmp"
|
|||
|
|
)
|
|||
|
|
tmp = Path(tmp_name)
|
|||
|
|
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
|||
|
|
json.dump(document, handle, ensure_ascii=False, indent=2)
|
|||
|
|
os.replace(tmp, self.path)
|
|||
|
|
except BaseException:
|
|||
|
|
if fd is not None:
|
|||
|
|
try:
|
|||
|
|
os.close(fd)
|
|||
|
|
except OSError:
|
|||
|
|
pass
|
|||
|
|
if tmp is not None and tmp.exists():
|
|||
|
|
try:
|
|||
|
|
tmp.unlink()
|
|||
|
|
except OSError:
|
|||
|
|
pass
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
def _sign(self, payload: dict[str, Any]) -> str:
|
|||
|
|
raw = json.dumps(payload, separators=(",", ":"), ensure_ascii=True).encode(
|
|||
|
|
"utf-8"
|
|||
|
|
)
|
|||
|
|
body = _b64encode(raw)
|
|||
|
|
signature = _b64encode(
|
|||
|
|
hmac.new(
|
|||
|
|
self._require_secret(), body.encode("ascii"), hashlib.sha256
|
|||
|
|
).digest()
|
|||
|
|
)
|
|||
|
|
return f"{body}.{signature}"
|
|||
|
|
|
|||
|
|
def _verify_token(self, token: str) -> dict[str, Any]:
|
|||
|
|
secret = self._require_secret()
|
|||
|
|
try:
|
|||
|
|
body, signature = token.split(".", 1)
|
|||
|
|
expected = _b64encode(
|
|||
|
|
hmac.new(secret, body.encode("ascii"), hashlib.sha256).digest()
|
|||
|
|
)
|
|||
|
|
if not hmac.compare_digest(signature, expected):
|
|||
|
|
raise ValueError("signature")
|
|||
|
|
payload = json.loads(_b64decode_text(body).decode("utf-8"))
|
|||
|
|
if (
|
|||
|
|
int(payload.get("v") or 0) != AGENT_TOKEN_VERSION
|
|||
|
|
or str(payload.get("kind") or "") != AGENT_TOKEN_KIND
|
|||
|
|
):
|
|||
|
|
raise ValueError("token shape")
|
|||
|
|
token_id = str(payload.get("tokenId") or "")
|
|||
|
|
if not token_id or len(token_id) > 128:
|
|||
|
|
raise ValueError("token id")
|
|||
|
|
return {
|
|||
|
|
"tokenId": token_id,
|
|||
|
|
"iat": int(payload.get("iat") or 0),
|
|||
|
|
"exp": int(payload.get("exp") or 0),
|
|||
|
|
}
|
|||
|
|
except AuthError:
|
|||
|
|
raise
|
|||
|
|
except Exception as exc:
|
|||
|
|
raise AuthError("AGENT_TOKEN_INVALID", "Pi Agent 凭证无效", 401) from exc
|
|||
|
|
|
|||
|
|
def _find(self, token_id: str) -> dict[str, Any] | None:
|
|||
|
|
document = self._load()
|
|||
|
|
for record in document.get("tokens") or []:
|
|||
|
|
if str(record.get("tokenId") or "") == token_id:
|
|||
|
|
return record
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def issue(
|
|||
|
|
self,
|
|||
|
|
principal: IdentityContext,
|
|||
|
|
label: str,
|
|||
|
|
*,
|
|||
|
|
ttl_days: int = 365,
|
|||
|
|
) -> tuple[IdentityContext, str]:
|
|||
|
|
self._require_secret()
|
|||
|
|
label = str(label or "").strip()
|
|||
|
|
if not label or len(label) > 80:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_LABEL_INVALID", "凭证标签必须为 1-80 字符", 400
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
user_id = int(principal.user_id)
|
|||
|
|
ttl_days = int(ttl_days)
|
|||
|
|
except (TypeError, ValueError) as exc:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_PRINCIPAL_INVALID", "签发主体标识无效", 400
|
|||
|
|
) from exc
|
|||
|
|
if user_id < 0 or user_id > 2**63 - 1:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_PRINCIPAL_INVALID", "签发主体 user_id 超出合法范围", 400
|
|||
|
|
)
|
|||
|
|
tenant_uuid = str(principal.tenant_uuid or "").strip()
|
|||
|
|
if not tenant_uuid or len(tenant_uuid) > 64:
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_PRINCIPAL_INVALID", "签发主体租户标识无效", 400
|
|||
|
|
)
|
|||
|
|
if not 1 <= ttl_days <= AGENT_TOKEN_MAX_TTL_DAYS:
|
|||
|
|
raise AuthError("AGENT_TOKEN_TTL_INVALID", "凭证有效期须为 1-3650 天", 400)
|
|||
|
|
|
|||
|
|
now = int(time.time())
|
|||
|
|
token_id = uuid.uuid4().hex
|
|||
|
|
payload = {
|
|||
|
|
"v": AGENT_TOKEN_VERSION,
|
|||
|
|
"kind": AGENT_TOKEN_KIND,
|
|||
|
|
"scope": [AGENT_TOKEN_SCOPE_READ],
|
|||
|
|
"tokenId": token_id,
|
|||
|
|
"iat": now,
|
|||
|
|
"exp": now + ttl_days * 24 * 60 * 60,
|
|||
|
|
}
|
|||
|
|
token = self._sign(payload)
|
|||
|
|
created_by = str(principal.username or principal.fullname or principal.user_id)
|
|||
|
|
record = {
|
|||
|
|
"tokenId": token_id,
|
|||
|
|
"label": label,
|
|||
|
|
"principalUserId": user_id,
|
|||
|
|
"principalUsername": str(principal.username or f"user-{user_id}"),
|
|||
|
|
"principalFullname": str(principal.fullname or ""),
|
|||
|
|
"tenantUuid": tenant_uuid,
|
|||
|
|
"scope": [AGENT_TOKEN_SCOPE_READ],
|
|||
|
|
"iat": now,
|
|||
|
|
"exp": now + ttl_days * 24 * 60 * 60,
|
|||
|
|
"revoked": False,
|
|||
|
|
"createdBy": created_by,
|
|||
|
|
"createdAt": now,
|
|||
|
|
}
|
|||
|
|
with self._lock:
|
|||
|
|
document = self._load()
|
|||
|
|
document["version"] = AGENT_TOKEN_VERSION
|
|||
|
|
document.setdefault("tokens", []).append(record)
|
|||
|
|
self._save(document)
|
|||
|
|
return self._identity_from_record(record), token
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _identity_from_record(record: dict[str, Any]) -> IdentityContext:
|
|||
|
|
return IdentityContext(
|
|||
|
|
user_id=int(record.get("principalUserId") or 0),
|
|||
|
|
username=str(
|
|||
|
|
record.get("principalUsername") or f"agent-{record.get('tokenId')}"
|
|||
|
|
),
|
|||
|
|
fullname=str(
|
|||
|
|
record.get("principalFullname") or record.get("label") or "Pi Agent"
|
|||
|
|
),
|
|||
|
|
tenant_uuid=str(record.get("tenantUuid") or "platform"),
|
|||
|
|
roles=("pi-agent",),
|
|||
|
|
expires_at=int(record.get("exp") or 0),
|
|||
|
|
auth_kind="agent",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def authenticate(self, token: str) -> IdentityContext:
|
|||
|
|
payload = self._verify_token(token)
|
|||
|
|
now = int(time.time())
|
|||
|
|
if payload["exp"] <= now:
|
|||
|
|
raise AuthError("AGENT_TOKEN_EXPIRED", "Pi Agent 凭证已到期", 401)
|
|||
|
|
record = self._find(str(payload["tokenId"]))
|
|||
|
|
if record is None:
|
|||
|
|
raise AuthError("AGENT_TOKEN_INVALID", "Pi Agent 凭证未登记或已删除", 401)
|
|||
|
|
if bool(record.get("revoked")):
|
|||
|
|
raise AuthError("AGENT_TOKEN_REVOKED", "Pi Agent 凭证已吊销", 401)
|
|||
|
|
if (
|
|||
|
|
int(record.get("exp") or 0) != payload["exp"]
|
|||
|
|
or int(record.get("iat") or 0) != payload["iat"]
|
|||
|
|
):
|
|||
|
|
raise AuthError("AGENT_TOKEN_INVALID", "Pi Agent 凭证登记不一致", 401)
|
|||
|
|
if AGENT_TOKEN_SCOPE_READ not in (record.get("scope") or []):
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_SCOPE_INVALID", "Pi Agent 凭证缺少只读作用域", 403
|
|||
|
|
)
|
|||
|
|
return self._identity_from_record(record)
|
|||
|
|
|
|||
|
|
def list_tokens(self, tenant_uuid: str | None = None) -> list[dict[str, Any]]:
|
|||
|
|
document = self._load()
|
|||
|
|
rows: list[dict[str, Any]] = []
|
|||
|
|
for record in document.get("tokens") or []:
|
|||
|
|
if tenant_uuid and str(record.get("tenantUuid") or "") != tenant_uuid:
|
|||
|
|
continue
|
|||
|
|
rows.append(
|
|||
|
|
{
|
|||
|
|
"tokenId": str(record.get("tokenId") or ""),
|
|||
|
|
"label": str(record.get("label") or ""),
|
|||
|
|
"principalUserId": int(record.get("principalUserId") or 0),
|
|||
|
|
"principalUsername": str(record.get("principalUsername") or ""),
|
|||
|
|
"tenantUuid": str(record.get("tenantUuid") or ""),
|
|||
|
|
"scope": list(record.get("scope") or []),
|
|||
|
|
"iat": int(record.get("iat") or 0),
|
|||
|
|
"exp": int(record.get("exp") or 0),
|
|||
|
|
"revoked": bool(record.get("revoked")),
|
|||
|
|
"createdBy": str(record.get("createdBy") or ""),
|
|||
|
|
"createdAt": int(record.get("createdAt") or 0),
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
return rows
|
|||
|
|
|
|||
|
|
def revoke(
|
|||
|
|
self,
|
|||
|
|
token_id: str,
|
|||
|
|
*,
|
|||
|
|
tenant_uuid: str | None = None,
|
|||
|
|
issuer: IdentityContext | None = None,
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
token_id = str(token_id or "").strip()
|
|||
|
|
if not token_id:
|
|||
|
|
raise AuthError("AGENT_TOKEN_ID_REQUIRED", "缺少凭证 ID", 400)
|
|||
|
|
self._require_secret()
|
|||
|
|
with self._lock:
|
|||
|
|
document = self._load()
|
|||
|
|
for record in document.get("tokens") or []:
|
|||
|
|
if str(record.get("tokenId") or "") != token_id:
|
|||
|
|
continue
|
|||
|
|
if tenant_uuid and str(record.get("tenantUuid") or "") != tenant_uuid:
|
|||
|
|
raise AuthError("AGENT_TOKEN_NOT_FOUND", "凭证不存在", 404)
|
|||
|
|
if not _can_manage(issuer, record):
|
|||
|
|
raise AuthError(
|
|||
|
|
"AGENT_TOKEN_FORBIDDEN", "当前身份无权吊销该凭证", 403
|
|||
|
|
)
|
|||
|
|
record["revoked"] = True
|
|||
|
|
self._save(document)
|
|||
|
|
return {
|
|||
|
|
"tokenId": token_id,
|
|||
|
|
"revoked": True,
|
|||
|
|
"tenantUuid": str(record.get("tenantUuid") or ""),
|
|||
|
|
}
|
|||
|
|
raise AuthError("AGENT_TOKEN_NOT_FOUND", "凭证不存在", 404)
|
|||
|
|
|
|||
|
|
|
|||
|
|
_service: AgentTokenService | None = None
|
|||
|
|
_service_signature: tuple[bytes | None, str] | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def reset_agent_token_service() -> None:
|
|||
|
|
"""丢弃进程级缓存，供环境变更后的测试与运维脚本调用。"""
|
|||
|
|
global _service, _service_signature
|
|||
|
|
_service = None
|
|||
|
|
_service_signature = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_service() -> AgentTokenService:
|
|||
|
|
global _service, _service_signature
|
|||
|
|
secret = (os.environ.get("APS_AGENT_TOKEN_SECRET") or "").strip().encode(
|
|||
|
|
"utf-8"
|
|||
|
|
) or None
|
|||
|
|
signature = (secret, str(_configured_path()))
|
|||
|
|
if _service is None or _service_signature != signature:
|
|||
|
|
_service = AgentTokenService(secret=secret, path=Path(signature[1]))
|
|||
|
|
_service_signature = signature
|
|||
|
|
return _service
|
|||
|
|
|
|||
|
|
|
|||
|
|
def issue(
|
|||
|
|
principal: IdentityContext, label: str, *, ttl_days: int = 365
|
|||
|
|
) -> tuple[IdentityContext, str]:
|
|||
|
|
return _get_service().issue(principal, label, ttl_days=ttl_days)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def authenticate(token: str) -> IdentityContext:
|
|||
|
|
return _get_service().authenticate(token)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def list_tokens(tenant_uuid: str | None = None) -> list[dict[str, Any]]:
|
|||
|
|
return _get_service().list_tokens(tenant_uuid=tenant_uuid)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def revoke(
|
|||
|
|
token_id: str,
|
|||
|
|
*,
|
|||
|
|
tenant_uuid: str | None = None,
|
|||
|
|
issuer: IdentityContext | None = None,
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
return _get_service().revoke(token_id, tenant_uuid=tenant_uuid, issuer=issuer)
|