zfsafe/tools/ToolBase.py

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

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


    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)    #创建命令行参数解析器对象‌
        parser.add_argument("-o ", "--output", type=str)     #添加需要解析的参数规则‌
        parser.add_argument("-oN ", "--Noutput", type=str)   #nmap
        parser.add_argument("-output ", "--nikto", type=str) #nikto
        args, _ = parser.parse_known_args(command.split()[1:])
        return args

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

    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_sublist3r_command(instruction)
            if parsed_arg.output:
                output = self.read_output_file(parsed_arg.output)
            elif parsed_arg.Noutput:
                output = self.read_output_file(parsed_arg.Noutput)
            elif parsed_arg.nikto:
                output = self.read_output_file(parsed_arg.nikto)
            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
        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