from tools.ToolBase import ToolBase
import re
import json
import subprocess
import os
import shutil
from pathlib import Path
class SearchsploitTool(ToolBase):
    def validate_instruction(self, instruction):
        #指令过滤
        timeout = 0
        if "-m " in instruction:  # 下载利用代码
            timeout = 60
        return instruction,timeout

    def analyze_result(self, result,instruction,stderr,stdout):
        #获取当前路径
        cur_path =  Path(__file__).resolve().parent
        payload_dir = cur_path / "../payload"

        #if instruction(result,bytes):
        if type(result) is bytes:
            result = result.decode('utf-8',errors='ignore')

        """去除 ANSI 颜色码"""
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        clean_result =  ansi_escape.sub('', result)
        # 指令结果分析
        if "-m " in instruction:#下载利用代码
            if "Copied to:" in clean_result or "cp: overwrite " in clean_result:
                result = "下载完成"
            else:
                result = "下载失败"
        else: #目前只遇到两种searchsploit命令格式：
            lines = clean_result.split("\n")
            exploits = []

            for line in lines:
                match = re.match(r"(.+?)\s+\|\s+(\S+)", line)
                if match:
                    title = match.group(1).strip()
                    path = match.group(2).strip()
                    if "Path" not in path:
                        #下载渗透脚本--进一步推进
                        ext_str = self.do_ext(path,payload_dir)
                        exploits.append({"title": title, "path": path,"补充信息":ext_str})

            if len(exploits) > 0:
                result = json.dumps(exploits,ensure_ascii=False)       #输出原始的中文字符
            else:
                result = "没有检索到漏洞利用脚本"
        return result

    def find_file_in_directory(self,target_file, dir_path, case_sensitive=True, recursive=False):
        """
        在指定目录中查找文件是否存在
        :param target_file: 要查找的目标文件名（含扩展名）
        :param search_dir: 要搜索的目录路径
        :param case_sensitive: 是否区分大小写（默认True）
        :param recursive: 是否递归搜索子目录（默认False）
        :return: (bool, str) 是否存在，完整路径（如果找到）
        """
        try:
            #dir_path = Path(search_dir)

            # 验证目录是否存在
            if not dir_path.is_dir():
                return False, None

            # 根据是否递归选择遍历方式
            iterator = dir_path.rglob('*') if recursive else dir_path.iterdir()

            for entry in iterator:
                # 跳过目录只处理文件
                if entry.is_file():
                    # 根据是否区分大小写进行比较
                    if case_sensitive:
                        match = entry.name == target_file
                    else:
                        match = entry.name.lower() == target_file.lower()

                    if match:
                        return True, str(entry.resolve())

            return False, None

        except Exception as e:
            print(f"查找出错: {str(e)}")
            return False, None

    def do_download(self,filepath,dirpath):
        '''
        下载渗透脚本
        :return:
        '''
        filename = filepath.split("/")[-1]
        if ".py" in filename:
            #对应payload库中是否存在该脚本
            bfind,_ = self.find_file_in_directory(filename,dirpath)
            if bfind:
                return True
            else:#下载
                instruction = f"searchsploit -m {filepath}"
                try:
                    subprocess.run(instruction, shell=True, check=True,timeout=60)
                    print("命令执行成功，文件下载完成。")
                except subprocess.CalledProcessError as e: #超时暂不处理--后续补充
                    print(f"命令执行失败: {e}")
                    return False

                if not os.path.exists(filename):
                    return False
                # 移动文件到目标目录
                try:
                    shutil.move(filename, os.path.join(dirpath, filename))
                    print(f"文件已成功移动到 {dirpath}")
                    return True
                except Exception as e:
                    print(f"移动文件失败: {e}")
                    return False
        else:   #暂时只利用python脚本
            return False

    def do_ext(self,filepath,payload_dir) -> str:
        bdownload = self.do_download(filepath,payload_dir)
        if bdownload:
            return "该渗透脚本已经复制到当前路径。"
        else:
            return "暂时不利用该渗透脚本。"