RuoYi-Vue3-FastAPI/ruoyi-fastapi-backend/cli/guards.py

import sys
from dataclasses import dataclass, field
from typing import Literal

import click
import typer

from cli.context import CliContext
from cli.exit_codes import GUARD_REJECTED
from cli.output import CommandResult

DangerousCommandRiskLevel = Literal['high', 'normal']


@dataclass(frozen=True)
class DangerousCommandRule:
    """
    危险命令保护规则。

    :param command_name: 命令唯一标识
    :param risk_level: 风险级别
    :param supports_dry_run: 是否支持演练执行
    """

    command_name: str
    risk_level: DangerousCommandRiskLevel
    supports_dry_run: bool


@dataclass(frozen=True)
class DangerousCommandRuleRegistry:
    """
    危险命令规则注册表。

    该注册表负责维护 CLI 内所有危险命令的风险元数据，并提供查询与
    强制获取能力，作为上下文工厂、风险元数据与测试的统一入口。

    :param rules: 按命令名索引的危险命令规则表
    """

    rules: dict[str, DangerousCommandRule]

    def get_rule(self, command_name: str) -> DangerousCommandRule | None:
        """
        获取指定命令的危险命令规则。

        :param command_name: 命令唯一标识
        :return: 命令规则，不存在时返回 None
        """
        return self.rules.get(command_name)

    def require_rule(self, command_name: str) -> DangerousCommandRule:
        """
        获取指定命令的危险命令规则，不存在时抛出异常。

        :param command_name: 命令唯一标识
        :return: 命令规则
        :raises ValueError: 命令未注册危险命令规则时抛出
        """
        rule = self.get_rule(command_name)
        if rule is None:
            raise ValueError(f'危险命令未注册保护规则：{command_name}')
        return rule


@dataclass(frozen=True)
class DangerousCommandResultBuilder:
    """
    危险命令结果构建器。

    该对象负责统一构建危险命令拒绝结果，避免保护服务内部继续拼装
    结构化负载细节。
    """

    @staticmethod
    def build_guard_reject_result(message: str, hint: str) -> CommandResult:
        """
        构建危险命令拒绝结果。

        :param message: 拒绝原因
        :param hint: 补充提示
        :return: 命令拒绝结果
        """
        return CommandResult(
            data={
                'ok': False,
                'message': message,
                'hint': hint,
            },
            exit_code=GUARD_REJECTED,
        )


@dataclass(frozen=True)
class DangerousCommandConfirmationService:
    """
    危险命令确认服务。

    该对象负责处理 TTY 检测、确认提示与交互取消异常收口。

    :param result_builder: 危险命令结果构建器
    """

    result_builder: DangerousCommandResultBuilder

    def confirm(self, ctx: CliContext, *, command_name: str) -> CommandResult | None:
        """
        执行危险命令交互确认。

        :param ctx: CLI 上下文
        :param command_name: 命令唯一标识
        :return: 拒绝结果或 None
        """
        if ctx.yes:
            return None

        if not sys.stdin.isatty():
            return self.result_builder.build_guard_reject_result(
                f'已取消危险命令执行：{command_name}',
                '当前命令需要交互确认；如需非交互执行，请传入 --yes',
            )

        try:
            confirmed = typer.confirm(
                f'确认执行危险命令 `{command_name}` 吗？ 当前环境：{ctx.env}{"（dry-run）" if ctx.dry_run else ""}',
                default=False,
            )
        except (click.Abort, EOFError, KeyboardInterrupt):
            return self.result_builder.build_guard_reject_result(
                f'已取消危险命令执行：{command_name}',
                '当前命令需要交互确认；如需非交互执行，请传入 --yes',
            )
        if confirmed:
            return None

        return self.result_builder.build_guard_reject_result(
            f'已取消危险命令执行：{command_name}',
            '如需跳过确认，请传入 --yes',
        )


@dataclass
class DangerousCommandGuardService:
    """
    危险命令保护执行服务。

    该服务负责根据命令规则和 CLI 上下文执行生产环境保护、交互确认
    与拒绝结果收口。

    :param rule_registry: 危险命令规则注册表
    :param result_builder: 危险命令结果构建器
    :param confirmation_service: 危险命令确认服务
    """

    rule_registry: DangerousCommandRuleRegistry
    result_builder: DangerousCommandResultBuilder = field(default_factory=DangerousCommandResultBuilder)
    confirmation_service: DangerousCommandConfirmationService = field(init=False)

    def __post_init__(self) -> None:
        """
        初始化危险命令保护服务依赖。

        :return: None
        """
        self.confirmation_service = DangerousCommandConfirmationService(self.result_builder)

    def guard(self, ctx: CliContext, *, rule: DangerousCommandRule) -> CommandResult | None:
        """
        执行危险命令保护与确认。

        :param ctx: CLI 上下文
        :param rule: 危险命令规则
        :return: 拒绝结果或 None
        """
        command_name = rule.command_name
        if ctx.env == 'prod' and not ctx.allow_prod:
            return self.result_builder.build_guard_reject_result(
                f'生产环境默认禁止直接执行危险命令：{command_name}',
                '如确认执行，请传入 --allow-prod；如需跳过确认，请同时传入 --yes',
            )

        return self.confirmation_service.confirm(ctx, command_name=command_name)


