56 changed files with 3135 additions and 743 deletions
@ -0,0 +1,38 @@ |
|||||
|
import queue |
||||
|
import threading |
||||
|
import time |
||||
|
from mycode.PythoncodeTool import PythoncodeTool |
||||
|
|
||||
|
|
||||
|
class PythonTManager: |
||||
|
def __init__(self,maxnum): |
||||
|
self.brun = True |
||||
|
self.cur_num = 0 |
||||
|
self.put_lock = threading.Lock() |
||||
|
# 构建进程池--Python 代码的执行在子进程完成-3个子进程 |
||||
|
self.maxnum = maxnum |
||||
|
self.python_tool = PythoncodeTool(maxnum) #python工具实例 |
||||
|
|
||||
|
def __del__(self): |
||||
|
self.python_tool.shutdown() |
||||
|
|
||||
|
def execute_instruction(self,instruction): |
||||
|
bwork = False |
||||
|
while self.brun: |
||||
|
with self.put_lock: |
||||
|
if self.cur_num < self.maxnum: |
||||
|
self.cur_num += 1 |
||||
|
bwork = True |
||||
|
if bwork:#还有空的子进程 |
||||
|
#提交给进程池执行 |
||||
|
_,instruction,analysis,_,ext_params = self.python_tool.execute_instruction(instruction) |
||||
|
#执行完成后，数量减一 |
||||
|
with self.put_lock: |
||||
|
self.cur_num -= 1 |
||||
|
#返回结果 |
||||
|
return instruction,analysis,analysis,ext_params |
||||
|
else: #如果没获取的许可，则等待N秒后再尝试---存在问题：多线程间没有先来先到的机制了，有可能第一个来排队的一直等到最后 |
||||
|
time.sleep(20) #休眠20秒 |
||||
|
|
||||
|
|
||||
|
|
||||
@ -0,0 +1,230 @@ |
|||||
|
#python代码动态执行 |
||||
|
import queue |
||||
|
import ast |
||||
|
import subprocess |
||||
|
import json |
||||
|
import builtins |
||||
|
import re |
||||
|
import paramiko |
||||
|
import impacket |
||||
|
import psycopg2 |
||||
|
import socket |
||||
|
import struct |
||||
|
import sys |
||||
|
import requests |
||||
|
import ssl |
||||
|
import mysql.connector |
||||
|
import telnetlib |
||||
|
import time |
||||
|
import uuid |
||||
|
import multiprocessing |
||||
|
import textwrap |
||||
|
from mycode.Result_merge import my_merge |
||||
|
from ftplib import FTP |
||||
|
from requests.auth import HTTPBasicAuth |
||||
|
|
||||
|
from myutils.ReturnParams import ReturnParams |
||||
|
from concurrent.futures import ProcessPoolExecutor, TimeoutError |
||||
|
|
||||
|
# -------------------------------------------- |
||||
|
# 1) 全局 helper：放在模块顶层，才能被子进程 picklable 调用 |
||||
|
# -------------------------------------------- |
||||
|
def _execute_dynamic(instruction_str): |
||||
|
""" |
||||
|
在子进程中执行 instruction_str 所描述的 dynamic_fun， |
||||
|
并返回 (status: bool, output: str)。 |
||||
|
""" |
||||
|
# 允许的内置函数白名单 |
||||
|
allowed_builtins = { |
||||
|
'__name__': __name__, |
||||
|
'__import__': builtins.__import__, |
||||
|
'abs': abs, 'all': all, 'any': any, 'bool': bool, |
||||
|
'chr': chr, 'dict': dict, 'enumerate': enumerate, |
||||
|
'float': float, 'int': int, 'len': len, 'list': list, |
||||
|
'max': max, 'min': min, 'print': print, 'range': range, |
||||
|
'set': set, 'str': str, 'sum': sum, 'type': type, |
||||
|
'open': open, 'Exception': Exception, 'locals': locals |
||||
|
} |
||||
|
# 构造安全的 globals |
||||
|
safe_globals = { |
||||
|
'__builtins__': allowed_builtins, |
||||
|
'subprocess': subprocess, |
||||
|
'json': json, |
||||
|
're': re, |
||||
|
'paramiko': paramiko, |
||||
|
'impacket': impacket, |
||||
|
'psycopg2': psycopg2, |
||||
|
'socket': socket, |
||||
|
'mysql': mysql, |
||||
|
'mysql.connector': mysql.connector, |
||||
|
'struct': struct, |
||||
|
'sys': sys, |
||||
|
'requests': requests, |
||||
|
'ssl': ssl, |
||||
|
'FTP': FTP, |
||||
|
'HTTPBasicAuth': HTTPBasicAuth, |
||||
|
'telnetlib': telnetlib, |
||||
|
'time': time, |
||||
|
'uuid':uuid, |
||||
|
} |
||||
|
safe_locals = {} |
||||
|
try: |
||||
|
# 编译并执行用户提供的 code 字符串 |
||||
|
compiled = compile(instruction_str, '<dynamic>', 'exec') |
||||
|
exec(compiled, safe_globals, safe_locals) |
||||
|
|
||||
|
# dynamic_fun 必须存在 |
||||
|
if 'dynamic_fun' not in safe_locals: |
||||
|
return False, "Function dynamic_fun() 未定义" |
||||
|
|
||||
|
# 调用它并返回结果 |
||||
|
res = safe_locals['dynamic_fun']() |
||||
|
if not (isinstance(res, tuple) and len(res) == 2 and isinstance(res[0], bool)): |
||||
|
return False, "dynamic_fun 返回值格式不对" |
||||
|
return res |
||||
|
except MemoryError: |
||||
|
return False, "内存溢出" |
||||
|
except RecursionError: |
||||
|
return False, "递归深度过深" |
||||
|
except Exception as e: |
||||
|
return False, f"子进程执行出错: {e}" |
||||
|
|
||||
|
class PythoncodeTool(): |
||||
|
def __init__(self,max_num): |
||||
|
self.proc_pool = ProcessPoolExecutor(max_workers=max_num) |
||||
|
|
||||
|
def preprocess(self,code: str) -> str: |
||||
|
# 去掉最外层空行 |
||||
|
code = code.strip('\n') |
||||
|
# 去除多余缩进 |
||||
|
return textwrap.dedent(code) |
||||
|
|
||||
|
def is_safe_code(self,code): |
||||
|
# List of high-risk functions to block (can be adjusted based on requirements) |
||||
|
# 只屏蔽这些“完整”函数调用 |
||||
|
HIGH_RISK = { |
||||
|
'eval', # eval(...) |
||||
|
'exec', # exec(...) |
||||
|
'os.system', # os.system(...) |
||||
|
'subprocess.call', # subprocess.call(...) |
||||
|
'subprocess.Popen' # subprocess.Popen(...) |
||||
|
} |
||||
|
try: |
||||
|
tree = ast.parse(code) |
||||
|
for node in ast.walk(tree): |
||||
|
if isinstance(node, ast.Call): |
||||
|
fn = node.func |
||||
|
# 1) 裸 exec/eval |
||||
|
if isinstance(fn, ast.Name): |
||||
|
if fn.id in ('exec', 'eval'): |
||||
|
return False,"有高风险函数，暂不执行！" |
||||
|
|
||||
|
# 2) 模块级别的 os.system、subprocess.call、subprocess.Popen |
||||
|
elif isinstance(fn, ast.Attribute): |
||||
|
# value 必须是 Name，才算"模块.方法" |
||||
|
if isinstance(fn.value, ast.Name): |
||||
|
fullname = f"{fn.value.id}.{fn.attr}" |
||||
|
if fullname in HIGH_RISK: |
||||
|
return False,"有高风险函数，暂不执行！" |
||||
|
|
||||
|
return True,"" |
||||
|
|
||||
|
except SyntaxError as se: |
||||
|
# 语法都不通过，也算不安全 |
||||
|
print("解析失败！", se, "第", se.lineno, "行") |
||||
|
print("出错的那行是：", code.splitlines()[se.lineno - 1]) |
||||
|
return False,str(se) |
||||
|
|
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 60*10 |
||||
|
instr = instruction.replace("python_code ","") |
||||
|
instr = instr.replace("python-code ", "") |
||||
|
instr = self.preprocess(instr) |
||||
|
# Safety check |
||||
|
bsafe,error = self.is_safe_code((instr)) |
||||
|
if not bsafe: |
||||
|
return "", timeout,error |
||||
|
return instr,timeout,"" |
||||
|
|
||||
|
def safe_import(self,name,*args,**kwargs): |
||||
|
ALLOWED_MODULES = ['subprocess', 'json','re'] |
||||
|
if name not in ALLOWED_MODULES: |
||||
|
raise ImportError(f"Import of '{name}' is not allowed") |
||||
|
return builtins.__import__(name, *args, **kwargs) |
||||
|
|
||||
|
def _run_dynamic(self, safe_locals, q): |
||||
|
"""子进程执行 dynamic_fun 并把结果放入队列""" |
||||
|
try: |
||||
|
fn = safe_locals['dynamic_fun'] |
||||
|
res = fn() |
||||
|
q.put(res) |
||||
|
except Exception as e: |
||||
|
q.put((False, f"执行出错: {e}")) |
||||
|
|
||||
|
def execute_instruction(self, instruction_old): |
||||
|
''' |
||||
|
执行指令：验证合法性 -> 执行 -> 分析结果 |
||||
|
:param instruction_old: |
||||
|
:return: |
||||
|
bool:true-正常返回给大模型，false-结果不返回给大模型 |
||||
|
str:执行的指令 |
||||
|
str:执行指令的结果 |
||||
|
''' |
||||
|
ext_params = ReturnParams() |
||||
|
ext_params["is_user"] = False # 是否要提交用户确认 -- 默认False |
||||
|
ext_params["is_vulnerability"] = False # 是否是脆弱点 |
||||
|
|
||||
|
# 第一步：验证指令合法性 |
||||
|
instruction,time_out,error = self.validate_instruction(instruction_old) |
||||
|
if not instruction: |
||||
|
return False, instruction_old, error,"",ext_params |
||||
|
# 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ |
||||
|
|
||||
|
# 第二步：执行指令 |
||||
|
future = self.proc_pool.submit(_execute_dynamic, instruction) |
||||
|
try: |
||||
|
# 在主进程中等待结果，超时则抛 TimeoutError |
||||
|
status, tmpout = future.result(timeout=time_out) #这里是阻塞的 |
||||
|
except TimeoutError: |
||||
|
# 超时处理 |
||||
|
future.cancel() |
||||
|
status, tmpout = False, f"执行超时（{time_out} 秒）" |
||||
|
except Exception as e: |
||||
|
# 其他异常 |
||||
|
status, tmpout = False, f"提交子进程运行出错: {e}" |
||||
|
output = f"status:{status},output:{tmpout}" |
||||
|
|
||||
|
# 第三步：分析执行结果 |
||||
|
analysis = self.analyze_result(output, instruction,"","") |
||||
|
# 指令和结果入数据库 |
||||
|
# ？ |
||||
|
if not analysis: # analysis为“” 不提交LLM |
||||
|
return False, instruction, analysis,"",ext_params |
||||
|
return True, instruction, analysis,"",ext_params |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 --- 要不要限定一个max_len? |
||||
|
if "enum4linux " in instruction: #存在指令包装成Python代码返回的情况 |
||||
|
result = my_merge("enum4linux",result) |
||||
|
else: |
||||
|
if len(result) > 3000: #超过2000长度时，尝试去重重复行 |
||||
|
lines = result.splitlines() |
||||
|
seen = set() |
||||
|
unique_lines = [] |
||||
|
for line in lines: |
||||
|
if line not in seen: |
||||
|
seen.add(line) |
||||
|
unique_lines.append(line) |
||||
|
return "\n".join(unique_lines) |
||||
|
return result |
||||
|
|
||||
|
#关闭进程池 |
||||
|
def shutdown(self): |
||||
|
self.proc_pool.shutdown(wait=True) |
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
llm_code = """ |
||||
|
def run_test(): |
||||
|
return 'Penetration test executed successfully!' |
||||
|
""" |
||||
@ -0,0 +1,5 @@ |
|||||
|
|
||||
|
class RsikManager: |
||||
|
def __init__(self): |
||||
|
pass |
||||
|
|
||||
@ -0,0 +1,11 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class ArpingTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,13 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class DirSearchTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
if "-o " not in instruction or "--output=" not in instruction: |
||||
|
instruction += " -o ds_result.txt" |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,11 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class GitdumperTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 60*2 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,11 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class MedusaTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,76 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
import os |
||||
|
import shlex |
||||
|
import subprocess |
||||
|
import tempfile |
||||
|
|
||||
|
class MsfvenomTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def do_worker_script(self,str_instruction,timeout,ext_params): |
||||
|
# 创建临时文件保存输出 |
||||
|
with tempfile.NamedTemporaryFile(delete=False) as tmpfile: |
||||
|
output_file = tmpfile.name |
||||
|
|
||||
|
# 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 |
||||
|
safe_instr = shlex.quote(str_instruction.strip()) |
||||
|
# 构建 script 命令 |
||||
|
# 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 |
||||
|
script_cmd = f"script -q -c {safe_instr} {output_file}" |
||||
|
# 选项 -q 表示静默（quiet），减少不必要的输出 |
||||
|
|
||||
|
# # 构建并执行 script 命令 |
||||
|
# script_cmd = f"script -c '{str_instruction}' {output_file}" |
||||
|
try: |
||||
|
if timeout ==0: |
||||
|
result = subprocess.run(script_cmd, shell=True, text=True) |
||||
|
else: |
||||
|
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): |
||||
|
ext_params = self.create_extparams() |
||||
|
# 第一步：验证指令合法性 |
||||
|
instruction,time_out = self.validate_instruction(instruction_old) |
||||
|
if not instruction: |
||||
|
return False, instruction_old, "该指令暂不执行！","",ext_params |
||||
|
# 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ |
||||
|
|
||||
|
# 第二步：执行指令---需要对ftp指令进行区分判断 |
||||
|
output = self.do_worker_script(instruction, time_out, ext_params) |
||||
|
|
||||
|
# 第三步：分析执行结果 |
||||
|
analysis = self.analyze_result(output,instruction,"","") |
||||
|
|
||||
|
return True, instruction, analysis,output,ext_params |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,75 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
import re |
||||
|
import os |
||||
|
import shlex |
||||
|
import subprocess |
||||
|
import tempfile |
||||
|
|
||||
|
|
||||
|
class OtherTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 60*2 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def do_worker_script(self,str_instruction,timeout,ext_params): |
||||
|
# 创建临时文件保存输出 |
||||
|
with tempfile.NamedTemporaryFile(delete=False) as tmpfile: |
||||
|
output_file = tmpfile.name |
||||
|
|
||||
|
# 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 |
||||
|
safe_instr = shlex.quote(str_instruction.strip()) |
||||
|
# 构建 script 命令 |
||||
|
# 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 |
||||
|
script_cmd = f"script -q -c {safe_instr} {output_file}" |
||||
|
# 选项 -q 表示静默（quiet），减少不必要的输出 |
||||
|
|
||||
|
# # 构建并执行 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): |
||||
|
ext_params = self.create_extparams() |
||||
|
# 第一步：验证指令合法性 |
||||
|
instruction,time_out = self.validate_instruction(instruction_old) |
||||
|
if not instruction: |
||||
|
return False, instruction_old, "该指令暂不执行！","",ext_params |
||||
|
# 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ |
||||
|
|
||||
|
# 第二步：执行指令---需要对ftp指令进行区分判断 |
||||
|
output = self.do_worker_script(instruction, time_out, ext_params) |
||||
|
|
||||
|
# 第三步：分析执行结果 |
||||
|
analysis = self.analyze_result(output,instruction,"","") |
||||
|
|
||||
|
return True, instruction, analysis,output,ext_params |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,11 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class PingTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -1,153 +0,0 @@ |
|||||
#python代码动态执行 |
|
||||
import ast |
|
||||
import subprocess |
|
||||
import json |
|
||||
import builtins |
|
||||
import re |
|
||||
import paramiko |
|
||||
import impacket |
|
||||
import psycopg2 |
|
||||
import socket |
|
||||
import struct |
|
||||
import sys |
|
||||
import requests |
|
||||
import ssl |
|
||||
import mysql.connector |
|
||||
from tools.ToolBase import ToolBase |
|
||||
from mycode.Result_merge import my_merge |
|
||||
|
|
||||
class PythoncodeTool(ToolBase): |
|
||||
|
|
||||
def is_safe_code(self,code): |
|
||||
# List of high-risk functions to block (can be adjusted based on requirements) |
|
||||
HIGH_RISK_FUNCTIONS = ['eval', 'exec', 'os.system', 'subprocess.call', 'subprocess.Popen'] |
|
||||
|
|
||||
"""Check if the code contains high-risk function calls.""" |
|
||||
try: |
|
||||
tree = ast.parse(code) |
|
||||
for node in ast.walk(tree): |
|
||||
if isinstance(node, ast.Call): |
|
||||
if isinstance(node.func, ast.Name) and node.func.id in HIGH_RISK_FUNCTIONS: |
|
||||
return False |
|
||||
elif isinstance(node.func, ast.Attribute) and node.func.attr in HIGH_RISK_FUNCTIONS: |
|
||||
return False |
|
||||
return True |
|
||||
except SyntaxError: |
|
||||
return False |
|
||||
|
|
||||
def validate_instruction(self, instruction): |
|
||||
#指令过滤 |
|
||||
timeout = 0 |
|
||||
instr = instruction.replace("python_code ","") |
|
||||
instr = instruction.replace("python-code ", "") |
|
||||
# Safety check |
|
||||
if not self.is_safe_code(instr): |
|
||||
return "", timeout |
|
||||
return instr,timeout |
|
||||
|
|
||||
def safe_import(self,name,*args,**kwargs): |
|
||||
ALLOWED_MODULES = ['subprocess', 'json','re'] |
|
||||
if name not in ALLOWED_MODULES: |
|
||||
raise ImportError(f"Import of '{name}' is not allowed") |
|
||||
return builtins.__import__(name, *args, **kwargs) |
|
||||
|
|
||||
def execute_instruction(self, instruction_old): |
|
||||
''' |
|
||||
执行指令：验证合法性 -> 执行 -> 分析结果 |
|
||||
:param instruction_old: |
|
||||
:return: |
|
||||
bool:true-正常返回给大模型，false-结果不返回给大模型 |
|
||||
str:执行的指令 |
|
||||
str:执行指令的结果 |
|
||||
''' |
|
||||
ext_params = self.create_extparams() |
|
||||
|
|
||||
# 定义允许的内置函数集合 --白名单 |
|
||||
allowed_builtins = { |
|
||||
'__name__': __name__, |
|
||||
'__import__': builtins.__import__, |
|
||||
"abs": abs, |
|
||||
"all": all, |
|
||||
"any": any, |
|
||||
"bool": bool, |
|
||||
"chr": chr, |
|
||||
"dict": dict, |
|
||||
"float": float, |
|
||||
"int": int, |
|
||||
"len": len, |
|
||||
"list": list, |
|
||||
"max": max, |
|
||||
"min": min, |
|
||||
"print": print, |
|
||||
"range": range, |
|
||||
"set": set, |
|
||||
"str": str, |
|
||||
"sum": sum, |
|
||||
"type": type, |
|
||||
'open':open, |
|
||||
'Exception':Exception, |
|
||||
# 根据需要可以添加其他安全的内置函数 |
|
||||
} |
|
||||
# 第一步：验证指令合法性 |
|
||||
instruction,time_out = self.validate_instruction(instruction_old) |
|
||||
if not instruction: |
|
||||
return False, instruction_old, "该指令暂不执行！","",ext_params |
|
||||
# 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ |
|
||||
|
|
||||
# 第二步：执行指令 |
|
||||
output = "" |
|
||||
try: |
|
||||
# 构造安全的全局命名空间，只包含我们允许的 __builtins__ |
|
||||
# 虽然动态代码中包含了import subprocess,但是还是需要在全局命名空间中添加subprocess这些库 |
|
||||
# 正常情况应该是不需要的，后续再研究 |
|
||||
safe_globals = {"__builtins__": allowed_builtins, |
|
||||
'subprocess':subprocess, |
|
||||
'json':json, |
|
||||
're':re, |
|
||||
'paramiko':paramiko, |
|
||||
'impacket':impacket, |
|
||||
'psycopg2':psycopg2, |
|
||||
'socket':socket, |
|
||||
'mysql':mysql, |
|
||||
'mysql.connector':mysql.connector, |
|
||||
'struct':struct, |
|
||||
'sys':sys, |
|
||||
'requests':requests, |
|
||||
'ssl':ssl} |
|
||||
safe_locals = {} #不需要预设局部参数 |
|
||||
# 在限制环境中执行代码 |
|
||||
exec(instruction, safe_globals,safe_locals) |
|
||||
# Check if check_samba_vuln is defined |
|
||||
if 'dynamic_fun' not in safe_locals: |
|
||||
analysis = "Function dynamic_fun() is not defined" |
|
||||
ext_params['is_use'] = True |
|
||||
return True,instruction,analysis,analysis,ext_params |
|
||||
# Get the function and call it |
|
||||
dynamic_fun = safe_locals['dynamic_fun'] |
|
||||
status, tmpout = dynamic_fun() #LLM存在status定义错误的情况（执行成功，却返回的是False） #重点要处理 |
|
||||
output = f"status:{status},output:{tmpout}" |
|
||||
except Exception as e: |
|
||||
analysis = f"执行动态代码时出错: {str(e)}" |
|
||||
ext_params['is_use'] = True |
|
||||
return True,instruction,analysis,analysis,ext_params |
|
||||
|
|
||||
|
|
||||
# 第三步：分析执行结果 |
|
||||
analysis = self.analyze_result(output, instruction,"","") |
|
||||
# 指令和结果入数据库 |
|
||||
# ？ |
|
||||
if not analysis: # analysis为“” 不提交LLM |
|
||||
return False, instruction, analysis,"",ext_params |
|
||||
return True, instruction, analysis,"",ext_params |
|
||||
|
|
||||
def analyze_result(self, result,instruction,stderr,stdout): |
|
||||
#指令结果分析 --- 要不要限定一个max_len? |
|
||||
if "enum4linux " in instruction: #存在指令包装成Python代码返回的情况 |
|
||||
result = my_merge("enum4linux",result) |
|
||||
return result |
|
||||
|
|
||||
if __name__ == "__main__": |
|
||||
llm_code = """ |
|
||||
def run_test(): |
|
||||
return 'Penetration test executed successfully!' |
|
||||
""" |
|
||||
@ -0,0 +1,11 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class RpcclientTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,13 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class SmbmapTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
if " grep " not in instruction: |
||||
|
instruction =instruction.strip() + " | grep -E 'READ|WRITE|Disk|path'" |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,77 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
import os |
||||
|
import shlex |
||||
|
import subprocess |
||||
|
import tempfile |
||||
|
|
||||
|
|
||||
|
class SshpassTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def do_worker_script(self,str_instruction,timeout,ext_params): |
||||
|
# 创建临时文件保存输出 |
||||
|
with tempfile.NamedTemporaryFile(delete=False) as tmpfile: |
||||
|
output_file = tmpfile.name |
||||
|
|
||||
|
# 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 |
||||
|
safe_instr = shlex.quote(str_instruction.strip()) |
||||
|
# 构建 script 命令 |
||||
|
# 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 |
||||
|
script_cmd = f"script -q -c {safe_instr} {output_file}" |
||||
|
# 选项 -q 表示静默（quiet），减少不必要的输出 |
||||
|
|
||||
|
# # 构建并执行 script 命令 |
||||
|
# script_cmd = f"script -c '{str_instruction}' {output_file}" |
||||
|
try: |
||||
|
if timeout ==0: |
||||
|
result = subprocess.run(script_cmd, shell=True, text=True) |
||||
|
else: |
||||
|
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): |
||||
|
ext_params = self.create_extparams() |
||||
|
# 第一步：验证指令合法性 |
||||
|
instruction,time_out = self.validate_instruction(instruction_old) |
||||
|
if not instruction: |
||||
|
return False, instruction_old, "该指令暂不执行！","",ext_params |
||||
|
# 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ |
||||
|
|
||||
|
# 第二步：执行指令---需要对ftp指令进行区分判断 |
||||
|
output = self.do_worker_script(instruction, time_out, ext_params) |
||||
|
|
||||
|
# 第三步：分析执行结果 |
||||
|
analysis = self.analyze_result(output,instruction,"","") |
||||
|
|
||||
|
return True, instruction, analysis,output,ext_params |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,73 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
import os |
||||
|
import shlex |
||||
|
import subprocess |
||||
|
import tempfile |
||||
|
|
||||
|
class TcpdumpTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 60*2 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def do_worker_script(self,str_instruction,timeout,ext_params): |
||||
|
# 创建临时文件保存输出 |
||||
|
with tempfile.NamedTemporaryFile(delete=False) as tmpfile: |
||||
|
output_file = tmpfile.name |
||||
|
|
||||
|
# 使用 shlex.quote 对 str_instruction 进行安全包装，确保整个命令作为一个参数传递 |
||||
|
safe_instr = shlex.quote(str_instruction.strip()) |
||||
|
# 构建 script 命令 |
||||
|
# 注意：此时 safe_instr 包含单引号，确保整个 -c 参数不被拆分 |
||||
|
script_cmd = f"script -q -c {safe_instr} {output_file}" |
||||
|
# 选项 -q 表示静默（quiet），减少不必要的输出 |
||||
|
|
||||
|
# # 构建并执行 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): |
||||
|
ext_params = self.create_extparams() |
||||
|
# 第一步：验证指令合法性 |
||||
|
instruction,time_out = self.validate_instruction(instruction_old) |
||||
|
if not instruction: |
||||
|
return False, instruction_old, "该指令暂不执行！","",ext_params |
||||
|
# 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？ |
||||
|
|
||||
|
# 第二步：执行指令---需要对ftp指令进行区分判断 |
||||
|
output = self.do_worker_script(instruction, time_out, ext_params) |
||||
|
|
||||
|
# 第三步：分析执行结果 |
||||
|
analysis = self.analyze_result(output,instruction,"","") |
||||
|
|
||||
|
return True, instruction, analysis,output,ext_params |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -0,0 +1,11 @@ |
|||||
|
from tools.ToolBase import ToolBase |
||||
|
|
||||
|
class XvfbrunTool(ToolBase): |
||||
|
def validate_instruction(self, instruction): |
||||
|
#指令过滤 |
||||
|
timeout = 0 |
||||
|
return instruction,timeout |
||||
|
|
||||
|
def analyze_result(self, result,instruction,stderr,stdout): |
||||
|
#指令结果分析 |
||||
|
return result |
||||
@ -1,4 +1,4 @@ |
|||||
from quart import Blueprint |
from quart import Blueprint |
||||
#定义模块 |
#定义模块 |
||||
api = Blueprint('api',__name__) |
api = Blueprint('api',__name__) |
||||
from . import user,task,wsm |
from . import user,task,wsm,system |
||||
|
|||||
@ -0,0 +1,23 @@ |
|||||
|
from . import api |
||||
|
from mycode.DBManager import app_DBM |
||||
|
from myutils.ReManager import mReM |
||||
|
from quart import Quart, render_template, redirect, url_for, request,jsonify |
||||
|
|
||||
|
@api.route('/system/getinfo',methods=['GET']) |
||||
|
async def get_system_info(): |
||||
|
data = app_DBM.getsystem_info() |
||||
|
return jsonify({"local_ip":data[0],"version":data[1]}) |
||||
|
|
||||
|
@api.route('/system/updateip',methods=['POST']) |
||||
|
async def update_local_ip(): |
||||
|
data = await request.get_json() |
||||
|
local_ip = data.get("local_ip") |
||||
|
if mReM.is_valid_ip(local_ip): |
||||
|
bsuccess = app_DBM.update_localip(local_ip) |
||||
|
error = "" |
||||
|
if not bsuccess: |
||||
|
error = "修改IP地址失败！" |
||||
|
else: |
||||
|
bsuccess = False |
||||
|
error = "IP不合法" |
||||
|
return jsonify({"bsuccess":bsuccess,"error":error}) |
||||
@ -0,0 +1,659 @@ |
|||||
|
// 全局变量，用于保存当前选中的节点数据
|
||||
|
let selectedNodeData = null; |
||||
|
|
||||
|
/** |
||||
|
* 根据节点数据递归生成树形结构（返回 <li> 元素） |
||||
|
* 假设节点数据格式： |
||||
|
* { |
||||
|
* "node_name":node.name, |
||||
|
* "node_path":node.path, |
||||
|
* "node_status":node.status, |
||||
|
* "node_bwork":node.bwork, |
||||
|
* "node_vultype":node.vul_type, |
||||
|
* "node_vulgrade":node.vul_grade, |
||||
|
* children: [ { ... }, { ... } ] |
||||
|
* } |
||||
|
*/ |
||||
|
function generateTreeHTML(nodeData) { |
||||
|
const li = document.createElement("li"); |
||||
|
const nodeSpan = document.createElement("span"); |
||||
|
nodeSpan.className = "tree-node"; |
||||
|
//设置data属性
|
||||
|
nodeSpan.setAttribute("data-node_name", nodeData.node_name); |
||||
|
nodeSpan.setAttribute("data-node_path", nodeData.node_path); |
||||
|
nodeSpan.setAttribute("data-node_status", nodeData.node_status); |
||||
|
nodeSpan.setAttribute("data-node_bwork", nodeData.node_bwork); |
||||
|
nodeSpan.setAttribute("data-node_vultype", nodeData.node_vultype); |
||||
|
nodeSpan.setAttribute("data-node_vulgrade", nodeData.node_vulgrade || ""); |
||||
|
nodeSpan.setAttribute("data-node_workstatus",nodeData.node_workstatus); |
||||
|
if(nodeData.node_workstatus ===0){ |
||||
|
nodeSpan.classList.add("no-work"); |
||||
|
}else { |
||||
|
nodeSpan.classList.remove("no-work"); |
||||
|
} |
||||
|
// 根据漏洞级别添加样式
|
||||
|
if (nodeData.node_vulgrade) { |
||||
|
nodeSpan.classList.remove("no-work"); |
||||
|
if (nodeData.node_vulgrade === "低危") { |
||||
|
nodeSpan.classList.add("vul-low"); |
||||
|
} else if (nodeData.node_vulgrade === "中危") { |
||||
|
nodeSpan.classList.add("vul-medium"); |
||||
|
} else if (nodeData.node_vulgrade === "高危") { |
||||
|
nodeSpan.classList.add("vul-high"); |
||||
|
} |
||||
|
} |
||||
|
// 创建容器用于存放切换图标与文本
|
||||
|
const container = document.createElement("div"); |
||||
|
container.className = "node-container"; |
||||
|
// 如果有子节点，则添加切换图标
|
||||
|
if (nodeData.children && nodeData.children.length > 0) { |
||||
|
const toggleIcon = document.createElement("span"); |
||||
|
toggleIcon.className = "toggle-icon"; |
||||
|
toggleIcon.textContent = "-"; // 默认展开时显示“-”
|
||||
|
container.appendChild(toggleIcon); |
||||
|
} |
||||
|
//节点文本
|
||||
|
const textSpan = document.createElement("span"); |
||||
|
textSpan.className = "node-text"; |
||||
|
textSpan.textContent = nodeData.node_name; |
||||
|
container.appendChild(textSpan); |
||||
|
nodeSpan.appendChild(container); |
||||
|
li.appendChild(nodeSpan); |
||||
|
//如果存在子节点，递归生成子节点列表
|
||||
|
if (nodeData.children && nodeData.children.length > 0) { |
||||
|
const ul = document.createElement("ul"); |
||||
|
nodeData.children.forEach((child) => { |
||||
|
ul.appendChild(generateTreeHTML(child)); |
||||
|
}); |
||||
|
li.appendChild(ul); |
||||
|
} |
||||
|
return li; |
||||
|
} |
||||
|
|
||||
|
// 绑定所有节点的点击事件
|
||||
|
function bindTreeNodeEvents() { |
||||
|
document.querySelectorAll(".tree-node").forEach((el) => { |
||||
|
el.addEventListener("click", (event) => { |
||||
|
// 阻止事件冒泡，避免点击时展开折叠影响
|
||||
|
event.stopPropagation(); |
||||
|
// 清除之前选中的节点样式
|
||||
|
document |
||||
|
.querySelectorAll(".tree-node.selected") |
||||
|
.forEach((node) => node.classList.remove("selected")); |
||||
|
// 当前节点标记为选中
|
||||
|
el.classList.add("selected"); |
||||
|
// 读取 data 属性更新右侧显示
|
||||
|
const nodeName = el.getAttribute("data-node_name"); |
||||
|
const status = el.getAttribute("data-node_status"); |
||||
|
const nodepath = el.getAttribute("data-node_path"); |
||||
|
const nodebwork = el.getAttribute("data-node_bwork"); |
||||
|
const vulType = el.getAttribute("data-node_vultype"); |
||||
|
const vulLevel = el.getAttribute("data-node_vulgrade"); |
||||
|
const workstatus = el.getAttribute("data-node_workstatus"); |
||||
|
//selectedNodeData = { nodeName, status, vulType, vulLevel,nodepath,nodebwork };
|
||||
|
// 示例中默认填充
|
||||
|
selectedNodeData = { |
||||
|
node_name: nodeName, |
||||
|
node_path: nodepath, |
||||
|
status: status, |
||||
|
node_bwork: nodebwork, |
||||
|
vul_type: vulType, |
||||
|
vul_grade: vulLevel || "-", |
||||
|
workstatus: workstatus |
||||
|
}; |
||||
|
//刷新界面内容
|
||||
|
update_select_node_data_show(nodeName,status,vulType,vulLevel,workstatus,nodebwork) |
||||
|
}); |
||||
|
// 双击事件：展开/收缩子节点区域
|
||||
|
el.addEventListener("dblclick", (event) => { |
||||
|
event.stopPropagation(); |
||||
|
// 找到该节点下的 <ul> 子节点列表
|
||||
|
const parentLi = el.parentElement; |
||||
|
const childUl = parentLi.querySelector("ul"); |
||||
|
if (childUl) { |
||||
|
// 切换 collapsed 类，控制 display
|
||||
|
childUl.classList.toggle("collapsed"); |
||||
|
// 更新切换图标
|
||||
|
const toggleIcon = el.querySelector(".toggle-icon"); |
||||
|
if (toggleIcon) { |
||||
|
toggleIcon.textContent = childUl.classList.contains("collapsed") |
||||
|
? "+" |
||||
|
: "-"; |
||||
|
} |
||||
|
} |
||||
|
}); |
||||
|
}); |
||||
|
} |
||||
|
|
||||
|
// 动态加载节点树数据
|
||||
|
async function loadNodeTree(task_id) { |
||||
|
// 清空选中状态
|
||||
|
selectedNodeData = null; |
||||
|
//刷新界面内容
|
||||
|
update_select_node_data_show("-","-","-","-","-",false) |
||||
|
try { |
||||
|
const res = await fetch("/api/task/gethistree", { |
||||
|
method: "POST", |
||||
|
headers: { "Content-Type": "application/json" }, |
||||
|
body: JSON.stringify({ task_id }), //task_id:task_id
|
||||
|
}); |
||||
|
if (!res.ok) { |
||||
|
const errorData = await res.json(); |
||||
|
throw new Error(errorData.error || `HTTP错误 ${res.status}`); |
||||
|
} |
||||
|
const data = await res.json(); |
||||
|
const treeData = data.tree; |
||||
|
if (!treeData) { |
||||
|
document.getElementById("treeContent").innerHTML = |
||||
|
"<p>无节点数据</p>"; |
||||
|
return; |
||||
|
} |
||||
|
|
||||
|
// 创建一个 <ul> 作为树的根容器
|
||||
|
const ul = document.createElement("ul"); |
||||
|
ul.className = "tree-root-ul"; |
||||
|
ul.appendChild(generateTreeHTML(treeData)); |
||||
|
// 替换节点树容器的内容
|
||||
|
const container = document.getElementById("treeContent"); |
||||
|
container.innerHTML = ""; |
||||
|
container.appendChild(ul); |
||||
|
// 绑定节点点击事件
|
||||
|
bindTreeNodeEvents(); |
||||
|
} catch (error) { |
||||
|
console.error("加载节点树失败：", error); |
||||
|
document.getElementById("treeContent").innerHTML = "<p>加载节点树失败</p>"; |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
function getWorkStatus_Str(workstatus){ |
||||
|
strworkstatus = "" |
||||
|
switch (workstatus){ |
||||
|
case 0: |
||||
|
strworkstatus = "无待执行任务"; |
||||
|
break; |
||||
|
case 1: |
||||
|
strworkstatus = "待执行指令中"; |
||||
|
break; |
||||
|
case 2: |
||||
|
strworkstatus = "指令执行中"; |
||||
|
break; |
||||
|
case 3: |
||||
|
strworkstatus = "待提交llm中"; |
||||
|
break; |
||||
|
case 4: |
||||
|
strworkstatus = "提交llm中"; |
||||
|
break; |
||||
|
default: |
||||
|
strworkstatus = "-" |
||||
|
} |
||||
|
return strworkstatus |
||||
|
} |
||||
|
|
||||
|
//根据web端过来的数据，更新节点的工作状态
|
||||
|
function updateTreeNode(node_path, node_workstatus) { |
||||
|
// 根据 node_path 查找对应节点（假设每个 .tree-node 上设置了 data-node_path 属性）
|
||||
|
const nodeEl = document.querySelector(`.tree-node[data-node_path="${node_path}"]`); |
||||
|
if (nodeEl) { |
||||
|
// 更新 DOM 属性（属性值均为字符串）
|
||||
|
nodeEl.setAttribute("data-node_workstatus", node_workstatus); |
||||
|
//判断是否需要更新界面
|
||||
|
if(selectedNodeData){ |
||||
|
if(node_path === selectedNodeData.node_path){ //只有是当前选中节点才更新数据
|
||||
|
selectedNodeData.workstatus = node_workstatus; |
||||
|
strnew = getWorkStatus_Str(node_workstatus); |
||||
|
document.getElementById("node_workstatus").textContent = strnew; |
||||
|
} |
||||
|
} |
||||
|
} else { |
||||
|
console.warn(`未找到节点 ${node_path}`); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
//刷新节点的数据显示
|
||||
|
function update_select_node_data_show(nodeName,testStatus,vulType,vulLevel,workStatus,nodebwork){ |
||||
|
document.getElementById("nodeName").textContent = nodeName; |
||||
|
document.getElementById("testStatus").textContent = testStatus; |
||||
|
document.getElementById("node_vulType").textContent = vulType; |
||||
|
document.getElementById("node_vulLevel").textContent = vulLevel; |
||||
|
str_workStatus = getWorkStatus_Str(Number(workStatus)); |
||||
|
document.getElementById("node_workstatus").textContent = str_workStatus; |
||||
|
if(nodebwork==="true"){ |
||||
|
document.getElementById("node_bwork").textContent = "执行中"; |
||||
|
}else { |
||||
|
document.getElementById("node_bwork").textContent = "暂停中"; |
||||
|
} |
||||
|
setNodeBtnStatus(); |
||||
|
} |
||||
|
|
||||
|
//节点按钮的状态控制
|
||||
|
function setNodeBtnStatus(){ |
||||
|
const btn_VI = document.getElementById("btnViewInstr"); |
||||
|
if(!selectedNodeData){ |
||||
|
//没有选择node，按钮全部置不可用
|
||||
|
btn_VI.disabled = true; |
||||
|
btn_VI.classList.add("disabled-btn"); |
||||
|
} |
||||
|
else{ |
||||
|
btn_VI.disabled = false; |
||||
|
btn_VI.classList.remove("disabled-btn"); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// // 刷新按钮事件绑定
|
||||
|
// document.getElementById("btnRefresh").addEventListener("click", () => {
|
||||
|
// // 重新加载节点树数据
|
||||
|
// loadNodeTree(cur_task_id);
|
||||
|
// });
|
||||
|
|
||||
|
// 按钮事件：当未选中节点时提示
|
||||
|
function checkSelectedNode() { |
||||
|
if (!selectedNodeData) { |
||||
|
alert("请先选择节点"); |
||||
|
return false; |
||||
|
} |
||||
|
return true; |
||||
|
} |
||||
|
|
||||
|
//----------------------查看节点--指令modal----------------------------
|
||||
|
let doneInstrs = []; // 已执行指令的所有数据
|
||||
|
let todoInstrs = []; // 待执行指令的所有数据
|
||||
|
let donePage = 1; // 已执行指令当前页
|
||||
|
let todoPage = 1; // 待执行指令当前页
|
||||
|
|
||||
|
document.getElementById("btnViewInstr").addEventListener("click", () => { |
||||
|
if (!checkSelectedNode()) return; |
||||
|
openInstrModal() |
||||
|
}); |
||||
|
// 打开对话框函数
|
||||
|
function openInstrModal() { |
||||
|
const instrCanvas = new bootstrap.Offcanvas(document.getElementById('instrCanvas')); |
||||
|
instrCanvas.show(); |
||||
|
|
||||
|
// const modalEl = document.getElementById("instrModal");
|
||||
|
// // 假设用 Bootstrap 5 的 Modal 组件
|
||||
|
// const instrModal = new bootstrap.Modal(modalEl, {keyboard: false});
|
||||
|
// 显示对话框
|
||||
|
//instrModal.show();
|
||||
|
|
||||
|
// 在打开 modal 时，先更新提示内容，将 loadingMsg 显示“请稍后，数据获取中…”
|
||||
|
const loadingMsg = document.getElementById("loadingMsg"); |
||||
|
if (loadingMsg) { |
||||
|
loadingMsg.textContent = "请稍后，数据获取中..."; |
||||
|
} |
||||
|
|
||||
|
// 加载指令数据
|
||||
|
loadInstrData(); |
||||
|
} |
||||
|
|
||||
|
// 调用后端接口，获取指令数据
|
||||
|
async function loadInstrData() { |
||||
|
task_id = cur_task_id; |
||||
|
node_path = selectedNodeData.node_path; |
||||
|
try { |
||||
|
const res = await fetch("/api/task/hisnodegetinstr", { |
||||
|
method: "POST", |
||||
|
headers: { "Content-Type": "application/json" }, |
||||
|
body: JSON.stringify({task_id,node_path}), |
||||
|
}); |
||||
|
if (!res.ok) { |
||||
|
const errorData = await res.json(); |
||||
|
throw new Error(errorData.error || `HTTP错误 ${res.status}`); |
||||
|
} |
||||
|
const data = await res.json(); |
||||
|
// 数据获取成功后，清除加载提示
|
||||
|
const loadingMsg = document.getElementById("loadingMsg"); |
||||
|
if (loadingMsg) { |
||||
|
loadingMsg.style.display = "none"; // 或者清空其 innerHTML
|
||||
|
} |
||||
|
doneInstrs = data.doneInstrs || []; |
||||
|
//todoInstrs = data.todoInstrs || [];
|
||||
|
donePage = 1; |
||||
|
todoPage = 1; |
||||
|
renderDoneInstrTable(donePage); |
||||
|
//renderTodoInstrTable(todoPage);
|
||||
|
} catch (error) { |
||||
|
console.error("加载指令数据异常:", error); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 渲染已执行指令表格
|
||||
|
function renderDoneInstrTable(page) { |
||||
|
const tbody = document.getElementById("doneInstrTbody"); |
||||
|
// 计算起始索引
|
||||
|
const startIndex = (page - 1) * pageSize; |
||||
|
const endIndex = startIndex + pageSize; |
||||
|
const pageData = doneInstrs.slice(startIndex, endIndex); |
||||
|
//select instruction,start_time,result from task_result where task_id=%s and node_path=%s;
|
||||
|
tbody.innerHTML = ""; |
||||
|
// 插入行
|
||||
|
pageData.forEach((item, i) => { |
||||
|
const tr = document.createElement("tr"); |
||||
|
|
||||
|
// 第一列：序号
|
||||
|
const tdIndex = document.createElement("td"); |
||||
|
tdIndex.textContent = startIndex + i + 1; |
||||
|
tr.appendChild(tdIndex); |
||||
|
|
||||
|
// 第二列：指令内容
|
||||
|
const tdInstr = document.createElement("td"); |
||||
|
tdInstr.textContent = item[0]; |
||||
|
tr.appendChild(tdInstr); |
||||
|
|
||||
|
// 第三列：开始时间（如果没有则显示空字符串）
|
||||
|
const tdStartTime = document.createElement("td"); |
||||
|
tdStartTime.textContent = item[1] || ""; |
||||
|
tr.appendChild(tdStartTime); |
||||
|
|
||||
|
// 第四列：执行结果
|
||||
|
const tdResult = document.createElement("td"); |
||||
|
tdResult.textContent = item[2] || ""; |
||||
|
tr.appendChild(tdResult); |
||||
|
|
||||
|
tbody.appendChild(tr); |
||||
|
}); |
||||
|
|
||||
|
// 若不足 10 行，补空行
|
||||
|
for (let i = pageData.length; i < pageSize; i++) { |
||||
|
const tr = document.createElement("tr"); |
||||
|
tr.innerHTML = ` |
||||
|
<td> </td> |
||||
|
<td> </td> |
||||
|
<td> </td> |
||||
|
<td> </td> |
||||
|
`;
|
||||
|
tbody.appendChild(tr); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 渲染待执行指令表格
|
||||
|
function renderTodoInstrTable(page) { |
||||
|
const tbody = document.getElementById("todoInstrTbody"); |
||||
|
const startIndex = (page - 1) * pageSize; |
||||
|
const endIndex = startIndex + pageSize; |
||||
|
const pageData = todoInstrs.slice(startIndex, endIndex); |
||||
|
|
||||
|
tbody.innerHTML = ""; |
||||
|
pageData.forEach((item, i) => { |
||||
|
const tr = document.createElement("tr"); |
||||
|
const idx = startIndex + i + 1; |
||||
|
// 第一列：序号
|
||||
|
const tdIndex = document.createElement("td"); |
||||
|
tdIndex.textContent = idx; |
||||
|
tr.appendChild(tdIndex); |
||||
|
|
||||
|
// 第二列：指令文本内容（直接使用 textContent）
|
||||
|
const tdItem = document.createElement("td"); |
||||
|
tdItem.textContent = item; // 使用 textContent 避免 HTML 解析
|
||||
|
tr.appendChild(tdItem); |
||||
|
|
||||
|
// 第三列：复制和删除按钮
|
||||
|
const tdAction = document.createElement("td"); |
||||
|
// const btn_cp = document.createElement("button");
|
||||
|
// btn_cp.className = "btn btn-primary btn-sm";
|
||||
|
// btn_cp.textContent = "复制";
|
||||
|
// btn_cp.style.marginRight = "2px"; // 设置间隔
|
||||
|
// btn_cp.onclick = () => confirmCopyTodoInstr(idx - 1);
|
||||
|
// tdAction.appendChild(btn_cp);
|
||||
|
const btn = document.createElement("button"); |
||||
|
btn.className = "btn btn-danger btn-sm"; |
||||
|
btn.textContent = "删除"; |
||||
|
btn.onclick = () => confirmDeleteTodoInstr(idx - 1); |
||||
|
tdAction.appendChild(btn); |
||||
|
tr.appendChild(tdAction); |
||||
|
|
||||
|
tbody.appendChild(tr); |
||||
|
}); |
||||
|
|
||||
|
// 补空行
|
||||
|
for (let i = pageData.length; i < pageSize; i++) { |
||||
|
const tr = document.createElement("tr"); |
||||
|
tr.innerHTML = ` |
||||
|
<td> </td> |
||||
|
<td> </td> |
||||
|
<td> </td> |
||||
|
`;
|
||||
|
tbody.appendChild(tr); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// 分页事件
|
||||
|
document.getElementById("doneInstrPrev").addEventListener("click", (e) => { |
||||
|
e.preventDefault(); |
||||
|
if (donePage > 1) { |
||||
|
donePage--; |
||||
|
renderDoneInstrTable(donePage); |
||||
|
} |
||||
|
}); |
||||
|
document.getElementById("doneInstrNext").addEventListener("click", (e) => { |
||||
|
e.preventDefault(); |
||||
|
const maxPage = Math.ceil(doneInstrs.length / pageSize); |
||||
|
if (donePage < maxPage) { |
||||
|
donePage++; |
||||
|
renderDoneInstrTable(donePage); |
||||
|
} |
||||
|
}); |
||||
|
document.getElementById("todoInstrPrev").addEventListener("click", (e) => { |
||||
|
e.preventDefault(); |
||||
|
if (todoPage > 1) { |
||||
|
todoPage--; |
||||
|
renderTodoInstrTable(todoPage); |
||||
|
} |
||||
|
}); |
||||
|
document.getElementById("todoInstrNext").addEventListener("click", (e) => { |
||||
|
e.preventDefault(); |
||||
|
const maxPage = Math.ceil(todoInstrs.length / pageSize); |
||||
|
if (todoPage < maxPage) { |
||||
|
todoPage++; |
||||
|
renderTodoInstrTable(todoPage); |
||||
|
} |
||||
|
}); |
||||
|
|
||||
|
// 导出当前页数据
|
||||
|
document.getElementById("btnExport").addEventListener("click", () => { |
||||
|
// 判断当前是哪个 Tab
|
||||
|
const activeTab = document.querySelector("#instrTab button.nav-link.active"); |
||||
|
if (activeTab.id === "doneInstrTab") { |
||||
|
exportCurrentPage(doneInstrs, donePage, ["序号", "执行指令", "执行时间", "执行结果"]); |
||||
|
} else { |
||||
|
exportCurrentPage(todoInstrs, todoPage, ["序号", "待执行指令"]); |
||||
|
} |
||||
|
}); |
||||
|
|
||||
|
function exportCurrentPage(dataArr, page, headerArr) { |
||||
|
const startIndex = (page - 1) * pageSize; |
||||
|
const endIndex = startIndex + pageSize; |
||||
|
const pageData = dataArr.slice(startIndex, endIndex); |
||||
|
|
||||
|
// 在 CSV 的开头加上 BOM，用于 Excel 识别 UTF-8 编码
|
||||
|
let csvContent = "\uFEFF" + headerArr.join(",") + "\n"; |
||||
|
pageData.forEach((item, i) => { |
||||
|
const rowIndex = startIndex + i + 1; |
||||
|
if (headerArr.length === 4) { |
||||
|
// 已执行：序号，执行指令，执行时间，执行结果
|
||||
|
csvContent += rowIndex + "," + |
||||
|
(item.command || "") + "," + |
||||
|
(item.execTime || "") + "," + |
||||
|
(item.result || "") + "\n"; |
||||
|
} else { |
||||
|
// 待执行：序号，待执行指令
|
||||
|
csvContent += rowIndex + "," + (item.command || "") + "\n"; |
||||
|
} |
||||
|
}); |
||||
|
// 如果不足 pageSize 行，补足空行（根据列数进行适当补全）
|
||||
|
for (let i = pageData.length; i < pageSize; i++) { |
||||
|
// 根据 headerArr.length 来设置空行的格式
|
||||
|
if (headerArr.length === 4) { |
||||
|
csvContent += ",,,\n"; |
||||
|
} else { |
||||
|
csvContent += ",\n"; |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
const blob = new Blob([csvContent], { type: "text/csv;charset=utf-8;" }); |
||||
|
const url = URL.createObjectURL(blob); |
||||
|
const link = document.createElement("a"); |
||||
|
link.href = url; |
||||
|
link.download = "指令导出.csv"; |
||||
|
link.click(); |
||||
|
URL.revokeObjectURL(url); |
||||
|
} |
||||
|
|
||||
|
//------------------测试数据和漏洞数据tab-------------------
|
||||
|
// 复用：根据返回的数据数组渲染表格 tbody，保证固定 10 行
|
||||
|
function renderTableRows(tbody, rowsData) { |
||||
|
tbody.innerHTML = ""; |
||||
|
// 遍历数据行，生成 <tr>
|
||||
|
rowsData.forEach((row, index) => { |
||||
|
const tr = document.createElement("tr"); |
||||
|
// 这里假设 row 为对象，包含各个字段；下标从1开始显示序号
|
||||
|
for (const cellData of Object.values(row)) { |
||||
|
const td = document.createElement("td"); |
||||
|
td.textContent = cellData; |
||||
|
tr.appendChild(td); |
||||
|
} |
||||
|
tbody.appendChild(tr); |
||||
|
}); |
||||
|
// 补足空行
|
||||
|
const rowCount = rowsData.length; |
||||
|
for (let i = rowCount; i < pageSize; i++) { |
||||
|
const tr = document.createElement("tr"); |
||||
|
for (let j = 0; j < tbody.parentElement.querySelectorAll("th").length; j++) { |
||||
|
const td = document.createElement("td"); |
||||
|
td.innerHTML = " "; |
||||
|
tr.appendChild(td); |
||||
|
} |
||||
|
tbody.appendChild(tr); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
//--------------------------测试指令-------------------------------
|
||||
|
let allInstrs = []; |
||||
|
let currentInstrPage = 1; |
||||
|
function renderInstrPage(page) { |
||||
|
currentInstrPage = page; |
||||
|
const start = (page - 1) * pageSize; |
||||
|
const end = start + pageSize; |
||||
|
const pageData = allInstrs.slice(start, end); |
||||
|
|
||||
|
const tbody = document.querySelector("#instrTable tbody"); |
||||
|
renderTableRows(tbody, pageData); |
||||
|
|
||||
|
// 更新分页按钮
|
||||
|
document.getElementById("instrPrev").dataset.page = page > 1 ? page - 1 : 1; |
||||
|
document.getElementById("instrNext").dataset.page = (end < allInstrs.length) ? page + 1 : page; |
||||
|
} |
||||
|
|
||||
|
// 查询测试指令
|
||||
|
async function searchInstructions(page = 1) { |
||||
|
if(cur_task_id === 0){ |
||||
|
return; |
||||
|
} |
||||
|
const nodeName = document.getElementById("instrNodeName").value.trim(); |
||||
|
try { |
||||
|
const res = await fetch("/api/task/getinstr", { |
||||
|
method: "POST", |
||||
|
headers: { |
||||
|
"Content-Type": "application/json", |
||||
|
}, |
||||
|
body: JSON.stringify({ |
||||
|
cur_task_id, |
||||
|
nodeName |
||||
|
}), |
||||
|
}); |
||||
|
if (!res.ok) { |
||||
|
const errorData = await res.json(); |
||||
|
throw new Error(errorData.error || `HTTP错误 ${res.status}`); |
||||
|
} |
||||
|
const data = await res.json(); |
||||
|
allInstrs = data.instrs; |
||||
|
renderInstrPage(1); //显示第一页数据
|
||||
|
} catch (error) { |
||||
|
console.error("获取测试指令失败：", error); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
//导出测试指令数据
|
||||
|
async function ExportInstructions(){ |
||||
|
alert("导出指令功能实现中。。。") |
||||
|
} |
||||
|
|
||||
|
// 绑定测试指令查询按钮事件
|
||||
|
document.getElementById("instrSearchBtn").addEventListener("click", () => { |
||||
|
searchInstructions(); |
||||
|
}); |
||||
|
// 绑定测试指令分页点击事件
|
||||
|
document.getElementById("instrPrev").addEventListener("click", (e) => { |
||||
|
const page = parseInt(e.target.dataset.page, 10); |
||||
|
renderInstrPage(page); |
||||
|
}); |
||||
|
document.getElementById("instrNext").addEventListener("click", (e) => { |
||||
|
const page = parseInt(e.target.dataset.page, 10); |
||||
|
renderInstrPage(page);; |
||||
|
}); |
||||
|
|
||||
|
//------------------漏洞数据---------------------------------
|
||||
|
let allVuls = []; |
||||
|
let currentVulPage = 1; |
||||
|
function renderVulPage(page) { |
||||
|
currentVulPage = page; |
||||
|
const start = (page - 1) * pageSize; |
||||
|
const end = start + pageSize; |
||||
|
const pageData = allVuls.slice(start, end); |
||||
|
|
||||
|
const tbody = document.querySelector("#vulTable tbody"); |
||||
|
renderTableRows(tbody, pageData); |
||||
|
|
||||
|
// 更新分页按钮
|
||||
|
document.getElementById("vulPrev").dataset.page = page > 1 ? page - 1 : 1; |
||||
|
document.getElementById("vulNext").dataset.page = (end < allVuls.length) ? page + 1 : page; |
||||
|
} |
||||
|
|
||||
|
// 查询漏洞数据
|
||||
|
async function searchVulnerabilities(page = 1) { |
||||
|
if(cur_task_id === 0){return;} |
||||
|
const nodeName = document.getElementById("vulNodeName").value.trim(); |
||||
|
const vulType = document.getElementById("vulType").value.trim(); |
||||
|
const vulLevel = document.getElementById("vulLevel").value; |
||||
|
try { |
||||
|
const res = await fetch("/api/task/getvul", { |
||||
|
method: "POST", |
||||
|
headers: { |
||||
|
"Content-Type": "application/json", |
||||
|
}, |
||||
|
body: JSON.stringify({ |
||||
|
cur_task_id, |
||||
|
nodeName, |
||||
|
vulType, |
||||
|
vulLevel |
||||
|
}), |
||||
|
}); |
||||
|
if (!res.ok) { |
||||
|
const errorData = await res.json(); |
||||
|
throw new Error(errorData.error || `HTTP错误 ${res.status}`); |
||||
|
} |
||||
|
const data = await res.json(); |
||||
|
allVuls = data.vuls; |
||||
|
renderVulPage(1) |
||||
|
} catch (error) { |
||||
|
console.error("获取漏洞数据失败：", error); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
//导出漏洞数据
|
||||
|
async function ExportVuls(){ |
||||
|
alert("导出漏洞功能实现中。。。") |
||||
|
} |
||||
|
|
||||
|
// 绑定漏洞数据查询按钮事件
|
||||
|
document.getElementById("vulSearchBtn").addEventListener("click", () => { |
||||
|
searchVulnerabilities(); |
||||
|
}); |
||||
|
// 绑定漏洞数据分页点击事件
|
||||
|
document.getElementById("vulPrev").addEventListener("click", (e) => { |
||||
|
const page = parseInt(e.target.dataset.page, 10); |
||||
|
renderVulPage(page); |
||||
|
}); |
||||
|
document.getElementById("vulNext").addEventListener("click", (e) => { |
||||
|
const page = parseInt(e.target.dataset.page, 10); |
||||
|
renderVulPage(page); |
||||
|
}); |
||||
Loading…
Reference in new issue