zfsafe/test.py

import re
import subprocess
import tempfile
import os
import pexpect
import struct
import sys
import mysql.connector
import requests
from mycode.LLMManager import LLMManager
from mycode.TaskObject import TaskObject
from myutils.PickleManager import g_PKM
from mycode.InstructionManager import g_instrM
from mycode.TaskManager import g_TaskM
from mycode.PythonTManager import PythonTManager
from myutils.ConfigManager import myCongif
from mycode.DBManager import app_DBM
import textwrap


class Mytest:
    def update_node_inter(self,attack_index):
        attack_tree = g_PKM.ReadData(attack_index)
        nodes = attack_tree.traverse_dfs()
        # 06-0=>p
        instr = nodes[6].get_instr_user().pop(0)
        nodes[6].parent.get_instr_user().append(instr)
        # 39-1
        instr = nodes[39].get_instr_user().pop(1)
        nodes[39].parent.get_instr_user().append(instr)
        # 49-0
        instr = nodes[49].get_instr_user().pop(0)
        nodes[49].parent.get_instr_user().append(instr)

        g_PKM.WriteData(attack_tree, attack_index)


    def dynamic_fun(self):
        import socket
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(20)  # 增加超时时间
            s.connect(("192.168.3.105", 52989))

            # 基于返回的 "99 -1 45973" 字符串构造特殊payload
            special_cmd = b'99\\x01\\x00\\x00\\x00'  # 模拟协议头
            s.sendall(special_cmd)

            response = s.recv(2048)
            s.close()

            return (True, f"SpecialCmd Response: {response.hex()}")
        except Exception as e:
            return (False, str(e))

    def do_test(self):
        pass

    def tmp_test(self):
        list_a = [0,1,2,3,4,5,6,7,8,9]

        isart = len(list_a) - 4  # 正常应该都是两个两个
        if isart % 2 != 0:
            print("c_msg数量不对称，需要检查逻辑！")
        for msg in list_a[isart:]:
            print(msg)


if __name__ == "__main__":
    # 示例使用
    mytest = Mytest()
    LLM = LLMManager(1)
    current_path = os.path.dirname(os.path.realpath(__file__))
    print(current_path)
    test_type = 1
    task_id = 67
    task_Object = TaskObject("test_target","cookie_info",1,1,1,"local_ip","",None)

    if test_type == 0:
        mytest.dynamic_fun()
    elif test_type == 1:
# # 获取所有自定义函数详情  HIGH_RISK_FUNCTIONS = ['eval', 'exec', 'os.system', 'subprocess.call', 'subprocess.Popen']
        instruction = '''python-code 
import requests
def dynamic_fun():
    try:
        r = requests.get('https://58.216.217.67/server-status', verify=False, timeout=5)
        return (1, f'HTTP:{r.status_code} Headers:{r.headers}') if r.status_code==200 else (0, '')
    except Exception as e:
        return (0, str(e))
'''
        task_Object.PythonM.start_pool()    #开个子进程池就行
        start_time, end_time, bsuccess, instr, reslut, source_result, ext_params = task_Object.do_instruction(instruction)
        # 暂存结果
        oneres = {'执行指令': instr, '结果': reslut}
        print("----执行结果----")
        print(reslut)
    elif test_type == 2:  #给节点添加指令
        node_path = "目标系统->192.168.3.108->80端口->PHP版本漏洞检测"
        instr_id = 3478
        g_TaskM.load_tasks()
        task = g_TaskM.tasks[task_id]
        nodes = task.attack_tree.traverse_dfs()
        cur_node = None
        for node in nodes:
            if node.path == node_path:
                cur_node = node
                break
        if cur_node:
            str_instr = app_DBM.get_one_instr(instr_id)
            if "import" in str_instr:
                str_instr = "python-code " + str_instr
            cur_node.test_add_instr(str_instr)
            cur_node.update_work_status(-4)
            #保存数据
            g_PKM.WriteData(task.attack_tree,str(task.task_id))
        else:
            print("没找到节点！")
    elif test_type ==3: #测试指令入节点
        strinstr = '''

'''
        strNodes = "执行系统命令探测，权限提升尝试，横向移动测试"
        nodes = strNodes.split(', ')
        unique_names = list(set(nodes))  # 去重
        for node_name in unique_names:
            print(node_name)
    elif test_type == 4:  # 修改Messages
        attact_tree = g_PKM.ReadData("88")
        # 创建一个新的节点
        from mycode.AttackMap import TreeNode
        testnode = TreeNode("test", 0,0)
        LLM.build_init_attact_prompt(testnode)  # 新的Message
        systems = testnode.parent_messages[0]["content"]
        # print(systems)
        # 遍历node，查看有instr的ndoe
        nodes = attact_tree.traverse_bfs()
        for node in nodes:
            node.parent_messages[0]["content"] = systems
        g_PKM.WriteData(attact_tree, "88")
        print("完成Messgae更新")
    elif test_type ==5:
        mytest.dynamic_fun()
    elif test_type == 6:
        import json
        strIPS = '''
[ {"action":"asset","URL":{"Domain":"www.czzfkjxx.cn","Subdomains":[],"Registrant":"","Email":"","Registrar":"","Creation_date":"","Expiration_date":""},"IPS":[{"IP":"58.216.217.67","IPtype":"IPv4","Ports":[ {"Port":"25","Service":"smtp?","Version":"","Protocol":"tcp","Status":"open"}, {"Port":"80","Service":"http","Version":"Apache httpd","Protocol":"tcp","Status":"open"}, {"Port":"110","Service":"pop3?","Version":"","Protocol":"tcp","Status":"open"}, {"Port":"443","Service":"ssl/http","Version":"Apache httpd","Protocol":"tcp","Status":"open"} ]}]} ]
        '''
        node_json = json.loads(strIPS)
        IPS = node_json[0]["IPS"]
        URL = node_json[0]["URL"]
        #task_Object.add_update_assets(URL,IPS,app_DBM)
        task_Object.update_attack_tree(URL, IPS, None)
    elif test_type == 7:
        task_Object.test(50)
    else:
        pass