Finally moved all class to its own file

src/modules/AmneziaWGPeer.py

@@ -0,0 +1,137 @@
+import os
+import random
+import re
+import subprocess
+import uuid
+
+from .Peer import Peer
+from .Utilities import ValidateIPAddressesWithRange, ValidateDNSAddress, GenerateWireguardPublicKey
+
+
+class AmneziaWGPeer(Peer):
+    def __init__(self, tableData, configuration):
+        self.advanced_security = tableData["advanced_security"]
+        super().__init__(tableData, configuration)
+
+    def downloadPeer(self) -> dict[str, str]:
+        filename = self.name
+        if len(filename) == 0:
+            filename = "UntitledPeer"
+        filename = "".join(filename.split(' '))
+        filename = f"{filename}_{self.configuration.Name}"
+        illegal_filename = [".", ",", "/", "?", "<", ">", "\\", ":", "*", '|' '\"', "com1", "com2", "com3",
+                            "com4", "com5", "com6", "com7", "com8", "com9", "lpt1", "lpt2", "lpt3", "lpt4",
+                            "lpt5", "lpt6", "lpt7", "lpt8", "lpt9", "con", "nul", "prn"]
+        for i in illegal_filename:
+            filename = filename.replace(i, "")
+
+        finalFilename = ""
+        for i in filename:
+            if re.match("^[a-zA-Z0-9_=+.-]$", i):
+                finalFilename += i
+
+        peerConfiguration = f'''[Interface]
+PrivateKey = {self.private_key}
+Address = {self.allowed_ip}
+MTU = {str(self.mtu)}
+Jc = {self.configuration.Jc}
+Jmin = {self.configuration.Jmin}
+Jmax = {self.configuration.Jmax}
+S1 = {self.configuration.S1}
+S2 = {self.configuration.S2}
+H1 = {self.configuration.H1}
+H2 = {self.configuration.H2}
+H3 = {self.configuration.H3}
+H4 = {self.configuration.H4}
+'''
+        if len(self.DNS) > 0:
+            peerConfiguration += f"DNS = {self.DNS}\n"
+        peerConfiguration += f'''
+[Peer]
+PublicKey = {self.configuration.PublicKey}
+AllowedIPs = {self.endpoint_allowed_ip}
+Endpoint = {self.configuration.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1]}:{self.configuration.ListenPort}
+PersistentKeepalive = {str(self.keepalive)}
+'''
+        if len(self.preshared_key) > 0:
+            peerConfiguration += f"PresharedKey = {self.preshared_key}\n"
+        return {
+            "fileName": finalFilename,
+            "file": peerConfiguration
+        }
+
+    def updatePeer(self, name: str, private_key: str,
+                   preshared_key: str,
+                   dns_addresses: str, allowed_ip: str, endpoint_allowed_ip: str, mtu: int,
+                   keepalive: int, advanced_security: str) -> tuple[bool, str] or tuple[bool, None]:
+        if not self.configuration.getStatus():
+            self.configuration.toggleConfiguration()
+
+        existingAllowedIps = [item for row in list(
+            map(lambda x: [q.strip() for q in x.split(',')],
+                map(lambda y: y.allowed_ip,
+                    list(filter(lambda k: k.id != self.id, self.configuration.getPeersList()))))) for item in row]
+
+        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"
+        if mtu < 0 or mtu > 1460:
+            return False, "MTU format is not correct"
+        if keepalive < 0:
+            return False, "Persistent Keepalive format is not correct"
+        if advanced_security != "on" and advanced_security != "off":
+            return False, "Advanced Security can only be on or off"
+        if len(private_key) > 0:
+            pubKey = GenerateWireguardPublicKey(private_key)
+            if not pubKey[0] or pubKey[1] != self.id:
+                return False, "Private key does not match with the public key"
+        try:
+            rd = random.Random()
+            uid = str(uuid.UUID(int=rd.getrandbits(128), version=4))
+            pskExist = len(preshared_key) > 0
+
+            if pskExist:
+                with open(uid, "w+") as f:
+                    f.write(preshared_key)
+            newAllowedIPs = allowed_ip.replace(" ", "")
+            updateAllowedIp = subprocess.check_output(
+                f"{self.configuration.Protocol} set {self.configuration.Name} peer {self.id} allowed-ips {newAllowedIPs} {f'preshared-key {uid}' if pskExist else 'preshared-key /dev/null'}",
+                shell=True, stderr=subprocess.STDOUT)
+
+            if pskExist: os.remove(uid)
+
+            if len(updateAllowedIp.decode().strip("\n")) != 0:
+                return False, "Update peer failed when updating Allowed IPs"
+            saveConfig = subprocess.check_output(f"{self.configuration.Protocol}-quick save {self.configuration.Name}",
+                                                 shell=True, stderr=subprocess.STDOUT)
+            if f"wg showconf {self.configuration.Name}" not in saveConfig.decode().strip('\n'):
+                return False, "Update peer failed when saving the configuration"
+            # sqlUpdate(
+            #     '''UPDATE '%s' SET name = ?, private_key = ?, DNS = ?, endpoint_allowed_ip = ?, mtu = ?, 
+            #     keepalive = ?, preshared_key = ?, advanced_security = ?  WHERE id = ?''' % self.configuration.Name,
+            #     (name, private_key, dns_addresses, endpoint_allowed_ip, mtu,
+            #      keepalive, preshared_key, advanced_security, self.id,)
+            # )
+
+            with self.configuration.engine.begin() as conn:
+                conn.execute(
+                    self.configuration.peersTable.update().values({
+                        "name": name,
+                        "private_key": private_key,
+                        "DNS": dns_addresses,
+                        "endpoint_allowed_ip": endpoint_allowed_ip,
+                        "mtu": mtu,
+                        "keepalive": keepalive,
+                        "preshared_key": preshared_key,
+                        "advanced_security": advanced_security
+                    }).where(
+                        self.configuration.peersTable.c.id == self.id
+                    )
+                )
+
+            return True, None
+        except subprocess.CalledProcessError as exc:
+            return False, exc.output.decode("UTF-8").strip()

