Merge branch 'master' into 2323-update-notification-with-new-messaging

backend/Dockerfile.pe

@@ -0,0 +1,24 @@
+FROM node:18-bullseye as build
+USER root
+
+WORKDIR /app
+
+COPY ./package* ./
+
+COPY src ./src
+
+RUN apt update && apt install git zlib1g-dev
+
+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
+
+RUN apt remove dav1d && apt autoclean && apt autoremove
+
+# 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"]

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:

backend/package.json

@@ -107,4 +107,4 @@
   },
   "author": "",
   "license": "ISC"
-}
+}

backend/serverless.yml

@@ -93,13 +93,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)}

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

backend/src/tasks/scanExecution.ts

@@ -1,57 +1,84 @@
 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.PE_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);
 
-    console.log(body);
+    if (message_body.scriptType! === 'shodan') {
+      // 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)
     };
   }
 };
+
+export 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();
+}

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 .

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

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 "$SHODAN_QUEUE_URL" --receipt-handle "$RECEIPT_HANDLE"
+done