refac: some WGDashboard code

src/modules/WireguardConfiguration.py

@@ -10,12 +10,15 @@ from datetime import datetime, timedelta
 from itertools import islice
 from flask import current_app
 
-from .ConnectionString import ConnectionString
+from .DatabaseConnection import ConnectionString
 from .DashboardConfig import DashboardConfig
 from .Peer import Peer
 from .PeerJobs import PeerJobs
 from .PeerShareLinks import PeerShareLinks
-from .Utilities import StringToBoolean, GenerateWireguardPublicKey, RegexMatch, ValidateDNSAddress, \
+from .Utilities import StringToBoolean, \
+    GenerateWireguardPublicKey, \
+    RegexMatch, \
+    ValidateDNSAddress, \
     ValidateEndpointAllowedIPs
 from .WireguardConfigurationInfo import WireguardConfigurationInfo, PeerGroupsClass
 from .DashboardWebHooks import DashboardWebHooks
@@ -61,13 +64,14 @@ class WireguardConfiguration:
         self.Protocol = "wg" if wg else "awg"
         self.AllPeerJobs = AllPeerJobs
         self.DashboardConfig = DashboardConfig
+        self.DashboardConfig.EnsureDatabaseIntegrity({self.Name: self})
         self.AllPeerShareLinks = AllPeerShareLinks
         self.DashboardWebHooks = DashboardWebHooks
         self.configPath = os.path.join(self.__getProtocolPath(), f'{self.Name}.conf')
         self.engine: sqlalchemy.Engine = sqlalchemy.create_engine(ConnectionString("wgdashboard"))
         self.metadata: sqlalchemy.MetaData = sqlalchemy.MetaData()
         self.dbType = self.DashboardConfig.GetConfig("Database", "type")[1]
-        
+
         if name is not None:
             if data is not None and "Backup" in data.keys():
                 db = self.__importDatabase(
@@ -150,7 +154,6 @@ class WireguardConfiguration:
         
         if self.Status:
             self.addAutostart()
-        
 
     def __getProtocolPath(self) -> str:
         _, path = self.DashboardConfig.GetConfig("Server", "wg_conf_path") if self.Protocol == "wg" \
@@ -242,54 +245,50 @@ class WireguardConfiguration:
         return True
 
     def createDatabase(self, dbName = None):
+        def generate_column_obj():
+            return [
+                sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False, primary_key=True),
+                sqlalchemy.Column('private_key', sqlalchemy.String(255)),
+                sqlalchemy.Column('DNS', sqlalchemy.Text),
+                sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
+                sqlalchemy.Column('name', sqlalchemy.Text),
+                sqlalchemy.Column('total_receive', sqlalchemy.Float),
+                sqlalchemy.Column('total_sent', sqlalchemy.Float),
+                sqlalchemy.Column('total_data', sqlalchemy.Float),
+                sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
+                sqlalchemy.Column('status', sqlalchemy.String(255)),
+                sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
+                sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
+                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('notes', sqlalchemy.Text),
+                sqlalchemy.Column('remote_endpoint', sqlalchemy.String(255)),
+                sqlalchemy.Column('preshared_key', sqlalchemy.String(255))
+            ]
+
         if dbName is None:
             dbName = self.Name
+
         self.peersTable = sqlalchemy.Table(
-            dbName, self.metadata,
-            sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False, primary_key=True),
-            sqlalchemy.Column('private_key', sqlalchemy.String(255)),
-            sqlalchemy.Column('DNS', sqlalchemy.Text),
-            sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
-            sqlalchemy.Column('name', sqlalchemy.Text),
-            sqlalchemy.Column('total_receive', sqlalchemy.Float),
-            sqlalchemy.Column('total_sent', sqlalchemy.Float),
-            sqlalchemy.Column('total_data', sqlalchemy.Float),
-            sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
-            sqlalchemy.Column('status', sqlalchemy.String(255)),
-            sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
-            sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
-            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(255)),
-            sqlalchemy.Column('preshared_key', sqlalchemy.String(255)),
-            extend_existing=True
+            f'{dbName}', self.metadata, *generate_column_obj(), extend_existing=True
         )
+
         self.peersRestrictedTable = sqlalchemy.Table(
-            f'{dbName}_restrict_access', self.metadata,
-            sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False, primary_key=True),
-            sqlalchemy.Column('private_key', sqlalchemy.String(255)),
-            sqlalchemy.Column('DNS', sqlalchemy.Text),
-            sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
-            sqlalchemy.Column('name', sqlalchemy.Text),
-            sqlalchemy.Column('total_receive', sqlalchemy.Float),
-            sqlalchemy.Column('total_sent', sqlalchemy.Float),
-            sqlalchemy.Column('total_data', sqlalchemy.Float),
-            sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
-            sqlalchemy.Column('status', sqlalchemy.String(255)),
-            sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
-            sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
-            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(255)),
-            sqlalchemy.Column('preshared_key', sqlalchemy.String(255)),
-            extend_existing=True
+            f'{dbName}_restrict_access', self.metadata, *generate_column_obj(), extend_existing=True
         )
