# Nmap工具类
import re
from typing import List, Dict, Any
from tools.ToolBase import ToolBase

# Nmap工具类
class NmapTool(ToolBase):
    def validate_instruction(self, instruction):
        #nmap过滤
        timeout = 0
        return instruction,timeout

    def parse_nmap_output(self,nmap_output: str) -> Dict[str, Any]:
        """
        分析 nmap 扫描输出，提取常见信息：
          - 端口信息（端口号、状态、服务、版本）
          - banner 信息（脚本输出部分）
          - OS 信息
          - 原始输出（作为补充）
        如果输出为空或分析过程中出现异常，则返回错误信息。
        """
        result: Dict[str, Any] = {}

        # 判断输出是否为空
        if not nmap_output.strip():
            result['error'] = "nmap 输出为空或无效"
            return result

        try:
            # 1. 提取端口信息
            # 这里匹配类似：
            #  10001/tcp open  softether-vpn-mgr  SoftEther VPN Server Manager 4.34 Build 9749
            port_pattern = re.compile(
                r"(?P<port>\d+/tcp)\s+(?P<state>\S+)\s+(?P<service>\S+)(?:\s+(?P<version>.+))?"
            )
            ports: List[Dict[str, str]] = []
            for match in port_pattern.finditer(nmap_output):
                port_info = match.groupdict()
                # 如果version为空，则设置为空字符串
                port_info["version"] = port_info.get("version") or ""
                ports.append(port_info)
            result['ports'] = ports if ports else "No port info found"

            # 2. 提取 banner 信息
            # 这里匹配以"| "开头的行
            banner_pattern = re.compile(r"^\|\s+(.*)", re.MULTILINE)
            banners = [line.strip() for line in banner_pattern.findall(nmap_output)]
            result['banners'] = banners if banners else "No banner info found"

            # 3. 提取 OS 信息
            # 常见格式：Service Info: OS: Linux; CPE: cpe:/o:linux:linux_kernel
            os_match = re.search(r"Service Info:\s+OS:\s*([^;]+);", nmap_output, re.IGNORECASE)
            result['os_info'] = os_match.group(1).strip() if os_match else "No OS info found"

            # 4. 其他信息：原始输出作为补充
            result['raw'] = nmap_output

        except Exception as e:
            result['error'] = f"Error parsing nmap output: {e}"

        return result

    def analyze_result(self, result,instruction,stderr,stdout):
        # 检查结果
        start_index = result.find("If you know the service/version")
        if start_index != -1:
            return result[:start_index]
        #result = self.parse_nmap_output(result)
        return result