[Update] 重构 LDAP/AD 同步功能，添加缓存机制

apps/authentication/signals_handlers.py

@@ -47,7 +47,7 @@ def on_openid_login_success(sender, user=None, request=None, **kwargs):
 
 @receiver(populate_user)
 def on_ldap_create_user(sender, user, ldap_user, **kwargs):
-    if user and user.username != 'admin':
+    if user and user.username not in ['admin']:
         user.source = user.SOURCE_LDAP
         user.save()
 

apps/settings/api.py

@@ -13,10 +13,16 @@ from django.core.mail import send_mail
 from django.utils.translation import ugettext_lazy as _
 
 from .models import Setting
-from .utils import LDAPUtil
+from .utils import (
+    LDAPServerUtil, LDAPCacheUtil, LDAPImportUtil, LDAPSyncUtil,
+    LDAP_USE_CACHE_FLAGS
+
+)
+from .tasks import sync_ldap_user_task
 from common.permissions import IsOrgAdmin, IsSuperUser
 from common.utils import get_logger
 from .serializers import MailTestSerializer, LDAPTestSerializer, LDAPUserSerializer
+from users.models import User
 
 
 logger = get_logger(__file__)
@@ -67,65 +73,107 @@ class LDAPTestingAPI(APIView):
     success_message = _("Test ldap success")
 
     @staticmethod
-    def get_ldap_util(serializer):
-        host = serializer.validated_data["AUTH_LDAP_SERVER_URI"]
+    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_ougroup = 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"]
-        try:
-            attr_map = json.loads(attr_map)
-        except json.JSONDecodeError:
-            return Response({"error": "AUTH_LDAP_USER_ATTR_MAP not valid"}, status=401)
-
-        util = LDAPUtil(
-            use_settings_config=False, server_uri=host, bind_dn=bind_dn,
-            password=password, use_ssl=use_ssl,
-            search_ougroup=search_ougroup, search_filter=search_filter,
-            attr_map=attr_map
-        )
-        return util
+        config = {
+            'server_uri': server_uri,
+            'bind_dn': bind_dn,
+            'password': password,
+            'use_ssl': use_ssl,
+            'search_ougroup': search_ougroup,
+            'search_filter': search_filter,
+            'attr_map': json.loads(attr_map),
+        }
+        return config
 
     def post(self, request):
         serializer = self.serializer_class(data=request.data)
         if not serializer.is_valid():
             return Response({"error": str(serializer.errors)}, status=401)
 
-        util = self.get_ldap_util(serializer)
+        attr_map = serializer.validated_data["AUTH_LDAP_USER_ATTR_MAP"]
+        try:
+            json.loads(attr_map)
+        except json.JSONDecodeError:
+            return Response({"error": "AUTH_LDAP_USER_ATTR_MAP not valid"}, status=401)
 
+        config = self.get_ldap_config(serializer)
+        util = LDAPServerUtil(config=config)
         try:
-            users = util.search_user_items()
+            users = util.search()
         except Exception as e:
             return Response({"error": str(e)}, status=401)
 
-        if len(users) > 0:
-            return Response({"msg": _("Match {} s users").format(len(users))})
-        else:
-            return Response({"error": "Have user but attr mapping error"}, status=401)
+        return Response({"msg": _("Match {} s users").format(len(users))})
 
 
 class LDAPUserListApi(generics.ListAPIView):
     permission_classes = (IsOrgAdmin,)
     serializer_class = LDAPUserSerializer
 
+    def get_queryset_from_cache(self):
+        search_value = self.request.query_params.get('search')
+        users = LDAPCacheUtil().search(search_value=search_value)
+        return users
+
+    def get_queryset_from_server(self):
+        search_value = self.request.query_params.get('search')
+        users = LDAPServerUtil().search(search_value=search_value)
+        return users
+
     def get_queryset(self):
         if hasattr(self, 'swagger_fake_view'):
             return []
