From 94124e010dd0149b7f04f5e419b717f2582644c5 Mon Sep 17 00:00:00 2001
From: aloftus23 <andylofti@gmail.com>
Date: Wed, 22 Nov 2023 12:40:00 -0500
Subject: [PATCH] Mkae sure the lambda checks 0 messages and sets desired count
 to 0

---
 backend/src/tasks/scanExecution.ts | 20 +++++++++++++++++++-
 infrastructure/pe_worker.tf        | 17 +++++++++++++++--
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/backend/src/tasks/scanExecution.ts b/backend/src/tasks/scanExecution.ts
index 5bff875c8..cb4fad528 100644
--- a/backend/src/tasks/scanExecution.ts
+++ b/backend/src/tasks/scanExecution.ts
@@ -29,6 +29,24 @@ export const handler: Handler = async (event) => {
     } else {
       console.log('Shodan is the only script type available right now.');
     }
+
+    // After processing each message, check if the SQS queue is empty
+    const sqsAttributes = await sqs
+      .getQueueAttributes({
+        QueueUrl: process.env.SHODAN_QUEUE_URL!,
+        AttributeNames: ['ApproximateNumberOfMessages']
+      })
+      .promise();
+
+    const approximateNumberOfMessages = parseInt(
+      sqsAttributes.Attributes?.ApproximateNumberOfMessages || '0',
+      10
+    );
+
+    // If the queue is empty, scale down to zero tasks
+    if (approximateNumberOfMessages === 0) {
+      await startFargateTask(clusterName, process.env.SHODAN_SERVICE_NAME!, 0);
+    }
   } catch (error) {
     console.error(error);
     return {
@@ -59,7 +77,7 @@ export async function startFargateTask(
       const service = serviceDescription.services[0];
 
       // Check if the desired task count is less than # provided
-      if (service.desiredCount! < desiredCountNum) {
+      if (service.desiredCount! !== desiredCountNum) {
         const updateServiceParams = {
           cluster: clusterName,
           service: serviceName,
diff --git a/infrastructure/pe_worker.tf b/infrastructure/pe_worker.tf
index cf039c9fd..714e1f438 100644
--- a/infrastructure/pe_worker.tf
+++ b/infrastructure/pe_worker.tf
@@ -177,7 +177,20 @@ resource "aws_ecs_service" "shodan_service" {
   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]
+    subnets          = aws_subnet.worker.id
+    security_groups  = [aws_security_group.worker.id]
+    assign_public_ip = true
+  }
+}
+
+# Create the  log group
+resource "aws_cloudwatch_log_group" "worker" {
+  name              = var.pe_worker_ecs_log_group_name # should match awslogs-group in service.json
+  retention_in_days = 3653
+  kms_key_id        = aws_kms_key.key.arn
+  tags = {
+    Project = var.project
+    Stage   = var.stage
+    Owner   = "Crossfeed managed resource"
   }
 }
\ No newline at end of file