jumpserver/apps/settings/ws.py

# -*- coding: utf-8 -*-
#
import json
import asyncio

from channels.generic.websocket import AsyncJsonWebsocketConsumer
from django.core.cache import cache
from django.conf import settings

from common.db.utils import close_old_connections
from common.utils import get_logger
from settings.serializers import (
    LDAPTestConfigSerializer,
    LDAPTestLoginSerializer
)
from settings.tasks import sync_ldap_user
from settings.utils import (
    LDAPSyncUtil, LDAPTestUtil
)
from .tools import (
    verbose_ping, verbose_telnet, verbose_nmap,
    verbose_tcpdump, verbose_traceroute
)

logger = get_logger(__name__)

CACHE_KEY_LDAP_TEST_CONFIG_MSG = 'CACHE_KEY_LDAP_TEST_CONFIG_MSG'
CACHE_KEY_LDAP_TEST_LOGIN_MSG = 'CACHE_KEY_LDAP_TEST_LOGIN_MSG'
CACHE_KEY_LDAP_SYNC_USER_MSG = 'CACHE_KEY_LDAP_SYNC_USER_MSG'
CACHE_KEY_LDAP_TEST_CONFIG_TASK_STATUS = 'CACHE_KEY_LDAP_TEST_CONFIG_TASK_STATUS'
CACHE_KEY_LDAP_TEST_LOGIN_TASK_STATUS = 'CACHE_KEY_LDAP_TEST_LOGIN_TASK_STATUS'
CACHE_KEY_LDAP_SYNC_USER_TASK_STATUS = 'CACHE_KEY_LDAP_SYNC_USER_TASK_STATUS'
TASK_STATUS_IS_RUNNING = 'RUNNING'
TASK_STATUS_IS_OVER = 'OVER'


class ToolsWebsocket(AsyncJsonWebsocketConsumer):

    async def connect(self):
        user = self.scope["user"]
        if user.is_authenticated:
            await self.accept()
        else:
            await self.close()

    async def send_msg(self, msg=''):
        await self.send_json({'msg': f'{msg}\r\n'})

    async def imitate_ping(self, dest_ips, timeout=3, count=5, psize=64):
        params = {
            'dest_ips': dest_ips, 'timeout': timeout,
            'count': count, 'psize': psize
        }
        logger.info(f'Receive request ping: {params}')
        await verbose_ping(display=self.send_msg, **params)

    async def imitate_telnet(self, dest_ips, dest_port=23, timeout=10):
        params = {
            'dest_ips': dest_ips, 'dest_port': dest_port, 'timeout': timeout,
        }
        logger.info(f'Receive request telnet: {params}')
        await verbose_telnet(display=self.send_msg, **params)

    async def imitate_nmap(self, dest_ips, dest_ports=None, timeout=None):
        params = {
            'dest_ips': dest_ips, 'dest_ports': dest_ports, 'timeout': timeout,
        }
        logger.info(f'Receive request nmap: {params}')
        await verbose_nmap(display=self.send_msg, **params)

    async def imitate_tcpdump(
            self, interfaces=None, src_ips='',
            src_ports='', dest_ips='', dest_ports=''
    ):
        params = {
            'interfaces': interfaces, 'src_ips': src_ips, 'src_ports': src_ports,
            'dest_ips': dest_ips, 'dest_ports': dest_ports
        }
        logger.info(f'Receive request tcpdump: {params}')
        await verbose_tcpdump(display=self.send_msg, **params)

    async def imitate_traceroute(self, dest_ips):
        params = {'dest_ips': dest_ips}
        await verbose_traceroute(display=self.send_msg, **params)

    async def receive(self, text_data=None, bytes_data=None, **kwargs):
        data = json.loads(text_data)
        tool_type = data.pop('tool_type', 'Ping')
        try:
            tool_func = getattr(self, f'imitate_{tool_type.lower()}')
            await tool_func(**data)
        except Exception as error:
            await self.send_msg('Exception: %s' % error)
        await self.send_msg()
        await self.close()

    async def disconnect(self, code):
        await self.close()
        close_old_connections()


