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