This repository has been archived on 2024-05-31. You can view files and clone it, but cannot push or open issues or pull requests.
authentik/passbook/outposts/channels.py

89 lines
2.8 KiB
Python
Raw Normal View History

2020-09-02 22:04:12 +00:00
"""Outpost websocket handler"""
from dataclasses import asdict, dataclass, field
from enum import IntEnum
from time import time
from typing import Any, Dict
from dacite import from_dict
from dacite.data import Data
from django.core.cache import cache
from guardian.shortcuts import get_objects_for_user
2020-09-02 22:04:12 +00:00
from structlog import get_logger
from passbook.core.channels import AuthJsonConsumer
2020-09-02 22:04:12 +00:00
from passbook.outposts.models import Outpost
LOGGER = get_logger()
class WebsocketMessageInstruction(IntEnum):
"""Commands which can be triggered over Websocket"""
# Simple message used by either side when a message is acknowledged
ACK = 0
# Message used by outposts to report their alive status
HELLO = 1
# Message sent by us to trigger an Update
TRIGGER_UPDATE = 2
@dataclass
class WebsocketMessage:
"""Complete Websocket Message that is being sent"""
instruction: int
args: Dict[str, Any] = field(default_factory=dict)
class OutpostConsumer(AuthJsonConsumer):
2020-09-02 22:04:12 +00:00
"""Handler for Outposts that connect over websockets for health checks and live updates"""
outpost: Outpost
def connect(self):
if not super().connect():
return
2020-09-02 22:04:12 +00:00
uuid = self.scope["url_route"]["kwargs"]["pk"]
outpost = get_objects_for_user(
self.user, "passbook_outposts.view_outpost"
).filter(pk=uuid)
2020-09-02 22:04:12 +00:00
if not outpost.exists():
self.close()
return
self.accept()
self.outpost = outpost.first()
self.outpost.channels.append(self.channel_name)
LOGGER.debug("added channel to outpost", channel_name=self.channel_name)
self.outpost.save()
# pylint: disable=unused-argument
def disconnect(self, close_code):
self.outpost.channels.remove(self.channel_name)
self.outpost.save()
LOGGER.debug("removed channel from outpost", channel_name=self.channel_name)
def receive_json(self, content: Data):
msg = from_dict(WebsocketMessage, content)
if msg.instruction == WebsocketMessageInstruction.HELLO:
cache.set(self.outpost.state_cache_prefix("health"), time(), timeout=60)
if "version" in msg.args:
cache.set(
self.outpost.state_cache_prefix("version"), msg.args["version"]
)
2020-09-02 22:04:12 +00:00
elif msg.instruction == WebsocketMessageInstruction.ACK:
return
response = WebsocketMessage(instruction=WebsocketMessageInstruction.ACK)
self.send_json(asdict(response))
# pylint: disable=unused-argument
def event_update(self, event):
"""Event handler which is called by post_save signals"""
self.send_json(
asdict(
WebsocketMessage(instruction=WebsocketMessageInstruction.TRIGGER_UPDATE)
)
)