Repository: lyft/cartography

Issue:
AWS IAM: Principals are not parsed correctly for multiple principal map entries
**AWS IAM: Principals are not parsed correctly for multiple principal map entries**

*Initial Fix + Unit Test*
roshinis78@0459cb317ee23f5b99b9668c19b283544db4594e

*Description*
In AWS policy statements, principals can have more than one entry (including "AWS", "Service", "Federated", or "CanonicalUser"). If more than one entry exists in the principal object, all entries should be ingested into the database as principals. For example, for a policy with entries for "AWS" and "Service", both entries should be ingested into the database. Currently, only the first entry (in this case, with the key "AWS") is ingested.

*To Reproduce*
Try ingesting a policy statement with the following principal. Only "test-role-1" is ingested into the database.
```python
SINGLE_PRINCIPAL = {
    "AWS": "test-role-1",
    "Service": ["test-service-1", "test-service-2"],
    "Federated": "test-provider-1"
}
```

*Notes*
When I ran the pre-commit lint with `make test_lint`, trailing commas were added to files that I didn't modify, so I included these in a separate commit: roshinis78@473b185d85f9086b9c24fb0bb586b45ec0e1d6dd

*Version Info*
>  - Cartography commit hash: fb68aac03f785a67990b5c31139a539c0798b255
>  - Python version: 3.8.1

*Additional Context*
> [AWS Policy Grammar](https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_grammar.html) 




Discussion:

Patch:
diff --git a/cartography/intel/aws/iam.py b/cartography/intel/aws/iam.py
index 07dfdd82..9635e65f 100644
--- a/cartography/intel/aws/iam.py
+++ b/cartography/intel/aws/iam.py
@@ -224,6 +224,21 @@ def load_groups(neo4j_session, groups, current_aws_account_id, aws_update_tag):
         )
 
 
+def _parse_principal_entries(principal):
+    """
+    Returns a list of tuples of the form (principal_type, principal_value)
+    e.g. [('AWS', 'example-role-name'), ('Service', 'example-service')]
+    """
+    principal_entries = []
+    for principal_type in principal:
+        principal_values = principal[principal_type]
+        if not isinstance(principal_values, list):
+            principal_values = [principal_values]
+        for principal_value in principal_values:
+            principal_entries.append((principal_type, principal_value))
+    return principal_entries
+
+
 @timeit
 def load_roles(neo4j_session, roles, current_aws_account_id, aws_update_tag):
     ingest_role = """
@@ -265,17 +280,8 @@ def load_roles(neo4j_session, roles, current_aws_account_id, aws_update_tag):
         )
 
         for statement in role["AssumeRolePolicyDocument"]["Statement"]:
-            principal = statement["Principal"]
-            principal_values = []
-            if 'AWS' in principal:
-                principal_type, principal_values = 'AWS', principal['AWS']
-            elif 'Service' in principal:
-                principal_type, principal_values = 'Service', principal['Service']
-            elif 'Federated' in principal:
-                principal_type, principal_values = 'Federated', principal['Federated']
-            if not isinstance(principal_values, list):
-                principal_values = [principal_values]
-            for principal_value in principal_values:
+            principal_entries = _parse_principal_entries(statement["Principal"])
+            for principal_type, principal_value in principal_entries:
                 neo4j_session.run(
                     ingest_policy_statement,
                     SpnArn=principal_value,


Test patch:
diff --git a/tests/integration/cartography/intel/aws/ec2/test_ec2_instances.py b/tests/integration/cartography/intel/aws/ec2/test_ec2_instances.py
index c8ec4061..9020953e 100644
--- a/tests/integration/cartography/intel/aws/ec2/test_ec2_instances.py
+++ b/tests/integration/cartography/intel/aws/ec2/test_ec2_instances.py
@@ -37,7 +37,7 @@ def test_load_ec2_instances(neo4j_session, *args):
     nodes = neo4j_session.run(
         """
         MATCH (i:EC2Instance) return i.id, i.instanceid