DEFAULT_DANGEROUS_COMMAND_RULES: dict[str, DangerousCommandRule] = {
    'cache clear': DangerousCommandRule(command_name='cache clear', risk_level='high', supports_dry_run=True),
    'cache warmup': DangerousCommandRule(command_name='cache warmup', risk_level='normal', supports_dry_run=False),
    'db upgrade': DangerousCommandRule(command_name='db upgrade', risk_level='high', supports_dry_run=True),
    'db init': DangerousCommandRule(command_name='db init', risk_level='high', supports_dry_run=True),
    'db downgrade': DangerousCommandRule(command_name='db downgrade', risk_level='high', supports_dry_run=True),
    'db revision': DangerousCommandRule(command_name='db revision', risk_level='high', supports_dry_run=True),
    'config set': DangerousCommandRule(command_name='config set', risk_level='high', supports_dry_run=True),
    'config sync-cache': DangerousCommandRule(
        command_name='config sync-cache',
        risk_level='normal',
        supports_dry_run=False,
    ),
    'crypto rotate': DangerousCommandRule(command_name='crypto rotate', risk_level='high', supports_dry_run=True),
    'job run-once': DangerousCommandRule(command_name='job run-once', risk_level='normal', supports_dry_run=False),
    'job pause': DangerousCommandRule(command_name='job pause', risk_level='normal', supports_dry_run=False),
    'job resume': DangerousCommandRule(command_name='job resume', risk_level='normal', supports_dry_run=False),
    'job sync': DangerousCommandRule(command_name='job sync', risk_level='normal', supports_dry_run=False),
    'gen import-table': DangerousCommandRule(command_name='gen import-table', risk_level='high', supports_dry_run=True),
    'gen create-table': DangerousCommandRule(command_name='gen create-table', risk_level='high', supports_dry_run=True),
    'gen export': DangerousCommandRule(command_name='gen export', risk_level='high', supports_dry_run=True),
    'gen sync-db': DangerousCommandRule(command_name='gen sync-db', risk_level='normal', supports_dry_run=False),
    'plugin install': DangerousCommandRule(command_name='plugin install', risk_level='high', supports_dry_run=True),
    'plugin install-deps': DangerousCommandRule(
        command_name='plugin install-deps',
        risk_level='high',
        supports_dry_run=True,
    ),
    'plugin upgrade': DangerousCommandRule(command_name='plugin upgrade', risk_level='high', supports_dry_run=True),
    'plugin batch': DangerousCommandRule(command_name='plugin batch', risk_level='high', supports_dry_run=True),
    'plugin enable': DangerousCommandRule(command_name='plugin enable', risk_level='normal', supports_dry_run=True),
    'plugin disable': DangerousCommandRule(command_name='plugin disable', risk_level='normal', supports_dry_run=True),
    'plugin config set': DangerousCommandRule(
        command_name='plugin config set',
        risk_level='normal',
        supports_dry_run=False,
    ),
    'plugin config import': DangerousCommandRule(
        command_name='plugin config import',
        risk_level='normal',
        supports_dry_run=False,
    ),
    'plugin config export': DangerousCommandRule(
        command_name='plugin config export',
        risk_level='normal',
        supports_dry_run=False,
    ),
    'plugin uninstall': DangerousCommandRule(
        command_name='plugin uninstall', risk_level='normal', supports_dry_run=True
    ),
    'plugin purge': DangerousCommandRule(command_name='plugin purge', risk_level='high', supports_dry_run=True),
    'plugin mark-success': DangerousCommandRule(
        command_name='plugin mark-success',
        risk_level='high',
        supports_dry_run=False,
    ),
    'plugin mark-failed': DangerousCommandRule(
        command_name='plugin mark-failed',
        risk_level='high',
        supports_dry_run=False,
    ),
}

DEFAULT_DANGEROUS_COMMAND_RULE_REGISTRY = DangerousCommandRuleRegistry(rules=DEFAULT_DANGEROUS_COMMAND_RULES)
DEFAULT_DANGEROUS_COMMAND_GUARD = DangerousCommandGuardService(rule_registry=DEFAULT_DANGEROUS_COMMAND_RULE_REGISTRY)