feat: 增加 DB Listen Port 映射规则

apps/applications/api/application.py

@@ -1,15 +1,19 @@
 # coding: utf-8
 #
 from orgs.mixins.api import OrgBulkModelViewSet
+from rest_framework import generics, status
 from rest_framework.decorators import action
 from rest_framework.response import Response
+from rest_framework.viewsets import GenericViewSet
+
 
 from common.tree import TreeNodeSerializer
 from common.mixins.api import SuggestionMixin
+from ..utils import db_port_manager
 from .. import serializers
 from ..models import Application
 
-__all__ = ['ApplicationViewSet']
+__all__ = ['ApplicationViewSet', 'DBListenPortViewSet']
 
 
 class ApplicationViewSet(SuggestionMixin, OrgBulkModelViewSet):
@@ -37,3 +41,27 @@ class ApplicationViewSet(SuggestionMixin, OrgBulkModelViewSet):
         tree_nodes = Application.create_tree_nodes(queryset, show_count=show_count)
         serializer = self.get_serializer(tree_nodes, many=True)
         return Response(serializer.data)
+
+
+class DBListenPortViewSet(GenericViewSet):
+    rbac_perms = {
+        'GET': 'applications.view_application',
+        'list': 'applications.view_application',
+        'db_info': 'applications.view_application',
+    }
+
+    http_method_names = ['get', 'post']
+
+    def list(self, request, *args, **kwargs):
+        ports = db_port_manager.get_already_use_ports()
+        return Response(data=ports, status=status.HTTP_200_OK)
+
+    @action(methods=['post'], detail=False, url_path='db-info')
+    def db_info(self, request, *args, **kwargs):
+        port = request.data.get("port")
+        db, msg = db_port_manager.get_db_by_port(port)
+        if db is None:
+            data = {'error': msg}
+            return Response(data=data, status=status.HTTP_404_NOT_FOUND)
+        serializer = serializers.AppSerializer(instance=db)
+        return Response(data=serializer.data, status=status.HTTP_201_CREATED)

apps/applications/apps.py

@@ -7,3 +7,7 @@ from django.apps import AppConfig
 class ApplicationsConfig(AppConfig):
     name = 'applications'
     verbose_name = _('Applications')
+
+    def ready(self):
+        from . import signal_handlers
+        super().ready()

apps/applications/models/application.py

@@ -12,7 +12,6 @@ from common.utils import is_uuid
 from assets.models import Asset, SystemUser
 from ..const import OracleVersion
 
-from ..utils import KubernetesTree
 from .. import const
 
 
@@ -175,6 +174,7 @@ class ApplicationTreeNodeMixin:
         return pid
 
     def as_tree_node(self, pid, k8s_as_tree=False):
+        from ..utils import KubernetesTree
         if self.type == const.AppType.k8s and k8s_as_tree:
             node = KubernetesTree(pid).as_tree_node(self)
         else:

apps/applications/signal_handlers.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+#
+from django.db.models.signals import post_save, post_delete
+
+from common.signals import django_ready
+from django.dispatch import receiver
+from common.utils import get_logger
+from .models import Application
+from .utils import db_port_manager
+
+
+logger = get_logger(__file__)
+
+
+@receiver(django_ready)
+def init_db_port_mapper(sender, **kwargs):
+    logger.info('Init db port mapper')
+    db_port_manager.init()
+
+
+@receiver(post_save, sender=Application)
+def on_db_app_created(sender, instance: Application, created, **kwargs):
+    if not instance.category_db:
+        return
+    if not created:
+        return
+    db_port_manager.add(instance)
+
+
+@receiver(post_delete, sender=Application)
+def on_db_app_delete(sender, instance, **kwargs):
+    if not instance.category_db:
+        return
+    db_port_manager.pop(instance)

apps/applications/urls/api_urls.py

@@ -13,6 +13,7 @@ router.register(r'applications', api.ApplicationViewSet, 'application')
 router.register(r'accounts', api.ApplicationAccountViewSet, 'application-account')
 router.register(r'system-users-apps-relations', api.SystemUserAppRelationViewSet, 'system-users-apps-relation')
 router.register(r'account-secrets', api.ApplicationAccountSecretViewSet, 'application-account-secret')
