120 lines
5.4 KiB
Python
120 lines
5.4 KiB
Python
import unittest
|
|
import warnings
|
|
from moto import mock_kms, mock_iam, mock_ecr
|
|
from endgame.shared.aws_login import get_boto3_client
|
|
import json
|
|
from endgame.shared import constants, utils
|
|
from endgame.shared.resource_results import ResourceResults, ServiceResourcesMultiRegion, ServiceResourcesSingleRegion
|
|
|
|
|
|
class ResourceResultsTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
self.mock_kms = mock_kms()
|
|
self.mock_iam = mock_iam()
|
|
self.mock_ecr = mock_ecr()
|
|
self.mock_kms.start()
|
|
self.mock_ecr.start()
|
|
self.mock_iam.start()
|
|
self.current_account_id = "123456789012"
|
|
# Set up KMS keys in 3 regions
|
|
self.ue1_kms_client = get_boto3_client(profile=None, service="kms", region="us-east-1")
|
|
self.ue2_kms_client = get_boto3_client(profile=None, service="kms", region="us-east-2")
|
|
self.euw1_kms_client = get_boto3_client(profile=None, service="kms", region="eu-west-1")
|
|
region = "us-east-1"
|
|
self.iam_client = get_boto3_client(profile=None, service="iam", region=region)
|
|
|
|
# self.iam_client = get_boto3_client(profile=None, service="iam", region="us-east-1")
|
|
self.uw1_ecr_client = get_boto3_client(profile=None, service="ecr", region="us-west-1")
|
|
|
|
# Get the ARNs so you can test them later
|
|
self.ue1_key_arn = self.ue1_kms_client.create_key()["KeyMetadata"]["Arn"]
|
|
self.ue2_key_arn = self.ue2_kms_client.create_key()["KeyMetadata"]["Arn"]
|
|
self.euw1_key_arn = self.euw1_kms_client.create_key()["KeyMetadata"]["Arn"]
|
|
iam_role = self.iam_client.create_role(
|
|
RoleName="yolo",
|
|
AssumeRolePolicyDocument=json.dumps(constants.EC2_ASSUME_ROLE_POLICY)
|
|
)
|
|
self.iam_role_arn = iam_role["Role"]["Arn"]
|
|
self.ecr_arn = self.uw1_ecr_client.create_repository(repositoryName="alpine")["repository"]["repositoryArn"]
|
|
|
|
# Create the objects that will store our results
|
|
self.service_resources_single_region = ServiceResourcesSingleRegion(
|
|
user_provided_service="kms",
|
|
region="us-east-1",
|
|
current_account_id=self.current_account_id,
|
|
profile=None,
|
|
cloak=False
|
|
)
|
|
self.service_resources_multi_region = ServiceResourcesMultiRegion(
|
|
user_provided_service="kms",
|
|
user_provided_region="all",
|
|
current_account_id=self.current_account_id,
|
|
profile=None,
|
|
cloak=False
|
|
)
|
|
# Let's exclude all the services that do not match the ones we specified above.
|
|
# That way, we can avoid an issue where moto tries to make API calls it does not support
|
|
excluded_services = []
|
|
for supported_service in constants.SUPPORTED_AWS_SERVICES:
|
|
if supported_service not in ["all", "iam", "ecr", "kms"]:
|
|
excluded_services.append(supported_service)
|
|
print(f"Excluded services: {excluded_services}")
|
|
self.resource_results = ResourceResults(
|
|
user_provided_service="all",
|
|
user_provided_region="all",
|
|
current_account_id=self.current_account_id,
|
|
profile=None,
|
|
cloak=False,
|
|
excluded_names=[],
|
|
excluded_services=excluded_services
|
|
)
|
|
|
|
def test_kms_single_region_arns(self):
|
|
results = self.service_resources_single_region.resources
|
|
for resource in results.resources:
|
|
print(resource.arn)
|
|
|
|
expected_results = [self.ue1_key_arn]
|
|
print(self.service_resources_single_region.arns)
|
|
self.assertListEqual(self.service_resources_single_region.arns, expected_results)
|
|
|
|
def test_kms_regions_multi_regions(self):
|
|
print("Test which regions the service appears in")
|
|
expected_regions = ['af-south-1', 'ap-east-1', 'ap-northeast-1', 'ap-northeast-2', 'ap-south-1',
|
|
'ap-southeast-1', 'ap-southeast-2', 'ca-central-1', 'eu-central-1', 'eu-north-1',
|
|
'eu-south-1', 'eu-west-1', 'eu-west-2', 'eu-west-3', 'me-south-1', 'sa-east-1', 'us-east-1',
|
|
'us-east-2', 'us-west-1', 'us-west-2']
|
|
kms_regions = self.service_resources_multi_region.regions
|
|
print(kms_regions)
|
|
# Let's future proof the tests, which is important once AWS adds more regions
|
|
for kms_region in kms_regions:
|
|
self.assertTrue(kms_region in expected_regions)
|
|
|
|
def test_kms_arns_multi_regions(self):
|
|
print("Test resource ARNs")
|
|
arns = self.service_resources_multi_region.arns
|
|
expected_arns = [
|
|
self.ue1_key_arn,
|
|
self.ue2_key_arn,
|
|
self.euw1_key_arn
|
|
]
|
|
print(arns)
|
|
for arn in expected_arns:
|
|
self.assertTrue(arn in arns)
|
|
|
|
def test_resource_results_arns(self):
|
|
print("Get ARNs from all the different services")
|
|
arns = self.resource_results.arns()
|
|
print(arns)
|
|
expected_arns = [
|
|
self.ue1_key_arn,
|
|
self.ue2_key_arn,
|
|
self.euw1_key_arn,
|
|
self.iam_role_arn,
|
|
self.ecr_arn
|
|
]
|
|
for arn in expected_arns:
|
|
self.assertTrue(arn in arns)
|