From e5f33ac3936281f2e9d03932fe3c60678249bd8b Mon Sep 17 00:00:00 2001 From: Nikita Melkozerov Date: Fri, 11 Oct 2024 18:16:57 +0200 Subject: [PATCH] Add scape_org_role_arn to account information (#635) --- .devcontainer/devcontainer.json | 2 +- fixbackend/collect/collect_queue.py | 1 + fixbackend/dispatcher/dispatcher_service.py | 16 ++++++++++++++-- tests/fixbackend/collect/collect_queue_test.py | 8 +++++++- .../dispatcher/dispatcher_service_test.py | 3 +++ 5 files changed, 26 insertions(+), 4 deletions(-) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 4ca19588..c50bc340 100644 --- 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 index 833a35fd..06a44b1e 100644 --- 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 index 7d312126..2e659e81 100644 --- 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,11 @@ 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/{role_name}") + if privileged_account_id + else None + ), external_id=external_id, ) case GcpCloudAccess(service_account_key_id): @@ -571,6 +577,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 +586,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 index 2c70a5a8..eb75169c 100644 --- 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,6 +84,7 @@ 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() == { @@ -87,6 +92,7 @@ def test_aws_account_info_json() -> None: "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 index 9e8c6fa5..9b992f5f 100644 --- 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,