Refactor ldap server

app/ioc.py

@@ -65,6 +65,7 @@
 from ldap_protocol.kerberos.service import KerberosService
 from ldap_protocol.kerberos.template_render import KRBTemplateRenderer
 from ldap_protocol.ldap_requests.contexts import (
+    AsyncSessionSearchRequest,
     LDAPAddRequestContext,
     LDAPBindRequestContext,
     LDAPDeleteRequestContext,
@@ -205,7 +206,7 @@ async def get_kadmin_http(
             yield KadminHTTPClient(client)
 
     @provide(scope=Scope.REQUEST)
-    async def get_kadmin(
+    def get_kadmin(
         self,
         client: KadminHTTPClient,
         kadmin_class: type[AbstractKadmin],
@@ -260,14 +261,14 @@ async def get_dns_http_client(
             yield DNSManagerHTTPClient(client)
 
     @provide(scope=Scope.REQUEST)
-    async def get_dns_mngr(
+    def get_dns_mngr(
         self,
         settings: DNSManagerSettings,
         dns_manager_class: type[AbstractDNSManager],
         http_client: DNSManagerHTTPClient,
-    ) -> AsyncIterator[AbstractDNSManager]:
+    ) -> AbstractDNSManager:
         """Get DNSManager class."""
-        yield dns_manager_class(settings=settings, http_client=http_client)
+        return dns_manager_class(settings=settings, http_client=http_client)
 
     @provide(scope=Scope.APP)
     async def get_redis_for_sessions(
@@ -284,7 +285,7 @@ async def get_redis_for_sessions(
         await client.aclose()
 
     @provide(scope=Scope.APP)
-    async def get_session_storage(
+    def get_session_storage(
         self,
         client: SessionStorageClient,
         settings: Settings,
@@ -297,7 +298,7 @@ async def get_session_storage(
         )
 
     @provide()
-    async def get_normalized_audit_event(
+    def get_normalized_audit_event(
         self,
     ) -> type[NormalizedAuditEvent]:
         """Get normalized audit event class."""
@@ -318,13 +319,13 @@ async def get_audit_redis_client(
         await client.aclose()
 
     @provide(scope=Scope.APP)
-    async def get_raw_audit_manager(
+    def get_raw_audit_manager(
         self,
         client: AuditRedisClient,
         settings: Settings,
-    ) -> AsyncIterator[RawAuditManager]:
+    ) -> RawAuditManager:
         """Get raw audit manager."""
-        yield RawAuditManager(
+        return RawAuditManager(
             client,
             settings.RAW_EVENT_STREAM_NAME,
             settings.EVENT_HANDLER_GROUP,
@@ -333,13 +334,13 @@ async def get_raw_audit_manager(
         )
 
     @provide(scope=Scope.APP)
-    async def get_normalized_audit_manager(
+    def get_normalized_audit_manager(
         self,
         client: AuditRedisClient,
         settings: Settings,
-    ) -> AsyncIterator[NormalizedAuditManager]:
+    ) -> NormalizedAuditManager:
         """Get raw audit manager."""
-        yield NormalizedAuditManager(
+        return NormalizedAuditManager(
             client,
             settings.NORMALIZED_EVENT_STREAM_NAME,
             settings.EVENT_SENDER_GROUP,
@@ -352,7 +353,7 @@ async def get_normalized_audit_manager(
     audit_destination_dao = provide(AuditDestinationDAO, scope=Scope.REQUEST)
 
     @provide(scope=Scope.REQUEST)
-    async def get_dhcp_manager_repository(
+    def get_dhcp_manager_repository(
         self,
         session: AsyncSession,
     ) -> DHCPManagerRepository:
@@ -368,20 +369,20 @@ async def get_dhcp_manager_state(
         return await dhcp_manager_repository.ensure_state()
 
     @provide(scope=Scope.REQUEST)
-    async def get_dhcp_mngr_class(
+    def get_dhcp_mngr_class(
         self,
         dhcp_state: DHCPManagerState,
     ) -> type[AbstractDHCPManager]:
         """Get DHCP manager type."""
-        return await get_dhcp_manager_class(dhcp_state)
+        return get_dhcp_manager_class(dhcp_state)
 
     @provide(scope=Scope.REQUEST)
-    async def get_dhcp_api_repository_class(
+    def get_dhcp_api_repository_class(
         self,
         dhcp_state: DHCPManagerState,
     ) -> type[DHCPAPIRepository]:
         """Get DHCP API repository type."""
-        return await get_dhcp_api_repository_class(dhcp_state)
+        return get_dhcp_api_repository_class(dhcp_state)
 
     @provide(scope=Scope.APP)
     async def get_dhcp_http_client(
@@ -395,7 +396,7 @@ async def get_dhcp_http_client(
             yield DHCPManagerHTTPClient(http_client)
 
     @provide(scope=Scope.REQUEST)
-    async def get_dhcp_api_repository(
+    def get_dhcp_api_repository(
         self,
         http_client: DHCPManagerHTTPClient,
         dhcp_api_repository_class: type[DHCPAPIRepository],
@@ -404,7 +405,7 @@ async def get_dhcp_api_repository(
         return dhcp_api_repository_class(http_client)
 
     @provide(scope=Scope.REQUEST)
-    async def get_dhcp_mngr(
+    def get_dhcp_mngr(
         self,
         dhcp_manager_class: type[AbstractDHCPManager],
         dhcp_api_repository: DHCPAPIRepository,
@@ -445,7 +446,7 @@ async def get_dhcp_mngr(
     )
     password_utils = provide(PasswordUtils, scope=Scope.RUNTIME)
 
-    access_manager = provide(AccessManager, scope=Scope.REQUEST)
+    access_manager = provide(AccessManager, scope=Scope.RUNTIME)
     role_dao = provide(RoleDAO, scope=Scope.REQUEST)
     ace_dao = provide(AccessControlEntryDAO, scope=Scope.REQUEST)
     role_use_case = provide(RoleUseCase, scope=Scope.REQUEST)
@@ -490,15 +491,37 @@ class LDAPContextProvider(Provider):
         LDAPModifyDNRequestContext,
         scope=Scope.REQUEST,
     )
-    search_request_context = provide(
-        LDAPSearchRequestContext,
-        scope=Scope.REQUEST,
-    )
     unbind_request_context = provide(
         LDAPUnbindRequestContext,
         scope=Scope.REQUEST,
     )
 
+    @provide(scope=Scope.SESSION)
+    async def create_search_session(
+        self,
+        async_session: async_sessionmaker[AsyncSession],
+    ) -> AsyncIterator[AsyncSessionSearchRequest]:
+        """Create session for request."""
+        async with async_session() as session:
+            yield session  # type: ignore
+
+    @provide(scope=Scope.SESSION, provides=LDAPSearchRequestContext)
+    def get_search_request_context(
+        self,
+        session: AsyncSessionSearchRequest,
+        ldap_session: LDAPSession,
+        settings: Settings,
+        access_manager: AccessManager,
+    ) -> LDAPSearchRequestContext:
+        """Get search request context."""
+        return LDAPSearchRequestContext(
+            session=session,
+            ldap_session=ldap_session,
+            settings=settings,
+            access_manager=access_manager,
+            rootdse_rd=RootDSEReader(settings, SADomainGateway(session)),
+        )
+
 
 class HTTPProvider(LDAPContextProvider):
     """HTTP LDAP session."""
@@ -508,7 +531,7 @@ class HTTPProvider(LDAPContextProvider):
     monitor_use_case = provide(AuditMonitorUseCase, scope=Scope.REQUEST)
 
     @provide()
-    async def get_audit_monitor(
+    def get_audit_monitor(
         self,
         session: AsyncSession,
         audit_use_case: "AuditUseCase",
@@ -568,7 +591,7 @@ def get_permissions_provider(
         return auth_provider
 
     @provide()
-    async def get_identity_provider(
+    def get_identity_provider(
         self,
         request: Request,
         session_storage: SessionStorage,
@@ -649,6 +672,23 @@ def get_krb_template_render(
     )
     network_policy_gateway = provide(NetworkPolicyGateway, scope=Scope.REQUEST)
 
+    @provide(scope=Scope.REQUEST, provides=LDAPSearchRequestContext)
+    async def get_search_request_context(
+        self,
+        session: AsyncSession,
+        ldap_session: LDAPSession,
+        settings: Settings,
+        access_manager: AccessManager,
+    ) -> LDAPSearchRequestContext:
+        """Get search request context."""
+        return LDAPSearchRequestContext(
+            session=session,  # type: ignore
+            ldap_session=ldap_session,
+            settings=settings,
+            access_manager=access_manager,
+            rootdse_rd=RootDSEReader(settings, SADomainGateway(session)),
+        )
+
 
 class LDAPServerProvider(LDAPContextProvider):
     """Provider with session scope."""
@@ -739,7 +779,7 @@ async def get_client(
             yield MFAHTTPClient(client)
 
     @provide(provides=MultifactorAPI)
-    async def get_http_mfa(
+    def get_http_mfa(
         self,
         credentials: MFA_HTTP_Creds,
         client: MFAHTTPClient,
@@ -761,7 +801,7 @@ async def get_http_mfa(
         )
 
     @provide(provides=LDAPMultiFactorAPI)
-    async def get_ldap_mfa(
+    def get_ldap_mfa(
         self,
         credentials: MFA_LDAP_Creds,
         client: MFAHTTPClient,

app/ldap_protocol/dhcp/__init__.py

@@ -26,7 +26,7 @@
 from .stub import StubDHCPAPIRepository, StubDHCPManager
 
 
-async def get_dhcp_manager_class(
+def get_dhcp_manager_class(
     dhcp_state: DHCPManagerState,
 ) -> type[AbstractDHCPManager]:
     """Get an instance of the DHCP manager."""
@@ -35,7 +35,7 @@ async def get_dhcp_manager_class(
     return StubDHCPManager
 
 
-async def get_dhcp_api_repository_class(
+def get_dhcp_api_repository_class(
     dhcp_state: DHCPManagerState,
 ) -> type[DHCPAPIRepository]:
     """Get an instance of the DHCP API repository."""

app/ldap_protocol/ldap_requests/abandon.py

@@ -27,7 +27,7 @@ def from_data(
         """Create structure from ASN1Row dataclass list."""
         return cls(message_id=1)
 
-    async def handle(self) -> AsyncGenerator:
+    async def handle(self, ctx: None) -> AsyncGenerator:  # noqa: ARG002
         """Handle message with current user."""
         await asyncio.sleep(0)
         return

app/ldap_protocol/ldap_requests/add.py

@@ -65,6 +65,7 @@ class AddRequest(BaseRequest):
     """
 
     PROTOCOL_OP: ClassVar[int] = ProtocolRequests.ADD
+    CONTEXT_TYPE: ClassVar[type] = LDAPAddRequestContext
 
     entry: str = Field(..., description="Any `DistinguishedName`")
     attributes: list[PartialAttribute]

app/ldap_protocol/ldap_requests/base.py

@@ -24,6 +24,7 @@
 from ldap_protocol.dependency import resolve_deps
 from ldap_protocol.dialogue import LDAPSession
 from ldap_protocol.ldap_responses import BaseResponse, LDAPResult
+from ldap_protocol.objects import ProtocolRequests
 from ldap_protocol.policies.audit.audit_use_case import AuditUseCase
 from ldap_protocol.policies.audit.events.factory import (
     RawAuditEventBuilderRedis,
@@ -62,6 +63,7 @@ class _APIProtocol: ...
 class BaseRequest(ABC, _APIProtocol, BaseModel):
     """Base request builder."""
 
+    CONTEXT_TYPE: ClassVar[type]
     handle: ClassVar[handler]
     from_data: ClassVar[serializer]
     __event_data: dict = {}
@@ -113,38 +115,42 @@ async def handle_tcp(
         container: AsyncContainer,
     ) -> AsyncIterator[BaseResponse]:
         """Hanlde response with tcp."""
-        kwargs = await resolve_deps(func=self.handle, container=container)
-        responses = []
+        if self.PROTOCOL_OP != ProtocolRequests.ABANDON:
+            ctx = await container.get(self.CONTEXT_TYPE)  # type: ignore
+        else:
+            ctx = None
 
-        async for response in self.handle(**kwargs):
+        responses = []
+        async for response in self.handle(ctx=ctx):
             responses.append(response)
             yield response
 
-        ldap_session = await container.get(LDAPSession)
-        settings = await container.get(Settings)
-        audit_use_case = await container.get(AuditUseCase)
-
-        if await audit_use_case.check_event_processing_enabled(
-            self.PROTOCOL_OP,
-        ):
-            username = getattr(
-                ldap_session.user,
-                "user_principal_name",
-                "ANONYMOUS",
-            )
-            event = RawAuditEventBuilderRedis.from_ldap_request(
-                self,
-                responses=responses,
-                username=username,
-                ip=ldap_session.ip,
-                protocol="TCP_LDAP",
-                settings=settings,
-                context=self.get_event_data(),
-            )
+        if self.PROTOCOL_OP != ProtocolRequests.SEARCH:
+            ldap_session = await container.get(LDAPSession)
+            settings = await container.get(Settings)
+            audit_use_case = await container.get(AuditUseCase)
+
+            if await audit_use_case.check_event_processing_enabled(
+                self.PROTOCOL_OP,
+            ):
+                username = getattr(
+                    ldap_session.user,
+                    "user_principal_name",
+                    "ANONYMOUS",
+                )
+                event = RawAuditEventBuilderRedis.from_ldap_request(
+                    self,
+                    responses=responses,
+                    username=username,
+                    ip=ldap_session.ip,
+                    protocol="TCP_LDAP",
+                    settings=settings,
+                    context=self.get_event_data(),
+                )
 
-            ldap_session.event_task_group.create_task(
-                audit_use_case.manager.send_event(event),
-            )
+                ldap_session.event_task_group.create_task(
+                    audit_use_case.manager.send_event(event),
+                )
 
     async def _handle_api(
         self,
@@ -156,7 +162,11 @@ async def _handle_api(
         :param AsyncSession session: db session
         :return list[BaseResponse]: list of handled responses
         """
-        kwargs = await resolve_deps(func=self.handle, container=container)
+        if self.PROTOCOL_OP != ProtocolRequests.ABANDON:
+            ctx = await container.get(self.CONTEXT_TYPE)  # type: ignore
+        else:
+            ctx = None
+
         ldap_session = await container.get(LDAPSession)
         settings = await container.get(Settings)
         audit_use_case = await container.get(AuditUseCase)
@@ -168,7 +178,7 @@ async def _handle_api(
         else:
             log_api.info(f"{get_class_name(self)}[{un}]")
 
-        responses = [response async for response in self.handle(**kwargs)]
+        responses = [response async for response in self.handle(ctx=ctx)]
 
         if settings.DEBUG:
             for response in responses: