wg-manager/wg_connection_manager/dj_wg_manager.py
2024-10-06 12:56:12 +02:00

45 lines
1.7 KiB
Python

from .wg_manager import WGManager
from django.conf import settings
from django.db import transaction
from .models import UserConnection
from pywireguard.factory import Peer
class DJWGManager:
wg_manager: WGManager
def __init__(self):
self.wg_manager = WGManager(settings.WG_INTERFACE)
def sync(self):
with transaction.atomic():
UserConnection.objects.filter(active=True).update(active=False)
for peer in self.wg_manager.get_peers():
pk = peer.public_key.decode("ascii")
psk = peer.preshared_key.decode("ascii")
if not peer.allowed_ips:
continue
connection = UserConnection.objects.filter(public_key=pk)
if len(connection) == 1:
connection = connection[0]
else:
connection = UserConnection()
connection.public_key = pk
connection.preshared_key = psk
connection.active = True
connection.vpn_ip = peer.allowed_ips[0]
connection.save()
def add_peer(self, user_connection: UserConnection):
self.wg_manager.add_peer(Peer(public_key=user_connection.public_key, preshared_key=user_connection.preshared_key, allowed_ips=[user_connection.vpn_ip]))
self.sync()
def remove_peer(self, user_connection: UserConnection):
pk = user_connection.public_key
peer = list(filter(lambda x: x.public_key.decode("ascii") == pk, self.wg_manager.get_peers()))
if not peer:
# TODO raise exception/ignore?
return
peer = peer[0]
self.wg_manager.remove_peer(peer)
self.sync()