Refactor: Introduce ObjectSIDUseCase and related gateways, enhancing RID management functionality

Author: m.shvets

app/alembic/versions/552b4eafb1aa_remove_objectsid_vals.py

@@ -6,6 +6,8 @@
 
 """
 
+import secrets
+
 import sqlalchemy as sa
 from alembic import op
 from dishka import AsyncContainer, Scope
@@ -16,19 +18,26 @@
 from enums import EntityTypeNames
 from ldap_protocol.ldap_schema.dto import EntityTypeDTO
 from ldap_protocol.ldap_schema.entity_type_use_case import EntityTypeUseCase
-from ldap_protocol.rid_manager.exceptions import RIDManagerNotFoundError
-from ldap_protocol.rid_manager.gateways import RIDManagerGateway
-from ldap_protocol.rid_manager.use_cases import (
-    RID_AVAILABLE_MAX,
+from ldap_protocol.rid_manager import (
+    RIDManagerGateway,
+    RIDManagerSetupGateway,
     RIDManagerSetupUseCase,
+    RIDManagerUseCase,
+    RIDSetUseCase,
+)
+from ldap_protocol.rid_manager.dtos import RIDSetAllocationParamsDTO
+from ldap_protocol.rid_manager.exceptions import (
+    RIDManagerNotFoundError,
+    RIDManagerRidSetNotFoundError,
 )
-from ldap_protocol.rid_manager.utils import create_qword
+from ldap_protocol.rid_manager.rid_set_gateway import RIDSetGateway
+from ldap_protocol.rid_manager.utils import from_qword, to_qword
 from ldap_protocol.utils.queries import get_base_directories
 from repo.pg.tables import queryable_attr as qa
 
 # revision identifiers, used by Alembic.
 revision: None | str = "552b4eafb1aa"
-down_revision: None | str = "19d86e660cf2"
+down_revision: None | str = "2dadf40c026a"
 branch_labels: None | list[str] = None
 depends_on: None | list[str] = None
 
@@ -78,8 +87,8 @@ async def _migrate_object_sids(
     ) -> None:
         """Move Directory.objectSid values into Attributes table.
 
-        Additionally, for domain directories move the domain SID prefix part
-        into the ``DomainIdentifier`` attribute.
+        Additionally, for domain directories create the ``DomainIdentifier``
+        attribute if it does not exist.
         """
         async with container(scope=Scope.REQUEST) as cnt:
             session = await cnt.get(AsyncSession)
@@ -106,16 +115,46 @@ async def _migrate_object_sids(
                     ),
                 )
 
-            if directory.name == "domain":
-                identifier = directory.object_sid.split("-")[
-                    -1
-                ]  # remove sid prefix
+        base_dn_list = await get_base_directories(session)
+        if base_dn_list:
+            domain = base_dn_list[0]
+
+            existing_identifier = await session.scalar(
+                select(Attribute).where(
+                    qa(Attribute.directory_id) == domain.id,
+                    qa(Attribute.name) == "DomainIdentifier",
+                ),
+            )
+
+            if not (existing_identifier and existing_identifier.value):
+                domain_object_sid = await session.scalar(
+                    select(Attribute).where(
+                        qa(Attribute.directory_id) == domain.id,
+                        qa(Attribute.name) == "objectSid",
+                    ),
+                )
+
+                identifier: str | None = None
+                if domain_object_sid and domain_object_sid.value:
+                    parts = domain_object_sid.value.split("-")
+                    # "S-1-5-21-AAA-BBB-CCC" -> "AAA-BBB-CCC"
+                    if len(parts) >= 7 and domain_object_sid.value.startswith(
+                        "S-1-5-21-",
+                    ):
+                        identifier = "-".join(parts[4:7])
+
+                if identifier is None:
+                    identifier = (
+                        f"{secrets.randbits(32)}-"
+                        f"{secrets.randbits(32)}-"
+                        f"{secrets.randbits(32)}"
+                    )
 
                 session.add(
                     Attribute(
                         name="DomainIdentifier",
                         value=identifier,
-                        directory_id=directory.id,
+                        directory_id=domain.id,
                     ),
                 )
 
@@ -129,27 +168,35 @@ async def _init_rid_manager(
         """Initialize RID Manager and RID Set for existing data."""
         async with container(scope=Scope.REQUEST) as cnt:
             session = await cnt.get(AsyncSession)
-        rid_setup_use_case = await cnt.get(RIDManagerSetupUseCase)
-        rid_gateway = await cnt.get(RIDManagerGateway)
+            rid_setup_gateway = await cnt.get(RIDManagerSetupGateway)
+            rid_gateway = await cnt.get(RIDManagerGateway)
+            rid_manager_use_case = await cnt.get(RIDManagerUseCase)
+            rid_set_gateway = await cnt.get(RIDSetGateway)
+            rid_set_use_case = await cnt.get(RIDSetUseCase)
 
         if not await get_base_directories(session):
             return
 
         try:
-            await rid_gateway.get_rid_manager()
+            rid_manager_dir = await rid_gateway.get_rid_manager()
         except RIDManagerNotFoundError:
-            await rid_setup_use_case.setup()
-            await rid_gateway.get_rid_manager()
+            rid_manager_dir = await rid_setup_gateway.set_rid_manager()
 
-        rid_set_dir = await rid_gateway.get_rid_set()
-        if not rid_set_dir:
+        base_dn_list = await get_base_directories(session)
+        if not base_dn_list:
             return
+        domain = base_dn_list[0]
 
-        base_domain = await rid_gateway.get_base_domain()
-        domain_identifier = await rid_gateway.get_domain_identifier(
-            base_domain,
+        domain_identifier = await session.scalar(
+            select(Attribute).where(
+                qa(Attribute.directory_id) == domain.id,
+                qa(Attribute.name) == "DomainIdentifier",
+            ),
         )
-        sid_prefix = f"S-1-5-21-{domain_identifier}-"
+        if not (domain_identifier and domain_identifier.value):
+            return
+
+        sid_prefix = f"S-1-5-21-{domain_identifier.value}-"
 
         sid_values = await session.scalars(
             select(Attribute).where(
@@ -172,25 +219,89 @@ async def _init_rid_manager(
 
         start_rid = max(max_rid, RIDManagerSetupUseCase.RID_USER_MIN)
 
-        qword = create_qword(start_rid, RID_AVAILABLE_MAX)
-        await rid_gateway.update_available_pool(qword)
+        qword = to_qword(start_rid, RIDManagerSetupUseCase.RID_AVAILABLE_MAX)
+        await rid_setup_gateway.set_rid_available_pool(rid_manager_dir, qword)
 
-        result = await session.execute(
-            update(Attribute)
-            .where(
+        domain_controller = await rid_gateway.get_domain_controller()
+        rid_set_dir: Directory | None = None
+        try:
+            rid_set_dir = await rid_set_gateway.get(domain_controller)
+        except RIDManagerRidSetNotFoundError:
+            rid_set_dir = None
+
+        if rid_set_dir is None:
+            previous_allocation_pool = (
+                await rid_manager_use_case.allocate_pool()
+            )
+            allocation_pool = await rid_manager_use_case.allocate_pool()
+            lower, _ = from_qword(previous_allocation_pool)
+
+            await rid_set_use_case.add(
+                domain_controller,
+                RIDSetAllocationParamsDTO(
+                    next_rid=lower,
+                    allocation_pool=allocation_pool,
+                    previous_allocation_pool=previous_allocation_pool,
+                ),
+            )
+            await session.commit()
+            return
+
+        existing_next_rid = await session.scalar(
+            select(Attribute).where(
                 qa(Attribute.directory_id) == rid_set_dir.id,
                 qa(Attribute.name) == "rIDNextRID",
-            )
-            .values(value=str(start_rid)),
+            ),
         )
-        if result.rowcount == 0:
-            session.add(
-                Attribute(
-                    directory_id=rid_set_dir.id,
-                    name="rIDNextRID",
-                    value=str(start_rid),
-                ),
+        existing_prev_pool = await session.scalar(
+            select(Attribute).where(
+                qa(Attribute.directory_id) == rid_set_dir.id,
+                qa(Attribute.name) == "rIDPreviousAllocationPool",
+            ),
+        )
+        existing_pool = await session.scalar(
+            select(Attribute).where(
+                qa(Attribute.directory_id) == rid_set_dir.id,
+                qa(Attribute.name) == "rIDAllocationPool",
+            ),
+        )
+
+        if (
+            existing_next_rid
+            and existing_next_rid.value
+            and existing_prev_pool
+            and existing_prev_pool.value
+            and existing_pool
+            and existing_pool.value
+        ):
+            await session.commit()
+            return
+
+        previous_allocation_pool = await rid_manager_use_case.allocate_pool()
+        allocation_pool = await rid_manager_use_case.allocate_pool()
+        lower, _ = from_qword(previous_allocation_pool)
+
+        for name, value in (
+            ("rIDNextRID", str(lower)),
+            ("rIDPreviousAllocationPool", str(previous_allocation_pool)),
+            ("rIDAllocationPool", str(allocation_pool)),
+        ):
+            result = await session.execute(
+                update(Attribute)
+                .where(
+                    qa(Attribute.directory_id) == rid_set_dir.id,
+                    qa(Attribute.name) == name,
+                )
+                .values(value=value),
             )
+            if result.rowcount == 0:
+                session.add(
+                    Attribute(
+                        directory_id=rid_set_dir.id,
+                        name=name,
+                        value=value,
+                    ),
+                )
 
         await session.commit()
 

app/extra/scripts/add_domain_controller.py

@@ -14,7 +14,7 @@
 from enums import SamAccountTypeCodes, SecurityPrincipalRid
 from ldap_protocol.ldap_schema.entity_type_dao import EntityTypeDAO
 from ldap_protocol.objects import UserAccountControlFlag
-from ldap_protocol.rid_manager.use_cases import RIDManagerUseCase
+from ldap_protocol.rid_manager import ObjectSIDUseCase
 from ldap_protocol.roles.role_use_case import RoleUseCase
 from repo.pg.tables import queryable_attr as qa
 
@@ -25,7 +25,7 @@ async def _add_domain_controller(
     entity_type_dao: EntityTypeDAO,
     settings: Settings,
     dc_ou_dir: Directory,
-    rid_manager_use_case: RIDManagerUseCase,
+    object_sid_use_case: ObjectSIDUseCase,
 ) -> None:
     dc_directory = Directory(
         object_class="",
@@ -37,7 +37,7 @@ async def _add_domain_controller(
     await session.flush()
 
     dc_directory.parent_id = dc_ou_dir.id
-    await rid_manager_use_case.set_object_sid(
+    await object_sid_use_case.add(
         directory=dc_directory,
         rid=SecurityPrincipalRid.DOMAIN_CONTROLLERS,
     )
@@ -103,7 +103,7 @@ async def add_domain_controller(
     settings: Settings,
     role_use_case: RoleUseCase,
     entity_type_dao: EntityTypeDAO,
-    rid_manager_use_case: RIDManagerUseCase,
+    object_sid_use_case: ObjectSIDUseCase,
 ) -> None:
     logger.info("Adding domain controller.")
 
@@ -137,7 +137,7 @@ async def add_domain_controller(
         entity_type_dao=entity_type_dao,
         settings=settings,
         dc_ou_dir=domain_controllers_ou,
-        rid_manager_use_case=rid_manager_use_case,
+        object_sid_use_case=object_sid_use_case,
     )
 
     logger.debug("Domain controller added.")

app/ioc.py

@@ -154,10 +154,14 @@
     UserPasswordHistoryUseCases,
 )
 from ldap_protocol.rid_manager import (
+    ObjectSIDGateway,
+    ObjectSIDUseCase,
     RIDManagerGateway,
     RIDManagerSetupGateway,
     RIDManagerSetupUseCase,
     RIDManagerUseCase,
+    RIDSetGateway,
+    RIDSetUseCase,
 )
 from ldap_protocol.roles.access_manager import AccessManager
 from ldap_protocol.roles.ace_dao import AccessControlEntryDAO
@@ -580,6 +584,10 @@ def get_dhcp_mngr(
         RIDManagerSetupUseCase,
         scope=Scope.REQUEST,
     )
+    object_sid_gateway = provide(ObjectSIDGateway, scope=Scope.REQUEST)
+    object_sid_use_case = provide(ObjectSIDUseCase, scope=Scope.REQUEST)
+    rid_set_gateway = provide(RIDSetGateway, scope=Scope.REQUEST)
+    rid_set_use_case = provide(RIDSetUseCase, scope=Scope.REQUEST)
 
 
 class LDAPContextProvider(Provider):

app/ldap_protocol/auth/setup_gateway.py

@@ -17,7 +17,7 @@
     AttributeValueValidator,
 )
 from ldap_protocol.ldap_schema.entity_type_dao import EntityTypeDAO
-from ldap_protocol.rid_manager.use_cases import RIDManagerUseCase
+from ldap_protocol.rid_manager import ObjectSIDUseCase
 from ldap_protocol.utils.async_cache import base_directories_cache
 from ldap_protocol.utils.queries import get_domain_object_class
 from password_utils import PasswordUtils
@@ -33,7 +33,7 @@ def __init__(
         password_utils: PasswordUtils,
         entity_type_dao: EntityTypeDAO,
         attribute_value_validator: AttributeValueValidator,
-        rid_manager_use_case: RIDManagerUseCase,
+        object_sid_use_case: ObjectSIDUseCase,
     ) -> None:
         """Initialize Setup use case.
 
