gchqgh-35 Restrict graph access to specific cognito user pool users.

lib/app-stack.ts

@@ -43,7 +43,8 @@ export class AppStack extends cdk.Stack {
         // REST API
         const kaiRest = new KaiRestApi(this, "KaiRestApi", {
             graphTable: database.table,
-            userPoolArn: userPool.userPoolArn
+            userPoolArn: userPool.userPoolArn,
+            userPoolId: userPool.userPoolId
         });
 
         // Kubectl Lambda layer

lib/rest-api/kai-rest-api-props.ts

@@ -19,4 +19,5 @@ import { Table } from "@aws-cdk/aws-dynamodb";
 export interface KaiRestApiProps {
     graphTable: Table;
     userPoolArn: string;
+    userPoolId: string;
 }

lib/rest-api/kai-rest-api.ts

@@ -56,11 +56,17 @@ export class KaiRestApi extends cdk.Construct {
             timeout: lambdaTimeout,
             environment: {
                 sqs_queue_url: this.addGraphQueue.queueUrl,
-                graph_table_name: props.graphTable.tableName
+                graph_table_name: props.graphTable.tableName,
+                user_pool_id: props.userPoolId
             }
         });
 
-        props.graphTable.grantWriteData(addGraphLambda);
+        addGraphLambda.addToRolePolicy(new PolicyStatement({
+            actions: [ "cognito-idp:ListUsers" ],
+            resources: [ props.userPoolArn ]
+        }));
+
+        props.graphTable.grantReadWriteData(addGraphLambda);
         this.addGraphQueue.grantSendMessages(addGraphLambda);
         graphsResource.addMethod("POST", new api.LambdaIntegration(addGraphLambda), methodOptions);
 
@@ -76,11 +82,12 @@ export class KaiRestApi extends cdk.Construct {
             timeout: lambdaTimeout,
             environment: {
                 sqs_queue_url: this.deleteGraphQueue.queueUrl,
-                graph_table_name: props.graphTable.tableName
+                graph_table_name: props.graphTable.tableName,
+                user_pool_id: props.userPoolId
             }
         });
 
-        props.graphTable.grantWriteData(deleteGraphLambda);
+        props.graphTable.grantReadWriteData(deleteGraphLambda);
         this.deleteGraphQueue.grantSendMessages(deleteGraphLambda);
         graph.addMethod("DELETE", new api.LambdaIntegration(deleteGraphLambda), methodOptions);
 
@@ -91,7 +98,8 @@ export class KaiRestApi extends cdk.Construct {
             handler: "get_graph_request.handler",
             timeout: lambdaTimeout,
             environment: {
-                graph_table_name: props.graphTable.tableName
+                graph_table_name: props.graphTable.tableName,
+                user_pool_id: props.userPoolId
             }
         });
 

lib/rest-api/lambdas/add_graph_request.py

@@ -1,8 +1,13 @@
 import boto3
 from botocore.exceptions import ClientError
+from graph import Graph
 import json
 import os
 import re
+from user import User
+
+graph = Graph()
+user = User()
 
 def is_graph_id_valid(graph_id):
     if graph_id is None:
@@ -30,22 +35,25 @@ def handler(event, context):
 
     # Get variables from env
     queue_url = os.getenv("sqs_queue_url")
-    graph_table_name = os.getenv("graph_table_name")
-
-    # Add Entry to table
-    dynamo = boto3.resource("dynamodb")
-    table = dynamo.Table(graph_table_name)
 
     initial_status = "DEPLOYMENT_QUEUED"
 
+    administrators = []
+    requesting_user = user.get_requesting_cognito_user(event)
+    if requesting_user is not None:
+        administrators.append(requesting_user)
+    if "administrators" in request_body:
+        administrators.extend(request_body["administrators"])
+    if user.contains_duplicates(administrators):
+        administrators = list(set(administrators))
+    if not user.valid_cognito_users(administrators):
+        return {
+            "statusCode": 400,
+            "body": "Not all of the supplied administrators are valid Cognito users: {}".format(str(administrators))
+        }
+
     try:
-        table.put_item(
-            Item={
-                "graphId": graph_id,
-                "currentState": initial_status
-            },
-            ConditionExpression=boto3.dynamodb.conditions.Attr("graphId").not_exists()
-        )
+        graph.create_graph(graph_id, initial_status, administrators)
     except ClientError as e:
         if e.response['Error']['Code']=='ConditionalCheckFailedException': 
             return {

lib/rest-api/lambdas/delete_graph_request.py

@@ -1,15 +1,15 @@
 import boto3
+from botocore.exceptions import ClientError
+from graph import Graph
 import json
 import os
-from botocore.exceptions import ClientError
+from user import User
 
 # Get variables from env
 queue_url = os.getenv("sqs_queue_url")
-graph_table_name = os.getenv("graph_table_name")
 
-# Dynamodb table
-dynamo = boto3.resource("dynamodb")
-table = dynamo.Table(graph_table_name)
+graph = Graph()
+user = User()
 
 def handler(event, context):
     params = event["pathParameters"]
@@ -23,20 +23,18 @@ def handler(event, context):
             body: "graphId is a required field"
         }
 
+    requesting_user = user.get_requesting_cognito_user(event)
+    if not user.is_authorized(requesting_user, graph_id):
+        return {
+            "statusCode": 403,
+            "body": "User: {} is not authorized to delete graph: {}".format(requesting_user, graph_id)
+        }
+
     initial_status = "DELETION_QUEUED"
 
     # Add Entry to table
     try:
