代理注册功能开发

3 files changed, 54 insertions, 28 deletions

diff --git a/agent/app.py b/agent/app.py
index d78c893..8216abb 100644
--- a/agent/app.py
+++ b/agent/app.py
@@ -1,20 +1,22 @@
 import argparse
 import ipaddress
+import os.path
 import platform
 import socket
 
 import geocoder
 import psutil
 import requests
+import yaml
 from apiflask import APIFlask
-from server.apps.datacheck import bp as datacheckbp
 
 from apps.delay import bp as scriptbp
 
 # 注册蓝图
 app = APIFlask(__name__, template_folder='./static/templates')
 app.register_blueprint(scriptbp)
-app.register_blueprint(datacheckbp)
+# 代理配置
+config = {}
 
 
 @app.get('/')
@@ -26,8 +28,8 @@ def hello():
 # 获取代理的部署环境信息
 def nodeinfo():
     # IP地址
-    ## IPv4地址通过向主控端发包时在主控端获取，从而定位最准确的IPv4地址（公网、内网）
-    ## IPv6地址获取本地的所有公网地址
+    # IPv4地址通过向主控端发包时在主控端获取，从而定位最准确的IPv4地址（公网、内网）
+    # IPv6地址获取本地的所有公网地址
     v6addr = []
     addr6_info_list = socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET6)
     for a in addr6_info_list:
@@ -54,21 +56,43 @@ def nodeinfo():
 
 
 # 注册代理
-def registernode(port=2525, atype="stgj", server="127.0.0.1:8888"):
+def registernode(proto="http", port=2525, atype="stgj", server="127.0.0.1:8888"):
     info = nodeinfo()
     info["port"] = port
     info["type"] = atype
-    requests.post("http://" + server)
+    r = requests.post(proto + "://" + server, json=info)
+    if r.status_code != 200:
+        print("注册失败，请检查参数配置是否正确")
+        return
+    return r.json()['id']
 
 
 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",
-                        help="代理的工作类型 {stgj(渗透攻击) / mbgz(目标感知) / ztgz(状态感知)}")
-    parser.add_argument("-s", "--server", type=str, default="127.0.0.1:8888", help="主控端访问地址+端口号")
-    args = parser.parse_args()
-    registernode(args.port, server=args.server)
-
-    app.run(host="0.0.0.0", debug=True, port=args.port)
+    # 判断是否已存在配置文件
+    if os.path.exists("config.yaml"):
+        # 存在则读取配置文件信息
+        with open("config.yaml", "r") as f:
+            config = yaml.safe_load(f)
+            app.run(host="0.0.0.0", port=config["port"], debug=True)
+    else:
+        # 不存在则解析命令行参数
+        # 命令行参数设置
+        parser = argparse.ArgumentParser()
+        parser.add_argument("-p", "--port", type=int, default=2525, help="代理的开放通信端口")
+        parser.add_argument("-t", "--type", type=str, default="stgj",
+                            help="代理的工作类型 {stgj(渗透攻击) / mbgz(目标感知) / ztgz(状态感知)}")
+        parser.add_argument("-s", "--server", type=str, default="127.0.0.1:8888", help="主控端访问地址+端口号")
+
+        # 解析参数
+        args = parser.parse_args()
+
+        # 注册代理，并获取主控分配的代理编号
+        id = registernode(port=args.port, server=args.server, atype=args.type)
+        config = {"id": id, "port": args.port, "atype": args.type, "server": args.server}
+
+        # 配置写入yaml文件存储
+        with open('config.yaml', 'w') as f:
+            yaml.dump(config, f)
+
+        # 运行
+        app.run(host="0.0.0.0", debug=True, port=config["port"])
diff --git a/agent/apps/util.py b/agent/apps/util.py
index 986918d..8cadf4d 100644
--- a/agent/apps/util.py
+++ b/agent/apps/util.py
@@ -1,9 +1,10 @@
 # 工具函数集合
 import requests
+
+
 # 代理输出回传
-def agent_echo(proto,server,level,info):
-    data={"level":level,"info":info}
-    r=requests.post(proto+"://"+server,json=data)
-    if r.status_code==200:
+def agent_echo(proto, server, level, info):
+    data = {"level": level, "info": info}
+    r = requests.post(proto + "://" + server, json=data)
+    if r.status_code == 200:
         print("ok")
-    
diff --git a/agent/requirements.txt b/agent/requirements.txt
index 5494f55..581b4bf 100644
--- a/agent/requirements.txt
+++ b/agent/requirements.txt
@@ -1,9 +1,10 @@
+requests~=2.31.0
+Flask~=3.0.0
+APIFlask~=2.1.0
 click~=8.1.7
-numpy~=1.25.0
+numpy~=1.26.3
 six~=1.16.0
-requests~=2.31.0
-pandas~=2.1.1
-APIFlask~=2.0.2
-flask~=3.0.0
-dnspython~=2.4.2
-icmplib~=3.0.4
\ No newline at end of file
+pandas~=2.1.4
+psutil~=5.9.0
+geocoder~=1.38.1
+pyyaml~=6.0.1
\ No newline at end of file