-        q = self.request.query_params.get('search')
-        try:
-            util = LDAPUtil()
-            extra_filter = util.construct_extra_filter(util.SEARCH_FIELD_ALL, q)
-            users = util.search_user_items(extra_filter)
-        except Exception as e:
-            users = []
-            logger.error(e)
-        # 前端data_table会根据row.id对table.selected值进行操作
-        for user in users:
-            user['id'] = user['username']
+        cache_police = self.request.query_params.get('cache_police', True)
+        if cache_police in LDAP_USE_CACHE_FLAGS:
+            users = self.get_queryset_from_cache()
+        else:
+            users = self.get_queryset_from_server()
         return users
 
+    def list(self, request, *args, **kwargs):
+        cache_police = self.request.query_params.get('cache_police', True)
+        # 不是用缓存
+        if cache_police not in LDAP_USE_CACHE_FLAGS:
+            return super().list(request, *args, **kwargs)
+
+        queryset = self.get_queryset()
+        # 缓存有数据
+        if queryset is not None:
+            return super().list(request, *args, **kwargs)
+
+        sync_util = LDAPSyncUtil()
+        # 还没有同步任务
+        if sync_util.task_no_start:
+            task = sync_ldap_user_task.delay()
+            data = {'msg': 'Cache no data, sync task {} started.'.format(task.id)}
+            return Response(data=data, status=409)
+        # 同步任务正在执行
+        if sync_util.task_is_running:
+            data = {'msg': 'synchronization is running.'}
+            return Response(data=data, status=409)
+        # 同步任务执行结束
+        if sync_util.task_is_over:
+            msg = sync_util.get_task_error_msg()
+            data = {'msg': 'Synchronization task report error: {}'.format(msg)}
+            return Response(data=data, status=400)
+
+        return super().list(request, *args, **kwargs)
+
+    @staticmethod
+    def processing_queryset(queryset):
+        db_username_list = User.objects.all().values_list('username', flat=True)
+        for q in queryset:
+            q['id'] = q['username']
+            q['existing'] = q['username'] in db_username_list
+        return queryset
+
     def sort_queryset(self, queryset):
         order_by = self.request.query_params.get('order')
         if not order_by:
@@ -138,32 +186,41 @@ class LDAPUserListApi(generics.ListAPIView):
         queryset = sorted(queryset, key=lambda x: x[order_by], reverse=reverse)
         return queryset
 
-    def list(self, request, *args, **kwargs):
-        queryset = self.get_queryset()
+    def filter_queryset(self, queryset):
+        queryset = self.processing_queryset(queryset)
         queryset = self.sort_queryset(queryset)
-        page = self.paginate_queryset(queryset)
-        if page is not None:
-            return self.get_paginated_response(page)
-        return Response(queryset)
+        return queryset
 
 
-class LDAPUserSyncAPI(APIView):
+class LDAPUserImportAPI(APIView):
     permission_classes = (IsOrgAdmin,)
 
+    def get_ldap_users(self):
+        username_list = self.request.data.get('username_list', [])
+        cache_police = self.request.query_params.get('cache_police', True)
+        if cache_police in LDAP_USE_CACHE_FLAGS:
+            users = LDAPCacheUtil().search(search_users=username_list)
+        else:
+            users = LDAPServerUtil().search(search_users=username_list)
+        return users
+
     def post(self, request):
-        username_list = request.data.get('username_list', [])
+        users = self.get_ldap_users()
+        errors = LDAPImportUtil().perform_import(users)
+        if errors:
+            return Response({'Error': errors}, status=401)
+        return Response({'msg': 'Imported {} users successfully'.format(len(users))})
 
-        util = LDAPUtil()
+
+class LDAPCacheRefreshAPI(generics.RetrieveAPIView):
+
+    def retrieve(self, request, *args, **kwargs):
         try:
-            result = util.sync_users(username_list)
+            LDAPSyncUtil().clear_cache()
         except Exception as e:
-            logger.error(e, exc_info=True)
-            return Response({'error': str(e)}, status=401)
-        else:
-            msg = _("succeed: {} failed: {} total: {}").format(
-                result['succeed'], result['failed'], result['total']
-            )
-            return Response({'msg': msg})
+            logger.error(str(e))
+            return Response(data={'msg': str(e)}, status=400)
+        return Response(data={'msg': 'success'})
 
 
 class ReplayStorageCreateAPI(APIView):

apps/settings/serializers.py

@@ -25,6 +25,7 @@ class LDAPTestSerializer(serializers.Serializer):
 class LDAPUserSerializer(serializers.Serializer):
     id = serializers.CharField()
     username = serializers.CharField()
+    name = serializers.CharField()
     email = serializers.CharField()
     existing = serializers.BooleanField(read_only=True)
 

apps/settings/tasks/__init__.py

+# coding: utf-8
+# 
+
+from .ldap import *

apps/settings/tasks/ldap.py

+# coding: utf-8
+# 
+
+from celery import shared_task
+
+from common.utils import get_logger
+from ..utils import LDAPSyncUtil
+
+__all__ = ['sync_ldap_user_task']
+
+
+logger = get_logger(__file__)
+
+
+@shared_task
+def sync_ldap_user_task():
+    LDAPSyncUtil().perform_sync()

apps/settings/templates/settings/_ldap_list_users_modal.html

@@ -23,6 +23,7 @@
    <div class="row">
        <div class="col-lg-12 animated fadeInRight" id="split-right">
            <div class="mail-box-header">
+               <div class="uc pull-left m-r-5"><a  id="id_refresh_cache" class="btn btn-sm btn-primary"> {% trans "Refresh cache" %} </a></div>
                <table class="table table-striped table-bordered table-hover " id="ldap_list_users_table" style="width: 100%">
                    <thead>
                        <tr>
@@ -43,8 +44,11 @@
 
 <script>
 var ldap_users_table = 0;
+var interval;
+
 function initLdapUsersTable() {
     if(ldap_users_table){
+        ldap_users_table.ajax.reload();
         return ldap_users_table
     }
     var options = {
@@ -68,21 +72,89 @@ function initLdapUsersTable() {
         ],
         pageLength: 15
     };
-
     ldap_users_table = jumpserver.initServerSideDataTable(options);
     return ldap_users_table
 }
 
+function testRequestLdapUser(){
+    var the_url = "{% url 'api-settings:ldap-user-list' %}";
+    var error = function (data, status) {
+        if (status === 409){
+            console.log(data);
+            return
+        }
+        if (status === 400){
+            toastr.error(data);
+            clearInterval(interval);
+            interval = undefined;
+            return
+        }
+        console.log(data, status)
+    };
+    var success = function() {
+        initLdapUsersTable();
+        clearInterval(interval);
+        interval = undefined
+    };
+    requestApi({
+        url: the_url,
+        method: 'GET',
+        flash_message: false,
+        error: error,
+        success: success
+    });
+}
+
+function timingTestRequestLdapUser(){
+    if (interval !== undefined){
+        return
+    }
+    interval = setInterval(testRequestLdapUser, 2000);
+}
 
 $(document).ready(function(){
+
 }).on('show.bs.modal', function () {
-    initLdapUsersTable();
+    timingTestRequestLdapUser()
 })
-.on('click','.close_btn1',function () {
-    window.location.reload()
+.on('click', '#id_refresh_cache', function () {
+    var the_url = "{% url "api-settings:ldap-cache-refresh" %}";
+    function error(data) {
+        toastr.error(data)
+    }
+    function success(){
+        timingTestRequestLdapUser();
+    }
+    requestApi({
+        url: the_url,
+        method: 'GET',
+        error: error,
+        success: success
+    })
 })
-.on('click','.close_btn2',function () {
-    window.location.reload()
+.on("click","#btn_ldap_modal_confirm",function () {
+    var username_list = ldap_users_table.selected;
+    if (username_list.length === 0){
+        var msg = "{% trans 'User is not currently selected, please check the user you want to import'%}";
+        toastr.error(msg);
+        return
+    }
+    var the_url =  "{% url "api-settings:ldap-user-import" %}";
+    function error(message) {
+        toastr.error(message)
+    }
+    function success(message) {
+        toastr.success(message.msg);
+        timingTestRequestLdapUser();
+    }
+    requestApi({
+        url: the_url,
+        body: JSON.stringify({'username_list':username_list}),
+        method: "POST",
+        flash_message: false,
+        success: success,
+        error: error
+    });
 })
 
 </script>

apps/settings/templates/settings/ldap_setting.html

@@ -109,33 +109,6 @@ $(document).ready(function () {
         error: error
     });
 })
-.on("click","#btn_ldap_modal_confirm",function () {
-    var username_list = ldap_users_table.selected;
-
-    if (username_list.length === 0){
-        var msg = "{% trans 'User is not currently selected, please check the user you want to import'%}";
-        toastr.error(msg);
-        return
-     }
-
-    var the_url =  "{% url "api-settings:ldap-user-sync" %}";
-
-    function error(message) {
-        toastr.error(message)
-    }
-
-    function success(message) {
-        toastr.success(message.msg)
-    }
-    requestApi({
-        url: the_url,
-        body: JSON.stringify({'username_list':username_list}),
-        method: "POST",
-        flash_message: false,
-        success: success,
-        error: error
-    });
- })
 
 </script>
 {% endblock %}

apps/settings/urls/api_urls.py

@@ -10,7 +10,8 @@ urlpatterns = [
     path('mail/testing/', api.MailTestingAPI.as_view(), name='mail-testing'),
     path('ldap/testing/', api.LDAPTestingAPI.as_view(), name='ldap-testing'),
     path('ldap/users/', api.LDAPUserListApi.as_view(), name='ldap-user-list'),
-    path('ldap/users/sync/', api.LDAPUserSyncAPI.as_view(), name='ldap-user-sync'),
+    path('ldap/users/import/', api.LDAPUserImportAPI.as_view(), name='ldap-user-import'),
+    path('ldap/cache/refresh/', api.LDAPCacheRefreshAPI.as_view(), name='ldap-cache-refresh'),
     path('terminal/replay-storage/create/', api.ReplayStorageCreateAPI.as_view(), name='replay-storage-create'),
     path('terminal/replay-storage/delete/', api.ReplayStorageDeleteAPI.as_view(), name='replay-storage-delete'),
     path('terminal/command-storage/create/', api.CommandStorageCreateAPI.as_view(), name='command-storage-create'),

apps/settings/utils.py

-# -*- coding: utf-8 -*-
-#
-
-from ldap3 import Server, Connection
-from django.utils.translation import ugettext_lazy as _
-
-from users.models import User
-from users.utils import construct_user_email
-from common.utils import get_logger
-from common.const import LDAP_AD_ACCOUNT_DISABLE
-
-from .models import settings
-
-
-logger = get_logger(__file__)
-
-
-class LDAPOUGroupException(Exception):
-    pass
-
-
-class LDAPUtil:
-    _conn = None
-
-    SEARCH_FIELD_ALL = 'all'
-    SEARCH_FIELD_USERNAME = 'username'
-
-    def __init__(self, use_settings_config=True, server_uri=None, bind_dn=None,
-                 password=None, use_ssl=None, search_ougroup=None,
-                 search_filter=None, attr_map=None, auth_ldap=None):
-        # config
-        self.paged_size = settings.AUTH_LDAP_SEARCH_PAGED_SIZE
-
-        if use_settings_config:
-            self._load_config_from_settings()
-        else:
-            self.server_uri = server_uri
-            self.bind_dn = bind_dn
-            self.password = password
-            self.use_ssl = use_ssl
-            self.search_ougroup = search_ougroup
-            self.search_filter = search_filter
-            self.attr_map = attr_map
-            self.auth_ldap = auth_ldap
-
-    def _load_config_from_settings(self):
-        self.server_uri = settings.AUTH_LDAP_SERVER_URI
-        self.bind_dn = settings.AUTH_LDAP_BIND_DN
-        self.password = settings.AUTH_LDAP_BIND_PASSWORD
-        self.use_ssl = settings.AUTH_LDAP_START_TLS
-        self.search_ougroup = settings.AUTH_LDAP_SEARCH_OU
-        self.search_filter = settings.AUTH_LDAP_SEARCH_FILTER
-        self.attr_map = settings.AUTH_LDAP_USER_ATTR_MAP
-        self.auth_ldap = settings.AUTH_LDAP
-
-    @property
-    def connection(self):
-        if self._conn is None:
-            server = Server(self.server_uri, use_ssl=self.use_ssl)
-            conn = Connection(server, self.bind_dn, self.password)
-            conn.bind()
-            self._conn = conn
-        return self._conn
-
-    @staticmethod
-    def get_user_by_username(username):
-        try:
-            user = User.objects.get(username=username)
-        except Exception as e:
-            return None
-        else:
-            return user
-
-    def _ldap_entry_to_user_item(self, entry):
-        user_item = {}
-        for attr, mapping in self.attr_map.items():
-            if not hasattr(entry, mapping):
-                continue
-            value = getattr(entry, mapping).value or ''
-            if mapping.lower() == 'useraccountcontrol' and attr == 'is_active'\
-                    and value:
-                value = int(value) & LDAP_AD_ACCOUNT_DISABLE \
-                        != LDAP_AD_ACCOUNT_DISABLE
-            user_item[attr] = value
-        return user_item
-
-    def _search_user_items_ou(self, search_ou, extra_filter=None, cookie=None):
-        search_filter = self.search_filter % {"user": "*"}
-        if extra_filter:
-            search_filter = '(&{}{})'.format(search_filter, extra_filter)
-
-        ok = self.connection.search(
-            search_ou, search_filter,
-            attributes=list(self.attr_map.values()),
-            paged_size=self.paged_size, paged_cookie=cookie
-        )
-        if not ok:
-            error = _("Search no entry matched in ou {}".format(search_ou))
-            raise LDAPOUGroupException(error)
-
-        user_items = []
-        for entry in self.connection.entries:
-            user_item = self._ldap_entry_to_user_item(entry)
-            user = self.get_user_by_username(user_item['username'])
-            user_item['existing'] = bool(user)
-            if user_item in user_items:
-                continue
-            user_items.append(user_item)
-        return user_items
-
-    def _cookie(self):
-        if self.paged_size is None:
-            cookie = None
-        else:
-            cookie = self.connection.result['controls']['1.2.840.113556.1.4.319']['value']['cookie']
-        return cookie
-
-    def search_user_items(self, extra_filter=None):
-        user_items = []
-        logger.info("Search user items")
-
-        for search_ou in str(self.search_ougroup).split("|"):
-            logger.info("Search user search ou: {}".format(search_ou))
-            _user_items = self._search_user_items_ou(search_ou, extra_filter=extra_filter)
-            user_items.extend(_user_items)
-            while self._cookie():
-                logger.info("Page Search user search ou: {}".format(search_ou))
-                _user_items = self._search_user_items_ou(search_ou, extra_filter, self._cookie())
-                user_items.extend(_user_items)
-        logger.info("Search user items end")
-        return user_items
-
-    def construct_extra_filter(self, field, q):
-        if not q:
-            return None
-        extra_filter = ''
-        if field == self.SEARCH_FIELD_ALL:
-            for attr in self.attr_map.values():
-                extra_filter += '({}={})'.format(attr, q)
-            extra_filter = '(|{})'.format(extra_filter)
-            return extra_filter
-
-        if field == self.SEARCH_FIELD_USERNAME and isinstance(q, list):
-            attr = self.attr_map.get('username')
-            for username in q:
-                extra_filter += '({}={})'.format(attr, username)
-            extra_filter = '(|{})'.format(extra_filter)
-            return extra_filter
-
-    def search_filter_user_items(self, username_list):
-        extra_filter = self.construct_extra_filter(
-            self.SEARCH_FIELD_USERNAME, username_list
-        )
-        user_items = self.search_user_items(extra_filter)
-        return user_items
-
-    @staticmethod
-    def save_user(user, user_item):
-        for field, value in user_item.items():
-            if not hasattr(user, field):
-                continue
-            if isinstance(getattr(user, field), bool):
-                if isinstance(value, str):
-                    value = value.lower()
-                value = value in ['true', 1, True]
-            setattr(user, field, value)
-        user.save()
-
-    def update_user(self, user_item):
-        user = self.get_user_by_username(user_item['username'])
-        if user.source != User.SOURCE_LDAP:
-            msg = _('The user source is not LDAP')
-            return False, msg
-        try:
-            self.save_user(user, user_item)
-        except Exception as e:
-            logger.error(e, exc_info=True)
-            return False, str(e)
-        else:
-            return True, None
-
-    def create_user(self, user_item):
-        user = User(source=User.SOURCE_LDAP)
-        try:
-            self.save_user(user, user_item)
-        except Exception as e:
-            logger.error(e, exc_info=True)
-            return False, str(e)
-        else:
-            return True, None
-
-    @staticmethod
-    def construct_user_email(user_item):
-        username = user_item['username']
-        email = user_item.get('email', '')
-        email = construct_user_email(username, email)
-        return email
-
-    def create_or_update_users(self, user_items):
-        succeed = failed = 0
-        for user_item in user_items:
-            exist = user_item.pop('existing', False)
-            user_item['email'] = self.construct_user_email(user_item)
-            if not exist:
-                ok, error = self.create_user(user_item)
-            else:
-                ok, error = self.update_user(user_item)
-            if not ok:
-                logger.info("Failed User: {}".format(user_item))
-                failed += 1
-            else:
-                succeed += 1
-        result = {'total': len(user_items), 'succeed': succeed, 'failed': failed}
-        return result
-
-    def sync_users(self, username_list=None):
-        user_items = self.search_filter_user_items(username_list)
-        result = self.create_or_update_users(user_items)
-        return result

apps/settings/utils/__init__.py

+# coding: utf-8
+# 
+
+from .ldap import *

apps/settings/views.py

@@ -5,6 +5,7 @@ from django.utils.translation import ugettext as _
 
 from common.permissions import PermissionsMixin, IsSuperUser
 from common import utils
+from .utils import LDAPSyncUtil
 from .forms import EmailSettingForm, LDAPSettingForm, BasicSettingForm, \
     TerminalSettingForm, SecuritySettingForm, EmailContentSettingForm
 
@@ -83,6 +84,7 @@ class LDAPSettingView(PermissionsMixin, TemplateView):
             form.save()
             msg = _("Update setting successfully")
             messages.success(request, msg)
+            LDAPSyncUtil().clear_cache()
             return redirect('settings:ldap-setting')
         else:
             context = self.get_context_data()

apps/users/tasks.py

@@ -11,7 +11,7 @@ from .models import User
 from .utils import (
     send_password_expiration_reminder_mail, send_user_expiration_reminder_mail
 )
-from settings.utils import LDAPUtil
+from settings.utils import LDAPServerUtil, LDAPImportUtil
 
 
 logger = get_logger(__file__)
@@ -70,16 +70,21 @@ def check_user_expired_periodic():
 
 
 @shared_task
-def sync_ldap_user():
-    logger.info("Start sync ldap user periodic task")
-    util = LDAPUtil()
-    result = util.sync_users()
-    logger.info("Result: {}".format(result))
+def import_ldap_user():
+    logger.info("Start import ldap user task")
+    util_server = LDAPServerUtil()
+    util_import = LDAPImportUtil()
+    users = util_server.search()
+    errors = util_import.perform_import(users)
+    if errors:
+        logger.error("Imported LDAP users errors: {}".format(errors))
+    else:
+        logger.info('Imported {} users successfully'.format(len(users)))
 
 
 @shared_task
 @after_app_ready_start
-def sync_ldap_user_periodic():
+def import_ldap_user_periodic():
     if not settings.AUTH_LDAP:
         return
     if not settings.AUTH_LDAP_SYNC_IS_PERIODIC:
@@ -91,10 +96,9 @@ def sync_ldap_user_periodic():
     else:
         interval = None
     crontab = settings.AUTH_LDAP_SYNC_CRONTAB
-
     tasks = {
-        'sync_ldap_user_periodic': {
-            'task': sync_ldap_user.name,
+        'import_ldap_user_periodic': {
+            'task': import_ldap_user.name,
             'interval': interval,
             'crontab': crontab,
             'enabled': True,