You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
85 lines
3.5 KiB
85 lines
3.5 KiB
1 month ago
|
import subprocess
|
||
|
|
import tempfile
|
||
|
|
import os
|
||
|
|
from tools.ToolBase import ToolBase
|
||
|
|
|
||
|
|
class PsqlTool(ToolBase):
|
||
|
|
def validate_instruction(self, instruction):
|
||
|
|
#指令过滤
|
||
|
|
timeout = 60
|
||
|
|
return instruction,timeout
|
||
|
|
|
||
|
|
def analyze_result(self, result,instruction,stderr,stdout):
|
||
|
|
#指令结果分析
|
||
|
|
return result
|
||
|
|
|
||
|
|
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_instruction1(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 #未
|
||
|
|
#过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？
|
||
|
|
|
||
|
|
# 第二步：执行指令
|
||
|
|
output = self.do_worker_script(instruction,timeout,ext_params)
|
||
|
|
|
||
|
|
# 第三步：分析执行结果
|
||
|
|
if isinstance(output,bytes):#若是bytes则转成str
|
||
|
|
output = output.decode('utf-8', errors='ignore')
|
||
|
|
|
||
|
|
analysis = self.analyze_result(output,instruction,"","")
|
||
|
|
|
||
|
|
if not analysis: #analysis为“” 不提交LLM
|
||
|
|
ext_params.is_user = True
|
||
|
|
return False,instruction,analysis,output,ext_params
|
||
|
|
return True,instruction, analysis,output,ext_params
|