xingrin/backend/apps/scan/utils/command_builder.py

"""
简化的命令构建工具
使用 Python 原生 f-string 和条件拼接，零依赖，性能更好。
"""

import logging
from typing import Dict, Any

logger = logging.getLogger(__name__)


def build_scan_command(
    tool_name: str,
    scan_type: str,
    command_params: Dict[str, Any],
    tool_config: Dict[str, Any]
) -> str:
    """
    构建扫描工具命令（使用 f-string）
    
    Args:
        tool_name: 工具名称（如 'subfinder'）
        scan_type: 扫描类型（如 'subdomain_discovery'）
        command_params: 命令占位符参数
            - domain: 目标域名
            - domains_file: 域名列表文件（用于端口扫描）
            - url_file: URL列表文件（用于站点扫描）
            - target_file: 目标文件路径（通用）
            - output_file: 输出文件路径
        tool_config: 工具配置参数（包含可选参数）
            - threads: 线程数
            - timeout: 超时时间（秒）
            - 其他可选参数...
    
    Returns:
        完整的命令字符串
    
    Example:
        >>> build_scan_command(
        ...     tool_name='subfinder',
        ...     scan_type='subdomain_discovery',
        ...     command_params={'domain': 'example.com', 'output_file': '/tmp/out.txt'},
        ...     tool_config={'threads': 10}
        ... )
        'subfinder -d example.com -o /tmp/out.txt -silent -t 10'
    """
    from apps.scan.configs.command_templates import get_command_template, SCAN_TOOLS_BASE_PATH
    
    # 获取命令模板
    template = get_command_template(scan_type, tool_name)
    if not template:
        raise ValueError(f"未找到工具 {tool_name} 的命令模板（扫描类型: {scan_type}）")
    
    # 合并所有参数，并将中划线统一转成下划线
    # 规范约定：
    #   - 配置文件（YAML）：参数名用中划线，贴近 CLI 原生参数（如 rate-limit, request-timeout）
    #   - 模板文件（Python）：参数名用下划线，符合 str.format() 占位符语法要求
    #   - 此处自动转换：rate-limit → rate_limit
    def normalize_key(k):
        return k.replace('-', '_') if isinstance(k, str) else k
    
    all_params = {
        'scan_tools_base': SCAN_TOOLS_BASE_PATH,
        **{normalize_key(k): v for k, v in command_params.items()},
        **{normalize_key(k): v for k, v in tool_config.items()}
    }

    # nuclei 特殊处理：要求 template_args 必填（支持多 -t），避免格式化缺失
    if tool_name == "nuclei":
        if not all_params.get("template_args"):
            raise ValueError("nuclei 命令构建缺少 template_args（请检查模板仓库列表配置）")
    
    try:
        # 1. 构建基础命令
        base_command = template['base'].format(**all_params)
        
        # 2. 拼接可选参数
        optional_parts = []
        for param_name, flag_template in template.get('optional', {}).items():
            # 检查参数是否存在且有值
            if param_name in all_params and all_params[param_name]:
                optional_parts.append(flag_template.format(**all_params))
        
        # 3. 组合完整命令
        full_command = base_command
        if optional_parts:
            full_command += ' ' + ' '.join(optional_parts)
        
        # 4. 清理多余空白
        import re
        cleaned_command = re.sub(r'\s+', ' ', full_command).strip()
        
        return cleaned_command
        
    except KeyError as e:
        raise ValueError(
            f"命令构建失败：缺少必需参数 {e}\n"
            f"模板: {template}\n"
            f"提供的参数: {list(all_params.keys())}"
        )
    except Exception as e:
        raise ValueError(
            f"命令构建失败: {e}\n"
            f"模板: {template}\n"
            f"提供的参数: {list(all_params.keys())}"
        )