# 工具基类
'''
工具基类----工具子类约束：
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,timeout=60*30)
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