src/modules/AmneziaWireguardConfiguration.py

@@ -0,0 +1,374 @@
+"""
+AmneziaWG Configuration
+"""
+import re
+
+import sqlalchemy
+
+from .PeerJobs import PeerJobs
+from .AmneziaWGPeer import AmneziaWGPeer
+from .PeerShareLinks import PeerShareLinks
+from .Utilities import RegexMatch
+from .WireguardConfiguration import WireguardConfiguration
+
+
+class AmneziaWireguardConfiguration(WireguardConfiguration):
+    def __init__(self, DashboardConfig,
+                 AllPeerJobs: PeerJobs,
+                 AllPeerShareLinks: PeerShareLinks,
+                 name: str = None, data: dict = None, backup: dict = None, startup: bool = False):
+        self.Jc = 0
+        self.Jmin = 0
+        self.Jmax = 0
+        self.S1 = 0
+        self.S2 = 0
+        self.H1 = 1
+        self.H2 = 2
+        self.H3 = 3
+        self.H4 = 4
+
+        super().__init__(DashboardConfig, AllPeerJobs, AllPeerShareLinks, name, data, backup, startup, wg=False)
+
+    def toJson(self):
+        self.Status = self.getStatus()
+        return {
+            "Status": self.Status,
+            "Name": self.Name,
+            "PrivateKey": self.PrivateKey,
+            "PublicKey": self.PublicKey,
+            "Address": self.Address,
+            "ListenPort": self.ListenPort,
+            "PreUp": self.PreUp,
+            "PreDown": self.PreDown,
+            "PostUp": self.PostUp,
+            "PostDown": self.PostDown,
+            "SaveConfig": self.SaveConfig,
+            "DataUsage": {
+                "Total": sum(list(map(lambda x: x.cumu_data + x.total_data, self.Peers))),
+                "Sent": sum(list(map(lambda x: x.cumu_sent + x.total_sent, self.Peers))),
+                "Receive": sum(list(map(lambda x: x.cumu_receive + x.total_receive, self.Peers)))
+            },
+            "ConnectedPeers": len(list(filter(lambda x: x.status == "running", self.Peers))),
+            "TotalPeers": len(self.Peers),
+            "Protocol": self.Protocol,
+            "Jc": self.Jc,
+            "Jmin": self.Jmin,
+            "Jmax": self.Jmax,
+            "S1": self.S1,
+            "S2": self.S2,
+            "H1": self.H1,
+            "H2": self.H2,
+            "H3": self.H3,
+            "H4": self.H4
+        }
+
+    def createDatabase(self, dbName = None):
+        if dbName is None:
+            dbName = self.Name
+
+
+        self.peersTable = sqlalchemy.Table(
+            dbName, self.metadata,
+            sqlalchemy.Column('id', sqlalchemy.String, nullable=False, primary_key=True),
+            sqlalchemy.Column('private_key', sqlalchemy.String),
+            sqlalchemy.Column('DNS', sqlalchemy.String),
+            sqlalchemy.Column('advanced_security', sqlalchemy.String),
+            sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.String),
+            sqlalchemy.Column('name', sqlalchemy.String),
+            sqlalchemy.Column('total_receive', sqlalchemy.Float),
+            sqlalchemy.Column('total_sent', sqlalchemy.Float),
+            sqlalchemy.Column('total_data', sqlalchemy.Float),
+            sqlalchemy.Column('endpoint', sqlalchemy.String),
+            sqlalchemy.Column('status', sqlalchemy.String),
+            sqlalchemy.Column('latest_handshake', sqlalchemy.String),
+            sqlalchemy.Column('allowed_ip', sqlalchemy.String),
+            sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_data', sqlalchemy.Float),
+            sqlalchemy.Column('mtu', sqlalchemy.Integer),
+            sqlalchemy.Column('keepalive', sqlalchemy.Integer),
+            sqlalchemy.Column('remote_endpoint', sqlalchemy.String),
+            sqlalchemy.Column('preshared_key', sqlalchemy.String),
+            extend_existing=True
+        )
+        self.peersRestrictedTable = sqlalchemy.Table(
+            f'{dbName}_restrict_access', self.metadata,
+            sqlalchemy.Column('id', sqlalchemy.String, nullable=False, primary_key=True),
+            sqlalchemy.Column('private_key', sqlalchemy.String),
+            sqlalchemy.Column('DNS', sqlalchemy.String),
+            sqlalchemy.Column('advanced_security', sqlalchemy.String),
+            sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.String),
+            sqlalchemy.Column('name', sqlalchemy.String),
+            sqlalchemy.Column('total_receive', sqlalchemy.Float),
+            sqlalchemy.Column('total_sent', sqlalchemy.Float),
+            sqlalchemy.Column('total_data', sqlalchemy.Float),
+            sqlalchemy.Column('endpoint', sqlalchemy.String),
+            sqlalchemy.Column('status', sqlalchemy.String),
+            sqlalchemy.Column('latest_handshake', sqlalchemy.String),
+            sqlalchemy.Column('allowed_ip', sqlalchemy.String),
+            sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_data', sqlalchemy.Float),
+            sqlalchemy.Column('mtu', sqlalchemy.Integer),
+            sqlalchemy.Column('keepalive', sqlalchemy.Integer),
+            sqlalchemy.Column('remote_endpoint', sqlalchemy.String),
+            sqlalchemy.Column('preshared_key', sqlalchemy.String),
+            extend_existing=True
+        )
+        self.peersTransferTable = sqlalchemy.Table(
+            f'{dbName}_transfer', self.metadata,
+            sqlalchemy.Column('id', sqlalchemy.String, nullable=False),
+            sqlalchemy.Column('total_receive', sqlalchemy.Float),
+            sqlalchemy.Column('total_sent', sqlalchemy.Float),
+            sqlalchemy.Column('total_data', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_data', sqlalchemy.Float),
+            sqlalchemy.Column('time', (sqlalchemy.DATETIME if self.DashboardConfig.GetConfig("Database", "type")[1] == 'sqlite' else sqlalchemy.TIMESTAMP),
+                              server_default=sqlalchemy.func.now()),
+            extend_existing=True
+        )
+        self.peersDeletedTable = sqlalchemy.Table(
+            f'{dbName}_deleted', self.metadata,
+            sqlalchemy.Column('id', sqlalchemy.String, nullable=False, primary_key=True),
+            sqlalchemy.Column('private_key', sqlalchemy.String),
+            sqlalchemy.Column('advanced_security', sqlalchemy.String),
+            sqlalchemy.Column('DNS', sqlalchemy.String),
+            sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.String),
+            sqlalchemy.Column('name', sqlalchemy.String),
+            sqlalchemy.Column('total_receive', sqlalchemy.Float),
+            sqlalchemy.Column('total_sent', sqlalchemy.Float),
+            sqlalchemy.Column('total_data', sqlalchemy.Float),
+            sqlalchemy.Column('endpoint', sqlalchemy.String),
+            sqlalchemy.Column('status', sqlalchemy.String),
+            sqlalchemy.Column('latest_handshake', sqlalchemy.String),
+            sqlalchemy.Column('allowed_ip', sqlalchemy.String),
+            sqlalchemy.Column('cumu_receive', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_sent', sqlalchemy.Float),
+            sqlalchemy.Column('cumu_data', sqlalchemy.Float),
+            sqlalchemy.Column('mtu', sqlalchemy.Integer),
+            sqlalchemy.Column('keepalive', sqlalchemy.Integer),
+            sqlalchemy.Column('remote_endpoint', sqlalchemy.String),
+            sqlalchemy.Column('preshared_key', sqlalchemy.String),
+            extend_existing=True
+        )
+
+        self.metadata.create_all(self.engine)
+
+        # existingTables = sqlSelect("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
+        # existingTables = [t['name'] for t in existingTables]
+        # if dbName not in existingTables:
+        #     sqlUpdate(
+        #         """
+        #         CREATE TABLE '%s'(
+        #             id VARCHAR NOT NULL, private_key VARCHAR NULL, DNS VARCHAR NULL, advanced_security VARCHAR NULL,
+        #             endpoint_allowed_ip VARCHAR NULL, name VARCHAR NULL, total_receive FLOAT NULL, 
+        #             total_sent FLOAT NULL, total_data FLOAT NULL, endpoint VARCHAR NULL, 
+        #             status VARCHAR NULL, latest_handshake VARCHAR NULL, allowed_ip VARCHAR NULL, 
+        #             cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, mtu INT NULL, 
+        #             keepalive INT NULL, remote_endpoint VARCHAR NULL, preshared_key VARCHAR NULL,
+        #             PRIMARY KEY (id)
+        #         )
+        #         """ % dbName
+        #     )
+        # 
+        # if f'{dbName}_restrict_access' not in existingTables:
+        #     sqlUpdate(
+        #         """
+        #         CREATE TABLE '%s_restrict_access' (
+        #             id VARCHAR NOT NULL, private_key VARCHAR NULL, DNS VARCHAR NULL, advanced_security VARCHAR NULL, 
+        #             endpoint_allowed_ip VARCHAR NULL, name VARCHAR NULL, total_receive FLOAT NULL, 
+        #             total_sent FLOAT NULL, total_data FLOAT NULL, endpoint VARCHAR NULL, 
+        #             status VARCHAR NULL, latest_handshake VARCHAR NULL, allowed_ip VARCHAR NULL, 
+        #             cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, mtu INT NULL, 
+        #             keepalive INT NULL, remote_endpoint VARCHAR NULL, preshared_key VARCHAR NULL,
+        #             PRIMARY KEY (id)
+        #         )
+        #         """ % dbName
+        #     )
+        # if f'{dbName}_transfer' not in existingTables:
+        #     sqlUpdate(
+        #         """
+        #         CREATE TABLE '%s_transfer' (
+        #             id VARCHAR NOT NULL, total_receive FLOAT NULL,
+        #             total_sent FLOAT NULL, total_data FLOAT NULL,
+        #             cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, time DATETIME
+        #         )
+        #         """ % dbName
+        #     )
+        # if f'{dbName}_deleted' not in existingTables:
+        #     sqlUpdate(
+        #         """
+        #         CREATE TABLE '%s_deleted' (
+        #             id VARCHAR NOT NULL, private_key VARCHAR NULL, DNS VARCHAR NULL, advanced_security VARCHAR NULL,
+        #             endpoint_allowed_ip VARCHAR NULL, name VARCHAR NULL, total_receive FLOAT NULL, 
+        #             total_sent FLOAT NULL, total_data FLOAT NULL, endpoint VARCHAR NULL, 
+        #             status VARCHAR NULL, latest_handshake VARCHAR NULL, allowed_ip VARCHAR NULL, 
+        #             cumu_receive FLOAT NULL, cumu_sent FLOAT NULL, cumu_data FLOAT NULL, mtu INT NULL, 
+        #             keepalive INT NULL, remote_endpoint VARCHAR NULL, preshared_key VARCHAR NULL,
+        #             PRIMARY KEY (id)
+        #         )
+        #         """ % dbName
+        #     )
+
+    def getPeers(self):
+        self.Peers.clear()
+        if self.configurationFileChanged():
+            with open(self.configPath, 'r') as configFile:
+                p = []
+                pCounter = -1
+                content = configFile.read().split('\n')
+                try:
+                    peerStarts = content.index("[Peer]")
+                    content = content[peerStarts:]
+                    for i in content:
+                        if not RegexMatch("#(.*)", i) and not RegexMatch(";(.*)", i):
+                            if i == "[Peer]":
+                                pCounter += 1
+                                p.append({})
+                                p[pCounter]["name"] = ""
+                            else:
+                                if len(i) > 0:
+                                    split = re.split(r'\s*=\s*', i, 1)
+                                    if len(split) == 2:
+                                        p[pCounter][split[0]] = split[1]
+
+                        if RegexMatch("#Name# = (.*)", i):
+                            split = re.split(r'\s*=\s*', i, 1)
+                            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(
+                                    self.peersTable.columns.id == i['PublicKey']
+                                )).mappings().fetchone()
+                                if tempPeer is None:
+                                    tempPeer = {
+                                        "id": i['PublicKey'],
+                                        "advanced_security": i.get('AdvancedSecurity', 'off'),
+                                        "private_key": "",
+                                        "DNS": self.DashboardConfig.GetConfig("Peers", "peer_global_DNS")[1],
+                                        "endpoint_allowed_ip": self.DashboardConfig.GetConfig("Peers", "peer_endpoint_allowed_ip")[
+                                            1],
+                                        "name": i.get("name"),
+                                        "total_receive": 0,
+                                        "total_sent": 0,
+                                        "total_data": 0,
+                                        "endpoint": "N/A",
+                                        "status": "stopped",
+                                        "latest_handshake": "N/A",
+                                        "allowed_ip": i.get("AllowedIPs", "N/A"),
+                                        "cumu_receive": 0,
+                                        "cumu_sent": 0,
+                                        "cumu_data": 0,
+                                        "mtu": self.DashboardConfig.GetConfig("Peers", "peer_mtu")[1],
+                                        "keepalive": self.DashboardConfig.GetConfig("Peers", "peer_keep_alive")[1],
+                                        "remote_endpoint": self.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1],
+                                        "preshared_key": i["PresharedKey"] if "PresharedKey" in i.keys() else ""
+                                    }
+                                    # sqlUpdate(
+                                    #     """
+                                    #     INSERT INTO '%s'
+                                    #         VALUES (:id, :private_key, :DNS, :advanced_security, :endpoint_allowed_ip, :name, :total_receive, :total_sent, 
+                                    #         :total_data, :endpoint, :status, :latest_handshake, :allowed_ip, :cumu_receive, :cumu_sent, 
+                                    #         :cumu_data, :mtu, :keepalive, :remote_endpoint, :preshared_key);
+                                    #     """ % self.Name
+                                    #     , newPeer)
+                                    conn.execute(
+                                        self.peersTable.insert().values(tempPeer)
+                                    )
+                                else:
+                                    # sqlUpdate("UPDATE '%s' SET allowed_ip = ? WHERE id = ?" % self.Name,
+                                    #           (i.get("AllowedIPs", "N/A"), i['PublicKey'],))
+                                    conn.execute(
+                                        self.peersTable.update().values({
+                                            "allowed_ip": i.get("AllowedIPs", "N/A")
+                                        }).where(
+                                            self.peersTable.columns.id == i['PublicKey']
+                                        )
+                                    )
+                                self.Peers.append(AmneziaWGPeer(tempPeer, self))
+                except Exception as e:
+                    if __name__ == '__main__':
+                        print(f"[WGDashboard] {self.Name} getPeers() Error: {str(e)}")
+        else:
+            # checkIfExist = sqlSelect("SELECT * FROM '%s'" % self.Name).fetchall()
+            with self.engine.connect() as conn:
+                existingPeers = conn.execute(self.peersTable.select()).mappings().fetchall()
+                for i in existingPeers:
+                    self.Peers.append(AmneziaWGPeer(i, self))
+
+    def addPeers(self, peers: list) -> tuple[bool, dict]:
+        result = {
+            "message": None,
+            "peers": []
+        }
+        try:
+            with self.engine.begin() as conn:
+                for i in peers:
+                    newPeer = {
+                        "id": i['id'],
+                        "private_key": i['private_key'],
+                        "DNS": i['DNS'],
+                        "endpoint_allowed_ip": i['endpoint_allowed_ip'],
+                        "name": i['name'],
+                        "total_receive": 0,
+                        "total_sent": 0,
+                        "total_data": 0,
+                        "endpoint": "N/A",
+                        "status": "stopped",
+                        "latest_handshake": "N/A",
+                        "allowed_ip": i.get("allowed_ip", "N/A"),
+                        "cumu_receive": 0,
+                        "cumu_sent": 0,
+                        "cumu_data": 0,
+                        "mtu": i['mtu'],
+                        "keepalive": i['keepalive'],
+                        "remote_endpoint": self.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1],
+                        "preshared_key": i["preshared_key"],
+                        "advanced_security": i['advanced_security']
+                    }
+                    conn.execute(
+                        self.peersTable.insert().values(newPeer)
+                    )
+                    # sqlUpdate(
+                    #     """
+                    #     INSERT INTO '%s'
+                    #         VALUES (:id, :private_key, :DNS, :advanced_security, :endpoint_allowed_ip, :name, :total_receive, :total_sent, 
+                    #         :total_data, :endpoint, :status, :latest_handshake, :allowed_ip, :cumu_receive, :cumu_sent, 
+                    #         :cumu_data, :mtu, :keepalive, :remote_endpoint, :preshared_key);
+                    #     """ % self.Name
+                    #     , newPeer)
+            for p in peers:
+                presharedKeyExist = len(p['preshared_key']) > 0
+                rd = random.Random()
+                uid = str(uuid.UUID(int=rd.getrandbits(128), version=4))
+                if presharedKeyExist:
+                    with open(uid, "w+") as f:
+                        f.write(p['preshared_key'])
+
+                subprocess.check_output(
+                    f"{self.Protocol} set {self.Name} peer {p['id']} allowed-ips {p['allowed_ip'].replace(' ', '')}{f' preshared-key {uid}' if presharedKeyExist else ''}",
+                    shell=True, stderr=subprocess.STDOUT)
+                if presharedKeyExist:
+                    os.remove(uid)
+            subprocess.check_output(
+                f"{self.Protocol}-quick save {self.Name}", shell=True, stderr=subprocess.STDOUT)
+            self.getPeersList()
+            for p in peers:
+                p = self.searchPeer(p['id'])
+                if p[0]:
+                    result['peers'].append(p[1])
+            return True, result
+        except Exception as e:
+            result['message'] = str(e)
+            return False, result
+
+    def getRestrictedPeers(self):
+        self.RestrictedPeers = []
+        # restricted = sqlSelect("SELECT * FROM '%s_restrict_access'" % self.Name).fetchall()
+        with self.engine.connect() as conn:
+            restricted = conn.execute(self.peersRestrictedTable.select()).mappings().fetchall()
+            for i in restricted:
+                self.RestrictedPeers.append(AmneziaWGPeer(i, self))

