45 lines
1.7 KiB
Python
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()
|