diff --git a/lib/app-stack.ts b/lib/app-stack.ts index 90d0449..9ab21fc 100644 --- a/lib/app-stack.ts +++ b/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 diff --git a/lib/rest-api/kai-rest-api-props.ts b/lib/rest-api/kai-rest-api-props.ts index 5552079..72c6a4d 100644 --- a/lib/rest-api/kai-rest-api-props.ts +++ b/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; } \ No newline at end of file diff --git a/lib/rest-api/kai-rest-api.ts b/lib/rest-api/kai-rest-api.ts index 6eb7ffa..d1a21f0 100644 --- a/lib/rest-api/kai-rest-api.ts +++ b/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 } }); diff --git a/lib/rest-api/lambdas/add_graph_request.py b/lib/rest-api/lambdas/add_graph_request.py index 6e48050..0353444 100644 --- a/lib/rest-api/lambdas/add_graph_request.py +++ b/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 { diff --git a/lib/rest-api/lambdas/delete_graph_request.py b/lib/rest-api/lambdas/delete_graph_request.py index 8d7ccf6..bf71ac1 100644 --- a/lib/rest-api/lambdas/delete_graph_request.py +++ b/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 { diff --git a/lib/rest-api/lambdas/get_graph_request.py b/lib/rest-api/lambdas/get_graph_request.py index 473351a..2e40f37 100644 --- a/lib/rest-api/lambdas/get_graph_request.py +++ b/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" diff --git a/lib/rest-api/lambdas/graph/__init__.py b/lib/rest-api/lambdas/graph/__init__.py new file mode 100644 index 0000000..8de4a70 --- /dev/null +++ b/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() + ) + diff --git a/lib/rest-api/lambdas/user/__init__.py b/lib/rest-api/lambdas/user/__init__.py new file mode 100644 index 0000000..d54d17b --- /dev/null +++ b/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 diff --git a/test/rest-api/kai-rest-api.test.ts b/test/rest-api/kai-rest-api.test.ts index 3fbd697..1d35788 100644 --- a/test/rest-api/kai-rest-api.test.ts +++ b/test/rest-api/kai-rest-api.test.ts @@ -28,7 +28,8 @@ function createRestAPI(stack: cdk.Stack, id = "Test"): rest.KaiRestApi { return new rest.KaiRestApi(stack, id, { "graphTable": table, - "userPoolArn": "userPoolArn" + "userPoolArn": "userPoolArn", + "userPoolId": "userPoolId" }); } @@ -227,7 +228,7 @@ test("Should create a lambda to send messages to the deleteGraph queue", () => { })); }); -test("Should allow the Delete Graph Lambda to write to the backend database and Send messages to the Queue", () => { +test("Should allow the Delete Graph Lambda to read and write to the backend database and Send messages to the Queue", () => { // Given const stack = new cdk.Stack(); @@ -240,6 +241,12 @@ test("Should allow the Delete Graph Lambda to write to the backend database and "Statement": [ { "Action": [ + "dynamodb:BatchGetItem", + "dynamodb:GetRecords", + "dynamodb:GetShardIterator", + "dynamodb:Query", + "dynamodb:GetItem", + "dynamodb:Scan", "dynamodb:BatchWriteItem", "dynamodb:PutItem", "dynamodb:UpdateItem", @@ -294,7 +301,8 @@ test("should create a queue for AddGraph messages to be sent to workers", () => // When new rest.KaiRestApi(stack, "Test", { "graphTable": table, - "userPoolArn": "userPoolArn" + "userPoolArn": "userPoolArn", + "userPoolId": "userPoolId" }); // Then @@ -313,7 +321,8 @@ test("should create lambda to write messages to the Add Graph Queue", () => { // When new rest.KaiRestApi(stack, "Test", { "graphTable": table, - "userPoolArn": "userPoolArn" + "userPoolArn": "userPoolArn", + "userPoolId": "userPoolId" }); // Then @@ -322,7 +331,7 @@ test("should create lambda to write messages to the Add Graph Queue", () => { })); }); -test("should allow AddGraphLambda to write messages to queue and write to Dynamodb", () => { +test("should allow AddGraphLambda to write messages to queue and read and write to Dynamodb", () => { // Given const stack = new cdk.Stack(); const table = new Table(stack, "test", { @@ -332,7 +341,8 @@ test("should allow AddGraphLambda to write messages to queue and write to Dynamo // When new rest.KaiRestApi(stack, "Test", { "graphTable": table, - "userPoolArn": "userPoolArn" + "userPoolArn": "userPoolArn", + "userPoolId": "userPoolId" }); // Then @@ -341,6 +351,12 @@ test("should allow AddGraphLambda to write messages to queue and write to Dynamo "Statement": [ { "Action": [ + "dynamodb:BatchGetItem", + "dynamodb:GetRecords", + "dynamodb:GetShardIterator", + "dynamodb:Query", + "dynamodb:GetItem", + "dynamodb:Scan", "dynamodb:BatchWriteItem", "dynamodb:PutItem", "dynamodb:UpdateItem",