src/modules/DashboardLogger.py

@@ -1,16 +1,13 @@
 """
 Dashboard Logger Class
 """
-import os, uuid
+import uuid
 import sqlalchemy as db
-from datetime import datetime
+
-from sqlalchemy_utils import database_exists, create_database
 
 class DashboardLogger:
     def __init__(self, DashboardConfig):
         self.engine = db.create_engine(DashboardConfig.getConnectionString("wgdashboard_log"))
-        if not database_exists(self.engine.url):
-            create_database(self.engine.url)
         self.metadata = db.MetaData()
         self.dashboardLoggerTable = db.Table('DashboardLog', self.metadata,
                                              db.Column('LogID', db.VARCHAR, nullable=False, primary_key=True),

src/modules/Peer.py

@@ -0,0 +1,203 @@
+"""
+Peer
+"""
+import os, subprocess, uuid, random, re
+from .PeerJob import PeerJob
+from .PeerShareLink import PeerShareLink
+from .Utilities import GenerateWireguardPublicKey, ValidateIPAddressesWithRange, ValidateDNSAddress
+
+
+class Peer:
+    def __init__(self, tableData, configuration):
+        self.configuration = configuration
+        self.id = tableData["id"]
+        self.private_key = tableData["private_key"]
+        self.DNS = tableData["DNS"]
+        self.endpoint_allowed_ip = tableData["endpoint_allowed_ip"]
+        self.name = tableData["name"]
+        self.total_receive = tableData["total_receive"]
+        self.total_sent = tableData["total_sent"]
+        self.total_data = tableData["total_data"]
+        self.endpoint = tableData["endpoint"]
+        self.status = tableData["status"]
+        self.latest_handshake = tableData["latest_handshake"]
+        self.allowed_ip = tableData["allowed_ip"]
+        self.cumu_receive = tableData["cumu_receive"]
+        self.cumu_sent = tableData["cumu_sent"]
+        self.cumu_data = tableData["cumu_data"]
+        self.mtu = tableData["mtu"]
+        self.keepalive = tableData["keepalive"]
+        self.remote_endpoint = tableData["remote_endpoint"]
+        self.preshared_key = tableData["preshared_key"]
+        self.jobs: list[PeerJob] = []
+        self.ShareLink: list[PeerShareLink] = []
+        self.getJobs()
+        self.getShareLink()
+
+    def toJson(self):
+        self.getJobs()
+        self.getShareLink()
+        return self.__dict__
+
+    def __repr__(self):
+        return str(self.toJson())
+
+    def updatePeer(self, name: str, private_key: str,
+                   preshared_key: str,
+                   dns_addresses: str, allowed_ip: str, endpoint_allowed_ip: str, mtu: int,
+                   keepalive: int) -> tuple[bool, str] or tuple[bool, None]:
+        if not self.configuration.getStatus():
+            self.configuration.toggleConfiguration()
+
+        existingAllowedIps = [item for row in list(
+            map(lambda x: [q.strip() for q in x.split(',')],
+                map(lambda y: y.allowed_ip,
+                    list(filter(lambda k: k.id != self.id, self.configuration.getPeersList()))))) for item in row]
+
+        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"
+        if mtu < 0 or mtu > 1460:
+            return False, "MTU format is not correct"
+        if keepalive < 0:
+            return False, "Persistent Keepalive format is not correct"
+        if len(private_key) > 0:
+            pubKey = GenerateWireguardPublicKey(private_key)
+            if not pubKey[0] or pubKey[1] != self.id:
+                return False, "Private key does not match with the public key"
+        try:
+            rd = random.Random()
+            uid = str(uuid.UUID(int=rd.getrandbits(128), version=4))
+            pskExist = len(preshared_key) > 0
+
+            if pskExist:
+                with open(uid, "w+") as f:
+                    f.write(preshared_key)
+            newAllowedIPs = allowed_ip.replace(" ", "")
+            updateAllowedIp = subprocess.check_output(
+                f"{self.configuration.Protocol} set {self.configuration.Name} peer {self.id} allowed-ips {newAllowedIPs} {f'preshared-key {uid}' if pskExist else 'preshared-key /dev/null'}",
+                shell=True, stderr=subprocess.STDOUT)
+
+            if pskExist: os.remove(uid)
+            if len(updateAllowedIp.decode().strip("\n")) != 0:
+                return False, "Update peer failed when updating Allowed IPs"
+            saveConfig = subprocess.check_output(f"{self.configuration.Protocol}-quick save {self.configuration.Name}",
+                                                 shell=True, stderr=subprocess.STDOUT)
+            if f"wg showconf {self.configuration.Name}" not in saveConfig.decode().strip('\n'):
+                return False, "Update peer failed when saving the configuration"
+            with self.configuration.engine.begin() as conn:
+                conn.execute(
+                    self.configuration.peersTable.update().values({
+                        "name": name,
+                        "private_key": private_key,
+                        "DNS": dns_addresses,
+                        "endpoint_allowed_ip": endpoint_allowed_ip,
+                        "mtu": mtu,
+                        "keepalive": keepalive,
+                        "preshared_key": preshared_key
+                    }).where(
+                        self.configuration.peersTable.c.id == self.id
+                    )
+                )
+            return True, None
+        except subprocess.CalledProcessError as exc:
+            return False, exc.output.decode("UTF-8").strip()
+
+    def downloadPeer(self) -> dict[str, str]:
+        filename = self.name
+        if len(filename) == 0:
+            filename = "UntitledPeer"
+        filename = "".join(filename.split(' '))
+        filename = f"{filename}"
+        illegal_filename = [".", ",", "/", "?", "<", ">", "\\", ":", "*", '|' '\"', "com1", "com2", "com3",
+                            "com4", "com5", "com6", "com7", "com8", "com9", "lpt1", "lpt2", "lpt3", "lpt4",
+                            "lpt5", "lpt6", "lpt7", "lpt8", "lpt9", "con", "nul", "prn"]
+        for i in illegal_filename:
+            filename = filename.replace(i, "")
+
+        finalFilename = ""
+        for i in filename:
+            if re.match("^[a-zA-Z0-9_=+.-]$", i):
+                finalFilename += i
+
+        peerConfiguration = f'''[Interface]
+PrivateKey = {self.private_key}
+Address = {self.allowed_ip}
+MTU = {str(self.mtu)}
+'''
+        if len(self.DNS) > 0:
+            peerConfiguration += f"DNS = {self.DNS}\n"
+
+        peerConfiguration += f'''
+[Peer]
+PublicKey = {self.configuration.PublicKey}
+AllowedIPs = {self.endpoint_allowed_ip}
+Endpoint = {self.configuration.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1]}:{self.configuration.ListenPort}
+PersistentKeepalive = {str(self.keepalive)}
+'''
+        if len(self.preshared_key) > 0:
+            peerConfiguration += f"PresharedKey = {self.preshared_key}\n"
+        return {
+            "fileName": finalFilename,
+            "file": peerConfiguration
+        }
+
+    def getJobs(self):
+        self.jobs = self.configuration.AllPeerJobs.searchJob(self.configuration.Name, self.id)
+
+    def getShareLink(self):
+        self.ShareLink = self.configuration.AllPeerShareLinks.getLink(self.configuration.Name, self.id)
+
+    def resetDataUsage(self, type):
+        try:
+            with self.configuration.engine.begin() as conn:
+                if type == "total":
+                    conn.execute(
+                        self.configuration.peersTable.update().values({
+                            "total_data": 0,
+                            "cumu_data": 0,
+                            "total_receive": 0,
+                            "cumu_receive": 0,
+                            "total_sent": 0,
+                            "cumu_sent": 0
+                        }).where(
+                            self.configuration.peersTable.c.id == self.id
+                        )
+                    )
+                    self.total_data = 0
+                    self.total_receive = 0
+                    self.total_sent = 0
+                    self.cumu_data = 0
+                    self.cumu_sent = 0
+                    self.cumu_receive = 0
+                elif type == "receive":
+                    conn.execute(
+                        self.configuration.peersTable.update().values({
+                            "total_receive": 0,
+                            "cumu_receive": 0,
+                        }).where(
+                            self.configuration.peersTable.c.id == self.id
+                        )
+                    )
+                    self.cumu_receive = 0
+                    self.total_receive = 0
+                elif type == "sent":
+                    conn.execute(
+                        self.configuration.peersTable.update().values({
+                            "total_sent": 0,
+                            "cumu_sent": 0
+                        }).where(
+                            self.configuration.peersTable.c.id == self.id
+                        )
+                    )
+                    self.cumu_sent = 0
+                    self.total_sent = 0
+                else:
+                    return False
+        except Exception as e:
+            print(e)
+            return False
+        return True