feat: 优化了节点启动和配置处理

fix: 修复无法更改DNS记录的问题
fix #3
Co-authored-by: cooollawf <[email protected]>

core/__init__.py

@@ -195,14 +195,14 @@ async def on_cluster_request_cert(sid, *args):
 ## 节点启动时
 @sio.on("enable")
 async def on_cluster_enable(sid, data: dict, *args):
-    # {'host': '127.0.0.1', 'port': 11451, 'version': '1.11.0', 'byoc': True, 'noFastEnable': True, 'flavor': {'runtime': 'python/3.12.2 python-openbmclapi/2.1.1', 'storage': 'file'}}
     session = await sio.get_session(sid)
     cluster = Cluster(str(session["cluster_id"]))
     cluster_is_exist = await cluster.initialize()
     if cluster_is_exist == False:
         return [{"message": "错误: 节点似乎并不存在，请检查配置文件"}]
     if oclm.include(cluster.id):
         return [{"message": "错误: 节点已经在线，请检查配置文件"}]
+    logger.debug(f"节点 {cluster.id} 请求启用")
     host = data.get("host", data.get("ip"))
     byoc = data.get("byoc", False)
     if byoc == False:
@@ -214,9 +214,9 @@ async def on_cluster_enable(sid, data: dict, *args):
             else:
                 cf_id = None
         if cf_id == None:
-            await cf_client.create_record(cluster.id, "A", data.get("host", data.get("ip")))
+            await cf_client.create_record(cluster.id, "A", host)
         else:
-            await cf_client.update_record(id, cluster.id, "A", data.get("host", data.get("ip")))
+            await cf_client.update_record(cf_id, cluster.id, "A", host)
         host = f"{cluster.id}.{config.get('cluster-certificate.domain')}"
 
     await cluster.edit(
@@ -237,8 +237,8 @@ async def on_cluster_enable(sid, data: dict, *args):
         await cluster.edit(measureBandwidth=int(bandwidth[1]))
         if cluster.trust < 0:
             await sio.emit("message", "节点信任度过低，请保持稳定在线。", sid)
-        oclm.append(cluster.id)
-        logger.debug(f"节点 {cluster.id} 上线: 测量带宽 = {bandwidth[1]}Mbps")
+        oclm.append(cluster.id, cluster.weight)
+        logger.debug(f"节点 {cluster.id} 启用: 测量带宽 = {bandwidth[1]}Mbps")
         return [None, True]
     elif bandwidth[0] and bandwidth[1] < 10:
         logger.debug(f"{cluster.id} 测速不合格: {bandwidth[1]}Mbps")
@@ -259,6 +259,7 @@ async def on_cluster_keep_alive(sid, data, *args):
     cluster_is_exist = await cluster.initialize()
     if cluster_is_exist == False or oclm.include(cluster.id) == False:
         return [None, False]
+    oclm.update(cluster.id, cluster.weight)
     logger.debug(
         f"节点 {cluster.id} 保活成功: 次数 = {data["hits"]}, 数据量 = {utils.hum_convert(data['bytes'])}"
     )

core/config.py

@@ -1,6 +1,10 @@
+# 第三方库
 import yaml
 from pathlib import Path
 
+# 本地库
+from core.logger import logger
+
 class Config:
     def __init__(self, config_file):
         self.config_file = config_file
@@ -11,14 +15,15 @@ def load_config(self):
             config = yaml.safe_load(file)
         return config
 
-    def get(self, path: str, default=None):
+    def get(self, path: str, default = None):
         keys = path.split(".")
         data = self.config
         try:
             for key in keys:
                 data = data[key]
             return data
         except (KeyError, TypeError):
+            logger.warning(f"{path} 未设置，已返回 None 作为默认值。")
             return default
 
 

core/dns/cloudflare.py

@@ -96,7 +96,15 @@ async def update_record(
     ):
         async with httpx.AsyncClient() as client:
             data = await client.patch(
-                f"https://api.cloudflare.com/client/v4/zones/{self.zone_id}/dns_records/{record_id}"
+                f"https://api.cloudflare.com/client/v4/zones/{self.zone_id}/dns_records/{record_id}",
+                headers=self.headers,
+                json={
+                    "type": type,
+                    "name": name,
+                    "content": content,
+                    "ttl": ttl,
+                    "proxied": proxied,
+                }
             )
         if data.status_code == 200:
             return data.json()

core/logger.py

@@ -1,3 +1,4 @@
+# 第三方库
 import os
 import sys
 from pathlib import Path

core/routes/agent.py

@@ -5,23 +5,25 @@
 from fastapi import APIRouter, Request, HTTPException, Response
 from fastapi.responses import JSONResponse, PlainTextResponse
 
-router = APIRouter()
-
 # 本地库
 import core.utils as utils
 import core.const as const
 from core.types import Cluster
 
 
-@router.get("/challenge", summary="颁发 challenge 码", tags=["nodes"])  # 颁发 challenge 验证码
+router = APIRouter()
+
+
+@router.get(
+    "/challenge", summary="颁发 challenge 码", tags=["nodes"]
+)  # 颁发 challenge 验证码
 async def get_challenge(response: Response, clusterId: str | None = ""):
     cluster = Cluster(clusterId)
     if await cluster.initialize():
         return {
             "challenge": utils.encode_jwt(
                 {
                     "cluster_id": clusterId,
-                    "cluster_secret": cluster.secret,
                     "iss": const.jwt_iss,
                     "exp": int(time.time()) + 1000 * 60 * 5,
                 }

core/routes/services.py

@@ -1,10 +1,10 @@
 # 第三方库
 from pathlib import Path
+from random import choices
 from fastapi import APIRouter, Request, HTTPException, Response
 from fastapi.responses import FileResponse, RedirectResponse
 
 # 本地库
-from random import choice
 import core.utils as utils
 from core.types import oclm, Cluster
 from core.logger import logger
@@ -22,7 +22,7 @@ async def download_path_file(hash: str):
         if len(oclm) == 0:
             return RedirectResponse(filedata["URL"], 302)
         else:
-            cluster = Cluster(choice(oclm.list))
+            cluster = Cluster(oclm.random())
             await cluster.initialize()
             sign = utils.get_sign(filedata["HASH"], cluster.secret)
             url = utils.get_url(

core/types.py

@@ -1,6 +1,10 @@
+# 第三方库
 import io
 import heapq
+from random import choices
 from typing import Optional
+
+# 本地库
 from core.mdb import cdb
 
 
@@ -88,21 +92,32 @@ def json(self):
 
 class OCLManager:
     def __init__(self):
-        self.list = []
+        self.id_list = []
+        self.weight_list = []
 
     def __len__(self):
-        return len(self.list)
+        return len(self.id_list)
 
-    def append(self, cluster_id: str):
-        if cluster_id not in self.list:
-            self.list.append(cluster_id)
+    def append(self, cluster_id: str, weight: float):
+        if cluster_id not in self.id_list:
+            self.id_list.append(cluster_id)
+            self.weight_list.append(weight)
 
     def remove(self, cluster_id: str):
-        if cluster_id in self.list:
-            self.list.remove(cluster_id)
+        if cluster_id in self.id_list:
+            self.weight_list.remove(self.weight_list[self.id_list.index(cluster_id)])
+            self.id_list.remove(cluster_id)
+
+    def update(self, cluster_id: str, weight: float):
+        if cluster_id in self.id_list:
+            self.weight_list[self.id_list.index(cluster_id)] = weight
 
     def include(self, cluster_id: str):
-        return cluster_id in self.list
+        return cluster_id in self.id_list
+
+    def random(self) -> str:
+        return choices(self.id_list, self.weight_list)[0]
+
 
 
 oclm = OCLManager()

core/utils.py

@@ -83,7 +83,7 @@ async def measure_cluster(size: int, cluster: Cluster):
     try:
         start_time = time.time()
         async with httpx.AsyncClient() as client:
-            response = await client.get(url, headers={"User-Agent": const.user_agent})
+            response = await client.get(url, headers={"User-Agent": const.user_agent}, timeout=10)
         end_time = time.time()
         elapsed_time = end_time - start_time
         # 计算测速时间