From d96b178a9c52c9759da9bb51ffd1a6230c9387c2 Mon Sep 17 00:00:00 2001 From: Donald Zou Date: Wed, 6 Aug 2025 17:27:14 +0800 Subject: [PATCH] Optimized loading speed for peer information --- src/dashboard.py | 6 +++--- src/modules/AmneziaWireguardConfiguration.py | 2 +- src/modules/Peer.py | 6 ++++-- src/modules/PeerJobs.py | 2 ++ src/modules/PeerShareLinks.py | 12 ++++++++---- src/modules/WireguardConfiguration.py | 11 ++++------- 6 files changed, 22 insertions(+), 17 deletions(-) diff --git a/src/dashboard.py b/src/dashboard.py index 0049328..bf221e1 100644 --- a/src/dashboard.py +++ b/src/dashboard.py @@ -609,7 +609,7 @@ def API_sharePeer_update(): ShareID: str = data.get("ShareID") ExpireDate: str = data.get("ExpireDate") - if ShareID is None: + if not all([ShareID, ExpireDate]): return ResponseObject(False, "Please specify ShareID") if len(AllPeerShareLinks.getLinkByID(ShareID)) == 0: @@ -1348,7 +1348,7 @@ def peerInformationBackgroundThread(): c.getPeersTransfer() c.getPeersLatestHandshake() c.getPeersEndpoint() - c.getPeersList() + c.getPeers() c.getRestrictedPeersList() time.sleep(10) @@ -1413,7 +1413,7 @@ _, WG_CONF_PATH = DashboardConfig.GetConfig("Server", "wg_conf_path") WireguardConfigurations: dict[str, WireguardConfiguration] = {} -AllPeerShareLinks: PeerShareLinks = PeerShareLinks(DashboardConfig) +AllPeerShareLinks: PeerShareLinks = PeerShareLinks(DashboardConfig, WireguardConfigurations) AllPeerJobs: PeerJobs = PeerJobs(DashboardConfig, WireguardConfigurations) DashboardLogger: DashboardLogger = DashboardLogger() diff --git a/src/modules/AmneziaWireguardConfiguration.py b/src/modules/AmneziaWireguardConfiguration.py index 9a8d890..53b5b44 100644 --- a/src/modules/AmneziaWireguardConfiguration.py +++ b/src/modules/AmneziaWireguardConfiguration.py @@ -353,7 +353,7 @@ class AmneziaWireguardConfiguration(WireguardConfiguration): os.remove(uid) subprocess.check_output( f"{self.Protocol}-quick save {self.Name}", shell=True, stderr=subprocess.STDOUT) - self.getPeersList() + self.getPeers() for p in peers: p = self.searchPeer(p['id']) if p[0]: diff --git a/src/modules/Peer.py b/src/modules/Peer.py index cb6b58a..37d6dd2 100644 --- a/src/modules/Peer.py +++ b/src/modules/Peer.py @@ -35,8 +35,8 @@ class Peer: self.getShareLink() def toJson(self): - self.getJobs() - self.getShareLink() + # self.getJobs() + # self.getShareLink() return self.__dict__ def __repr__(self): @@ -56,8 +56,10 @@ class Peer: if allowed_ip in existingAllowedIps: return False, "Allowed IP already taken by another peer" + if not ValidateIPAddressesWithRange(endpoint_allowed_ip): return False, f"Endpoint Allowed IPs format is incorrect" + if len(dns_addresses) > 0 and not ValidateDNSAddress(dns_addresses): return False, f"DNS format is incorrect" diff --git a/src/modules/PeerJobs.py b/src/modules/PeerJobs.py index c80294a..8acad45 100644 --- a/src/modules/PeerJobs.py +++ b/src/modules/PeerJobs.py @@ -95,6 +95,7 @@ class PeerJobs: ) self.JobLogger.log(Job.JobID, Message=f"Job is updated from if {currentJob[0].Field} {currentJob[0].Operator} {currentJob[0].Value} then {currentJob[0].Action}; to if {Job.Field} {Job.Operator} {Job.Value} then {Job.Action}") self.__getJobs() + self.WireguardConfigurations.get(Job.Configuration).searchPeer(Job.Peer)[1].getJobs() return True, list( filter(lambda x: x.Configuration == Job.Configuration and x.Peer == Job.Peer and x.JobID == Job.JobID, self.Jobs)) @@ -116,6 +117,7 @@ class PeerJobs: ) self.JobLogger.log(Job.JobID, Message=f"Job is removed due to being deleted or finshed.") self.__getJobs() + self.WireguardConfigurations.get(Job.Configuration).searchPeer(Job.Peer)[1].getJobs() return True, None except Exception as e: return False, str(e) diff --git a/src/modules/PeerShareLinks.py b/src/modules/PeerShareLinks.py index 5123d84..206e2fd 100644 --- a/src/modules/PeerShareLinks.py +++ b/src/modules/PeerShareLinks.py @@ -8,7 +8,7 @@ import uuid Peer Share Links """ class PeerShareLinks: - def __init__(self, DashboardConfig): + def __init__(self, DashboardConfig, WireguardConfigurations): self.Links: list[PeerShareLink] = [] self.engine = db.create_engine(ConnectionString("wgdashboard")) self.metadata = db.MetaData() @@ -23,6 +23,7 @@ class PeerShareLinks: ) self.metadata.create_all(self.engine) self.__getSharedLinks() + self.wireguardConfigurations = WireguardConfigurations def __getSharedLinks(self): self.Links.clear() with self.engine.connect() as conn: @@ -68,18 +69,21 @@ class PeerShareLinks: ) ) self.__getSharedLinks() + self.wireguardConfigurations.get(Configuration).searchPeer(Peer)[1].getShareLink() except Exception as e: return False, str(e) return True, newShareID def updateLinkExpireDate(self, ShareID, ExpireDate: datetime = None) -> tuple[bool, str]: with self.engine.begin() as conn: - conn.execute( + updated = conn.execute( self.peerShareLinksTable.update().values( { "ExpireDate": ExpireDate } - ).where(db.and_(self.peerShareLinksTable.columns.ShareID == ShareID)) - ) + ).returning(self.peerShareLinksTable.c.Configuration, self.peerShareLinksTable.c.Peer) + .where(self.peerShareLinksTable.columns.ShareID == ShareID) + ).mappings().fetchone() self.__getSharedLinks() + self.wireguardConfigurations.get(updated.Configuration).searchPeer(updated.Peer)[1].getShareLink() return True, "" \ No newline at end of file diff --git a/src/modules/WireguardConfiguration.py b/src/modules/WireguardConfiguration.py index 4b78816..3da07bf 100644 --- a/src/modules/WireguardConfiguration.py +++ b/src/modules/WireguardConfiguration.py @@ -73,9 +73,6 @@ class WireguardConfiguration: self.__parseConfigurationFile() self.__initPeersList() - - - else: self.Name = data["ConfigurationName"] self.configPath = os.path.join(self.__getProtocolPath(), f'{self.Name}.conf') @@ -132,7 +129,7 @@ class WireguardConfiguration: def __initPeersList(self): self.Peers: list[Peer] = [] - self.getPeersList() + self.getPeers() self.getRestrictedPeersList() def getRawConfigurationFile(self): @@ -350,7 +347,6 @@ class WireguardConfiguration: return changed def getPeers(self): - # self.Peers = [] tmpList = [] if self.configurationFileChanged(): with open(self.configPath, 'r') as configFile: @@ -377,6 +373,8 @@ class WireguardConfiguration: if len(split) == 2: p[pCounter]["name"] = split[1] with self.engine.begin() as conn: + + for i in p: if "PublicKey" in i.keys(): tempPeer = conn.execute(self.peersTable.select().where( @@ -475,7 +473,7 @@ class WireguardConfiguration: os.remove(uid) subprocess.check_output( f"{self.Protocol}-quick save {self.Name}", shell=True, stderr=subprocess.STDOUT) - self.getPeersList() + self.getPeers() for p in peers: p = self.searchPeer(p['id']) if p[0]: @@ -758,7 +756,6 @@ class WireguardConfiguration: return True, None def getPeersList(self): - self.getPeers() return self.Peers def getRestrictedPeersList(self) -> list: