From 337338ebf3872d3ec52b5f3466c0c33b124fa87f Mon Sep 17 00:00:00 2001
From: pufer <puferyu@gmail.com>
Date: Tue, 7 Aug 2018 18:11:46 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=A4=9AOU=E6=94=AF=E6=8C=81?=
 =?UTF-8?q?=20(#1262)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 增加多OU支持

* Update models.py

need to import LDAPSearchUnion
---
 apps/common/api.py          | 25 +++++++++++++------------
 apps/common/forms.py        |  3 ++-
 apps/common/models.py       | 11 ++++++-----
 apps/jumpserver/settings.py | 10 ++++++----
 4 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/apps/common/api.py b/apps/common/api.py
index b6658b665..e09cf9726 100644
--- a/apps/common/api.py
+++ b/apps/common/api.py
@@ -48,7 +48,7 @@ class LDAPTestingAPI(APIView):
             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_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"]
 
@@ -64,18 +64,19 @@ class LDAPTestingAPI(APIView):
             except Exception as e:
                 return Response({"error": str(e)}, status=401)
 
-            ok = conn.search(search_ou, search_filter % ({"user": "*"}),
-                             attributes=list(attr_map.values()))
-            if not ok:
-                return Response({"error": "Search no entry matched"}, status=401)
-
             users = []
-            for entry in conn.entries:
-                user = {}
-                for attr, mapping in attr_map.items():
-                    if hasattr(entry, mapping):
-                        user[attr] = getattr(entry, mapping)
-                users.append(user)
+            for search_ou in str(search_ougroup).split("|"):
+                ok = conn.search(search_ou, search_filter % ({"user": "*"}),
+                                 attributes=list(attr_map.values()))
+                if not ok:
+                    return Response({"error": _("Search no entry matched in ou {}").format(search_ou)}, status=401)
+
+                for entry in conn.entries:
+                    user = {}
+                    for attr, mapping in attr_map.items():
+                        if hasattr(entry, mapping):
+                            user[attr] = getattr(entry, mapping)
+                    users.append(user)
             if len(users) > 0:
                 return Response({"msg": _("Match {} s users").format(len(users))})
             else:
diff --git a/apps/common/forms.py b/apps/common/forms.py
index 8667aa128..3ee553016 100644
--- a/apps/common/forms.py
+++ b/apps/common/forms.py
@@ -114,7 +114,8 @@ class LDAPSettingForm(BaseForm):
         widget=forms.PasswordInput, required=False
     )
     AUTH_LDAP_SEARCH_OU = forms.CharField(
-        label=_("User OU"), initial='ou=tech,dc=jumpserver,dc=org'
+        label=_("User OU"), initial='ou=tech,dc=jumpserver,dc=org',
+        help_text=_("Use | split User OUs")
     )
     AUTH_LDAP_SEARCH_FILTER = forms.CharField(
         label=_("User search filter"), initial='(cn=%(user)s)',
diff --git a/apps/common/models.py b/apps/common/models.py
index c90458985..007683917 100644
--- a/apps/common/models.py
+++ b/apps/common/models.py
@@ -5,7 +5,7 @@ from django.db import models
 from django.db.utils import ProgrammingError, OperationalError
 from django.utils.translation import ugettext_lazy as _
 from django.conf import settings
-from django_auth_ldap.config import LDAPSearch
+from django_auth_ldap.config import LDAPSearch, LDAPSearchUnion
 
 
 class SettingQuerySet(models.QuerySet):
@@ -72,10 +72,11 @@ class Setting(models.Model):
                 settings.AUTHENTICATION_BACKENDS.remove(settings.AUTH_LDAP_BACKEND)
 
         if self.name == "AUTH_LDAP_SEARCH_FILTER":
-            settings.AUTH_LDAP_USER_SEARCH = LDAPSearch(
-                settings.AUTH_LDAP_SEARCH_OU, ldap.SCOPE_SUBTREE,
-                settings.AUTH_LDAP_SEARCH_FILTER,
-            )
+            settings.AUTH_LDAP_USER_SEARCH_UNION = [
+                LDAPSearch(USER_SEARCH, ldap.SCOPE_SUBTREE, settings.AUTH_LDAP_SEARCH_FILTER)
+                for USER_SEARCH in str(settings.AUTH_LDAP_SEARCH_OU).split("|")
+            ]
+            settings.AUTH_LDAP_USER_SEARCH = LDAPSearchUnion(*settings.AUTH_LDAP_USER_SEARCH_UNION)
 
     class Meta:
         db_table = "settings"
diff --git a/apps/jumpserver/settings.py b/apps/jumpserver/settings.py
index e4a932775..37c14dbf1 100644
--- a/apps/jumpserver/settings.py
+++ b/apps/jumpserver/settings.py
@@ -14,7 +14,7 @@ import os
 import sys
 
 import ldap
-from django_auth_ldap.config import LDAPSearch
+from django_auth_ldap.config import LDAPSearch, LDAPSearchUnion
 from django.urls import reverse_lazy
 
 # Build paths inside the project like this: os.path.join(BASE_DIR, ...)
@@ -356,9 +356,11 @@ AUTH_LDAP_SEARCH_OU = CONFIG.AUTH_LDAP_SEARCH_OU
 AUTH_LDAP_SEARCH_FILTER = CONFIG.AUTH_LDAP_SEARCH_FILTER
 AUTH_LDAP_START_TLS = CONFIG.AUTH_LDAP_START_TLS
 AUTH_LDAP_USER_ATTR_MAP = CONFIG.AUTH_LDAP_USER_ATTR_MAP
-AUTH_LDAP_USER_SEARCH = LDAPSearch(
-    AUTH_LDAP_SEARCH_OU, ldap.SCOPE_SUBTREE, AUTH_LDAP_SEARCH_FILTER,
-)
+AUTH_LDAP_USER_SEARCH_UNION = [
+    LDAPSearch(USER_SEARCH, ldap.SCOPE_SUBTREE, AUTH_LDAP_SEARCH_FILTER)
+    for USER_SEARCH in str(AUTH_LDAP_SEARCH_OU).split("|")
+]
+AUTH_LDAP_USER_SEARCH = LDAPSearchUnion(*AUTH_LDAP_USER_SEARCH_UNION)
 AUTH_LDAP_GROUP_SEARCH_OU = CONFIG.AUTH_LDAP_GROUP_SEARCH_OU
 AUTH_LDAP_GROUP_SEARCH_FILTER = CONFIG.AUTH_LDAP_GROUP_SEARCH_FILTER
 AUTH_LDAP_GROUP_SEARCH = LDAPSearch(