Enhance: Introduce HostMachineShortName type and refactor RID Manager components to utilize it, improving clarity and consistency in domain controller and RID set management. Update various functions to streamline object SID handling and enhance overall data integrity.

Author: m.shvets

app/alembic/versions/01f3f05a5b11_add_primary_group_id.py

@@ -68,7 +68,6 @@ async def _add_domain_computers_group(connection: AsyncConnection) -> None:  # n
 
             dir_, group_ = await create_group(
                 name=DOMAIN_COMPUTERS_GROUP_NAME,
-                sid=515,
                 attribute_value_validator=AttributeValueValidator(),
                 session=session,
             )

app/alembic/versions/552b4eafb1aa_remove_objectsid_vals.py

@@ -6,8 +6,6 @@
 
 """
 
-import secrets
-
 import sqlalchemy as sa
 from alembic import op
 from dishka import AsyncContainer, Scope
@@ -180,26 +178,19 @@ async def _migrate_object_sids(
             ),
         )
 
-        identifier: str | None = None
+        domain_identifier: str | None = None
         if domain_sid_from_column:
             parts = domain_sid_from_column.split("-")
             # "S-1-5-21-AAA-BBB-CCC" -> "AAA-BBB-CCC"
             if len(parts) >= 7 and domain_sid_from_column.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)}"
-            )
+                domain_identifier = "-".join(parts[4:7])
 
         session.add(
             Attribute(
                 name="DomainIdentifier",
-                value=identifier,
+                value=domain_identifier,
                 directory_id=domain.id,
             ),
         )
@@ -228,7 +219,7 @@ async def _migrate_object_sids(
                 ),
             )
 
-        built_in_sid_prefix = "S-1-5-32"
+        sid_prefix = "S-1-5-21"
         for dir_name, rid in (
             (DOMAIN_ADMIN_GROUP_NAME, SecurityPrincipalRid.DOMAIN_ADMINS),
             (DOMAIN_USERS_GROUP_NAME, SecurityPrincipalRid.DOMAIN_USERS),
@@ -249,7 +240,7 @@ async def _migrate_object_sids(
                     ),
                 )
                 .values(
-                    value=f"{built_in_sid_prefix}-{int(rid)}",
+                    value=f"{sid_prefix}-{domain_identifier}-{int(rid)}",
                 ),
             )
 
@@ -263,8 +254,8 @@ async def _migrate_object_sids(
             )
             .values(
                 value=(
-                    f"{built_in_sid_prefix}"
-                    f"-{int(SecurityPrincipalRid.ADMINISTRATOR)}"
+                    f"{sid_prefix}"
+                    f"-{domain_identifier}-{int(SecurityPrincipalRid.ADMINISTRATOR)}"
                 ),
             ),
         )
@@ -286,19 +277,16 @@ async def _init_rid_manager(
             rid_set_use_case = await cnt.get(RIDSetUseCase)
             role_use_case = await cnt.get(RoleUseCase)
 
-        if not await get_base_directories(session):
+        base_dn_list = await get_base_directories(session)
+        if not base_dn_list:
             return
+        domain = base_dn_list[0]
 
         try:
             rid_manager_dir = await rid_gateway.get_rid_manager()
         except RIDManagerNotFoundError:
             rid_manager_dir = await rid_setup_gateway.set_rid_manager()
 
-        base_dn_list = await get_base_directories(session)
-        if not base_dn_list:
-            return
-        domain = base_dn_list[0]
-
         domain_identifier = await session.scalar(
             select(Attribute).where(
                 qa(Attribute.directory_id) == domain.id,
@@ -341,21 +329,17 @@ async def _init_rid_manager(
             directory=rid_manager_dir,
         )
 
-        domain_controller = await rid_gateway.get_domain_controller()
-        rid_set_dir: Directory | None = None
+        domain_controller = await rid_setup_gateway.get_domain_controller()
         try:
-            rid_set_dir = await rid_set_gateway.get(domain_controller)
+            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)
 
-            rid_set_dir = await rid_set_use_case.add(
+            await rid_set_use_case.add(
                 domain_controller,
                 RIDSetAllocationParamsDTO(
                     next_rid=lower,

app/alembic/versions/fafc3d0b11ec_.py

@@ -61,7 +61,6 @@ async def _create_readonly_grp_and_plcy(
             if not group_dir:
                 dir_, _ = await create_group(
                     name="readonly domain controllers",
-                    sid=521,
                     attribute_value_validator=attribute_value_validator,
                     session=session,
                 )

app/ioc.py

@@ -188,6 +188,7 @@
     RIDSetGateway,
     RIDSetUseCase,
 )
+from ldap_protocol.rid_manager.types import HostMachineShortName
 from ldap_protocol.roles.access_manager import AccessManager
 from ldap_protocol.roles.ace_dao import AccessControlEntryDAO
 from ldap_protocol.roles.migrations_ace_dao import (
@@ -216,6 +217,13 @@ class MainProvider(Provider):
     scope = Scope.APP
     settings = from_context(provides=Settings, scope=Scope.APP)
 
+    @provide(scope=Scope.APP)
+    def host_machine_short_name(
+        self,
+        settings: Settings,
+    ) -> HostMachineShortName:
+        return HostMachineShortName(settings.HOST_MACHINE_SHORT_NAME)
+
     @provide(scope=Scope.APP)
     def get_engine_registry(self, settings: Settings) -> EngineRegistry:
         return EngineRegistry(

app/ldap_protocol/auth/setup_gateway.py

@@ -12,7 +12,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from entities import Attribute, Directory, Group, NetworkPolicy, User
-from enums import EntityTypeNames, SidPrefix
+from enums import EntityTypeNames
 from ldap_protocol.ldap_schema.attribute_value_validator import (
     AttributeValueValidator,
 )
@@ -168,7 +168,6 @@ async def create_dir(
             await self._object_sid_use_case.add(
                 directory_id=dir_.id,
                 rid=int(data["objectSid"]),
-                sid_prefix=SidPrefix.BUILT_IN_DOMAIN,
             )
 
         if dir_.object_class == "group":

app/ldap_protocol/auth/use_cases.py

@@ -253,10 +253,6 @@ async def _create(self, dto: SetupDTO, data: list) -> None:
             await self._role_use_case.create_read_only_role()
             await self._audit_use_case.create_policies()
             await self._rid_manager_setup_use_case.setup()
-            dc = await self._rid_manager_use_case.get_domain_controller()
-            await self._object_sid_use_case.add(
-                directory_id=dc.id,
-            )
 
             await self._session.commit()
         except IntegrityError:

app/ldap_protocol/rid_manager/object_sid_use_case.py

@@ -28,9 +28,9 @@ def __init__(
         self._session = session
         self._rid_manager_use_case = rid_manager_use_case
 
-    async def get(self, directory_id: int) -> str:
-        """Get object SID."""
-        return await self._gateway.get(directory_id)
+    async def get_domain_identifier(self) -> str:
+        """Get domain identifier."""
+        return await self._gateway.get_domain_identifier()
 
     async def add(
         self,
@@ -40,8 +40,7 @@ async def add(
     ) -> None:
         """Add object SID."""
         if rid is None:
-            rid_set_id = await self._rid_set_use_case.get_rid_set_id()
-            rid = await self._rid_set_use_case.allocate_next_rid(rid_set_id)
+            rid = await self._rid_set_use_case.allocate_next_rid()
 
         if sid_prefix == SidPrefix.BUILT_IN_DOMAIN:
             object_sid = f"{sid_prefix}-{rid}"

app/ldap_protocol/rid_manager/rid_manager_gateway.py

@@ -7,24 +7,26 @@
 from sqlalchemy import select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from config import Settings
-from constants import DOMAIN_CONTROLLERS_OU_NAME
 from entities import Attribute, Directory
 from ldap_protocol.rid_manager.exceptions import (
     RIDManagerAvailablePoolNotFoundError,
-    RIDManagerDomainControllerNotFoundError,
     RIDManagerNotFoundError,
 )
+from ldap_protocol.rid_manager.types import HostMachineShortName
 from repo.pg.tables import queryable_attr as qa
 
 
 class RIDManagerGateway:
     """RID Manager gateway."""
 
-    def __init__(self, session: AsyncSession, settings: Settings) -> None:
+    def __init__(
+        self,
+        session: AsyncSession,
+        host_machine_short_name: HostMachineShortName,
+    ) -> None:
         """Initialize RID Manager gateway."""
         self._session = session
-        self._settings = settings
+        self._host_machine_short_name = host_machine_short_name
 
     async def get_rid_manager(self) -> Directory:
         """Get RID Manager directory."""
@@ -55,41 +57,3 @@ async def update_rid_available_pool(self, available_pool: int) -> None:
             .where(qa(Attribute.name) == "rIDAvailablePool")
             .values(value=str(available_pool)),
         )
-
-    async def get_domain_controller(
-        self,
-    ) -> Directory:
-        """Get domain controller."""
-        domain = await self._session.scalar(
-            select(Directory).where(
-                qa(Directory.object_class) == "domain",
-                qa(Directory.parent_id).is_(None),
-            ),
-        )
-        if not domain:
-            raise RIDManagerDomainControllerNotFoundError(
-                "Domain controller not found",
-            )
-
-        domain_controllers_ou = await self._session.scalar(
-            select(Directory).where(
-                qa(Directory.name) == DOMAIN_CONTROLLERS_OU_NAME,
-                qa(Directory.parent_id) == domain.id,
-            ),
-        )
-        if not domain_controllers_ou:
-            raise RIDManagerDomainControllerNotFoundError(
-                "Domain controllers OU not found",
-            )
-
-        domain_controller = await self._session.scalar(
-            select(Directory).where(
-                qa(Directory.name) == self._settings.HOST_MACHINE_SHORT_NAME,
-                qa(Directory.parent_id) == domain_controllers_ou.id,
-            ),
-        )
-        if not domain_controller:
-            raise RIDManagerDomainControllerNotFoundError(
-                "Domain controller not found",
-            )
-        return domain_controller

app/ldap_protocol/rid_manager/rid_manager_use_case.py

@@ -6,7 +6,6 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from entities import Directory
 from ldap_protocol.rid_manager.exceptions import RIDManagerPoolExceededError
 from ldap_protocol.rid_manager.rid_manager_gateway import RIDManagerGateway
 from ldap_protocol.rid_manager.utils import from_qword, to_qword
@@ -41,7 +40,3 @@ async def allocate_pool(self) -> int:
             await self._gateway.update_rid_available_pool(new_available_pool)
 
         return to_qword(lower, lower + self.RID_BLOCK_SIZE)
-
-    async def get_domain_controller(self) -> Directory:
-        """Locate best Domain Controller via DNS SRV records."""
-        return await self._gateway.get_domain_controller()

app/ldap_protocol/rid_manager/rid_set_gateway.py

@@ -8,7 +8,6 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import aliased
 
-from config import Settings
 from constants import DOMAIN_CONTROLLERS_OU_NAME
 from entities import Attribute, Directory
 from ldap_protocol.rid_manager.dtos import RIDSetAllocationParamsDTO
@@ -18,17 +17,22 @@
     RIDManagerRidPreviousAllocationPoolNotFoundError,
     RIDManagerRidSetNotFoundError,
 )
+from ldap_protocol.rid_manager.types import HostMachineShortName
 from ldap_protocol.utils.async_cache import rid_set_id_cache
 from repo.pg.tables import queryable_attr as qa
 
 
 class RIDSetGateway:
     """RID Set gateway."""
 
-    def __init__(self, session: AsyncSession, settings: Settings) -> None:
+    def __init__(
+        self,
+        session: AsyncSession,
+        host_machine_short_name: HostMachineShortName,
+    ) -> None:
         """Initialize RID Set gateway."""
         self._session = session
-        self._settings = settings
+        self._host_machine_short_name = host_machine_short_name
 
     @rid_set_id_cache
     async def get_rid_set_id(self) -> int:
@@ -63,7 +67,7 @@ async def get_rid_set_id_value(self) -> int:
                 qa(domain.parent_id).is_(None),
                 qa(domain_controllers_ou.name) == DOMAIN_CONTROLLERS_OU_NAME,
                 qa(domain_controller.name)
-                == self._settings.HOST_MACHINE_SHORT_NAME,
+                == self._host_machine_short_name,
                 qa(rid_set.name) == "RID Set",
             ),
         )
@@ -84,13 +88,15 @@ async def get(self, domain_controller: Directory) -> Directory:
 
         return rid_set
 
-    async def add(self, domain_controller: Directory) -> Directory:
-        """Add RID Set directory."""
+    async def create_rid_set_directory(
+        self,
+        domain_controller: Directory,
+    ) -> Directory:
+        """Create RID Set directory."""
         rid_set_dir = Directory(
             is_system=True,
             name="RID Set",
         )
-        rid_set_dir.create_path(domain_controller, "cn")
 
         self._session.add(rid_set_dir)
         await self._session.flush()
@@ -163,10 +169,12 @@ async def set_allocation_attrs(
     async def get_rid_allocation_pool(self, rid_set_id: int) -> int:
         """Get RID allocation pool from RID Set directory."""
         allocation_pool = await self._session.scalar(
-            select(Attribute).where(
+            select(Attribute)
+            .where(
                 qa(Attribute.name) == "rIDAllocationPool",
                 qa(Attribute.directory_id) == rid_set_id,
-            ),
+            )
+            .with_for_update(),
         )
         if not (allocation_pool and allocation_pool.value):
             raise RIDManagerRidAllocationPoolNotFoundError(
@@ -193,7 +201,7 @@ async def get_rid_previous_allocation_pool(
             )
         return int(previous_allocation_pool.value)
 
-    async def get_rid_next_rid(self, rid_set_id: int) -> int:
+    async def get_next_rid_value(self, rid_set_id: int) -> int:
         """Get next RID from RID Set directory."""
         next_rid = await self._session.scalar(
             select(Attribute)