-        """
+        """,
     )
     actual_nodes = {
         (
@@ -80,7 +80,7 @@ def test_ec2_reservations_to_instances(neo4j_session, *args):
     nodes = neo4j_session.run(
         """
     MATCH (r:EC2Reservation)<-[:MEMBER_OF_EC2_RESERVATION]-(i:EC2Instance) RETURN r.reservationid, i.id
-    """
+    """,
     )
     actual_nodes = {
         (
diff --git a/tests/integration/cartography/intel/aws/ec2/test_ec2_key_pairs.py b/tests/integration/cartography/intel/aws/ec2/test_ec2_key_pairs.py
index f4ea542e..0c62d2ce 100644
--- a/tests/integration/cartography/intel/aws/ec2/test_ec2_key_pairs.py
+++ b/tests/integration/cartography/intel/aws/ec2/test_ec2_key_pairs.py
@@ -34,7 +34,7 @@ def test_load_ec2_key_pairs(neo4j_session, *args):
     nodes = neo4j_session.run(
         """
         MATCH (k:EC2KeyPair) return k.arn, k.keyfingerprint
-        """
+        """,
     )
     actual_nodes = {
         (
diff --git a/tests/integration/cartography/intel/aws/ec2/test_ec2_load_balancers.py b/tests/integration/cartography/intel/aws/ec2/test_ec2_load_balancers.py
index 135b746f..fb8d4f73 100644
--- a/tests/integration/cartography/intel/aws/ec2/test_ec2_load_balancers.py
+++ b/tests/integration/cartography/intel/aws/ec2/test_ec2_load_balancers.py
@@ -234,7 +234,7 @@ def test_load_load_balancer_v2_subnets(neo4j_session, *args):
     nodes = neo4j_session.run(
         """
         MATCH (subnet:EC2Subnet) return subnet.subnetid, subnet.region, subnet.lastupdated
-        """
+        """,
     )
     actual_nodes = {
         (
diff --git a/tests/integration/cartography/intel/aws/test_dynamodb.py b/tests/integration/cartography/intel/aws/test_dynamodb.py
index 3b723364..6356b745 100644
--- a/tests/integration/cartography/intel/aws/test_dynamodb.py
+++ b/tests/integration/cartography/intel/aws/test_dynamodb.py
@@ -27,7 +27,7 @@ def test_load_dynamodb(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (d:DynamoDBTable) return d.arn, d.rows
-        """
+        """,
     )
     actual_nodes = {(n['d.arn'], n['d.rows']) for n in nodes}
     assert actual_nodes == expected_nodes
diff --git a/tests/integration/cartography/intel/aws/test_eks.py b/tests/integration/cartography/intel/aws/test_eks.py
index c3f465bf..cf358342 100644
--- a/tests/integration/cartography/intel/aws/test_eks.py
+++ b/tests/integration/cartography/intel/aws/test_eks.py
@@ -24,7 +24,7 @@ def test_load_eks_clusters_nodes(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (r:EKSCluster) RETURN r.arn;
-        """
+        """,
     )
     actual_nodes = {n['r.arn'] for n in nodes}
 
@@ -61,7 +61,7 @@ def test_load_eks_clusters_relationships(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:AWSAccount)-[:RESOURCE]->(n2:EKSCluster) RETURN n1.id, n2.arn;
-        """
+        """,
     )
     actual = {
         (r['n1.id'], r['n2.arn']) for r in result
diff --git a/tests/integration/cartography/intel/aws/test_iam.py b/tests/integration/cartography/intel/aws/test_iam.py
index a0ebbe03..9e40f9db 100644
--- a/tests/integration/cartography/intel/aws/test_iam.py
+++ b/tests/integration/cartography/intel/aws/test_iam.py
@@ -59,7 +59,7 @@ def test_load_roles_creates_trust_relationships(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:AWSRole)-[:TRUSTS_AWS_PRINCIPAL]->(n2:AWSPrincipal) RETURN n1.arn, n2.arn;
-        """
+        """,
     )
 
     # Define the relationships we expect in terms of role ARN and principal ARN.
diff --git a/tests/integration/cartography/intel/aws/test_redshift.py b/tests/integration/cartography/intel/aws/test_redshift.py
index fb8cc962..a213ac82 100644
--- a/tests/integration/cartography/intel/aws/test_redshift.py
+++ b/tests/integration/cartography/intel/aws/test_redshift.py
@@ -16,7 +16,7 @@ def test_load_redshift_cluster_data(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (n:RedshiftCluster) RETURN n.id;
-        """
+        """,
     )
     actual_nodes = {n['n.id'] for n in nodes}
     assert actual_nodes == expected_nodes
@@ -42,7 +42,7 @@ def test_load_redshift_cluster_and_aws_account(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:AWSAccount)-[:RESOURCE]->(n2:RedshiftCluster) RETURN n1.id, n2.id;
-        """
+        """,
     )
     actual = {
         (r['n1.id'], r['n2.id']) for r in result
@@ -70,7 +70,7 @@ def test_load_redshift_cluster_and_security_group(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:EC2SecurityGroup)<-[:MEMBER_OF_EC2_SECURITY_GROUP]-(n2:RedshiftCluster) RETURN n1.id, n2.id;
-        """
+        """,
     )
     actual = {
         (r['n1.id'], r['n2.id']) for r in result
@@ -98,7 +98,7 @@ def test_load_redshift_cluster_and_iam_role(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:AWSPrincipal)<-[:STS_ASSUMEROLE_ALLOW]-(n2:RedshiftCluster) RETURN n1.arn, n2.id;
-        """
+        """,
     )
     actual = {
         (r['n1.arn'], r['n2.id']) for r in result
@@ -126,7 +126,7 @@ def test_load_redshift_cluster_and_vpc(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:AWSVpc)<-[:MEMBER_OF_AWS_VPC]-(n2:RedshiftCluster) RETURN n1.id, n2.id;
-        """
+        """,
     )
     actual = {
         (r['n1.id'], r['n2.id']) for r in result
diff --git a/tests/integration/cartography/intel/aws/test_resourcegroupstaggingapi.py b/tests/integration/cartography/intel/aws/test_resourcegroupstaggingapi.py
index 7f89a71c..0238f59f 100644
--- a/tests/integration/cartography/intel/aws/test_resourcegroupstaggingapi.py
+++ b/tests/integration/cartography/intel/aws/test_resourcegroupstaggingapi.py
@@ -38,7 +38,7 @@ def test_transform_and_load_ec2_tags(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:EC2Instance)-[:TAGGED]->(n2:AWSTag) RETURN n1.id, n2.id;
-        """
+        """,
     )
     actual = {
         (r['n1.id'], r['n2.id']) for r in result
diff --git a/tests/integration/cartography/intel/aws/test_route53.py b/tests/integration/cartography/intel/aws/test_route53.py
index a7b89362..720fb237 100644
--- a/tests/integration/cartography/intel/aws/test_route53.py
+++ b/tests/integration/cartography/intel/aws/test_route53.py
@@ -77,7 +77,7 @@ def test_load_and_cleanup_dnspointsto_relationships(neo4j_session):
         """
         MERGE (n1:AWSDNSRecord{id:"/hostedzone/HOSTED_ZONE/example.com/NS"})
         -[:DNS_POINTS_TO]->(:NewTestAsset{name:"hello"})
-        """
+        """,
     )
 
     # Verify that the expected AWS DNS records point to each other
@@ -85,7 +85,7 @@ def test_load_and_cleanup_dnspointsto_relationships(neo4j_session):
         """
         MATCH (n1:AWSDNSRecord{id:"/hostedzone/HOSTED_ZONE/example.com/NS"})-[:DNS_POINTS_TO]->(n2:AWSDNSRecord)
         RETURN n1.name, n2.id
-        """
+        """,
     )
     expected = {("example.com", "/hostedzone/HOSTED_ZONE/example.com/A")}
     actual = {(r['n1.name'], r['n2.id']) for r in result}
@@ -109,7 +109,7 @@ def test_load_and_cleanup_dnspointsto_relationships(neo4j_session):
         """
         MATCH (n1:AWSDNSRecord{id:"/hostedzone/HOSTED_ZONE/example.com/NS"})-[:DNS_POINTS_TO]->(n2:AWSDNSRecord)
         RETURN count(n2) as recordcount
-        """
+        """,
     )
     for r in result:
         assert r["recordcount"] == 0
@@ -119,7 +119,7 @@ def test_load_and_cleanup_dnspointsto_relationships(neo4j_session):
         """
         MATCH (n1:AWSDNSRecord{id:"/hostedzone/HOSTED_ZONE/example.com/NS"})-[:DNS_POINTS_TO]->(n2:NewTestAsset)
         RETURN n1.id, n2.name
-        """
+        """,
     )
     actual = {(r['n1.id'], r['n2.name']) for r in result}
     expected = {("/hostedzone/HOSTED_ZONE/example.com/NS", "hello")}
diff --git a/tests/integration/cartography/intel/aws/test_s3.py b/tests/integration/cartography/intel/aws/test_s3.py
index 92d5bf5c..cf29bcb5 100644
--- a/tests/integration/cartography/intel/aws/test_s3.py
+++ b/tests/integration/cartography/intel/aws/test_s3.py
@@ -32,7 +32,7 @@ def test_load_s3_buckets(neo4j_session, *args):
     nodes = neo4j_session.run(
         """
         MATCH (s:S3Bucket) return s.id, s.name
-        """
+        """,
     )
     actual_nodes = {
         (
diff --git a/tests/integration/cartography/intel/gcp/test_crm.py b/tests/integration/cartography/intel/gcp/test_crm.py
index 7b335930..6692fc96 100644
--- a/tests/integration/cartography/intel/gcp/test_crm.py
+++ b/tests/integration/cartography/intel/gcp/test_crm.py
@@ -32,7 +32,7 @@ def test_load_gcp_projects(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (d:GCPProject) return d.id
-        """
+        """,
     )
     actual_nodes = {(n['d.id']) for n in nodes}
     assert actual_nodes == expected_nodes
@@ -80,7 +80,7 @@ def test_load_gcp_projects_without_parent(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (d:GCPProject) WHERE NOT (d)<-[:RESOURCE]-() RETURN d.id
-        """
+        """,
     )
     actual_nodes = {(n['d.id']) for n in nodes}
     assert actual_nodes == expected_nodes
diff --git a/tests/integration/cartography/intel/gcp/test_gke.py b/tests/integration/cartography/intel/gcp/test_gke.py
index c92ad2ff..7e80af61 100644
--- a/tests/integration/cartography/intel/gcp/test_gke.py
+++ b/tests/integration/cartography/intel/gcp/test_gke.py
@@ -22,7 +22,7 @@ def test_load_gke_clusters(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (r:GKECluster) RETURN r.id;
-        """
+        """,
     )
 
     actual_nodes = {n['r.id'] for n in nodes}
@@ -59,7 +59,7 @@ def test_load_eks_clusters_relationships(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n1:GCPProject)-[:RESOURCE]->(n2:GKECluster) RETURN n1.id, n2.id;
-        """
+        """,
     )
 
     actual = {
diff --git a/tests/integration/cartography/intel/jamf/test_jamf.py b/tests/integration/cartography/intel/jamf/test_jamf.py
index 214d038e..b29bc187 100644
--- a/tests/integration/cartography/intel/jamf/test_jamf.py
+++ b/tests/integration/cartography/intel/jamf/test_jamf.py
@@ -16,7 +16,7 @@ def test_load_jamf_computer_group_data(neo4j_session):
     nodes = neo4j_session.run(
         """
         MATCH (n:JamfComputerGroup) RETURN n.id;
-        """
+        """,
     )
     actual_nodes = {n['n.id'] for n in nodes}
     assert actual_nodes == expected_nodes
diff --git a/tests/unit/cartography/intel/aws/test_iam.py b/tests/unit/cartography/intel/aws/test_iam.py
index 411b026b..30beb9a2 100644
--- a/tests/unit/cartography/intel/aws/test_iam.py
+++ b/tests/unit/cartography/intel/aws/test_iam.py
@@ -5,6 +5,14 @@ SINGLE_STATEMENT = {
     "Action": "*",
 }
 
+# Example principal field in an AWS policy statement
+# see: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_principal.html
+SINGLE_PRINCIPAL = {
+    "AWS": "test-role-1",
+    "Service": ["test-service-1", "test-service-2"],
+    "Federated": "test-provider-1",
+}
+
 
 def test__generate_policy_statements():
     statements = iam._transform_policy_statements(SINGLE_STATEMENT, "test_policy_id")
@@ -12,3 +20,13 @@ def test__generate_policy_statements():
     assert isinstance(statements[0]["Action"], list)
     assert isinstance(statements[0]["Resource"], list)
     assert statements[0]["id"] == "test_policy_id/statement/1"
+
+
+def test__parse_principal_entries():
+    principal_entries = iam._parse_principal_entries(SINGLE_PRINCIPAL)
+    assert isinstance(principal_entries, list)
+    assert len(principal_entries) == 4
+    assert principal_entries[0] == ("AWS", "test-role-1")
+    assert principal_entries[1] == ("Service", "test-service-1")
+    assert principal_entries[2] == ("Service", "test-service-2")
+    assert principal_entries[3] == ("Federated", "test-provider-1")