worldquant-alpha-system/backend/app/cli.py

"""Administrative commands run inside the backend container."""

import argparse
import asyncio
import getpass

from sqlalchemy import delete, update

from .config import Settings
from .db import create_database
from .models import Admin, LoginSession, MCPToken, now
from .security import password_hasher


async def reset_password():
    password = getpass.getpass("New admin password (12+ characters): ")
    if len(password) < 12 or password != getpass.getpass("Confirm password: "):
        raise SystemExit("Password too short or confirmation does not match")
    engine, sessions = create_database(Settings().database_url)
    async with sessions() as db:
        admin = await db.get(Admin, 1)
        admin.password_hash = password_hasher.hash(password)
        await db.execute(delete(LoginSession))
        await db.execute(update(MCPToken).where(MCPToken.revoked_at.is_(None)).values(revoked_at=now()))
        await db.commit()
    await engine.dispose()
    print("Admin password updated; all system sessions and MCP tokens revoked.")


async def token_command(args):
    import json

    from sqlalchemy import select

    from .mcp_api.auth import create_token
    from .models import MCPToken, now
    from .research.serialization import encode_snapshot

    engine, sessions = create_database(Settings().database_url)
    try:
        async with sessions.begin() as db:
            if args.command == "mcp-token-create":
                row, secret = await create_token(db, args.name, args.scope, args.days)
                result = {"id": row.id, "name": row.name, "scopes": row.scopes,
                          "expires_at": row.expires_at, "token": secret}
            elif args.command == "mcp-token-revoke":
                row = await db.get(MCPToken, args.token_id)
                if not row:
                    raise ValueError("令牌不存在")
                row.revoked_at = row.revoked_at or now()
                result = {"id": row.id, "revoked": True}
            else:
                rows = list(await db.scalars(select(MCPToken).order_by(MCPToken.created_at.desc())))
                result = [{k: getattr(row, k) for k in
                           ("id", "name", "scopes", "created_at", "expires_at", "revoked_at")} for row in rows]
        # Reveal only after the transaction has committed successfully.
        print(json.dumps(encode_snapshot(result), ensure_ascii=False, indent=2))
    finally:
        await engine.dispose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("reset-password")
    create = commands.add_parser("mcp-token-create")
    create.add_argument("--name", required=True)
    create.add_argument("--scope", action="append", default=None)
    create.add_argument("--days", type=int, default=90)
    commands.add_parser("mcp-token-list")
    revoke = commands.add_parser("mcp-token-revoke")
    revoke.add_argument("token_id")
    args = parser.parse_args()
    try:
        asyncio.run(reset_password() if args.command == "reset-password" else token_command(args))
    except ValueError as exc:
        parser.error(str(exc))