主控端数据库贯通实现尝试

6 files changed, 293 insertions, 39 deletions

diff --git a/agent/app.py b/agent/app.py
index 86e5c87..ad85985 100644
--- a/agent/app.py
+++ b/agent/app.py
@@ -64,7 +64,7 @@ def nodeinfo():
 def registernode(proto="http", port=2525, atype="stgj", server="127.0.0.1:8888"):
     info = nodeinfo()
     info["port"] = port
-    info["type"] = atype
+    info["atype"] = atype
     print(info)
     r = requests.post(proto + "://" + server + "/agent/register", json=info)
     if r.status_code != 200:
@@ -86,7 +86,7 @@ if __name__ == '__main__':
         # 命令行参数设置
         parser = argparse.ArgumentParser()
         parser.add_argument("-p", "--port", type=int, default=2525, help="代理的开放通信端口")
-        parser.add_argument("-t", "--type", type=str, default="stgj",
+        parser.add_argument("-t", "--atype", type=str, default="stgj",
                             help="代理的工作类型 {stgj(渗透攻击) / mbgz(目标感知) / ztgz(状态感知)}")
         parser.add_argument("-s", "--server", type=str, default="127.0.0.1:8888", help="主控端访问地址+端口号")
 
diff --git a/server/app.py b/server/app.py
index a53a829..4f02203 100644
--- a/server/app.py
+++ b/server/app.py
@@ -1,7 +1,6 @@
-import argparse
-
 from apiflask import APIFlask
 
+import settings
 from apps.agentcomm import bp as agentbp
 from apps.sysinfo import bp as sysbp
 from apps.sysmange import bp as mangbp
@@ -29,7 +28,4 @@ def hello():
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-p", "--port", type=int, default=2525, help="主控服务的开放通信端口")
-    args = parser.parse_args()
-    app.run(host="0.0.0.0", debug=True, port=args.port)
+    app.run(host="0.0.0.0", debug=settings.DEBUG, port=settings.SERVER_PORT)
diff --git a/server/apps/agentcomm.py b/server/apps/agentcomm.py
index 3364915..d053c2b 100644
--- a/server/apps/agentcomm.py
+++ b/server/apps/agentcomm.py
@@ -1,11 +1,13 @@
 # 代理通信与注册接口
+import datetime
 import random
 
 from apiflask import APIBlueprint, Schema
 from apiflask.fields import String, Integer, List, Nested, Boolean, DateTime
 from apiflask.validators import OneOf
-from apps.util import fake
+from flask import request
 
+from util import da
 
 bp = APIBlueprint("代理管理接口集合", __name__, url_prefix="/agent")
 
@@ -23,19 +25,65 @@ class agent(Schema):
     start_time = DateTime()
 
 
+# 数据库列与返回值的键对应关系
+agent_map = {
+    "AGENT_ID": "id",
+    "IPADDR": "ipaddr",
+    "START_TIME": "start_time",
+    "LAT": "lat",
+    "LNG": "lng",
+    "AGENT_TYPE": "atype",
+    "SYS": "sys",
+    "PORT": "port",
+    "CPU_NUM": "cpu_num",
+    "STATUS": "status",
+    "IDLE": "idle",
+    "MEM": "mem"
+}
+
+
 # 代理注册接口
[email protected]("/register", )
[email protected]("/register")
 @bp.doc("代理注册接口", "返回分配给代理的编号值", hide=True)
[email protected]({
-    # 代理类型
-    "type": String(required=True, validate=OneOf(["gjst", "mbgz", "ztgz"])),
-    # 代理的通信端口
-    "port": Integer(required=True),
-    # 代理所在的操作系统
-    "sys": String()
-})
-def register_agent(query_data):
-    print(query_data)
+# 参数说明
+# 代理类型
+# "type": String(required=True, validate=OneOf(["stgj", "mbgz", "ztgz"])),
+# 代理的通信端口
+# "port": Integer(required=True),
+# # 代理所在的操作系统
+# "sys": String(),
+# # IPv6地址
+# # "v6addr": List(String()),
+# # 经度
+# "lat": String(),
+# # 纬度
+# "lng": String(),
+# # cpu核心数
+# "cpu_num": Integer(),
+# # 内存大小
+# "ram_size": Integer(),
+# # 已用内存比例
+# "ram_per": Float()
+def register_agent():
+    data = dict(request.json)
+    id = "".join(random.sample('1234567890zyxwvutsrqponmlkjihgfedcba', 8))
+    # 代理所有参数
+    param = {"id": id,
+             "ipaddr": (request.remote_addr).join(("|" + i) for i in data["v6addr"]),
+             "start_time": str(datetime.datetime.now()),
+             "lat": data["lat"],
+             "lng": data["lng"],
+             "cpunum": data["cpu_num"],
+             "mem": data["ram_size"] + "GB" + "(" + {1 - data["ram_per"]} + "%" + "可用)",
+             "sys": data["sys"],
+             "status": True,
+             "idle": True,
+             "port": data["port"],
+             "atype": data["atype"],
+             }
+    err = insert_agent(param)
+    if not err:
+        return {"id": id}
 
 
 # 代理任务下发
@@ -43,7 +91,6 @@ def deliver_task():
     pass
 
 
-
 @bp.post("/res")
 @bp.doc("代理输出消息接收接口", hide=True)
 # TODO:代理输出接收接口实现
@@ -74,19 +121,40 @@ def agent_info(query_data):
     per_page = query_data["per_page"]
     page = query_data["page"]
     agent_list = []
-    # TODO:接口数据库贯通实现
-    for i in range(per_page):
-        agent_list.append({
-            "id": fake.unique.random_int(),
-            "ipaddr": [fake.ipv4_public(), fake.ipv6()],
-            "atype": fake.word(ext_word_list=["参数感知", "状态感知", "攻击渗透"]),
-            "status": random.choice([True, False]),
-            "idle": random.choice([True, False]),
-            "port": random.randint(1000, 65534),
-            "sys": random.choice(["Linux", "Windows", "Macos"]),
-            "cpu_num": 4,
-            "mem": random.choice(["4", "8", "16", "32"]) + "GB(" + str(
-                '{:.2f}'.format(100 * random.random())) + "%可用)",
-            "start_time": fake.date_time_between(start_date="-1y")})
-
+    res = da.get_data(data_type="agent", offset=(page - 1) * per_page, limit=per_page)
+    for r in res:
+        agent_list.append({agent_map[key]: value for key, value in r.items()})
     return {"agent_data": agent_list}
+
+
+# 代理信息存储到数据库
+def insert_agent(param={}):
+    sql = """REPLACE INTO %s(AGENT_ID,ADDRv4,
+            ADDRv6,
+            CREATED_TIME,
+            LAT,
+            LNG ,
+            AGENT_TYPE,
+            SYS ,
+            PORT ,
+            CPU_NUM ,
+            STATUS ,
+            RAM_SIZE ,
+            WORK_STATUS) VALUES(
+                %(id)s,
+                %(v4addr)s,
+                %(v6addr)s,
+                %(create_time)s,
+                %(lat)s,
+                %(lng)s,
+                %(atype)s,
+                %(sys)s,
+                %(port)s,
+                %(cpunum)s,
+                %(status)s,
+                %(mem)s,
+                %(work_status)s)
+            )""" % param
+    da.cursor.execute(sql)
+    da.conn.commit()
+    return None
diff --git a/server/apps/util.py b/server/apps/util.py
index 9dff80f..ad95002 100644
--- a/server/apps/util.py
+++ b/server/apps/util.py
@@ -1,8 +1,197 @@
 # 工具包
+import pymysql
 from faker import Faker
-from faker.providers import date_time
 from faker.providers import company
+from faker.providers import date_time
+from pymysql.cursors import DictCursor
+
+from settings import *
 
 fake = Faker("zh_CN")
 fake.add_provider(date_time)
 fake.add_provider(company)
+
+
+# 数据库访问类
+class DataHandler:
+    # 参数与数据表映射关系
+    tabmapping = {
+        "agent": MYSQL_TAB_AGENT,
+        "task": MYSQL_TAB_TASK,
+        "user": MYSQL_TAB_USER}
+
+    # 数据库链接及数据库初始化
+    def __init__(self):
+        # mysql连接,采用字典游标，返回一系列字典值
+        self.conn = pymysql.connect(cursorclass=DictCursor, host=MYSQL_HOST, user='root',
+                                    password=MYSQL_PAWD, port=MYSQL_PORT)
+        self.cursor = self.conn.cursor()
+
+        # 初始化sql语句
+        # 创建数据库
+        dbsql = "CREATE DATABASE IF NOT EXISTS %s" % MYSQL_DATADB
+
+        agentsql = """CREATE TABLE IF NOT EXISTS %s (
+            `AGENT_ID` varchar(255) NOT NULL PRIMARY KEY,
+            `IPADDR` varchar(255),,
+            `START_TIME` datetime,
+            `LAT` varchar(255),
+            `LNG` varchar(255),
+            `AGENT_TYPE` varchar(255),
+            `SYS` varchar(255),
+            `PORT` int,
+            `CPU_NUM` int,
+            `STATUS` bool,
+            `MEM` varchar(255),
+            `IDLE` bool)ENGINE=innodb DEFAULT CHARSET=utf8; """ % MYSQL_TAB_AGENT
+
+        policysql = """CREATE TABLE IF NOT EXISTS %s(
+            `P_ID` int NOT NULL AUTO_INCREMENT PRIMARY KEY,
+            `P_EXE` varchar(255),
+            `P_TYPE` varchar(255),
+            `P_DESC` varchar(255),
+            `P_PAYLOAD` varchar(255),
+            `P_NAME` varchar(255),
+        `P_PROTO` varchar(255))ENGINE=innodb DEFAULT CHARSET=utf8; """ % MYSQL_TAB_POLICY
+
+        syslogsql = """CREATE TABLE IF NOT EXISTS %s(
+            `LOG_ID` int NOT NULL AUTO_INCREMENT PRIMARY KEY,
+            `LOG_LEVEL` varchar(255),  
+            `LOG_INFO` varchar(255),
+            `LOG_TIME` datetime,
+            `S_IP` varchar(255),
+            `USER_ID` varchar(255))ENGINE=innodb DEFAULT CHARSET=utf8;  """ % MYSQL_TAB_SYSLOG
+
+        targetsql = """CREATE TABLE IF NOT EXISTS %s(
+            `TARGET_ID` varchar(255) NOT NULL PRIMARY KEY,
+        `ADDRv4` varchar(255),
+        `ADDRv6` varchar(255),
+        `IPv6` bool,
+        `DNSESEC` bool,
+        `DOT` bool,
+        `DOH` bool,
+        `COU` varchar(255),
+        `ISP` varchar(255),
+        `LAT` varchar(255),
+        `LNG` varchar(255),
+        `UPDATED_TIME` datetime,
+        `PROTECT` varchar(255),
+        `DOH_DOMAIN` varchar(255))ENGINE=innodb DEFAULT CHARSET=utf8; """ % MYSQL_TAB_TARGETDATA
+
+        tasksql = """CREATE TABLE IF NOT EXISTS %s(
+            `TASK_ID` varchar(255) NOT NULL PRIMARY KEY,
+        `TASK_NAME` varchar(255),
+        `AGENT_ID` varchar(255),
+        `CREATED_BY` varchar(255),
+        `TARGET_IP` varchar(255),
+        `CREATED_TIME` datetime,
+        `POLICY` varchar(255),
+        `STATUS` varchar(255),
+        `POLICY_DELAY` varchar(255),
+        `TASK_DELAY` varchar(255),
+        `TARGET_SCAN` varchar(255),
+        `TARGET_DOMAIN` varchar(255),
+        `TARGET_RTYPE` varchar(255),
+        `TARGET_RR` varchar(255))ENGINE=innodb DEFAULT CHARSET=utf8; """ % MYSQL_TAB_TASK
+
+        tasklogsql = """CREATE TABLE IF NOT EXISTS %s(
+            `TLOG_ID` int NOT NULL AUTO_INCREMENT PRIMARY KEY,
+        `CREATED_BY_AGENT` varchar(255),
+        `CREATED_TIME` datetime,
+        `TLOG_LEVEL` varchar(255),
+        `TLOG_INFO` varchar(255),
+        `TLOG_TP` varchar(255))ENGINE=innodb DEFAULT CHARSET=utf8;""" % MYSQL_TAB_TASK_LOG
+
+        taskpolicysql = """CREATE TABLE IF NOT EXISTS %s(
+            `TP_ID` int NOT NULL AUTO_INCREMENT PRIMARY KEY,
+        `TP_TIME` datetime,
+        `POLICY` varchar(255),
+        `POLICY_PARAM` varchar(255),
+        `FOR_TASK` varchar(255)) ENGINE=innodb DEFAULT CHARSET=utf8;""" % MYSQL_TAB_TASKPOLICY
+
+        usersql = """CREATE TABLE IF NOT EXISTS %s(
+            `USER_ID` varchar(255) NOT NULL PRIMARY KEY,
+        `USER_NAME` varchar(255),
+        `CREATED_BY` varchar(255),
+        `CREATED_TIME` datetime,
+        `USER_PWD_HASH` varchar(255),
+        `USER_GROUP` varchar(255)) ENGINE=innodb DEFAULT CHARSET=utf8;""" % MYSQL_TAB_USER
+
+        # 外键
+        fkey1 = """ALTER TABLE %s ADD CONSTRAINT `fk_SYSLOG_USER_1`
+        FOREIGN KEY(`USER_ID`) REFERENCES %s(`USER_ID`);""" % (MYSQL_TAB_SYSLOG, MYSQL_TAB_USER)
+
+        fkey2 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_AGENT_1`
+        FOREIGN KEY(`AGENT_ID`) REFERENCES %s(`AGENT_ID`);""" % (MYSQL_TAB_TASK, MYSQL_TAB_AGENT)
+
+        fkey3 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_USER_1`
+        FOREIGN KEY(`CREATED_BY`) REFERENCES %s(`USER_ID`);""" % (MYSQL_TAB_TASK, MYSQL_TAB_USER)
+
+        fkey4 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_TASK_POLICY_1`
+        FOREIGN KEY(`POLICY`) REFERENCES %s(`TP_ID`);""" % (MYSQL_TAB_TASK, MYSQL_TAB_TASKPOLICY)
+
+        fkey5 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_LOG_AGENT_1`
+                FOREIGN KEY(`CREATED_BY_AGENT`) REFERENCES %s(`AGENT_ID`);""" % (MYSQL_TAB_TASK_LOG, MYSQL_TAB_AGENT)
+
+        fkey6 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_LOG_TASK_POLICY_1`
+                        FOREIGN KEY(`TLOG_TP`) REFERENCES %s(`TP_ID`);""" % (
+            MYSQL_TAB_TASK_LOG, MYSQL_TAB_TASKPOLICY)
+
+        fkey7 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_POLICY_POLICY_1`
+                                FOREIGN KEY(`POLICY`) REFERENCES %s(`P_ID`);""" % (
+            MYSQL_TAB_TASKPOLICY, MYSQL_TAB_POLICY)
+
+        fkey8 = """ALTER TABLE %s ADD CONSTRAINT `fk_TASK_POLICY_TASK_1`
+                                        FOREIGN KEY(`FOR_TASK`) REFERENCES %s(`TASK_ID`);""" % (
+            MYSQL_TAB_TASKPOLICY, MYSQL_TAB_TASK)
+
+        # 执行sql语句
+        try:
+            # 创建数据库
+            self.cursor.execute(dbsql)
+            self.conn.commit()
+            self.conn.select_db(MYSQL_DATADB)
+
+            # 创建表格
+            self.cursor.execute(agentsql)
+            self.cursor.execute(tasksql)
+            self.cursor.execute(tasklogsql)
+            self.cursor.execute(targetsql)
+            self.cursor.execute(syslogsql)
+            self.cursor.execute(policysql)
+            self.cursor.execute(taskpolicysql)
+            self.cursor.execute(usersql)
+            self.conn.commit()
+
+            # 创建外键
+            self.cursor.execute(fkey1)
+            self.cursor.execute(fkey2)
+            self.cursor.execute(fkey3)
+            self.cursor.execute(fkey4)
+            self.cursor.execute(fkey5)
+            self.cursor.execute(fkey6)
+            self.cursor.execute(fkey7)
+            self.cursor.execute(fkey8)
+            self.conn.commit()
+        except Exception as e:
+            print(e)
+
+    # 获取信息(代理、任务)
+    # data_type可选范围参照DataHandler.tabmapping的键
+    # 若需要按条件检索，则以将检索维度与检索值以字典形式传入search
+    def get_data(self, data_type="agent", offset=0, limit=10, search=None):
+        # 参数映射到表名
+        tabname = self.tabmapping[data_type]
+        if search is None:
+            sql = """SELECT * FROM %s LIMIT %s, %s""" % (tabname, offset, limit)
+            self.cursor.execute(sql)
+            return self.cursor.fetchall()
+        else:
+            search = dict(search)
+            key, value = search.popitem()
+            sql = """SELECT * FROM %s WHERE %s=%s LIMIT %s, %s""" % (tabname, key, value, offset, limit)
+            self.cursor.execute(sql)
+            return self.cursor.fetchall()
+
+
+da = DataHandler()
diff --git a/server/requirements.txt b/server/requirements.txt
index af9f81b..0da63cc 100644
--- a/server/requirements.txt
+++ b/server/requirements.txt
@@ -6,4 +6,5 @@ numpy~=1.26.3
 six~=1.16.0
 pandas~=2.1.4
 faker~=18.9.0
-dnspython~=2.6.1
\ No newline at end of file
+dnspython~=2.6.1
+pymysql~=1.1.0
\ No newline at end of file
diff --git a/server/settings.py b/server/settings.py
index 5cc1536..4791c0f 100644
--- a/server/settings.py
+++ b/server/settings.py
@@ -15,4 +15,4 @@ MYSQL_TAB_TARGETDATA = "TARGETDATA"
 MYSQL_TAB_TASK = "TASK"
 MYSQL_TAB_POLICY = "POLICY"
 MYSQL_TAB_TASKPOLICY = "TASK_POLICY"
-MYSQL_TASK_LOG = "TASK_LOG"
+MYSQL_TAB_TASK_LOG = "TASK_LOG"