class LdapWebsocket(AsyncJsonWebsocketConsumer):
    async def connect(self):
        user = self.scope["user"]
        if user.is_authenticated:
            await self.accept()
        else:
            await self.close()

    async def receive(self, text_data=None, bytes_data=None, **kwargs):
        data = json.loads(text_data)
        msg_type = data.pop('msg_type', 'testing_config')
        try:
            tool_func = getattr(self, f'run_{msg_type.lower()}')
            await asyncio.to_thread(tool_func, data)
            if msg_type == 'testing_config':
                ok, msg = cache.get(CACHE_KEY_LDAP_TEST_CONFIG_MSG)
            elif msg_type == 'sync_user':
                ok, msg = cache.get(CACHE_KEY_LDAP_SYNC_USER_MSG)
            else:
                ok, msg = cache.get(CACHE_KEY_LDAP_TEST_LOGIN_MSG)
            await self.send_msg(ok, msg)
        except Exception as error:
            await self.send_msg(msg='Exception: %s' % error)

    async def send_msg(self, ok=True, msg=''):
        await self.send_json({'ok': ok, 'msg': f'{msg}'})

    async def disconnect(self, code):
        await self.close()
        close_old_connections()

    @staticmethod
    def get_ldap_config(serializer):
        server_uri = serializer.validated_data["AUTH_LDAP_SERVER_URI"]
        bind_dn = serializer.validated_data["AUTH_LDAP_BIND_DN"]
        password = serializer.validated_data["AUTH_LDAP_BIND_PASSWORD"]
        use_ssl = serializer.validated_data.get("AUTH_LDAP_START_TLS", False)
        search_ou = serializer.validated_data["AUTH_LDAP_SEARCH_OU"]
        search_filter = serializer.validated_data["AUTH_LDAP_SEARCH_FILTER"]
        attr_map = serializer.validated_data["AUTH_LDAP_USER_ATTR_MAP"]
        auth_ldap = serializer.validated_data.get('AUTH_LDAP', False)

        if not password:
            password = settings.AUTH_LDAP_BIND_PASSWORD

        config = {
            'server_uri': server_uri,
            'bind_dn': bind_dn,
            'password': password,
            'use_ssl': use_ssl,
            'search_ou': search_ou,
            'search_filter': search_filter,
            'attr_map': attr_map,
            'auth_ldap': auth_ldap
        }
        return config

    @staticmethod
    def task_is_over(task_key):
        return cache.get(task_key) == TASK_STATUS_IS_OVER

    @staticmethod
    def set_task_status_over(task_key, ttl=120):
        cache.set(task_key, TASK_STATUS_IS_OVER, ttl)

    @staticmethod
    def set_task_msg(task_key, ok, msg):
        cache.set(task_key, (ok, msg), 120)

    def run_testing_config(self, data):
        while True:
            if self.task_is_over(CACHE_KEY_LDAP_TEST_CONFIG_TASK_STATUS):
                break
            else:
                serializer = LDAPTestConfigSerializer(data=data)
                if not serializer.is_valid():
                    self.send_msg(msg=f'error: {str(serializer.errors)}')
                config = self.get_ldap_config(serializer)
                ok, msg = LDAPTestUtil(config).test_config()
                self.set_task_status_over(CACHE_KEY_LDAP_TEST_CONFIG_TASK_STATUS)
            self.set_task_msg(CACHE_KEY_LDAP_TEST_CONFIG_MSG, ok, msg)

    def run_testing_login(self, data):
        while True:
            if self.task_is_over(CACHE_KEY_LDAP_TEST_LOGIN_TASK_STATUS):
                break
            else:
                serializer = LDAPTestLoginSerializer(data=data)
                if not serializer.is_valid():
                    self.send_msg(msg=f'error: {str(serializer.errors)}')
                username = serializer.validated_data['username']
                password = serializer.validated_data['password']
                ok, msg = LDAPTestUtil().test_login(username, password)
                self.set_task_status_over(CACHE_KEY_LDAP_TEST_LOGIN_TASK_STATUS, 3)
            self.set_task_msg(CACHE_KEY_LDAP_TEST_LOGIN_MSG, ok, msg)

    def run_sync_user(self, data):
        while True:
            if self.task_is_over(CACHE_KEY_LDAP_SYNC_USER_TASK_STATUS):
                break
            else:
                sync_util = LDAPSyncUtil()
                sync_util.clear_cache()
                sync_ldap_user()
                msg = sync_util.get_task_error_msg()
                ok = False if msg else True
                self.set_task_status_over(CACHE_KEY_LDAP_SYNC_USER_TASK_STATUS)
            self.set_task_msg(CACHE_KEY_LDAP_SYNC_USER_MSG, ok, msg)