From 37e920d5edc7e782a10d397b9e5dd15429278317 Mon Sep 17 00:00:00 2001
From: handingkang <18791985373@163.com>
Date: Thu, 23 May 2024 22:46:11 +0800
Subject: 1. 在settings.py中新增RUN_DEV选项，模拟代理运行正常的情况 2.
 修复/task/create接口存在问题 3. 更新/target/delay接口
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/apps/agentcomm.py | 22 +++++++++++++++++-----
 server/apps/target.py    | 43 ++++++++++++++++++++++++++++++++-----------
 server/apps/task.py      | 12 ++++++++----
 server/settings.py       |  2 ++
 4 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/server/apps/agentcomm.py b/server/apps/agentcomm.py
index b5c55ad..e5adbca 100644
--- a/server/apps/agentcomm.py
+++ b/server/apps/agentcomm.py
@@ -118,14 +118,26 @@ def deliver_task(agent_id, policy, policy_param):
         """ % (MYSQL_TAB_AGENT, agent_id)
         da.cursor.execute(sql)
         data = da.cursor.fetchall()
-        ip, port = data[0]['IPADDR'], data[0]['PORT']
+
+        # 解析代理的ipv4地址和通信端口
+        # 如果代理有多个地址，则会以"|"分隔，详见register_agent函数
+        # 第一个永远是ipv4地址（在代理向ipv4地址的端口注册信息的前提下），详见register_agent函数
+        ip, port = str(data[0]['IPADDR']).split("|")[0], data[0]['PORT']
+
         # 将policy和policy_param作为参数传递到对应接口
-        res = requests.post("http://%s:%s/script" % (ip, port), json={'policy': policy, 'param': policy_param})
-        if res.status_code == 200:
-            return None
-        else:
+        try:
+            res = requests.post("http://%s:%s/script" % (ip, port), timeout=3,
+                                json={'policy': policy, 'param': policy_param})
+            if res.status_code == 200:
+                return None
             error("任务分发错误，code: " + str(res.status_code) + " err: " + str(res.text))
             return res.text
+        except Exception as e:
+            if RUN_DEV:
+                return None
+            else:
+                error(str(e))
+                return str(e)
 
 
 @bp.post("/res")
diff --git a/server/apps/target.py b/server/apps/target.py
index 154e9d7..600e08a 100644
--- a/server/apps/target.py
+++ b/server/apps/target.py
@@ -6,6 +6,7 @@ import threading
 import requests
 from apiflask import APIBlueprint, Schema
 from apiflask.fields import Integer, String, List, Nested, IP, DateTime, Dict
+from apiflask.validators import OneOf
 from requests.exceptions import Timeout
 
 from settings import *
@@ -99,24 +100,44 @@ def get_nodes(query_data):
     return {"code": 200, "nodes": node_list}
 
 
-@bp.get("/delay/<string:type>")
+@bp.get("/delay")
 @bp.doc("获取每个节点的时延数据", "type参数为{icmp,dns,tcp}中的一个")
-@bp.input({"ip": IP(required=False)}, location="query")
+@bp.input({
+    "ip": IP(required=True),
+    "taskid": String(required=True),
+    "type": String(required=True, validate=OneOf(['icmp', 'dns', 'tcp']))}, location="query")
 @bp.output(DelayOut)
-def get_pernode_delay(query_data, scan_type):
-    # TODO:节点选择
-    addr = ""
-    if 'ip' in query_data.keys():
-        addr = query_data['ip']
+# TODO:和实际节点联调测试
+def get_pernode_delay(query_data):
+    # TODO:DoH处理
+    # 探测地址
+    addr = query_data['ip']
+    taskid = query_data['taskid']
+    scan_type = query_data['type']
+
+    # 响应数据
     ans = []
-    # threads = []
-    # df = pd.read_csv("./server.csv", encoding="utf-8")
-    # for index, row in df.iterrows():
-    #     mythread = threading.Thread(target=task, args=[ans, addr, row, type])
+    # 线程池
+    threads = []
+    # # 检索探测节点信息
+    # sql = """
+    # SELECT SCAN_AGENT_ID_LIST as node_info
+    # FROM %s
+    # WHERE TASK_ID='%s'
+    # """ % (MYSQL_TAB_TASK, taskid)
+    # da.cursor.execute(sql)
+    # # 探测节点ID与地址
+    # nodes = json.loads(da.cursor.fetchall()[0]["node_info"].replace('"', "\""))
+    #
+    # for id, ip in nodes.items():
+    #     mythread = threading.Thread(target=task, args=[ans, addr, {'id': id, 'ip': ip}, type])
     #     mythread.start()
     #     threads.append(mythread)
     # for t in threads:
     #     t.join()
+
+    # 暂未部署实际代理节点，以假数据返回
+    ans = []
     for i in range(10):
         ans.append({"Id": str(i), "Type": scan_type, "CurrDelay": random.randint(1, 1000)})
     return {"code": 200, 'delay_data': ans}
diff --git a/server/apps/task.py b/server/apps/task.py
index 4d69bf2..74be644 100644
--- a/server/apps/task.py
+++ b/server/apps/task.py
@@ -94,13 +94,16 @@ def valid_task_info(task_param: dict):
     "policy": String(validate=OneOf(["auto", "ddos", "sjqp"])),
     # 状态感知方式
     "scan": String(validate=OneOf(["auto", "icmp", "dns", "tcp", "record"])),
-    # 策略切换时限
+    # 策略切换时限,单位分钟
     "policy_time": Integer(),
-    # 任务执行时限
+    # 任务执行时限，单位分钟
     "run_time": Integer(),
     # 运行配置
     "run_flag": String(validate=OneOf(["now", "man"]))
-})
+}, example={'name': "test_task", 'target': "1.2.3.4", 'agent': "8a9ces", 'target_domain': "www.google.com",
+            'target_rr': "NS ns.ourattack.com",
+            'policy': "auto", 'scan': "auto",
+            'policy_time': 60, "run_time": 600, 'run_flag': "now"})
 @bp.output({
     "code": Integer(),
     "msg": String()
@@ -144,7 +147,8 @@ def make_task(json_data):
     if warn is not None:
         return {"code": 400, "msg": warn}
     if task_info["target_rr"] != "":
-        task_info["target_rtype"], task_info["target_rr"] = tuple(task_info["target_rr"].split(" "))
+        rr = task_info["target_rr"].split(" ")
+        task_info["target_rtype"], task_info["target_rr"] = rr[0], rr[1]
 
     # 添加表名
     task_info["tab"] = MYSQL_TAB_TASK
diff --git a/server/settings.py b/server/settings.py
index a67923f..039a4c2 100644
--- a/server/settings.py
+++ b/server/settings.py
@@ -1,6 +1,8 @@
 import datetime
 
 DEBUG = True
+# 是否以运行效果模拟状态运行，用于控制和代理通信的处理
+RUN_DEV = True
 SERVER_PORT = 12526
 START_DAY = datetime.date.today()
 
-- 
cgit v1.2.3