+router.register(r'db-listen-ports', api.DBListenPortViewSet, 'db-listen-ports')
 
 
 urlpatterns = [

apps/applications/utils/__init__.py

@@ -2,3 +2,4 @@
 #
 
 from .kubernetes_util import *
+from .db_port_mapper import *

apps/applications/utils/db_port_mapper.py

@@ -0,0 +1,88 @@
+from common.decorator import Singleton
+from django.core.cache import cache
+from django.conf import settings
+from applications.const import AppCategory
+from applications.models import Application
+from common.utils import get_logger
+from common.utils import get_object_or_none
+
+
+logger = get_logger(__file__)
+
+
+@Singleton
+class DBPortManager(object):
+    """ 管理端口-数据库ID的映射, Magnus 要使用 """
+
+    CACHE_KEY = 'PORT_DB_MAPPER'
+
+    def __init__(self):
+        self.port_start = settings.MAGNUS_DB_PORTS_START
+        self.port_limit = settings.MAGNUS_DB_PORTS_LIMIT_COUNT
+        self.port_end = self.port_start + self.port_limit + 1
+        # 可以使用的端口列表
+        self.all_usable_ports = [i for i in range(self.port_start, self.port_end)]
+
+    def init(self):
+        db_ids = Application.objects.filter(category=AppCategory.db).values_list('id', flat=True)
+        db_ids = [str(i) for i in db_ids]
+        mapper = dict(zip(self.all_usable_ports, list(db_ids)))
+        self.set_mapper(mapper)
+
+    def get_db_by_port(self, port):
+        mapper = self.get_mapper()
+        db_id = mapper.get(port, None)
+        if db_id:
+            db = get_object_or_none(Application, id=db_id)
+            if not db:
+                msg = 'Database not exists, database id: {}'.format(db_id)
+            else:
+                msg = ''
+        else:
+            db = None
+            msg = 'Port not in port-db mapper, port: {}'.format(port)
+        return db, msg
+
+    def add(self, db: Application):
+        mapper = self.get_mapper()
+        usable_port = self.get_next_usable_port()
+        if not usable_port:
+            return False
+        mapper.update({usable_port: str(db.id)})
+        self.set_mapper(mapper)
+        return True
+
+    def pop(self, db: Application):
+        mapper = self.get_mapper()
+        to_delete_port = None
+        for port, db_id in mapper.items():
+            if db_id == str(db.id):
+                to_delete_port = port
+                break
+        mapper.pop(to_delete_port, None)
+        self.set_mapper(mapper)
+
+    def get_next_usable_port(self):
+        already_use_ports = self.get_already_use_ports()
+        usable_ports = list(set(self.all_usable_ports) - set(already_use_ports))
+        if len(usable_ports) > 1:
+            return usable_ports[0]
+        else:
+            already_use_ports = self.get_already_use_ports()
+            msg = 'No port is usable, All usable port count: {}, Already use port count: {}'.format(
+                len(self.all_usable_ports), len(already_use_ports)
+            )
+            logger.warning(msg)
+
+    def get_already_use_ports(self):
+        mapper = self.get_mapper()
+        return list(mapper.keys())
+
+    def get_mapper(self):
+        return cache.get(self.CACHE_KEY, {})
+
+    def set_mapper(self, value):
+        cache.set(self.CACHE_KEY, value, timeout=None)
+
+
+db_port_manager = DBPortManager()

apps/jumpserver/conf.py

@@ -484,6 +484,9 @@ class Config(dict):
         'SERVER_REPLAY_STORAGE': {},
         'SECURITY_DATA_CRYPTO_ALGO': None,
         'GMSSL_ENABLED': False,
+        # Magnus 组件需要监听的端口范围
+        'MAGNUS_DB_PORTS_START': 30000,
+        'MAGNUS_DB_PORTS_LIMIT_COUNT': 1000,
 
         # 记录清理清理
         'LOGIN_LOG_KEEP_DAYS': 200,

apps/jumpserver/settings/custom.py

@@ -177,3 +177,7 @@ HELP_SUPPORT_URL = CONFIG.HELP_SUPPORT_URL
 
 SESSION_RSA_PRIVATE_KEY_NAME = 'jms_private_key'
 SESSION_RSA_PUBLIC_KEY_NAME = 'jms_public_key'
+
+# Magnus DB Port
+MAGNUS_DB_PORTS_START = CONFIG.MAGNUS_DB_PORTS_START
+MAGNUS_DB_PORTS_LIMIT_COUNT = CONFIG.MAGNUS_DB_PORTS_LIMIT_COUNT

apps/terminal/models/endpoint.py

@@ -14,6 +14,8 @@ class Endpoint(JMSModel):
     http_port = PortField(default=80, verbose_name=_('HTTP Port'))
     ssh_port = PortField(default=2222, verbose_name=_('SSH Port'))
     rdp_port = PortField(default=3389, verbose_name=_('RDP Port'))
+
+    # Todo: Delete
     mysql_port = PortField(default=33060, verbose_name=_('MySQL Port'))
     mariadb_port = PortField(default=33061, verbose_name=_('MariaDB Port'))
     postgresql_port = PortField(default=54320, verbose_name=_('PostgreSQL Port'))