From 1619c7db28fd0f8c665f26c5821209680f41bb3d Mon Sep 17 00:00:00 2001 From: Surbhi Date: Wed, 4 Dec 2024 16:23:33 -0500 Subject: [PATCH 1/2] feat: adding capture for 2fa enforcements for everyone in the org --- posthog/api/organization.py | 24 +++++++++++++++++++++++- posthog/api/test/test_organization.py | 17 ++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/posthog/api/organization.py b/posthog/api/organization.py index c522ca164c0b9..b5d9ccf9f28c1 100644 --- a/posthog/api/organization.py +++ b/posthog/api/organization.py @@ -5,6 +5,8 @@ from django.shortcuts import get_object_or_404 from rest_framework import exceptions, permissions, serializers, viewsets from rest_framework.request import Request +from rest_framework.response import Response +import posthoganalytics from posthog import settings from posthog.api.routing import TeamAndOrgViewSetMixin @@ -12,7 +14,7 @@ from posthog.auth import PersonalAPIKeyAuthentication from posthog.cloud_utils import is_cloud from posthog.constants import INTERNAL_BOT_EMAIL_SUFFIX, AvailableFeature -from posthog.event_usage import report_organization_deleted +from posthog.event_usage import report_organization_deleted, groups from posthog.models import Organization, User from posthog.models.async_deletion import AsyncDeletion, DeletionType from posthog.rbac.user_access_control import UserAccessControlSerializerMixin @@ -240,3 +242,23 @@ def get_serializer_context(self) -> dict[str, Any]: **super().get_serializer_context(), "user_permissions": UserPermissions(cast(User, self.request.user)), } + + def update(self, request: Request, *args: Any, **kwargs: Any) -> Response: + if "enforce_2fa" in request.data: + enforce_2fa_value = request.data["enforce_2fa"] + organization = self.get_object() + + # Add capture event for 2FA enforcement change + posthoganalytics.capture( + request.user.distinct_id, + "organization 2fa enforcement toggled", + properties={ + "enabled": enforce_2fa_value, + "organization_id": str(organization.id), + "organization_name": organization.name, + "user_role": request.user.organization_memberships.get(organization=organization).level, + }, + groups=groups(organization), + ) + + return super().update(request, *args, **kwargs) diff --git a/posthog/api/test/test_organization.py b/posthog/api/test/test_organization.py index 2396f78e3c557..143fbe3f524b9 100644 --- a/posthog/api/test/test_organization.py +++ b/posthog/api/test/test_organization.py @@ -1,4 +1,5 @@ from rest_framework import status +from unittest.mock import patch, ANY from posthog.models import Organization, OrganizationMembership, Team from posthog.models.personal_api_key import PersonalAPIKey, hash_key_value @@ -128,7 +129,8 @@ def test_cant_update_plugins_access_level(self): self.organization.refresh_from_db() self.assertEqual(self.organization.plugins_access_level, 3) - def test_enforce_2fa_for_everyone(self): + @patch("posthoganalytics.capture") + def test_enforce_2fa_for_everyone(self, mock_capture): # Only admins should be able to enforce 2fa response = self.client.patch(f"/api/organizations/{self.organization.id}/", {"enforce_2fa": True}) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) @@ -142,6 +144,19 @@ def test_enforce_2fa_for_everyone(self): self.organization.refresh_from_db() self.assertEqual(self.organization.enforce_2fa, True) + # Verify the capture event was called correctly + mock_capture.assert_any_call( + self.user.distinct_id, + "organization 2fa enforcement toggled", + properties={ + "enabled": True, + "organization_id": str(self.organization.id), + "organization_name": self.organization.name, + "user_role": OrganizationMembership.Level.ADMIN, + }, + groups={"instance": ANY, "organization": str(self.organization.id)}, + ) + def test_projects_outside_personal_api_key_scoped_organizations_not_listed(self): other_org, _, _ = Organization.objects.bootstrap(self.user) personal_api_key = generate_random_token_personal() From ad8e07838f361d1dd486e2a3264581707a61bec7 Mon Sep 17 00:00:00 2001 From: Surbhi Date: Wed, 4 Dec 2024 16:39:20 -0500 Subject: [PATCH 2/2] adding type casts to resolve type check errors --- posthog/api/organization.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/posthog/api/organization.py b/posthog/api/organization.py index b5d9ccf9f28c1..6fe798479dd7b 100644 --- a/posthog/api/organization.py +++ b/posthog/api/organization.py @@ -247,16 +247,17 @@ def update(self, request: Request, *args: Any, **kwargs: Any) -> Response: if "enforce_2fa" in request.data: enforce_2fa_value = request.data["enforce_2fa"] organization = self.get_object() + user = cast(User, request.user) # Add capture event for 2FA enforcement change posthoganalytics.capture( - request.user.distinct_id, + str(user.distinct_id), "organization 2fa enforcement toggled", properties={ "enabled": enforce_2fa_value, "organization_id": str(organization.id), "organization_name": organization.name, - "user_role": request.user.organization_memberships.get(organization=organization).level, + "user_role": user.organization_memberships.get(organization=organization).level, }, groups=groups(organization), )