WGDashboard/src/modules/DashboardClientsPeerAssignment.py

import uuid

from .ConnectionString import ConnectionString
from .DashboardLogger import DashboardLogger
import sqlalchemy as db

from .WireguardConfiguration import WireguardConfiguration


class DashboardClientsPeerAssignment:
    def __init__(self, wireguardConfigurations: dict[str, WireguardConfiguration]):
        self.logger = DashboardLogger()
        self.engine = db.create_engine(ConnectionString("wgdashboard"))
        self.metadata = db.MetaData()
        self.wireguardConfigurations = wireguardConfigurations
        self.dashboardClientsPeerAssignmentTable = db.Table(
            'DashboardClientsPeerAssignment', self.metadata,
            db.Column('AssignmentID', db.String(255), nullable=False, primary_key=True),
            db.Column('ClientID', db.String(255), nullable=False, index=True),
            db.Column('ConfigurationName', db.String(255)),
            db.Column('PeerID', db.String(500)),
            db.Column('AssignedDate',
                      (db.DATETIME if 'sqlite:///' in ConnectionString("wgdashboard") else db.TIMESTAMP),
                      server_default=db.func.now()),
            db.Column('UnassignedDate',
                      (db.DATETIME if 'sqlite:///' in ConnectionString("wgdashboard") else db.TIMESTAMP)),
            extend_existing=True
        )
        self.metadata.create_all(self.engine)
        self.assignments = []
        self.__getAssignments()
    
        self.AssignClient("0117a895-bd8b-4ba2-9116-6658372417fb", "wg0", "3kv6Bo46u7ULT07B3I1VHw/rYomVnrCD5TFU369jRSc=")
        self.GetAssignedPeers("0117a895-bd8b-4ba2-9116-6658372417fb")
        
    def __getAssignments(self):
        with self.engine.connect() as conn:
            self.assignments = conn.execute(
                self.dashboardClientsPeerAssignmentTable.select().where(
                    self.dashboardClientsPeerAssignmentTable.c.UnassignedDate == db.null()
                )
            ).mappings().fetchall()
            
    def AssignClient(self, ClientID, ConfigurationName, PeerID):
        existing = list(
            filter(lambda e: 
                   e['ClientID'] == ClientID and 
                   e['ConfigurationName'] == ConfigurationName and
                   e['PeerID'] == PeerID, self.assignments)
        )
        if len(existing) == 0:
            if ConfigurationName in self.wireguardConfigurations.keys():
                config = self.wireguardConfigurations.get(ConfigurationName)
                peer = list(filter(lambda x : x.id == PeerID, config.Peers))
                if len(peer) == 1:
                    with self.engine.begin() as conn:
                        data = {
                            "AssignmentID": uuid.uuid4(),
                            "ClientID": ClientID,
                            "ConfigurationName": ConfigurationName,
                            "PeerID": PeerID
                        }
                        conn.execute(
                            self.dashboardClientsPeerAssignmentTable.insert().values(data)
                        )
                    self.__getAssignments()
                    return True, data
        return False, None
    
    def UnassignClient(self, AssignmentID):
        pass
    
    def GetAssignedClient(self, ConfigurationName, PeerID):
        pass
    
    def GetAssignedPeers(self, ClientID):
        self.__getAssignments()
        
        peers = []
        assigned = filter(lambda e:
                          e['ClientID'] == ClientID, self.assignments)
        
        for a in assigned:
            peer = filter(lambda e : e.id == a['PeerID'], 
                          self.wireguardConfigurations[a['ConfigurationName']].Peers)
            for p in peer:
                peers.append({
                    'protocol': self.wireguardConfigurations[a['ConfigurationName']].Protocol,
                    'id': p.id,
                    'private_key': p.private_key,
                    'name': p.name,
                    'received_data': p.total_receive + p.cumu_receive,
                    'sent_data': p.total_sent + p.cumu_sent,
                    'data': p.total_data + p.cumu_data,
                    'status': p.status,
                    'latest_handshake': p.latest_handshake,
                    'allowed_ip': p.allowed_ip,
                    'jobs': p.jobs,
                    'configuration_name': a['ConfigurationName'],
                    'peer_configuration_data': p.downloadPeer()
                })
        return peers