@@ -45,7 +45,7 @@ def __init__(
         self._password_utils = password_utils
         self._entity_type_dao = entity_type_dao
         self._attribute_value_validator = attribute_value_validator
-        self._rid_manager_use_case = rid_manager_use_case
+        self._object_sid_use_case = object_sid_use_case
 
     async def is_setup(self) -> bool:
         """Check if setup is performed.
@@ -165,7 +165,7 @@ async def create_dir(
         )
 
         if "objectSid" in data:
-            await self._rid_manager_use_case.set_object_sid(
+            await self._object_sid_use_case.add(
                 directory=dir_,
                 rid=int(data["objectSid"]),
                 sid_prefix=SidPrefix.BUILT_IN_DOMAIN,

app/ldap_protocol/auth/use_cases.py

@@ -27,7 +27,7 @@
 from ldap_protocol.objects import UserAccountControlFlag
 from ldap_protocol.policies.audit.audit_use_case import AuditUseCase
 from ldap_protocol.policies.password import PasswordPolicyUseCases
-from ldap_protocol.rid_manager.use_cases import RIDManagerSetupUseCase
+from ldap_protocol.rid_manager import RIDManagerSetupUseCase
 from ldap_protocol.roles.role_use_case import RoleUseCase
 from ldap_protocol.utils.helpers import create_integer_hash, ft_now
 

app/ldap_protocol/ldap_requests/add.py

@@ -218,7 +218,7 @@ async def handle(  # noqa: C901
             ctx.session.add(new_dir)
 
             await ctx.session.flush()
-            await ctx.rid_manager_use_case.set_object_sid(
+            await ctx.object_sid_use_case.add(
                 directory=new_dir,
             )
             await ctx.session.flush()