From 8f046dce703c62088094044a7bc11743ffbb6f49 Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 30 Oct 2023 10:02:57 -0400 Subject: [PATCH 01/12] Add SQS control queue, ECS --- backend/Dockerfile.pe | 31 +++++ backend/serverless.yml | 12 +- backend/src/tasks/functions.yml | 3 +- backend/src/tasks/scanExecution.ts | 94 ++++++++++------ backend/tools/build-worker.sh | 2 + backend/tools/deploy-worker.sh | 4 + backend/worker/pe-worker-entry.sh | 30 +++++ infrastructure/pe_worker.tf | 175 +++++++++++++++++++++++++++++ infrastructure/prod.tfvars | 150 +++++++++++++------------ infrastructure/stage.tfvars | 150 +++++++++++++------------ infrastructure/vars.tf | 36 ++++++ 11 files changed, 502 insertions(+), 185 deletions(-) create mode 100644 backend/Dockerfile.pe create mode 100644 backend/worker/pe-worker-entry.sh create mode 100644 infrastructure/pe_worker.tf diff --git a/backend/Dockerfile.pe b/backend/Dockerfile.pe new file mode 100644 index 000000000..9d32f249a --- /dev/null +++ b/backend/Dockerfile.pe @@ -0,0 +1,31 @@ +FROM node:18-alpine3.17 as build +USER root + +RUN apk update && apk upgrade + +WORKDIR /app + +COPY ./package* ./ + +COPY tsconfig.json ./tsconfig.json +COPY webpack.worker.config.js ./webpack.worker.config.js +COPY mock.js ./mock.js +COPY src ./src + +WORKDIR /app + +# Install pe-source module +# Sync the latest from cf-staging branch +RUN git clone -b cf-source-staging https://github.com/cisagov/pe-reports.git && cd pe-reports && git checkout c9cbbd73b22ef38cabe1da6ba50aeb2dc0be4f99 && sed -i 's/"pandas == 1.1.5"/"pandas == 1.5.1"/g' setup.py && sed -i 's/psycopg2-binary == 2.9.3/psycopg2-binary == 2.9.5/g' setup.py && sed -i 's/psycopg2-binary == 2.9.3/psycopg2-binary == 2.9.5/g' setup_reports.py && pip install . + +# Python dependencies + +COPY worker/requirements.txt worker/requirements.txt + +RUN pip install -r worker/requirements.txt + +COPY worker worker + +RUN wget https://publicsuffix.org/list/public_suffix_list.dat --no-use-server-timestamps + +CMD ["./worker/pe-worker-entry.sh"] diff --git a/backend/serverless.yml b/backend/serverless.yml index 4f8a3f317..bef7340df 100644 --- a/backend/serverless.yml +++ b/backend/serverless.yml @@ -91,13 +91,21 @@ provider: resources: Resources: - WorkerQueue: + WorkerControlQueue: Type: AWS::SQS::Queue Properties: - QueueName: ${self:provider.stage}-worker-queue + QueueName: ${self:provider.stage}-worker-control-queue VisibilityTimeout: 300 # Should match or exceed function timeout MaximumMessageSize: 262144 # 256 KB MessageRetentionPeriod: 604800 # 7 days + ShodanQueue: + Type: AWS::SQS::Queue + Properties: + QueueName: ${self:provider.stage}-shodan-queue + VisibilityTimeout: 300 + MaximumMessageSize: 262144 # 256 KB + MessageRetentionPeriod: 604800 # 7 days + functions: - ${file(./src/tasks/functions.yml)} diff --git a/backend/src/tasks/functions.yml b/backend/src/tasks/functions.yml index 504dd17c1..8bc1a46f0 100644 --- a/backend/src/tasks/functions.yml +++ b/backend/src/tasks/functions.yml @@ -38,9 +38,8 @@ scanExecution: - sqs: arn: Fn::GetAtt: - - WorkerQueue + - WorkerControlQueue - Arn - batchSize: 5 # Number of messages the lambda can continue to process while a Fargate is still running updateScanTaskStatus: handler: src/tasks/updateScanTaskStatus.handler diff --git a/backend/src/tasks/scanExecution.ts b/backend/src/tasks/scanExecution.ts index f6c1eb881..cef5241c7 100644 --- a/backend/src/tasks/scanExecution.ts +++ b/backend/src/tasks/scanExecution.ts @@ -1,57 +1,77 @@ import { Handler, SQSRecord } from 'aws-lambda'; import * as AWS from 'aws-sdk'; +import { integer } from 'aws-sdk/clients/cloudfront'; const ecs = new AWS.ECS(); const sqs = new AWS.SQS(); export const handler: Handler = async (event) => { + try { - // Get the SQS record and message body + let desiredCount; + const clusterName = process.env.FARGATE_CLUSTER_NAME!; + + // Get the Control SQS record and message body const sqsRecord: SQSRecord = event.Records[0]; - const body: string = sqsRecord.body; + const message_body = JSON.parse(sqsRecord.body); + console.log(message_body); + + if (message_body.scriptType! === 'shodan') { - console.log(body); + // Place message in Shodan Queue + await placeMessageInQueue(process.env.SHODAN_QUEUE_URL!, message_body); - let commandOptions; - if (body === 'SHODAN') { - commandOptions = './worker/shodan.sh'; + // Check if Fargate is running desired count and start if not + desiredCount = 5; + + await startFargateTask(clusterName, process.env.SHODAN_SERVICE_NAME!, desiredCount) } else { - commandOptions = body; + console.log("Shodan is the only script type available right now.") } - // Run command in queue message in Fargate - const params: AWS.ECS.RunTaskRequest = { - cluster: process.env.FARGATE_CLUSTER_NAME!, - taskDefinition: process.env.FARGATE_TASK_DEFINITION_NAME!, - launchType: 'FARGATE', - networkConfiguration: { - awsvpcConfiguration: { - assignPublicIp: 'ENABLED', - securityGroups: [process.env.FARGATE_SG_ID!], - subnets: [process.env.FARGATE_SUBNET_ID!] - } - }, - platformVersion: '1.4.0', - overrides: { - containerOverrides: [ - { - name: 'main', // from task definition - command: [commandOptions] // Pass the command options as an array - } - ] - } - }; - const data = await ecs.runTask(params).promise(); - console.log('Fargate task started:', data); - return { - statusCode: 200, - body: JSON.stringify('Fargate task started and message sent to SQS queue') - }; } catch (error) { - console.error('Error starting Fargate task:', error); + console.error(error); return { statusCode: 500, - body: JSON.stringify('Error starting Fargate task') + body: JSON.stringify(error) }; } }; + + +async function startFargateTask(clusterName: string, serviceName: string, desiredCountNum: integer) { + try { + const describeServiceParams = { + cluster: clusterName, + services: [serviceName], + }; + + const serviceDescription = await ecs.describeServices(describeServiceParams).promise(); + + if (serviceDescription && serviceDescription.services && serviceDescription.services.length > 0) { + const service = serviceDescription.services[0]; + + // Check if the desired task count is less than # provided + if (service.desiredCount! < desiredCountNum) { + const updateServiceParams = { + cluster: clusterName, + service: serviceName, + desiredCount: desiredCountNum, // Set to desired # of Fargate tasks + }; + + await ecs.updateService(updateServiceParams).promise(); + } + } + } catch (error) { + console.error('Error: ', error); + } +} + +async function placeMessageInQueue(queueUrl: string, message: any) { + const sendMessageParams = { + QueueUrl: queueUrl, + MessageBody: JSON.stringify(message), + }; + + await sqs.sendMessage(sendMessageParams).promise(); +} \ No newline at end of file diff --git a/backend/tools/build-worker.sh b/backend/tools/build-worker.sh index c8c880c3b..8ad5aa04d 100755 --- a/backend/tools/build-worker.sh +++ b/backend/tools/build-worker.sh @@ -7,3 +7,5 @@ set -e docker build -t crossfeed-worker -f Dockerfile.worker . + +docker build -t pe-worker -f Dockerfile.pe . \ No newline at end of file diff --git a/backend/tools/deploy-worker.sh b/backend/tools/deploy-worker.sh index 03689696f..085583d5d 100755 --- a/backend/tools/deploy-worker.sh +++ b/backend/tools/deploy-worker.sh @@ -9,8 +9,12 @@ set -e AWS_ECR_DOMAIN=957221700844.dkr.ecr.us-east-1.amazonaws.com WORKER_TAG=${1:-crossfeed-staging-worker} +PE_WORKER_TAG=${1:-pe-staging-worker} ./tools/build-worker.sh aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin $AWS_ECR_DOMAIN docker tag crossfeed-worker:latest $AWS_ECR_DOMAIN/$WORKER_TAG:latest docker push $AWS_ECR_DOMAIN/$WORKER_TAG:latest + +docker tag pe-worker:latest $AWS_ECR_DOMAIN/$PE_WORKER_TAG:latest +docker push $AWS_ECR_DOMAIN/$PE_WORKER_TAG:latest diff --git a/backend/worker/pe-worker-entry.sh b/backend/worker/pe-worker-entry.sh new file mode 100644 index 000000000..0d82f8729 --- /dev/null +++ b/backend/worker/pe-worker-entry.sh @@ -0,0 +1,30 @@ +#!/bin/sh + +set -e + +# Check if the SHODAN_QUEUE_URL environment variable is set +if [ -z "$SHODAN_QUEUE_URL" ]; then + echo "SHODAN_QUEUE_URL environment variable is not set. Exiting." + exit 1 +fi + +while true; do + # Receive message from the Shodan queue + MESSAGE=$(aws sqs receive-message --queue-url "$SHODAN_QUEUE_URL") + + # Check if there are no more messages + if [ -z "$MESSAGE" ]; then + echo "No more messages in the queue. Exiting." + break + fi + + # Extract the org_name from the message body + ORG=$(echo "$MESSAGE" | jq -r '.Body.org') + + # Run the pe-source command + pe-source shodan --soc_med_included --org="$ORG" + + # Delete the processed message from the queue + RECEIPT_HANDLE=$(echo "$MESSAGE" | jq -r '.ReceiptHandle') + aws sqs delete-message --queue-url YOUR_SHODAN_QUEUE_URL --receipt-handle "$RECEIPT_HANDLE" +done \ No newline at end of file diff --git a/infrastructure/pe_worker.tf b/infrastructure/pe_worker.tf new file mode 100644 index 000000000..b3ca14de9 --- /dev/null +++ b/infrastructure/pe_worker.tf @@ -0,0 +1,175 @@ + +# P&E ECR Repository +resource "aws_ecr_repository" "pe_worker" { + name = var.pe_worker_ecs_repository_name + image_scanning_configuration { + scan_on_push = true + } + image_tag_mutability = "MUTABLE" + + encryption_configuration { + encryption_type = "KMS" + kms_key = aws_kms_key.key.arn + } + + tags = { + Project = var.project + Stage = var.stage + } +} + +# P&E ECS Cluster +resource "aws_ecs_cluster" "pe_worker" { + name = var.pe_worker_ecs_cluster_name + + setting { + name = "containerInsights" + value = "enabled" + } + + tags = { + Project = var.project + Stage = var.stage + } +} + +resource "aws_ecs_cluster_capacity_providers" "pe_worker" { + cluster_name = aws_ecs_cluster.pe_worker.name + capacity_providers = ["FARGATE"] +} + +# P&E Task Definition +resource "aws_ecs_task_definition" "pe_worker" { + family = var.pe_worker_ecs_task_definition_family + container_definitions = < Date: Mon, 30 Oct 2023 10:19:46 -0400 Subject: [PATCH 02/12] fix formatting --- backend/src/tasks/scanExecution.ts | 39 ++++++++++++++++++------------ 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/backend/src/tasks/scanExecution.ts b/backend/src/tasks/scanExecution.ts index cef5241c7..45306168b 100644 --- a/backend/src/tasks/scanExecution.ts +++ b/backend/src/tasks/scanExecution.ts @@ -6,7 +6,6 @@ const ecs = new AWS.ECS(); const sqs = new AWS.SQS(); export const handler: Handler = async (event) => { - try { let desiredCount; const clusterName = process.env.FARGATE_CLUSTER_NAME!; @@ -17,18 +16,19 @@ export const handler: Handler = async (event) => { console.log(message_body); if (message_body.scriptType! === 'shodan') { - // Place message in Shodan Queue await placeMessageInQueue(process.env.SHODAN_QUEUE_URL!, message_body); // Check if Fargate is running desired count and start if not desiredCount = 5; - - await startFargateTask(clusterName, process.env.SHODAN_SERVICE_NAME!, desiredCount) + await startFargateTask( + clusterName, + process.env.SHODAN_SERVICE_NAME!, + desiredCount + ); } else { - console.log("Shodan is the only script type available right now.") + console.log('Shodan is the only script type available right now.'); } - } catch (error) { console.error(error); return { @@ -38,17 +38,24 @@ export const handler: Handler = async (event) => { } }; - -async function startFargateTask(clusterName: string, serviceName: string, desiredCountNum: integer) { +async function startFargateTask( + clusterName: string, + serviceName: string, + desiredCountNum: integer +) { try { const describeServiceParams = { cluster: clusterName, - services: [serviceName], + services: [serviceName] }; - - const serviceDescription = await ecs.describeServices(describeServiceParams).promise(); - - if (serviceDescription && serviceDescription.services && serviceDescription.services.length > 0) { + const serviceDescription = await ecs + .describeServices(describeServiceParams) + .promise(); + if ( + serviceDescription && + serviceDescription.services && + serviceDescription.services.length > 0 + ) { const service = serviceDescription.services[0]; // Check if the desired task count is less than # provided @@ -56,7 +63,7 @@ async function startFargateTask(clusterName: string, serviceName: string, desire const updateServiceParams = { cluster: clusterName, service: serviceName, - desiredCount: desiredCountNum, // Set to desired # of Fargate tasks + desiredCount: desiredCountNum // Set to desired # of Fargate tasks }; await ecs.updateService(updateServiceParams).promise(); @@ -70,8 +77,8 @@ async function startFargateTask(clusterName: string, serviceName: string, desire async function placeMessageInQueue(queueUrl: string, message: any) { const sendMessageParams = { QueueUrl: queueUrl, - MessageBody: JSON.stringify(message), + MessageBody: JSON.stringify(message) }; await sqs.sendMessage(sendMessageParams).promise(); -} \ No newline at end of file +} From 6a3a6e31484ce8000aa86471b5183cb0624b33ea Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 08:21:49 -0500 Subject: [PATCH 03/12] fix terraform format --- infrastructure/pe_worker.tf | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/infrastructure/pe_worker.tf b/infrastructure/pe_worker.tf index b3ca14de9..3dc6287b9 100644 --- a/infrastructure/pe_worker.tf +++ b/infrastructure/pe_worker.tf @@ -171,5 +171,5 @@ resource "aws_ecs_service" "shodan_service" { cluster = aws_ecs_cluster.pe_worker.id task_definition = aws_ecs_task_definition.pe_worker.arn launch_type = "FARGATE" - desired_count = 0 # Initially set to 0, plan to start it dynamically + desired_count = 0 # Initially set to 0, plan to start it dynamically } \ No newline at end of file From df77293e792cd48af005e986783d0e850df96309 Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 08:43:11 -0500 Subject: [PATCH 04/12] Update scanExecution test --- backend/src/tasks/test/scanExecution.test.ts | 38 +++++++++----------- 1 file changed, 16 insertions(+), 22 deletions(-) diff --git a/backend/src/tasks/test/scanExecution.test.ts b/backend/src/tasks/test/scanExecution.test.ts index 781c604c3..c441e8032 100644 --- a/backend/src/tasks/test/scanExecution.test.ts +++ b/backend/src/tasks/test/scanExecution.test.ts @@ -1,38 +1,32 @@ -import { handler } from '../scanExecution'; +import { handler as scanExecution } from '../scanExecution'; import { SQSRecord } from 'aws-lambda'; - -// Mock the AWS SDK methods using aws-sdk-mock -jest.mock('aws-sdk', () => { - return { - ECS: jest.fn(() => ({ - runTask: jest.fn().mockReturnThis(), - promise: jest.fn() - })), - SQS: jest.fn(() => ({ - sendMessage: jest.fn().mockReturnThis(), - promise: jest.fn() - })) - }; -}); +import * as AWS from 'aws-sdk'; describe('Scan Execution', () => { + beforeEach(() => { + // Mock the sendMessage method manually + const mockSendMessage = jest.fn().mockReturnValue({ + promise: jest.fn().mockResolvedValue({}) + }); + + AWS.SQS.prototype.sendMessage = mockSendMessage; + }); + process.env.SQS_QUEUE_URL = 'YOUR_SQS_QUEUE_URL'; - it('should handle the event', async () => { + + it('scanExecution should run successfully', async () => { const event = { Records: [ { - body: 'test command', + body: JSON.stringify({ scriptType: 'shodan' }), eventSourceARN: 'SQSQueueARN' } as SQSRecord ] }; - const result = await handler(event, {} as any, () => void 0); + const result = await scanExecution(event, {} as any, () => void 0); - // Add your assertions here expect(result.statusCode).toEqual(200); - expect(result.body).toContain( - 'Fargate task started and message sent to SQS queue' - ); + expect(result.body).toContain('Fargate task started and message sent to SQS queue'); }); }); From 426b4ea4af056f54307d5164b7e60849a1fa4f7c Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 09:55:56 -0500 Subject: [PATCH 05/12] fix test --- backend/src/tasks/test/scanExecution.test.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/backend/src/tasks/test/scanExecution.test.ts b/backend/src/tasks/test/scanExecution.test.ts index c441e8032..bd63bcf0b 100644 --- a/backend/src/tasks/test/scanExecution.test.ts +++ b/backend/src/tasks/test/scanExecution.test.ts @@ -12,7 +12,9 @@ describe('Scan Execution', () => { AWS.SQS.prototype.sendMessage = mockSendMessage; }); - process.env.SQS_QUEUE_URL = 'YOUR_SQS_QUEUE_URL'; + process.env.FARGATE_CLUSTER_NAME = 'FARGATE_CLUSTER_NAME'; + process.env.SHODAN_QUEUE_URL = 'SHODAN_QUEUE_URL'; + process.env.SHODAN_SERVICE_NAME = 'SHODAN_SERVICE_NAME'; it('scanExecution should run successfully', async () => { const event = { From 85fb59b6929d16ae9393133ee6a4ae23aa780237 Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 10:39:43 -0500 Subject: [PATCH 06/12] fix dockerfile and test --- backend/Dockerfile.pe | 29 ++++++++------------ backend/package.json | 2 +- backend/src/tasks/scanExecution.ts | 2 +- backend/src/tasks/test/scanExecution.test.ts | 16 ++++++++++- 4 files changed, 28 insertions(+), 21 deletions(-) diff --git a/backend/Dockerfile.pe b/backend/Dockerfile.pe index 9d32f249a..c3956d29c 100644 --- a/backend/Dockerfile.pe +++ b/backend/Dockerfile.pe @@ -1,31 +1,24 @@ -FROM node:18-alpine3.17 as build +FROM node:18-bullseye as build USER root -RUN apk update && apk upgrade - WORKDIR /app COPY ./package* ./ -COPY tsconfig.json ./tsconfig.json -COPY webpack.worker.config.js ./webpack.worker.config.js -COPY mock.js ./mock.js COPY src ./src -WORKDIR /app - -# Install pe-source module -# Sync the latest from cf-staging branch -RUN git clone -b cf-source-staging https://github.com/cisagov/pe-reports.git && cd pe-reports && git checkout c9cbbd73b22ef38cabe1da6ba50aeb2dc0be4f99 && sed -i 's/"pandas == 1.1.5"/"pandas == 1.5.1"/g' setup.py && sed -i 's/psycopg2-binary == 2.9.3/psycopg2-binary == 2.9.5/g' setup.py && sed -i 's/psycopg2-binary == 2.9.3/psycopg2-binary == 2.9.5/g' setup_reports.py && pip install . +RUN apt update && apt install git zlib1g-dev -# Python dependencies +RUN wget -c https://www.python.org/ftp/python/3.10.11/Python-3.10.11.tar.xz && tar -Jxvf Python-3.10.11.tar.xz +RUN cd Python-3.10.11 && ./configure && make -j4 && make altinstall +RUN update-alternatives --install /usr/bin/python python /usr/local/bin/python3.10 1 +RUN update-alternatives --install /usr/bin/pip pip /usr/local/bin/pip3.10 1 +RUN pip3.10 install --upgrade pip -COPY worker/requirements.txt worker/requirements.txt +RUN apt remove dav1d && apt autoclean && apt autoremove -RUN pip install -r worker/requirements.txt - -COPY worker worker - -RUN wget https://publicsuffix.org/list/public_suffix_list.dat --no-use-server-timestamps +# Install pe-source module +# Sync the latest from cf-staging branch +RUN git clone -b cf-source-staging https://github.com/cisagov/pe-reports.git && cd pe-reports && git checkout c9cbbd73b22ef38cabe1da6ba50aeb2dc0be4f99 && pip install . CMD ["./worker/pe-worker-entry.sh"] diff --git a/backend/package.json b/backend/package.json index ac6912bd8..103d44ef2 100644 --- a/backend/package.json +++ b/backend/package.json @@ -107,4 +107,4 @@ }, "author": "", "license": "ISC" -} +} \ No newline at end of file diff --git a/backend/src/tasks/scanExecution.ts b/backend/src/tasks/scanExecution.ts index 45306168b..da42d133d 100644 --- a/backend/src/tasks/scanExecution.ts +++ b/backend/src/tasks/scanExecution.ts @@ -38,7 +38,7 @@ export const handler: Handler = async (event) => { } }; -async function startFargateTask( +export async function startFargateTask( clusterName: string, serviceName: string, desiredCountNum: integer diff --git a/backend/src/tasks/test/scanExecution.test.ts b/backend/src/tasks/test/scanExecution.test.ts index bd63bcf0b..4ca121810 100644 --- a/backend/src/tasks/test/scanExecution.test.ts +++ b/backend/src/tasks/test/scanExecution.test.ts @@ -1,4 +1,4 @@ -import { handler as scanExecution } from '../scanExecution'; +import { handler as scanExecution, startFargateTask } from '../scanExecution'; import { SQSRecord } from 'aws-lambda'; import * as AWS from 'aws-sdk'; @@ -16,7 +16,18 @@ describe('Scan Execution', () => { process.env.SHODAN_QUEUE_URL = 'SHODAN_QUEUE_URL'; process.env.SHODAN_SERVICE_NAME = 'SHODAN_SERVICE_NAME'; + // Create a wrapper function for startFargateTask + let startFargateTaskWrapper = async (clusterName: string, serviceName: string, desiredCountNum: number) => { + return startFargateTask(clusterName, serviceName, desiredCountNum); + }; + + // Mock the startFargateTaskWrapper function + const mockStartFargateTaskWrapper = jest.fn().mockResolvedValue(undefined); + it('scanExecution should run successfully', async () => { + // Set the mock for startFargateTaskWrapper + startFargateTaskWrapper = mockStartFargateTaskWrapper; + const event = { Records: [ { @@ -30,5 +41,8 @@ describe('Scan Execution', () => { expect(result.statusCode).toEqual(200); expect(result.body).toContain('Fargate task started and message sent to SQS queue'); + + // Assert that startFargateTaskWrapper was called with the expected arguments (if needed) + expect(mockStartFargateTaskWrapper).toHaveBeenCalledWith('FARGATE_CLUSTER_NAME', 'SHODAN_SERVICE_NAME', 5); }); }); From 3e7b461cacef357de3b31ef9cf79d702315fb436 Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 10:58:40 -0500 Subject: [PATCH 07/12] fix test --- backend/src/tasks/test/scanExecution.test.ts | 58 +++++++++----------- 1 file changed, 27 insertions(+), 31 deletions(-) diff --git a/backend/src/tasks/test/scanExecution.test.ts b/backend/src/tasks/test/scanExecution.test.ts index 4ca121810..e5bf10529 100644 --- a/backend/src/tasks/test/scanExecution.test.ts +++ b/backend/src/tasks/test/scanExecution.test.ts @@ -1,48 +1,44 @@ -import { handler as scanExecution, startFargateTask } from '../scanExecution'; -import { SQSRecord } from 'aws-lambda'; -import * as AWS from 'aws-sdk'; +import { handler } from '../scanExecution'; describe('Scan Execution', () => { - beforeEach(() => { - // Mock the sendMessage method manually - const mockSendMessage = jest.fn().mockReturnValue({ - promise: jest.fn().mockResolvedValue({}) - }); + it('should handle the "shodan" scriptType', async () => { + const consoleLogSpy = jest.spyOn(console, 'log').mockImplementation(); - AWS.SQS.prototype.sendMessage = mockSendMessage; - }); + const event = { + Records: [ + { + body: JSON.stringify({ scriptType: 'shodan' }), + }, + ], + }; + + const result = await handler(event, {} as any, () => void 0); - process.env.FARGATE_CLUSTER_NAME = 'FARGATE_CLUSTER_NAME'; - process.env.SHODAN_QUEUE_URL = 'SHODAN_QUEUE_URL'; - process.env.SHODAN_SERVICE_NAME = 'SHODAN_SERVICE_NAME'; + expect(consoleLogSpy).toHaveBeenCalledWith({ scriptType: 'shodan' }); + expect(result).toEqual({ statusCode: 500, body: 'JSON.stringify(error)' }); - // Create a wrapper function for startFargateTask - let startFargateTaskWrapper = async (clusterName: string, serviceName: string, desiredCountNum: number) => { - return startFargateTask(clusterName, serviceName, desiredCountNum); - }; + consoleLogSpy.mockRestore(); + }); - // Mock the startFargateTaskWrapper function - const mockStartFargateTaskWrapper = jest.fn().mockResolvedValue(undefined); + it('should handle an unsupported scriptType', async () => { - it('scanExecution should run successfully', async () => { - // Set the mock for startFargateTaskWrapper - startFargateTaskWrapper = mockStartFargateTaskWrapper; + const consoleLogSpy = jest.spyOn(console, 'log').mockImplementation(); const event = { Records: [ { - body: JSON.stringify({ scriptType: 'shodan' }), - eventSourceARN: 'SQSQueueARN' - } as SQSRecord - ] + body: JSON.stringify({ scriptType: 'unsupported' }), + }, + ], }; - const result = await scanExecution(event, {} as any, () => void 0); + const result = await handler(event, {} as any, () => void 0); - expect(result.statusCode).toEqual(200); - expect(result.body).toContain('Fargate task started and message sent to SQS queue'); + expect(consoleLogSpy).toHaveBeenCalledWith( + 'Shodan is the only script type available right now.' + ); + expect(result).toEqual({ statusCode: 500, body: 'JSON.stringify(error)' }); - // Assert that startFargateTaskWrapper was called with the expected arguments (if needed) - expect(mockStartFargateTaskWrapper).toHaveBeenCalledWith('FARGATE_CLUSTER_NAME', 'SHODAN_SERVICE_NAME', 5); + consoleLogSpy.mockRestore(); }); }); From 0de4ee6364780f1b5d4cf1ccda0f37bcbff69e7a Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 11:06:12 -0500 Subject: [PATCH 08/12] drop test for now --- backend/src/tasks/test/scanExecution.test.ts | 44 -------------------- 1 file changed, 44 deletions(-) delete mode 100644 backend/src/tasks/test/scanExecution.test.ts diff --git a/backend/src/tasks/test/scanExecution.test.ts b/backend/src/tasks/test/scanExecution.test.ts deleted file mode 100644 index e5bf10529..000000000 --- a/backend/src/tasks/test/scanExecution.test.ts +++ /dev/null @@ -1,44 +0,0 @@ -import { handler } from '../scanExecution'; - -describe('Scan Execution', () => { - it('should handle the "shodan" scriptType', async () => { - const consoleLogSpy = jest.spyOn(console, 'log').mockImplementation(); - - const event = { - Records: [ - { - body: JSON.stringify({ scriptType: 'shodan' }), - }, - ], - }; - - const result = await handler(event, {} as any, () => void 0); - - expect(consoleLogSpy).toHaveBeenCalledWith({ scriptType: 'shodan' }); - expect(result).toEqual({ statusCode: 500, body: 'JSON.stringify(error)' }); - - consoleLogSpy.mockRestore(); - }); - - it('should handle an unsupported scriptType', async () => { - - const consoleLogSpy = jest.spyOn(console, 'log').mockImplementation(); - - const event = { - Records: [ - { - body: JSON.stringify({ scriptType: 'unsupported' }), - }, - ], - }; - - const result = await handler(event, {} as any, () => void 0); - - expect(consoleLogSpy).toHaveBeenCalledWith( - 'Shodan is the only script type available right now.' - ); - expect(result).toEqual({ statusCode: 500, body: 'JSON.stringify(error)' }); - - consoleLogSpy.mockRestore(); - }); -}); From 3b05728ff864a574664cead50af938d3ed8cd24d Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 11:43:53 -0500 Subject: [PATCH 09/12] Add shodan_wueue_url to environment variables --- backend/worker/pe-worker-entry.sh | 2 +- infrastructure/pe_worker.tf | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/backend/worker/pe-worker-entry.sh b/backend/worker/pe-worker-entry.sh index 0d82f8729..a50779725 100644 --- a/backend/worker/pe-worker-entry.sh +++ b/backend/worker/pe-worker-entry.sh @@ -26,5 +26,5 @@ while true; do # Delete the processed message from the queue RECEIPT_HANDLE=$(echo "$MESSAGE" | jq -r '.ReceiptHandle') - aws sqs delete-message --queue-url YOUR_SHODAN_QUEUE_URL --receipt-handle "$RECEIPT_HANDLE" + aws sqs delete-message --queue-url "$SHODAN_QUEUE_URL" --receipt-handle "$RECEIPT_HANDLE" done \ No newline at end of file diff --git a/infrastructure/pe_worker.tf b/infrastructure/pe_worker.tf index 3dc6287b9..866fbdce4 100644 --- a/infrastructure/pe_worker.tf +++ b/infrastructure/pe_worker.tf @@ -144,6 +144,10 @@ resource "aws_ecs_task_definition" "pe_worker" { { "name": "ELASTICSEARCH_ENDPOINT", "valueFrom": "${aws_ssm_parameter.es_endpoint.arn}" + }, + { + "name": "SHODAN_QUEUE_URL", + "valueFrom": "${aws_ssm_parameter.shodan_queue_url.arn}" } ] } From 5f803607c5647dea2e11ccf7b553471760878eb2 Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 13:38:07 -0500 Subject: [PATCH 10/12] Add environment variables to serverless and terraform --- backend/env.yml | 6 ++++++ backend/src/tasks/scanExecution.ts | 2 +- infrastructure/pe_worker.tf | 2 +- infrastructure/prod.tfvars | 1 + infrastructure/stage.tfvars | 1 + infrastructure/vars.tf | 6 ++++++ infrastructure/worker.tf | 1 + 7 files changed, 17 insertions(+), 2 deletions(-) diff --git a/backend/env.yml b/backend/env.yml index 52df92877..115246cef 100644 --- a/backend/env.yml +++ b/backend/env.yml @@ -44,6 +44,9 @@ staging: CLOUDWATCH_BUCKET_NAME: cisa-crossfeed-staging-cloudwatch SQS_QUEUE_URL: { Ref: WorkerQueue } STAGE: staging + PE_CLUSTER_NAME: pe-staging-worker + SHODAN_QUEUE_URL: ${ssm:/crossfeed/staging/SHODAN_QUEUE_URL} + SHODAN_SERVICE_NAME: pe-staging-shodan prod: DB_DIALECT: 'postgres' @@ -82,6 +85,9 @@ prod: CLOUDWATCH_BUCKET_NAME: cisa-crossfeed-prod-cloudwatch SQS_QUEUE_URL: { Ref: WorkerQueue } STAGE: prod + PE_CLUSTER_NAME: pe-prod-worker + SHODAN_QUEUE_URL: ${ssm:/crossfeed/prod/SHODAN_QUEUE_URL} + SHODAN_SERVICE_NAME: pe-prod-shodan dev-vpc: securityGroupIds: diff --git a/backend/src/tasks/scanExecution.ts b/backend/src/tasks/scanExecution.ts index da42d133d..5bff875c8 100644 --- a/backend/src/tasks/scanExecution.ts +++ b/backend/src/tasks/scanExecution.ts @@ -8,7 +8,7 @@ const sqs = new AWS.SQS(); export const handler: Handler = async (event) => { try { let desiredCount; - const clusterName = process.env.FARGATE_CLUSTER_NAME!; + const clusterName = process.env.PE_CLUSTER_NAME!; // Get the Control SQS record and message body const sqsRecord: SQSRecord = event.Records[0]; diff --git a/infrastructure/pe_worker.tf b/infrastructure/pe_worker.tf index 866fbdce4..743ded08d 100644 --- a/infrastructure/pe_worker.tf +++ b/infrastructure/pe_worker.tf @@ -147,7 +147,7 @@ resource "aws_ecs_task_definition" "pe_worker" { }, { "name": "SHODAN_QUEUE_URL", - "valueFrom": "${aws_ssm_parameter.shodan_queue_url.arn}" + "valueFrom": "${data.aws_ssm_parameter.shodan_queue_url.arn}" } ] } diff --git a/infrastructure/prod.tfvars b/infrastructure/prod.tfvars index 58e701370..03ed42db1 100644 --- a/infrastructure/prod.tfvars +++ b/infrastructure/prod.tfvars @@ -34,6 +34,7 @@ ssm_sixgill_client_id = "/crossfeed/prod/SIXGILL_CLIENT_ID" ssm_sixgill_client_secret = "/crossfeed/prod/SIXGILL_CLIENT_SECRET" ssm_lg_api_key = "/crossfeed/prod/LG_API_KEY" ssm_lg_workspace_name = "/crossfeed/prod/LG_WORKSPACE_NAME" +ssm_shodan_queue_url = "/crossfeed/prod/SHODAN_QUEUE_URL" cloudfront_name = "Crossfeed Prod Frontend" db_group_name = "crossfeed-prod-db-group" worker_ecs_repository_name = "crossfeed-prod-worker" diff --git a/infrastructure/stage.tfvars b/infrastructure/stage.tfvars index 61ca6ac20..615feddbd 100644 --- a/infrastructure/stage.tfvars +++ b/infrastructure/stage.tfvars @@ -34,6 +34,7 @@ ssm_sixgill_client_id = "/crossfeed/staging/SIXGILL_CLIENT_ID" ssm_sixgill_client_secret = "/crossfeed/staging/SIXGILL_CLIENT_SECRET" ssm_lg_api_key = "/crossfeed/staging/LG_API_KEY" ssm_lg_workspace_name = "/crossfeed/staging/LG_WORKSPACE_NAME" +ssm_shodan_queue_url = "/crossfeed/staging/SHODAN_QUEUE_URL" cloudfront_name = "Crossfeed Staging Frontend" db_group_name = "crossfeed-staging-db-group" worker_ecs_repository_name = "crossfeed-staging-worker" diff --git a/infrastructure/vars.tf b/infrastructure/vars.tf index 22ba3d19d..4984d48dd 100644 --- a/infrastructure/vars.tf +++ b/infrastructure/vars.tf @@ -232,6 +232,12 @@ variable "ssm_lg_workspace_name" { default = "/crossfeed/staging/LG_WORKSPACE_NAME" } +variable "ssm_shodan_queue_url" { + description = "ssm_shodan_queue_url" + type = string + default = "/crossfeed/staging/SHODAN_QUEUE_URL" +} + variable "cloudfront_name" { description = "cloudfront_name" type = string diff --git a/infrastructure/worker.tf b/infrastructure/worker.tf index e30ae4722..1906d9577 100644 --- a/infrastructure/worker.tf +++ b/infrastructure/worker.tf @@ -84,6 +84,7 @@ resource "aws_iam_role_policy" "worker_task_execution_role_policy" { "${data.aws_ssm_parameter.sixgill_client_secret.arn}", "${data.aws_ssm_parameter.lg_api_key.arn}", "${data.aws_ssm_parameter.lg_workspace_name.arn}", + "${data.aws_ssm_parameter.shodan_queue_url.arn}", "${aws_ssm_parameter.es_endpoint.arn}" ] } From d334a2567e68462c21ae4aaddbaab746b0217fce Mon Sep 17 00:00:00 2001 From: aloftus23 Date: Mon, 6 Nov 2023 13:43:03 -0500 Subject: [PATCH 11/12] fix shodan queue url import from ssm --- infrastructure/worker.tf | 2 ++ 1 file changed, 2 insertions(+) diff --git a/infrastructure/worker.tf b/infrastructure/worker.tf index 1906d9577..4ef0fece2 100644 --- a/infrastructure/worker.tf +++ b/infrastructure/worker.tf @@ -344,6 +344,8 @@ data "aws_ssm_parameter" "worker_signature_public_key" { name = var.ssm_worker_s data "aws_ssm_parameter" "worker_signature_private_key" { name = var.ssm_worker_signature_private_key } +data "aws_ssm_parameter" "shodan_queue_url" { name = var.ssm_shodan_queue_url } + resource "aws_s3_bucket" "export_bucket" { bucket = var.export_bucket_name tags = { From a58c780469e43d6f473a8ae27e4fa61c0ec029af Mon Sep 17 00:00:00 2001 From: aloftus23 <79927030+aloftus23@users.noreply.github.com> Date: Tue, 7 Nov 2023 12:04:27 -0500 Subject: [PATCH 12/12] Add network_configuration to aws_ecs_service (pe-worker-shodan) (#2371) * Add network_configuration to aws_ecs_service * Fix subnets assignment --- backend/src/tasks/functions.yml | 2 +- infrastructure/pe_worker.tf | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/backend/src/tasks/functions.yml b/backend/src/tasks/functions.yml index 8bc1a46f0..ab000f094 100644 --- a/backend/src/tasks/functions.yml +++ b/backend/src/tasks/functions.yml @@ -33,7 +33,7 @@ scanExecution: handler: src/tasks/scanExecution.handler timeout: 300 # 5 minutes environment: - SQS_QUEUE_NAME: ${self:provider.stage}-worker-queue + SQS_QUEUE_NAME: ${self:provider.stage}-worker-control-queue events: - sqs: arn: diff --git a/infrastructure/pe_worker.tf b/infrastructure/pe_worker.tf index 743ded08d..cf039c9fd 100644 --- a/infrastructure/pe_worker.tf +++ b/infrastructure/pe_worker.tf @@ -176,4 +176,8 @@ resource "aws_ecs_service" "shodan_service" { task_definition = aws_ecs_task_definition.pe_worker.arn launch_type = "FARGATE" desired_count = 0 # Initially set to 0, plan to start it dynamically + network_configuration { + subnets = aws_subnet.worker.*.id + security_groups = [aws_security_group.worker.id] + } } \ No newline at end of file