jumpserver/apps/terminal/ws.py

import datetime

from channels.generic.websocket import JsonWebsocketConsumer
from django.utils import timezone
from rest_framework.renderers import JSONRenderer

from common.db.utils import safe_db_connection
from common.utils import get_logger
from common.utils.connection import Subscription
from terminal.models import Terminal
from terminal.models import Session
from terminal.serializers import TaskSerializer, StatSerializer
from .signal_handlers import component_event_chan

logger = get_logger(__name__)


class TerminalTaskWebsocket(JsonWebsocketConsumer):
    sub: Subscription = None
    terminal: Terminal = None

    def connect(self):
        user = self.scope["user"]
        if user.is_authenticated and user.terminal:
            self.accept()
            self.terminal = user.terminal
            self.sub = self.watch_component_event()
        else:
            self.close()

    def receive_json(self, content, **kwargs):
        req_type = content.get('type')
        if req_type == "status":
            payload = content.get('payload')
            self.handle_status(payload)

    def handle_status(self, content):
        serializer = StatSerializer(data=content)
        if not serializer.is_valid():
            logger.error('Invalid status data: {}'.format(serializer.errors))
            return
        serializer.validated_data["terminal"] = self.terminal
        session_ids = serializer.validated_data.pop('sessions', [])
        Session.set_sessions_active(session_ids)
        with safe_db_connection():
            serializer.save()

    def send_kill_tasks_msg(self, task_id=None):
        content = self.get_terminal_tasks(task_id)
        self.send(bytes_data=content)

    def get_terminal_tasks(self, task_id=None):
        with safe_db_connection():
            critical_time = timezone.now() - datetime.timedelta(minutes=10)
            tasks = self.terminal.task_set.filter(is_finished=False, date_created__gte=critical_time)
            if task_id:
                tasks = tasks.filter(id=task_id)
            serializer = TaskSerializer(tasks, many=True)
            return JSONRenderer().render(serializer.data)

    def watch_component_event(self):
        # 先发一次已有的任务
        self.send_kill_tasks_msg()

        ws = self

        def handle_task_msg_recv(msg):
            logger.debug('New component task msg recv: {}'.format(msg))
            msg_type = msg.get('type')
            payload = msg.get('payload')
            if msg_type == "kill_session":
                ws.send_kill_tasks_msg(payload.get('id'))

        return component_event_chan.subscribe(handle_task_msg_recv)

    def disconnect(self, code):
        if self.sub is None:
            return
        self.sub.unsubscribe()