mirror of
https://github.com/yyhuni/xingrin.git
synced 2026-01-31 19:53:11 +08:00
107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
"""
|
||
简化的命令构建工具
|
||
使用 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())}"
|
||
)
|