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

# Nmap工具类
class NmapTool(ToolBase):
    def validate_instruction(self, instruction):
        #nmap过滤
        timeout = 0
        if "&& nc " in instruction or "&& ftp " in instruction:
            timeout = 60
        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
            '''常见格式
            Running: Linux 2.6.X
            OS CPE: cpe:/o:linux:linux_kernel:2.6
            OS details: Linux 2.6.9 - 2.6.33
            Network Distance: 1 hop
            Service Info: Hosts:  metasploitable.localdomain, irc.Metasploitable.LAN; OSs: Unix, 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}"

        result = json.dumps(result,ensure_ascii=False)
        return result

    def extract_key_info(self,nmap_text):
        """
            利用正则表达式提取每一行的端口、状态、服务及版本/额外信息。
            只处理形如 "端口/tcp   状态   服务   版本/额外信息" 的行，达到字符缩减的效果。
            """
        # 正则模式：匹配开头数字、/tcp、状态、服务名称，后续内容为可选的额外信息
        pattern = re.compile(r'^(\d+)/tcp\s+(\w+)\s+(\S+)(?:\s+(.*))?$')
        key_info = []
        for line in nmap_text.splitlines():
            match = pattern.match(line)
            if match:
                port, state, service, extra = match.groups()
                extra = extra.strip() if extra else ""
                key_info.append({
                    "port": port,
                    "state": state,
                    "service": service,
                    "info": extra
                })
        return key_info

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

        #result = self.parse_nmap_output(result)
        # tmpstr = self.extract_key_info(result)
        # result = json.dumps(tmpstr)

        return result