Fix #348 - Connect AWS Lambda functions to their associated AWSAccoun…

…t nodes (#359) * add account owner for lambda * updated integration test and clean up * added test_account_id
cartography-cncf · Jul 19, 2020 · 1b9b86b · 1b9b86b
1 parent cfa33fc
commit 1b9b86b
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 0 deletions.
diff --git a/cartography/data/jobs/cleanup/aws_import_lambda_cleanup.json b/cartography/data/jobs/cleanup/aws_import_lambda_cleanup.json
@@ -5,6 +5,11 @@
             "iterative": true,
             "iterationsize": 100
         },
+        {
+            "query": "MATCH (:AWSAccount{id: {AWS_ID}})-[r:RESOURCE]->(:AWSLambda) WHERE r.lastupdated <> {UPDATE_TAG} WITH r LIMIT {LIMIT_SIZE} DELETE (r) return COUNT(*) as TotalCompleted",
+            "iterative": true,
+            "iterationsize": 100
+        },
         {
             "query": "MATCH (:AWSLambda)-[r:STS_ASSUME_ROLE_ALLOW]->(:AWSPrincipal) WITH r LIMIT {LIMIT_SIZE} DELETE r return COUNT(*) as TotalCompleted",
             "iterative": true,

diff --git a/cartography/intel/aws/lambda_function.py b/cartography/intel/aws/lambda_function.py
@@ -36,6 +36,11 @@ def load_lambda_functions(neo4j_session, data, region, current_aws_account_id, a
     lambda.memory = {MemorySize},
     lambda.lastupdated = {aws_update_tag}
     WITH lambda
+    MATCH (owner:AWSAccount{id: {AWS_ACCOUNT_ID}})
+    MERGE (owner)-[r:RESOURCE]->(lambda)
+    ON CREATE SET r.firstseen = timestamp()
+    SET r.lastupdated = {aws_update_tag}
+    WITH lambda
     MATCH (role:AWSPrincipal{arn: {Role}})
     MERGE (lambda)-[r:STS_ASSUME_ROLE_ALLOW]->(role)
     ON CREATE SET r.firstseen = timestamp()

diff --git a/tests/integration/cartography/intel/aws/test_lambda.py b/tests/integration/cartography/intel/aws/test_lambda.py
@@ -39,3 +39,52 @@ def test_load_lambda_functions(neo4j_session):
     actual_nodes = {n['r.id'] for n in nodes}
 
     assert actual_nodes == expected_nodes
+
+
+def test_load_lambda_relationships(neo4j_session):
+    # Create Test AWSAccount
+    neo4j_session.run(
+        """
+        MERGE (aws:AWSAccount{id: {aws_account_id}})
+        ON CREATE SET aws.firstseen = timestamp()
+        SET aws.lastupdated = {aws_update_tag}
+        """,
+        aws_account_id=TEST_ACCOUNT_ID,
+        aws_update_tag=TEST_UPDATE_TAG,
+    )
+
+    # Load Test Lambda Functions
+    data = tests.data.aws.lambda_function.LIST_LAMBDA_FUNCTIONS
+
+    cartography.intel.aws.lambda_function.load_lambda_functions(
+        neo4j_session,
+        data,
+        TEST_REGION,
+        TEST_ACCOUNT_ID,
+        TEST_UPDATE_TAG,
+    )
+
+    expected_nodes = {
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-1"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-2"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-3"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-4"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-5"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-6"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-7"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-8"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-9"),
+        (TEST_ACCOUNT_ID, "arn:aws:lambda:us-west-2:000000000000:function:sample-function-10"),
+    }
+
+    # Fetch relationships
+    result = neo4j_session.run(
+        """
+        MATCH (n1:AWSAccount)-[:RESOURCE]->(n2:AWSLambda) RETURN n1.id, n2.id;
+        """,
+    )
+    actual = {
+        (r['n1.id'], r['n2.id']) for r in result
+    }
+
+    assert actual == expected_nodes