aps-agent/scripts/pi_agent_token.py

111 lines
3.2 KiB
Python

from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
sys.path.insert(0, str(_REPO_ROOT))
from server.auth import agent_tokens
from server.auth.context import IdentityContext
from server.auth.providers import AuthError
def _principal(args: argparse.Namespace) -> IdentityContext:
user_id = int(args.user_id or os.environ.get("APS_AGENT_PRINCIPAL_USER_ID") or "1")
username = (
args.username or os.environ.get("APS_AGENT_PRINCIPAL_USERNAME") or "local-admin"
)
fullname = args.fullname or username
tenant_uuid = args.tenant or os.environ.get("APS_AGENT_TENANT") or "platform"
return IdentityContext(
user_id=user_id,
username=username,
fullname=fullname,
tenant_uuid=tenant_uuid,
roles=("admin", "system"),
auth_kind="local-admin",
)
def _print(payload: dict) -> None:
print(json.dumps(payload, ensure_ascii=False, indent=2))
def _cmd_issue(args: argparse.Namespace) -> int:
identity, token = agent_tokens.issue(
_principal(args),
args.label,
ttl_days=args.ttl_days,
)
rows = agent_tokens.list_tokens(tenant_uuid=identity.tenant_uuid)
_print(
{
"ok": True,
"token": token,
"tokenId": rows[-1]["tokenId"] if rows else None,
"label": args.label,
"scope": ["read"],
"actor": f"pi-agent:{identity.username}",
"tenantUuid": identity.tenant_uuid,
}
)
return 0
def _cmd_list(args: argparse.Namespace) -> int:
_print({"ok": True, "tokens": agent_tokens.list_tokens(tenant_uuid=args.tenant)})
return 0
def _cmd_revoke(args: argparse.Namespace) -> int:
result = agent_tokens.revoke(args.token_id, tenant_uuid=args.tenant)
_print({"ok": True, **result})
return 0
def main() -> int:
parser = argparse.ArgumentParser(description="Pi Agent 长期只读凭证管理")
sub = parser.add_subparsers(dest="command", required=True)
issue = sub.add_parser("issue", help="签发长期只读 token")
issue.add_argument("--label", default="pi-agent-read")
issue.add_argument("--user-id", type=int)
issue.add_argument("--username")
issue.add_argument("--fullname")
issue.add_argument("--tenant")
issue.add_argument("--ttl-days", type=int, default=365)
issue.set_defaults(handler=_cmd_issue)
listed = sub.add_parser("list", help="列出已登记 token")
listed.add_argument("--tenant")
listed.set_defaults(handler=_cmd_list)
revoked = sub.add_parser("revoke", help="吊销 token")
revoked.add_argument("token_id")
revoked.add_argument("--tenant")
revoked.set_defaults(handler=_cmd_revoke)
args = parser.parse_args()
try:
return int(args.handler(args))
except (AuthError, ValueError) as exc:
_print(
{
"ok": False,
"error": {
"code": getattr(exc, "code", "ARG_INVALID"),
"message": str(exc),
},
}
)
return 2
if __name__ == "__main__":
sys.exit(main())