完善目标信息获取接口

4 files changed, 76 insertions, 33 deletions

diff --git a/server/apps/model.py b/server/apps/model.py
index 45ea54d..45cf61a 100644
--- a/server/apps/model.py
+++ b/server/apps/model.py
@@ -19,12 +19,21 @@ timemapping = {
     "target": "UPDATED_TIME",
 }
 
-#
-agent_keymapping = bidict({
-    "atype": "AGENT_TYPE",
-    "status": "STATUS",
-    "idle": "IDLE"
-})
+# 不同接口的参数（或值）和对应数据表字段的映射关系
+keymapping = {
+    "agent": bidict({
+        "atype": "AGENT_TYPE",
+        "status": "STATUS",
+        "idle": "IDLE"
+    }),
+    "target": bidict({
+        "cou": "COU",
+        "isp": "ISP",
+        "IPv6": "IPv6",
+        "DNSSEC": "DNSSEC",
+        "DoH": "DOH",
+        "DoT": "DOT",
+    })}
 
 typemapping = {
     "atype": "str",
@@ -34,5 +43,6 @@ typemapping = {
 
 # 默认参数
 default_data = {
-    "agent": {"atype": "all", "status": 2, "idle": 2}
+    "agent": {"atype": "all", "status": 2, "idle": 2},
+    "target": {"proto": None, "isp": None, "cou": None}
 }
diff --git a/server/apps/target.py b/server/apps/target.py
index 0a47599..05c1263 100644
--- a/server/apps/target.py
+++ b/server/apps/target.py
@@ -217,9 +217,9 @@ def record(query_data):
 @bp.get("/")
 @bp.doc("(表格)目标信息获取接口", "返回目标信息")
 @bp.input({
-    "protocol": String(validate=OneOf(['IPv6', 'DNSSEC', "DoH", "DoT"])),
-    "cou": String(),
-    "isp": String(),
+    "protocol": String(load_default=None, validate=OneOf(['IPv6', 'DNSSEC', "DoH", "DoT"])),
+    "cou": String(load_default=None),
+    "isp": String(load_default=None),
     "per_page": Integer(load_default=10),
     "page": Integer(load_default=1),
     "ip": IP(load_default=None)
@@ -232,16 +232,24 @@ def record(query_data):
 def target_info(query_data):
     per_page = query_data["per_page"]
     page = query_data["page"]
+    proto = query_data["protocol"]
+    cou = query_data["cou"]
+    isp = query_data["isp"]
+    ip = query_data["ip"]
+
+    # 目标信息列表
     target_list = []
 
-    if query_data["ip"] is None:
+    if ip is None:
         # 普通检索
-        res = da.get_data(data_type="target", offset=(page - 1) * per_page, limit=per_page)
+        res = da.get_data(data_type="target", search={"proto": proto, "cou": cou, "isp": isp},
+                          offset=(page - 1) * per_page, limit=per_page)
+        res_count = da.count_data(data_type="target", search={"proto": proto, "cou": cou, "isp": isp})
     else:
         # 查询目标
         res = da.get_data(data_type="target", offset=(page - 1) * per_page, limit=per_page,
                           search={"ip": query_data["ip"]})
-
+        res_count = 1
     # 结果转换
     for r in res:
         target = {}
@@ -260,8 +268,7 @@ def target_info(query_data):
             # 原数据的值赋给经过映射之后的键
             target[target_map[k]] = r[k]
         target_list.append(target)
-    # TODO:total
-    return {"code": 200, "data": target_list}
+    return {"code": 200, "data": target_list, "total": res_count}
 
 
 @bp.get("/filter")
diff --git a/server/apps/task.py b/server/apps/task.py
index c3d36ce..751a4a1 100644
--- a/server/apps/task.py
+++ b/server/apps/task.py
@@ -6,10 +6,23 @@ from apiflask import APIBlueprint, Schema
 from apiflask.fields import String, Integer, IP, DateTime, List, Nested
 from apiflask.validators import OneOf
 
-from .util import fake
+from .util import fake, da
 
 bp = APIBlueprint("任务管理接口集合", __name__, url_prefix="/task")
 
+# 数据库列与返回值的键对应关系
+task_response_map = {
+    "TASK_ID": "id",
+    "TARGET_IP": "target",
+    "TASK_NAME": "name",
+    "AGENT_ID": "agent",
+    "TARGET_DOMAIN": "target_domain",
+    "TARGET_RR": "target_rr",
+    "POLICY": "policy",
+    "CREATE_TIME": "create_time",
+    "STATUS": "status",
+}
+
 
 class Task(Schema):
     id = Integer()
@@ -122,20 +135,20 @@ def ops_task(json_data):
 })
 # TODO：查询任务状态接口具体实现
 def tasks_state(query_data):
+    per_page = query_data["per_page"]
+    page = query_data["page"]
+
+    # 任务列表
     task_list = []
-    for _ in range(query_data["per_page"]):
-        task_list.append({
-            "id": fake.random.randint(1, 10000),
-            "target": fake.ipv4(),
-            "name": "示例任务",
-            "agent": fake.random.randint(1, 10000),
-            "target_domain": fake.domain_name(),
-            "target_rr": "NS attack.com",
-            "policy": random.choice(["auto", "ddos", "sjqp"]),
-            "create_time": fake.date_time_between(start_date="-1y"),
-            "status": random.choice(["working", "stop", "finish"])
-        })
-    return {"code": 200, "data": task_list}
+    res = da.get_data(data_type="task",
+                      offset=(page - 1) * per_page, limit=per_page)
+    res_count = da.count_data(data_type="task")
+    for r in res:
+        task = {}
+        for key, value in r.items():
+            task[task_response_map[key]] = value
+        task_list.append(task)
+    return {"code": 200, "data": task_list, "total": res_count}
 
 
 # 任务详情接口
diff --git a/server/apps/util.py b/server/apps/util.py
index 96711ec..a608782 100644
--- a/server/apps/util.py
+++ b/server/apps/util.py
@@ -192,11 +192,24 @@ class DataHandler:
             condition = {}
             for _ in range(l):
                 key, val = differ.pop()
-                tab_key = model.agent_keymapping[key]
-                if model.typemapping[key] == "str":
-                    condition[tab_key] = "\"" + val + "\""
+                # target表的协议参数和 ip 单独处理
+                if data_type == "target" and key == "proto":
+                    tab_key = model.keymapping[data_type][val]
+                    condition[tab_key] = True
+                if data_type == "target" and key == "ip":
+                    if "." in val:
+                        condition["ADDRv4"] = "\"" + val + "\""
+                    elif ":" in val:
+                        condition["ADDRv6"] = "\"" + val + "\""
+                    else:
+                        error("错误的地址输入：  " + str(val))
                 else:
-                    condition[tab_key] = str(val)
+                    # 参数在数据表中对应的字段名
+                    tab_key = model.keymapping[data_type][key]
+                    if model.typemapping[key] == "str":
+                        condition[tab_key] = "\"" + val + "\""
+                    else:
+                        condition[tab_key] = str(val)
             if not count:
                 sql = """SELECT * FROM %s WHERE %s LIMIT %s, %s""" % (
                     tabname, " AND ".join(["=".join(condition.popitem()) for _ in range(l)]), offset, offset + limit)