# 工具基类 ''' 工具基类----工具子类约束： 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