122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
@File : ext_ldap.py
|
|||
|
|
@Time : 2025/3/5 {TIME}
|
|||
|
|
@Author : xxlaila
|
|||
|
|
@Software: dify
|
|||
|
|
"""
|
|||
|
|
from flask_ldap3_login import LDAP3LoginManager
|
|||
|
|
from configs import dify_config
|
|||
|
|
from dify_app import DifyApp
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
from queue import Queue
|
|||
|
|
from ldap3 import Server, Connection, ALL
|
|||
|
|
|
|||
|
|
def is_enabled():
|
|||
|
|
return getattr(dify_config, 'LDAP_ENABLED', False)
|
|||
|
|
|
|||
|
|
def create_ldap_pool():
|
|||
|
|
pool = Queue(maxsize=dify_config.LDAP_POOL_SIZE)
|
|||
|
|
for i in range(dify_config.LDAP_POOL_SIZE):
|
|||
|
|
server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL)
|
|||
|
|
conn = Connection(
|
|||
|
|
server,
|
|||
|
|
user=dify_config.AUTH_LDAP_BIND_DN,
|
|||
|
|
password=dify_config.AUTH_LDAP_BIND_PASSWORD,
|
|||
|
|
receive_timeout=dify_config.LDAP_CONN_TIMEOUT
|
|||
|
|
)
|
|||
|
|
if conn.bind():
|
|||
|
|
logging.info(f"LDAP connection {i} bound successfully")
|
|||
|
|
pool.put(conn)
|
|||
|
|
else:
|
|||
|
|
logging.error(f"LDAP connection {i} failed to bind")
|
|||
|
|
return pool
|
|||
|
|
|
|||
|
|
def get_ldap_connection():
|
|||
|
|
"""Get a connection from the connection pool"""
|
|||
|
|
conn = LDAP_POOL.get()
|
|||
|
|
|
|||
|
|
# Add connection validity check
|
|||
|
|
if not conn.bound or not conn.server:
|
|||
|
|
try:
|
|||
|
|
if not conn.bind():
|
|||
|
|
logging.error("LDAP connection is unbound, rebinding...")
|
|||
|
|
conn.unbind()
|
|||
|
|
server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL)
|
|||
|
|
new_conn = Connection(
|
|||
|
|
server,
|
|||
|
|
user=dify_config.AUTH_LDAP_BIND_DN,
|
|||
|
|
password=dify_config.AUTH_LDAP_BIND_PASSWORD,
|
|||
|
|
receive_timeout=dify_config.LDAP_CONN_TIMEOUT
|
|||
|
|
)
|
|||
|
|
if new_conn.bind():
|
|||
|
|
return new_conn
|
|||
|
|
raise Exception("LDAP connection reconstruction failed")
|
|||
|
|
except Exception as e:
|
|||
|
|
logging.error(f"LDAP connection recovery failed: {str(e)}")
|
|||
|
|
raise
|
|||
|
|
return conn
|
|||
|
|
|
|||
|
|
def release_ldap_connection(conn):
|
|||
|
|
"""Return the connection to the connection pool"""
|
|||
|
|
try:
|
|||
|
|
# Reset connection status
|
|||
|
|
if conn.bound:
|
|||
|
|
conn.unbind()
|
|||
|
|
conn.open() # Reopen connection without binding
|
|||
|
|
LDAP_POOL.put(conn)
|
|||
|
|
except Exception as e:
|
|||
|
|
logging.error(f"Failed to recycle LDAP connection: {str(e)}")
|
|||
|
|
conn.unbind()
|
|||
|
|
|
|||
|
|
def init_app(app: DifyApp):
|
|||
|
|
"""Initialize LDAP authentication integration"""
|
|||
|
|
if not is_enabled():
|
|||
|
|
app.ldap_manager = None # Explicitly set the manager to None
|
|||
|
|
logging.info("LDAP authentication is disabled")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# Global LDAP connection pool
|
|||
|
|
global LDAP_POOL
|
|||
|
|
LDAP_POOL = create_ldap_pool()
|
|||
|
|
|
|||
|
|
# Parsing User Attribute Mapping
|
|||
|
|
if isinstance(dify_config.AUTH_LDAP_USER_ATTR_MAP, str):
|
|||
|
|
ldap_user_attr_map = json.loads(dify_config.AUTH_LDAP_USER_ATTR_MAP)
|
|||
|
|
else:
|
|||
|
|
ldap_user_attr_map = dify_config.AUTH_LDAP_USER_ATTR_MAP
|
|||
|
|
|
|||
|
|
# Setting up LDAP configuration
|
|||
|
|
app.config.update({
|
|||
|
|
"LDAP_HOST": dify_config.AUTH_LDAP_SERVER_URI,
|
|||
|
|
"LDAP_BASE_DN": dify_config.AUTH_LDAP_SEARCH_BASE_DN,
|
|||
|
|
"LDAP_BIND_DN": dify_config.AUTH_LDAP_BIND_DN,
|
|||
|
|
"LDAP_BIND_PASSWORD": dify_config.AUTH_LDAP_BIND_PASSWORD,
|
|||
|
|
"LDAP_USER_FILTER": dify_config.AUTH_LDAP_USER_FILTER,
|
|||
|
|
"LDAP_USER_RDN_ATTR": "uid",
|
|||
|
|
"LDAP_USER_LOGIN_ATTR": "uid",
|
|||
|
|
"LDAP_USER_SEARCH_SCOPE": "SUBTREE",
|
|||
|
|
"LDAP_USER_MAPPING": ldap_user_attr_map,
|
|||
|
|
"LDAP_DEFAULT_ROLE": dify_config.LDAP_DEFAULT_ROLE,
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
|
|||
|
|
# Initializing the LDAP Manager
|
|||
|
|
ldap_manager = LDAP3LoginManager()
|
|||
|
|
ldap_manager.init_app(app)
|
|||
|
|
|
|||
|
|
# Mount the LDAP manager into the app
|
|||
|
|
app.ldap_manager = ldap_manager
|
|||
|
|
# Confirm that the mount was successful
|
|||
|
|
logging.info(f"LDAP manager mounted: {hasattr(app, 'ldap_manager')}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# Configuring Logging
|
|||
|
|
if app.debug:
|
|||
|
|
app.logger.info("LDAP configuration loaded：")
|
|||
|
|
app.logger.info(f"Server: {app.config['LDAP_HOST']}")
|
|||
|
|
app.logger.info(f"Base DN: {app.config['LDAP_BASE_DN']}")
|
|||
|
|
|
|||
|
|
return ldap_manager
|