Add scape_org_role_arn to account information

someengineering · Oct 11, 2024 · 3bbc814 · 3bbc814
1 parent 38dd9f7
commit 3bbc814
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 4 deletions.
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -8,7 +8,7 @@
 	// Features to add to the dev container. More info: https://containers.dev/features.
 	"features": {
 		"ghcr.io/devcontainers-contrib/features/poetry:2": {},
-		"ghcr.io/devcontainers/features/node:1": {},
+		"ghcr.io/devcontainers/features/node:1.5": {}
 	},
 	// Use 'forwardPorts' to make a list of ports inside the container available locally.
 	// This can be used to network with other containers or the host.

diff --git a/fixbackend/collect/collect_queue.py b/fixbackend/collect/collect_queue.py
@@ -54,6 +54,7 @@ class AwsAccountInformation(AccountInformation):
     aws_account_id: CloudAccountId
     aws_account_name: Optional[str]
     aws_role_arn: AwsARN
+    scrape_org_role_arn: Optional[AwsARN]
     external_id: ExternalId
 
 

diff --git a/fixbackend/dispatcher/dispatcher_service.py b/fixbackend/dispatcher/dispatcher_service.py
@@ -488,6 +488,7 @@ async def trigger_collect(
         reason: str,
         defer_by: Optional[timedelta] = None,
         retry_failed_for: Optional[timedelta] = None,
+        privileged_account_id: Optional[CloudAccountId] = None,
         **kwargs: Any,
     ) -> None:
         set_cloud_account_id(account.account_id)
@@ -507,6 +508,9 @@ async def account_information() -> Optional[AccountInformation]:
                                 aws_account_id=account.account_id,
                                 aws_account_name=account.final_name(),
                                 aws_role_arn=AwsARN(f"arn:aws:iam::{account.account_id}:role/{role_name}"),
+                                scrape_org_role_arn=AwsARN(
+                                    f"arn:aws:iam::{privileged_account_id}:role/FixCrossAccountAccessRole"
+                                ),
                                 external_id=external_id,
                             )
                         case GcpCloudAccess(service_account_key_id):
@@ -571,6 +575,7 @@ async def schedule_next_runs(self) -> None:
             accounts = healthy_accounts + degraded_accounts
             product_tier = await self.workspace_repository.get_product_tier(workspace_id)
             log.info(f"scheduling next run for workspace {workspace_id}, {len(accounts)} accounts")
+            priveleged_account_id = next((acc.account_id for acc in accounts if acc.privileged), None)
             for account in accounts:
                 reason = "regular_collect"
 
@@ -579,9 +584,14 @@ async def schedule_next_runs(self) -> None:
 
                 if account.cloud == CloudNames.Azure and not azure_graph_scheduled:
                     azure_graph_scheduled = True
-                    await self.trigger_collect(account, reason=reason, collect_microsoft_graph=True)
+                    await self.trigger_collect(
+                        account,
+                        reason=reason,
+                        collect_microsoft_graph=True,
+                        privileged_account_id=priveleged_account_id,
+                    )
                 else:
-                    await self.trigger_collect(account, reason=reason)
+                    await self.trigger_collect(account, reason=reason, privileged_account_id=priveleged_account_id)
 
             for account in degraded_accounts:
                 await self.cloud_account_repo.update(

diff --git a/tests/fixbackend/collect/collect_queue_test.py b/tests/fixbackend/collect/collect_queue_test.py
@@ -49,7 +49,11 @@ async def test_redis_collect_queue(
     assert set(await arq_redis.keys()) == set()
     # enqueue new job
     aws_account = AwsAccountInformation(
-        CloudAccountId("123"), CloudAccountName("test"), AwsARN("arn"), ExternalId(uuid4())
+        CloudAccountId("123"),
+        CloudAccountName("test"),
+        AwsARN("arn"),
+        AwsARN("scarpe_arn"),
+        ExternalId(uuid4()),
     )
     await collect_queue.enqueue(graph_db_access, aws_account, job_id="test")
     assert set(await arq_redis.keys()) == {b"arq:queue", b"arq:job:test"}
@@ -80,13 +84,15 @@ def test_aws_account_info_json() -> None:
         aws_account_id=CloudAccountId("123456789012"),
         aws_account_name=CloudAccountName("test"),
         aws_role_arn=AwsARN("arn:aws:iam::123456789012:role/test"),
+        scrape_org_role_arn=AwsARN("arn:aws:iam::123456789012:role/scrape-test"),
         external_id=external_id,
     )
     assert aws_account_info.to_json() == {
         "kind": "aws_account_information",
         "aws_account_id": "123456789012",
         "aws_account_name": "test",
         "aws_role_arn": "arn:aws:iam::123456789012:role/test",
+        "scrape_org_role_arn": "arn:aws:iam::123456789012:role/scrape-test",
         "external_id": str(external_id),
     }
 

diff --git a/tests/fixbackend/dispatcher/dispatcher_service_test.py b/tests/fixbackend/dispatcher/dispatcher_service_test.py
@@ -241,6 +241,7 @@ async def jobs_mapping_hash_len() -> int:
             aws_account_id=aws_account_id,
             aws_account_name=CloudAccountName("test"),
             aws_role_arn=AwsARN("arn"),
+            scrape_org_role_arn=AwsARN("scrape_arn"),
             external_id=external_id,
         ),
         job_id,
@@ -391,6 +392,7 @@ async def jobs_mapping_hash_len() -> int:
             aws_account_id=aws_account_id,
             aws_account_name=CloudAccountName("test"),
             aws_role_arn=AwsARN("arn"),
+            scrape_org_role_arn=AwsARN("scrape_arn"),
             external_id=external_id,
         ),
         job_id,
@@ -478,6 +480,7 @@ async def jobs_mapping_hash_len() -> int:
             aws_account_id=aws_account_id,
             aws_account_name=CloudAccountName("test"),
             aws_role_arn=AwsARN("arn"),
+            scrape_org_role_arn=AwsARN("scrape_arn"),
             external_id=external_id,
         ),
         job_id,