zfsafe/tools/ToolBase.py

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

class ToolBase(abc.ABC):
    def __init__(self,TM):
        #self.res_type = 0   #补充一个结果类型 0-初始状态,1-有安全问题,
        #由于工具类会被多个线程调用，全局变量不能修改，只能读取
        self.TM = TM


    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 ""

        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 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
        except subprocess.TimeoutExpired as e:
            stdout = e.stdout if e.stdout is not None else ""
            stderr = e.stderr if e.stderr is not None else ""
            ext_params.is_user = True   #对于超时的也需要人工进行确认，是否是预期的超时
        except Exception as e:
            ext_params.is_user = True
            return False,instruction,f"执行失败：{str(e)}","",ext_params #执行失败，提交给人工确认指令的正确性

        # 第三步：分析执行结果
        output = stdout
        if stderr:
            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 False,instruction,analysis,output,ext_params -- 单节点后都要提交结果
        return True,instruction, analysis,output,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