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)
|