fix(edge-identities): prevent unauthorised identity access (#5135)

api/edge_api/identities/models.py

@@ -27,41 +27,52 @@ class EdgeIdentity:
     dynamo_wrapper = DynamoIdentityWrapper()
 
     def __init__(self, engine_identity_model: IdentityModel):
-        self._engine_identity_model = engine_identity_model
+        self.engine_identity_model = engine_identity_model
         self._reset_initial_state()  # type: ignore[no-untyped-call]
 
     @classmethod
     def from_identity_document(cls, identity_document: dict) -> "EdgeIdentity":  # type: ignore[type-arg]
         return EdgeIdentity(IdentityModel.model_validate(identity_document))
 
-    @property
-    def django_id(self) -> int:
-        return self._engine_identity_model.django_id  # type: ignore[return-value]
-
     @property
     def environment_api_key(self) -> str:
-        return self._engine_identity_model.environment_api_key
+        return self.engine_identity_model.environment_api_key
 
     @property
     def feature_overrides(self) -> IdentityFeaturesList:
-        return self._engine_identity_model.identity_features
+        return self.engine_identity_model.identity_features
 
     @property
     def id(self) -> typing.Union[int, str]:
-        return self._engine_identity_model.django_id or str(
-            self._engine_identity_model.identity_uuid
+        return self.engine_identity_model.django_id or str(
+            self.engine_identity_model.identity_uuid
         )
 
     @property
     def identifier(self) -> str:
-        return self._engine_identity_model.identifier
+        return self.engine_identity_model.identifier
 
     @property
     def identity_uuid(self) -> str:
-        return str(self._engine_identity_model.identity_uuid)
+        return str(self.engine_identity_model.identity_uuid)
+
+    @property
+    def environment(self) -> Environment:
+        environment: Environment = Environment.objects.get(
+            api_key=self.environment_api_key
+        )
+        return environment
+
+    @property
+    def dashboard_alias(self) -> str | None:
+        return self.engine_identity_model.dashboard_alias
+
+    @dashboard_alias.setter
+    def dashboard_alias(self, dashboard_alias: str) -> None:
+        self.engine_identity_model.dashboard_alias = dashboard_alias
 
     def add_feature_override(self, feature_state: FeatureStateModel) -> None:
-        self._engine_identity_model.identity_features.append(feature_state)
+        self.engine_identity_model.identity_features.append(feature_state)
 
     def get_all_feature_states(
         self,
@@ -77,9 +88,9 @@ def get_all_feature_states(
             for the identity specifically)
         """
         segment_ids = self.dynamo_wrapper.get_segment_ids(
-            identity_model=self._engine_identity_model
+            identity_model=self.engine_identity_model
         )
-        django_environment = Environment.objects.get(api_key=self.environment_api_key)
+        django_environment = self.environment
 
         # since identity overrides are included in the document retrieved from dynamo,
         # we only want to retrieve the environment default and (relevant) segment overrides
@@ -138,7 +149,7 @@ def match_feature_state(fs):  # type: ignore[no-untyped-def]
         feature_state = next(
             filter(
                 match_feature_state,
-                self._engine_identity_model.identity_features,
+                self.engine_identity_model.identity_features,
             ),
             None,
         )
@@ -151,19 +162,19 @@ def get_feature_state_by_featurestate_uuid(
         return next(
             filter(
                 lambda fs: str(fs.featurestate_uuid) == featurestate_uuid,  # type: ignore[arg-type,union-attr]
-                self._engine_identity_model.identity_features,
+                self.engine_identity_model.identity_features,
             ),
             None,
         )
 
     def get_hash_key(self, use_identity_composite_key_for_hashing: bool) -> str:
-        return self._engine_identity_model.get_hash_key(
+        return self.engine_identity_model.get_hash_key(
             use_identity_composite_key_for_hashing
         )
 
     def remove_feature_override(self, feature_state: FeatureStateModel) -> None:
         with suppress(ValueError):  # ignore if feature state didn't exist
-            self._engine_identity_model.identity_features.remove(feature_state)
+            self.engine_identity_model.identity_features.remove(feature_state)
 
     def save(self, user: FFAdminUser | APIKeyUser = None):  # type: ignore[no-untyped-def,assignment]
         self.dynamo_wrapper.put_item(self.to_document())
@@ -175,8 +186,8 @@ def save(self, user: FFAdminUser | APIKeyUser = None):  # type: ignore[no-untype
         self._reset_initial_state()  # type: ignore[no-untyped-call]
 
     def delete(self, user: FFAdminUser | APIKeyUser = None) -> None:  # type: ignore[assignment]
-        self.dynamo_wrapper.delete_item(self._engine_identity_model.composite_key)
-        self._engine_identity_model.identity_features.clear()
+        self.dynamo_wrapper.delete_item(self.engine_identity_model.composite_key)
+        self.engine_identity_model.identity_features.clear()
         changeset = self._get_changes()
         self._update_feature_overrides(
             changeset=changeset,
@@ -186,14 +197,14 @@ def delete(self, user: FFAdminUser | APIKeyUser = None) -> None:  # type: ignore
 
     def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None:
         identity_feature_names = {
-            fs.feature.name for fs in self._engine_identity_model.identity_features
+            fs.feature.name for fs in self.engine_identity_model.identity_features
         }
         if not identity_feature_names.issubset(valid_feature_names):
-            self._engine_identity_model.prune_features(list(valid_feature_names))
+            self.engine_identity_model.prune_features(list(valid_feature_names))
             sync_identity_document_features.delay(args=(str(self.identity_uuid),))
 
     def to_document(self) -> dict:  # type: ignore[type-arg]
-        return map_engine_identity_to_identity_document(self._engine_identity_model)
+        return map_engine_identity_to_identity_document(self.engine_identity_model)
 
     def _update_feature_overrides(
         self, changeset: IdentityChangeset, user: FFAdminUser | APIKeyUser

api/edge_api/identities/serializers.py

@@ -3,9 +3,7 @@
 
 from django.utils import timezone
 from flag_engine.features.models import FeatureModel as EngineFeatureModel
-from flag_engine.features.models import (
-    FeatureStateModel as EngineFeatureStateModel,
-)
+from flag_engine.features.models import FeatureStateModel as EngineFeatureStateModel
 from flag_engine.features.models import (
     MultivariateFeatureOptionModel as EngineMultivariateFeatureOptionModel,
 )
@@ -23,7 +21,9 @@
 from environments.models import Environment
 from features.models import Feature, FeatureState, FeatureStateValue
 from features.multivariate.models import MultivariateFeatureOption
-from features.serializers import FeatureStateValueSerializer  # type: ignore[attr-defined]
+from features.serializers import (  # type: ignore[attr-defined]
+    FeatureStateValueSerializer,
+)
 from util.mappers import (
     map_engine_identity_to_identity_document,
     map_feature_to_engine,
@@ -84,19 +84,13 @@ def get_fields(self):  # type: ignore[no-untyped-def]
         return fields
 
     def update(
-        self, instance: dict[str, typing.Any], validated_data: dict[str, typing.Any]
-    ) -> EngineIdentity:
-        engine_identity = EngineIdentity.model_validate(instance)
-
-        engine_identity.dashboard_alias = (
-            self.validated_data.get("dashboard_alias")
-            or engine_identity.dashboard_alias
+        self, instance: EdgeIdentity, validated_data: dict[str, typing.Any]
+    ) -> EdgeIdentity:
+        instance.dashboard_alias = (
+            self.validated_data.get("dashboard_alias") or instance.dashboard_alias
         )
-
-        edge_identity = EdgeIdentity(engine_identity)
-        edge_identity.save()
-
-        return engine_identity
+        instance.save()
+        return instance
 
 
 class EdgeMultivariateFeatureOptionField(serializers.IntegerField):

api/edge_api/identities/views.py

@@ -3,11 +3,14 @@
 import typing
 
 import pydantic
-from common.environments.permissions import MANAGE_IDENTITIES, VIEW_IDENTITIES  # type: ignore[import-untyped]
+from common.environments.permissions import (  # type: ignore[import-untyped]
+    MANAGE_IDENTITIES,
+    VIEW_IDENTITIES,
+)
 from django.shortcuts import get_object_or_404
 from django.utils.decorators import method_decorator
 from drf_yasg.utils import swagger_auto_schema  # type: ignore[import-untyped]
-from flag_engine.identities.models import IdentityFeaturesList, IdentityModel
+from flag_engine.identities.models import IdentityFeaturesList
 from flag_engine.identities.traits.models import TraitModel
 from pyngo import drf_error_details
 from rest_framework import status, viewsets
@@ -24,7 +27,7 @@
     RetrieveModelMixin,
     UpdateModelMixin,
 )
-from rest_framework.permissions import IsAuthenticated
+from rest_framework.permissions import BasePermission, IsAuthenticated
 from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.views import APIView
@@ -56,7 +59,6 @@
 from features.models import FeatureState
 from features.permissions import IdentityFeatureStatePermissions
 from projects.exceptions import DynamoNotEnabledError
-from util.mappers import map_engine_identity_to_identity_document
 
 from . import edge_identity_service
 from .exceptions import TraitPersistenceError
@@ -87,7 +89,7 @@ class EdgeIdentityViewSet(
     lookup_field = "identity_uuid"
 
     def initial(self, request, *args, **kwargs):  # type: ignore[no-untyped-def]
-        environment = self.get_environment_from_request()  # type: ignore[no-untyped-call]
+        environment = self.get_environment_from_request()
         if not environment.project.enable_dynamo_db:
             raise DynamoNotEnabledError()
 
@@ -98,11 +100,13 @@ def get_serializer_class(self):  # type: ignore[no-untyped-def]
             return EdgeIdentityUpdateSerializer
         return EdgeIdentitySerializer
 
-    def get_object(self):  # type: ignore[no-untyped-def]
-        # TODO: should this return an EdgeIdentity object instead of a dict?
-        return EdgeIdentity.dynamo_wrapper.get_item_from_uuid_or_404(
+    def get_object(self) -> EdgeIdentity:
+        identity_document = EdgeIdentity.dynamo_wrapper.get_item_from_uuid_or_404(
             self.kwargs["identity_uuid"]
         )
+        edge_identity = EdgeIdentity.from_identity_document(identity_document)
+        self.check_object_permissions(self.request, edge_identity)
+        return edge_identity
 
     def get_queryset(self):  # type: ignore[no-untyped-def]
         page_size = self.pagination_class().get_page_size(self.request)
@@ -131,40 +135,42 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
             start_key=start_key,  # type: ignore[arg-type]
         )
 
-    def get_permissions(self):  # type: ignore[no-untyped-def]
+    def get_permissions(self) -> list[BasePermission]:
         return [
             IsAuthenticated(),
             NestedEnvironmentPermissions(
                 action_permission_map={
                     "retrieve": VIEW_IDENTITIES,
                     "list": VIEW_IDENTITIES,
                     "create": MANAGE_IDENTITIES,
+                    "destroy": MANAGE_IDENTITIES,
                     "get_traits": VIEW_IDENTITIES,
                     "update_traits": MANAGE_IDENTITIES,
-                    "perform_destroy": MANAGE_IDENTITIES,
                 },
             ),
         ]
 
-    def get_environment_from_request(self):  # type: ignore[no-untyped-def]
+    def get_environment_from_request(self) -> Environment:
         """
         Get environment object from URL parameters in request.
         """
         return get_object_or_404(
             Environment, api_key=self.kwargs["environment_api_key"]
         )
 
-    def perform_destroy(self, instance):  # type: ignore[no-untyped-def]
-        edge_identity = EdgeIdentity.from_identity_document(instance)
-        edge_identity.delete(user=self.request.user)  # type: ignore[arg-type]
+    def perform_destroy(self, instance: EdgeIdentity) -> None:
+        instance.delete(user=self.request.user)  # type: ignore[arg-type]
 
     @swagger_auto_schema(
         responses={200: EdgeIdentityTraitsSerializer(many=True)},
     )
     @action(detail=True, methods=["get"], url_path="list-traits")
     def get_traits(self, request, *args, **kwargs):  # type: ignore[no-untyped-def]
-        identity = IdentityModel.model_validate(self.get_object())  # type: ignore[no-untyped-call]
-        data = [trait.dict() for trait in identity.identity_traits]
+        edge_identity = self.get_object()
+        data = [
+            trait.dict()
+            for trait in edge_identity.engine_identity_model.identity_traits
+        ]
         return Response(data=data, status=status.HTTP_200_OK)
 
     @swagger_auto_schema(
@@ -174,21 +180,19 @@ def get_traits(self, request, *args, **kwargs):  # type: ignore[no-untyped-def]
     )
     @action(detail=True, methods=["put"], url_path="update-traits")
     def update_traits(self, request, *args, **kwargs):  # type: ignore[no-untyped-def]
-        environment = self.get_environment_from_request()  # type: ignore[no-untyped-call]
+        environment = self.get_environment_from_request()
         if not environment.project.organisation.persist_trait_data:
             raise TraitPersistenceError()
-        identity = IdentityModel.model_validate(self.get_object())  # type: ignore[no-untyped-call]
+        edge_identity = self.get_object()
         try:
             trait = TraitModel(**request.data)
         except pydantic.ValidationError as validation_error:
             raise ValidationError(
                 drf_error_details(validation_error)
             ) from validation_error
-        _, traits_updated = identity.update_traits([trait])
+        _, traits_updated = edge_identity.engine_identity_model.update_traits([trait])
         if traits_updated:
-            EdgeIdentity.dynamo_wrapper.put_item(
-                map_engine_identity_to_identity_document(identity)
-            )
+            edge_identity.save()
 
         data = trait.dict()
         return Response(data, status=status.HTTP_200_OK)

api/tests/conftest.py

@@ -1,3 +1,5 @@
+import typing
+
 import pytest
 from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table
 from pytest_django.fixtures import SettingsWrapper
@@ -8,7 +10,10 @@
     DynamoIdentityWrapper,
 )
 from environments.models import Environment
-from util.mappers import map_environment_to_environment_document
+from util.mappers import (
+    map_environment_to_environment_document,
+    map_environment_to_environment_v2_document,
+)
 
 
 @pytest.fixture()
@@ -75,15 +80,22 @@ def dynamodb_wrapper_v2(
 @pytest.fixture()
 def dynamo_enabled_project_environment_one_document(
     flagsmith_environment_table: Table,
+    flagsmith_environments_v2_table: Table,
     dynamo_enabled_project_environment_one: Environment,
-) -> dict:  # type: ignore[type-arg]
+) -> dict[str, typing.Any]:
     environment_dict = map_environment_to_environment_document(
         dynamo_enabled_project_environment_one
     )
-
     flagsmith_environment_table.put_item(
         Item=environment_dict,
     )
+
+    flagsmith_environments_v2_table.put_item(
+        Item=map_environment_to_environment_v2_document(
+            dynamo_enabled_project_environment_one
+        )
+    )
+
     return environment_dict
 
 

api/tests/types.py

@@ -4,6 +4,8 @@
 from organisations.permissions.models import UserOrganisationPermission
 from projects.models import UserProjectPermission
 
+# TODO: these type aliases aren't strictly correct according to mypy
+#  See here for more details: https://github.com/Flagsmith/flagsmith/issues/5140
 WithProjectPermissionsCallable = Callable[
     [list[str] | None, int | None, bool], UserProjectPermission
 ]

api/tests/unit/edge_api/identities/conftest.py

@@ -3,6 +3,13 @@
 from edge_api.identities.models import EdgeIdentity
 from environments.models import Environment
 from features.models import Feature
+from projects.models import Project
+
+
+@pytest.fixture(autouse=True)
+def enable_dynamodb_on_default_project(project: Project) -> None:
+    project.enable_dynamo_db = True
+    project.save()
 
 
 @pytest.fixture()