-        table.update_item(
-            Key={
-                "graphId": graph_id
-            },
-            UpdateExpression="SET currentState = :state",
-            ExpressionAttributeValues={
-                ":state": initial_status
-            },
-            ConditionExpression=boto3.dynamodb.conditions.Attr("graphId").exists()
-        )
+        graph.update_graph(graph_id, initial_status)
     except ClientError as e:
         if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
             return {

lib/rest-api/lambdas/get_graph_request.py

@@ -1,36 +1,11 @@
 import os
 import boto3
+from graph import Graph
 import json
+from user import User
 
-graph_table_name = os.getenv("graph_table_name")
-
-class NotFoundException(Exception):
-    pass
-
-
-# Dynamodb table
-dynamo = boto3.resource("dynamodb")
-table = dynamo.Table(graph_table_name)
-
-def get_all_graphs():
-    """
-    Gets all graphs from Dynamodb table
-    """
-    return table.scan()["Items"]
-
-
-def get_graph(graph_id):
-    """
-    Gets a specific graph from Dynamodb table
-    """
-    response = table.get_item(
-        Key={
-            "graphId": graph_id
-        }
-    )
-    if "Item" in response:
-        return response["Item"]
-    raise NotFoundException
+graph = Graph()
+user = User()
 
 
 def handler(event, context):
@@ -48,18 +23,26 @@ def handler(event, context):
     else:
         graph_id = path_params["graphId"]
 
+    requesting_user = user.get_requesting_cognito_user(event)
+
     if return_all:
         return {
             "statusCode": 200,
-            "body": json.dumps(get_all_graphs())
+            "body": json.dumps(graph.get_all_graphs(requesting_user))
         }
     else:
+        if not user.is_authorized(requesting_user, graph_id):
+            return {
+                "statusCode": 403,
+                "body": "User: {} is not authorized to retrieve graph: {}".format(requesting_user, graph_id)
+            }
+
         try:
             return {
                 "statusCode": 200,
-                "body": json.dumps(get_graph(graph_id))
+                "body": json.dumps(graph.get_graph(graph_id))
             }
-        except NotFoundException as e:
+        except Exception as e:
             return {
                 "statusCode": 404,
                 "body": graph_id + " was not found"

lib/rest-api/lambdas/graph/__init__.py

@@ -0,0 +1,59 @@
+import boto3
+import os
+
+
+class Graph:
+
+    def __init__(self):
+        dynamodb = boto3.resource("dynamodb")
+        graph_table_name = os.getenv("graph_table_name")
+        self.table = dynamodb.Table(graph_table_name)
+
+
+    def get_all_graphs(self, requesting_user):
+        """
+        Gets all graphs from Dynamodb table
+        """
+        graphs = self.table.scan()["Items"]
+        if requesting_user is None:
+            return graphs
+        else:
+            return list(filter(lambda graph: requesting_user in graph["administrators"], graphs))
+
+
+    def get_graph(self, graph_id):
+        """
+        Gets a specific graph from Dynamodb table
+        """
+        response = self.table.get_item(
+            Key={
+                "graphId": graph_id
+            }
+        )
+        if "Item" in response:
+            return response["Item"]
+        raise Exception
+
+
+    def update_graph(self, graph_id, status):
+        self.table.update_item(
+            Key={
+                "graphId": graph_id
+            },
+            UpdateExpression="SET currentState = :state",
+            ExpressionAttributeValues={
+                ":state": status
+            },
+            ConditionExpression=boto3.dynamodb.conditions.Attr("graphId").exists()
+        )
+
+    def create_graph(self, graph_id, status, administrators):
+        self.table.put_item(
+            Item={
+                "graphId": graph_id,
+                "currentState": status,
+                "administrators": administrators
+            },
+            ConditionExpression=boto3.dynamodb.conditions.Attr("graphId").not_exists()
+        )
+

lib/rest-api/lambdas/user/__init__.py

@@ -0,0 +1,40 @@
+import boto3
+from graph import Graph
+import os
+
+
+class User:
+
+    def __init__(self):
+        self.cognito_client = boto3.client('cognito-idp')
+        self.user_pool_id = os.getenv("user_pool_id")
+        self.graph = Graph()
+
+    def valid_cognito_users(self, users):
+        response = self.cognito_client.list_users(UserPoolId=self.user_pool_id)
+        cognito_users = map(self.to_user_name, response["Users"])
+        return set(users).issubset(set(cognito_users))
+
+    def to_user_name(self, user):
+        return user["Username"]
+
+    def contains_duplicates(self, items):
+        return set([item for item in items if items.count(item) > 1])
+
+    def get_requesting_cognito_user(self, request):
+        if ("authorizer" not in request["requestContext"]
+            or "claims" not in request["requestContext"]["authorizer"]
+            or "cognito:username" not in request["requestContext"]["authorizer"]["claims"]):
+            return None
+        return request["requestContext"]["authorizer"]["claims"]["cognito:username"]
+
+    def is_authorized(self, user, graphId):
+        # If Authenticated through AWS account treat as admin for all graphs
+        if (user is None):
+            return True
+        # Otherwise check the list of administrators configured on the graph
+        try:
+            graph_record = self.graph.get_graph(graphId)
+            return user in graph_record["administrators"]
+        except Exception as e:
+            return False