zfsafe/tools/ToolBase.py

# 工具基类
'''
工具基类----工具子类约束：
1.一个工具子类一个py文件
2.工具类文件名必须为：[指令名]+Tool.py
3.未完成的工具子类以__开头
4.
'''
import abc
import subprocess
import argparse
import shlex
import subprocess
import tempfile
import os
from myutils.ReturnParams import ReturnParams

class ToolBase(abc.ABC):
    def get_time_out(self):
        return 60*30

    def create_extparams(self):
        ext_params = ReturnParams()
        ext_params["is_user"] = False  # 是否要提交用户确认 -- 默认False
        ext_params["is_vulnerability"] = False  # 是否是脆弱点
        return ext_params

    def parse_sublist3r_command(self,command):
        parser = argparse.ArgumentParser(add_help=False,allow_abbrev=False)    #创建命令行参数解析器对象‌
        parser.add_argument("-o", "--output", type=str)     #添加需要解析的参数规则‌
        parser.add_argument("-oN", "--Noutput", type=str)   #nmap
        parser.add_argument("-oG","--NMG",type=str)         #nmap
        parser.add_argument("-output", "--nikto", type=str) #nikto
        args, _ = parser.parse_known_args(command.split()[1:])
        return args

    def parse_output_params(self,command, valid_flags=None):
        """
        解析 shell 命令中用于输出文件的参数。
        只检查完全匹配的 token，忽略诸如 "-oKexAlgorithms" 这样的参数。

        :param command: 完整的命令字符串
        :param valid_flags: 合法的输出参数列表，默认支持 -o, -oN, -oG, -output
        :return: dict，键为输出参数，值为对应的文件路径（如果有的话）
        """
        if command.strip().startswith("mkdir"):
            return ""
        #额外指令排除:ssh -o 不是输出结果到文件
        if "ssh" in command:
            return ""

        if valid_flags is None:
            valid_flags = ['-o', '-oN', '-oG', '-output']

        tokens = shlex.split(command)
        output_params = {}
        output_file = ""
        i = 0
        while i < len(tokens):
            token = tokens[i]
            # 如果 token 完全匹配有效的参数，则认为它是输出参数
            if token in valid_flags:
                # 如果紧跟着有值，则把它作为输出文件路径，否则为 None
                if i + 1 < len(tokens):
                    #output_params[token] = tokens[i + 1]
                    #i += 2
                    output_file = tokens[i+1]
                # else:
                #     output_params[token] = None
                #     i += 1
                break  # 一般只有一个输出参数
            else:
                i += 1
        return output_file

    def read_output_file(self,filename):
        """
        读取 -o 参数指定的文件内容
        """
        content = ""
        try:
            with open(filename, "r") as f:
                content = f.read()
            with open(filename,'w') as f:
                f.write("")
        except FileNotFoundError:
            print(f"错误: 文件 {filename} 不存在")
        except PermissionError:
            print(f"错误: 无权限读取文件 {filename}")
        return content

    def do_worker_script(self,str_instruction,timeout,ext_params):
        # 创建临时文件保存输出
        with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
            output_file = tmpfile.name

        # 构建并执行 script 命令
        script_cmd = f"script -c '{str_instruction}' {output_file}"
        try:
            result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout)
            # 读取输出文件内容
            with open(output_file, 'r') as f:
                output = f.read()
                lines = output.splitlines()
                # 跳过第一行（Script started）和最后一行（Script done）
                ftp_output = lines[1:-1]
                output = '\n'.join(ftp_output)
        except subprocess.TimeoutExpired:
            output = "命令超时返回"
            try:
                with open(output_file, 'r') as f:
                    partial_output = f.read()
                if partial_output:
                    output += f"\n部分输出:\n{partial_output}"
            except FileNotFoundError:
                pass  # 文件可能未创建
        except subprocess.CalledProcessError as e:
            output = f"错误: {e}"
        finally:
            # 删除临时文件
            try:
                os.remove(output_file)
            except FileNotFoundError:
                pass  # 文件可能未创建
        return output

    def execute_instruction(self, instruction_old):
        '''
        执行指令：验证合法性 -> 执行 -> 分析结果
        *****如果指令要做验证，只做白名单，所有逻辑不是全放开就是白名单*****
        :param instruction_old:
        :return:
            bool:true-正常返回给大模型处理下一步，false-结果不返回给大模型,2--需要人工确认的指令
            str:执行的指令
            str:执行指令的结果-解析过滤后的结果--也是提交给LLM的结果
            str:执行指令的结果-原结果
            object:补充参数-封装一个对象： 0-不知是否攻击成功，1-明确存在漏洞，2-明确不存在漏洞
        '''

        ext_params = self.create_extparams()
        # 第一步：验证指令合法性
        instruction,timeout = self.validate_instruction(instruction_old)
        if not instruction:
            ext_params.is_user= True
            return False,instruction_old,"该指令暂不执行！由用户确认是否要兼容支持","",ext_params      #未
        #过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？

        # 第二步：执行指令
        #print(f"执行指令：{instruction}")
        output = ""
        stdout = ""
        stderr = ""
        try:
            if timeout == 0:
                result = subprocess.run(instruction, shell=True, capture_output=True, text=True)
            elif timeout >0:
                result = subprocess.run(instruction, shell=True, capture_output=True, text=True,  timeout=timeout)
            else:
                print("timeout参数错误,需要自查程序逻辑！")
            #output = result.stdout if result.stdout else result.stderr
            #-o 的命令需要处理
            parsed_arg = self.parse_output_params(instruction)
            if parsed_arg:
                fole_out = self.read_output_file(parsed_arg)
                stderr =""
                stdout = fole_out
            else:
                stderr = result.stderr
                stdout = result.stdout

            # 第三步：分析执行结果
            # if result.returncode == 0:  # 执行正确
            #     output = stdout
            if stdout:
                output = stdout
            else:
                if result.returncode ==28:    #curl 命令超时
                    output = "指令执行超时！"
                else:   #执行错误只拿错误信息
                    output = stderr
            if isinstance(output, bytes):  # 若是bytes则转成str
                output = output.decode('utf-8', errors='ignore')

            analysis = self.analyze_result(output, instruction, stderr, stdout)
            if not analysis:  # analysis为“” 不提交LLM
                ext_params.is_user = True
            return True, instruction, analysis, output, ext_params

        except subprocess.TimeoutExpired:
            ext_params.is_user = True   #对于超时的也需要人工进行确认，是否是预期的超时
            return False, instruction, f"执行超时：{str(timeout)}秒", "", ext_params  # 执行失败，提交给人工确认指令的正确性
        except Exception as e:
            ext_params.is_user = True
            return False,instruction,f"执行失败：{str(e)}","",ext_params #执行失败，提交给人工确认指令的正确性

    @abc.abstractmethod
    def validate_instruction(self, instruction:str):
        """指令合法性分析
        :return:
        str：非空代表合法（函数内会涉及修改指令），为空代表非法已取消指令
        int：表示超时时间，0不设置超时时间，正值具体的超时时间设置。单位秒
        """
        pass

    @abc.abstractmethod
    def analyze_result(self, result:str,instruction:str,stderr:str,stdout:str) -> str:
        """分析执行结果"""
        pass

def read_output_file_test(filename):
    """
    读取 -o 参数指定的文件内容
    """
    try:
        with open(filename, "r") as f:
            return f.read()
    except FileNotFoundError:
        print(f"错误: 文件 {filename} 不存在")
    except PermissionError:
        print(f"错误: 无权限读取文件 {filename}")
    return ""

if __name__ =="__main__":
    pass