+
+        self.peersDeletedTable = sqlalchemy.Table(
+            f'{dbName}_deleted', self.metadata, *generate_column_obj(), extend_existing=True
+        )
+
+        if self.DashboardConfig.GetConfig("Database", "type")[1] == 'sqlite':
+            time_col_type = sqlalchemy.DATETIME
+        else:
+            time_col_type = sqlalchemy.TIMESTAMP
+
         self.peersTransferTable = sqlalchemy.Table(
             f'{dbName}_transfer', self.metadata,
             sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False),
@@ -299,8 +298,7 @@ class WireguardConfiguration:
             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()),
+            sqlalchemy.Column('time', time_col_type, server_default=sqlalchemy.func.now()),
             extend_existing=True
         )
         
@@ -308,34 +306,9 @@ class WireguardConfiguration:
             f'{dbName}_history_endpoint', self.metadata,
             sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False),
             sqlalchemy.Column('endpoint', sqlalchemy.String(255), nullable=False),
-            sqlalchemy.Column('time', 
-                              (sqlalchemy.DATETIME if self.DashboardConfig.GetConfig("Database", "type")[1] == 'sqlite' else sqlalchemy.TIMESTAMP)),
-            extend_existing=True
+            sqlalchemy.Column('time', time_col_type)
         )
         
-        self.peersDeletedTable = sqlalchemy.Table(
-            f'{dbName}_deleted', self.metadata,
-            sqlalchemy.Column('id', sqlalchemy.String(255), nullable=False, primary_key=True),
-            sqlalchemy.Column('private_key', sqlalchemy.String(255)),
-            sqlalchemy.Column('DNS', sqlalchemy.Text),
-            sqlalchemy.Column('endpoint_allowed_ip', sqlalchemy.Text),
-            sqlalchemy.Column('name', sqlalchemy.Text),
-            sqlalchemy.Column('total_receive', sqlalchemy.Float),
-            sqlalchemy.Column('total_sent', sqlalchemy.Float),
-            sqlalchemy.Column('total_data', sqlalchemy.Float),
-            sqlalchemy.Column('endpoint', sqlalchemy.String(255)),
-            sqlalchemy.Column('status', sqlalchemy.String(255)),
-            sqlalchemy.Column('latest_handshake', sqlalchemy.String(255)),
-            sqlalchemy.Column('allowed_ip', sqlalchemy.String(255)),
-            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(255)),
-            sqlalchemy.Column('preshared_key', sqlalchemy.String(255)),
-            extend_existing=True
-        )
         self.infoTable = sqlalchemy.Table(
             'ConfigurationsInfo', self.metadata,
             sqlalchemy.Column('ID', sqlalchemy.String(255), primary_key=True),
@@ -450,8 +423,7 @@ class WireguardConfiguration:
                                     "id": i['PublicKey'],
                                     "private_key": "",
                                     "DNS": self.DashboardConfig.GetConfig("Peers", "peer_global_DNS")[1],
-                                    "endpoint_allowed_ip": self.DashboardConfig.GetConfig("Peers", "peer_endpoint_allowed_ip")[
-                                        1],
+                                    "endpoint_allowed_ip": self.DashboardConfig.GetConfig("Peers", "peer_endpoint_allowed_ip")[1],
                                     "name": i.get("name"),
                                     "total_receive": 0,
                                     "total_sent": 0,
@@ -465,6 +437,7 @@ class WireguardConfiguration:
                                     "cumu_data": 0,
                                     "mtu": self.DashboardConfig.GetConfig("Peers", "peer_mtu")[1] if len(self.DashboardConfig.GetConfig("Peers", "peer_mtu")[1]) > 0 else None,
                                     "keepalive": self.DashboardConfig.GetConfig("Peers", "peer_keep_alive")[1] if len(self.DashboardConfig.GetConfig("Peers", "peer_keep_alive")[1]) > 0 else None,
+                                    "notes": "",
                                     "remote_endpoint": self.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1],
                                     "preshared_key": i["PresharedKey"] if "PresharedKey" in i.keys() else ""
                                 }
@@ -557,6 +530,7 @@ class WireguardConfiguration:
                         "cumu_data": 0,
                         "mtu": i['mtu'],
                         "keepalive": i['keepalive'],
+                        "notes": i.get("notes", ""),
                         "remote_endpoint": self.DashboardConfig.GetConfig("Peers", "remote_endpoint")[1],
                         "preshared_key": i["preshared_key"]
                     }