wg-manager/wg_connection_manager_worker/tasks.py

import environ
from celery import Celery
from pywireguard.factory import Peer
from .wg_manager import WGManager


env = environ.Env(WG_INTERFACE=(str, "wg0"), CELERY_BROKER=(str, "redis://localhost"), CELERY_BACKEND=(str, "redis://localhost"))
environ.Env.read_env(".env")
WG_INTERFACE = env("WG_INTERFACE")
CELERY_BROKER = env("CELERY_BROKER")
CELERY_BACKEND = env("CELERY_BACKEND")
app = Celery("wg_manager_tasks", broker=CELERY_BROKER, backend=CELERY_BACKEND)
# app.conf.event_serializer = "pickle"  # this event_serializer is optional. somehow i missed this when writing this solution and it still worked without.
# app.conf.task_serializer = "pickle"
# app.conf.result_serializer = "pickle"
# app.conf.accept_content = ["application/json", "application/x-python-serialize"]


@app.task()
def add_peer(peer: dict):
    p = Peer(**peer)
    wg = WGManager(WG_INTERFACE)
    wg.add_peer(p)


@app.task()
def remove_peer(peer: dict):
    p = Peer(**peer)
    wg = WGManager(WG_INTERFACE)
    wg.remove_peer(p)


@app.task
def get_peers() -> list[Peer]:
    wg = WGManager(WG_INTERFACE)
    return [{"public_key": x.public_key, "preshared_key": x.preshared_key, "allowed_ips": x.allowed_ips} for x in wg.get_peers()]


@app.task
def ping() -> str:
    return "pong"