Merge pull request #208 from cisagov/SQS-fixes

Clean-up SQS code and fix visibility timeout

backend/env.yml

@@ -54,11 +54,7 @@ staging:
   STAGE: staging
   PE_FARGATE_CLUSTER_NAME: pe-staging-worker
   PE_FARGATE_TASK_DEFINITION_NAME: pe-staging-worker
-  SHODAN_QUEUE_URL: ${ssm:/crossfeed/staging/SHODAN_QUEUE_URL}
-  DNSTWIST_QUEUE_URL: ${ssm:/crossfeed/staging/DNSTWIST_QUEUE_URL}
-  HIBP_QUEUE_URL: ${ssm:/crossfeed/staging/HIBP_QUEUE_URL}
-  INTELX_QUEUE_URL: ${ssm:/crossfeed/staging/INTELX_QUEUE_URL}
-  CYBERSIXGILL_QUEUE_URL: ${ssm:/crossfeed/staging/CYBERSIXGILL_QUEUE_URL}
+  QUEUE_URL: ${ssm:/crossfeed/staging/QUEUE_URL}
   EMAIL_BUCKET_NAME: cisa-crossfeed-staging-html-email
 
 prod:
@@ -105,11 +101,7 @@ prod:
   STAGE: prod
   PE_FARGATE_CLUSTER_NAME: pe-prod-worker
   PE_FARGATE_TASK_DEFINITION_NAME: pe-prod-worker
-  SHODAN_QUEUE_URL: ${ssm:/crossfeed/prod/SHODAN_QUEUE_URL}
-  DNSTWIST_QUEUE_URL: ${ssm:/crossfeed/prod/DNSTWIST_QUEUE_URL}
-  HIBP_QUEUE_URL: ${ssm:/crossfeed/prod/HIBP_QUEUE_URL}
-  INTELX_QUEUE_URL: ${ssm:/crossfeed/prod/INTELX_QUEUE_URL}
-  CYBERSIXGILL_QUEUE_URL: ${ssm:/crossfeed/prod/CYBERSIXGILL_QUEUE_URL}
+  QUEUE_URL: ${ssm:/crossfeed/staging/QUEUE_URL}
   EMAIL_BUCKET_NAME: cisa-crossfeed-staging-html-email
 
 dev-vpc:

backend/package.json

@@ -112,11 +112,10 @@
     "lint:fix": "eslint '**/*.{ts,tsx,js,jsx}' --fix",
     "pesyncdb": "docker-compose exec -T backend npx ts-node src/tools/run-pesyncdb.ts",
     "scan-exec": "docker-compose exec -T backend npx ts-node src/tools/run-scanExecution.ts",
-    "send-message": "node sendMessage.js",
     "syncdb": "docker-compose exec -T backend npx ts-node src/tools/run-syncdb.ts",
     "syncmdl": "docker-compose exec -T backend npx ts-node src/tools/run-syncmdl.ts",
     "test": "jest --detectOpenHandles",
     "test-python": "pytest"
   },
   "version": "1.0.0"
-}
+}

backend/serverless.yml

@@ -112,35 +112,35 @@ resources:
       Type: AWS::SQS::Queue
       Properties:
         QueueName: ${self:provider.stage}-shodan-queue
-        VisibilityTimeout: 300
+        VisibilityTimeout: 18000  # 5 hours
         MaximumMessageSize: 262144  # 256 KB
         MessageRetentionPeriod: 604800  # 7 days
     DnstwistQueue:
       Type: AWS::SQS::Queue
       Properties:
         QueueName: ${self:provider.stage}-dnstwist-queue
-        VisibilityTimeout: 300
+        VisibilityTimeout: 18000  # 5 hours
         MaximumMessageSize: 262144  # 256 KB
         MessageRetentionPeriod: 604800  # 7 days
     HibpQueue:
       Type: AWS::SQS::Queue
       Properties:
         QueueName: ${self:provider.stage}-hibp-queue
-        VisibilityTimeout: 300
+        VisibilityTimeout: 18000  # 5 hours
         MaximumMessageSize: 262144  # 256 KB
         MessageRetentionPeriod: 604800  # 7 days
     IntelxQueue:
       Type: AWS::SQS::Queue
       Properties:
         QueueName: ${self:provider.stage}-intelx-queue
-        VisibilityTimeout: 300
+        VisibilityTimeout: 18000  # 5 hours
         MaximumMessageSize: 262144  # 256 KB
         MessageRetentionPeriod: 604800  # 7 days
     CybersixgillQueue:
       Type: AWS::SQS::Queue
       Properties:
         QueueName: ${self:provider.stage}-cybersixgill-queue
-        VisibilityTimeout: 300
+        VisibilityTimeout: 18000  # 5 hours
         MaximumMessageSize: 262144  # 256 KB
         MessageRetentionPeriod: 604800  # 7 days
 

backend/src/tasks/pesyncdb.ts

@@ -6576,6 +6576,9 @@ VALUES ('WhoisXML', 'DNS lookpus', '2022-03-14');
 INSERT INTO public.data_source(name, description, last_run)
 VALUES ('findomain', 'Domain enumerator', '2022-03-14');
 
+INSERT INTO public.data_source(name, description, last_run)
+VALUES ('IntelX', 'Credentials', '2022-03-14');
+
 INSERT INTO public.data_source(name, description, last_run)
 VALUES ('Sublist3r', 'Domain Permutations', '2022-03-14');
 

backend/src/tasks/scanExecution.ts

@@ -4,6 +4,8 @@ import { integer } from 'aws-sdk/clients/cloudfront';
 
 const ecs = new AWS.ECS();
 let docker: any;
+const QUEUE_URL = process.env.QUEUE_URL!;
+const SCAN_LIST = ['dnstwist', 'hibp', 'intelx', 'cybersixgill', 'shodan'];
 
 if (process.env.IS_LOCAL) {
   const Docker = require('dockerode');
@@ -15,19 +17,33 @@ const toSnakeCase = (input) => input.replace(/ /g, '-');
 async function startDesiredTasks(
   scanType: string,
   desiredCount: integer,
-  queueUrl: string
+  shodanApiKeyList: string[] = []
 ) {
+  const queueUrl = QUEUE_URL + `${scanType}-queue`;
   try {
     // ECS can only start 10 tasks at a time. Split up into batches
-    const batchSize = 10;
+    let batchSize = 10;
+    if (scanType == 'shodan') {
+      batchSize = 1;
+    }
     let remainingCount = desiredCount;
     while (remainingCount > 0) {
+      let shodan_api_key = '';
+      if (shodanApiKeyList.length > 0) {
+        shodan_api_key = shodanApiKeyList[remainingCount - 1];
+      }
       const currentBatchCount = Math.min(remainingCount, batchSize);
 
       if (process.env.IS_LOCAL) {
         // If running locally, use RabbitMQ and Docker instead of SQS and ECS
         console.log('Starting local containers');
-        await startLocalContainers(currentBatchCount, scanType, queueUrl);
+        console.log(queueUrl);
+        await startLocalContainers(
+          currentBatchCount,
+          scanType,
+          queueUrl,
+          shodan_api_key
+        );
       } else {
         await ecs
           .runTask({
@@ -55,6 +71,10 @@ async function startDesiredTasks(
                     {
                       name: 'SERVICE_QUEUE_URL',
                       value: queueUrl
+                    },
+                    {
+                      name: 'PE_SHODAN_API_KEYS',
+                      value: shodan_api_key
                     }
                   ]
                 }
@@ -75,7 +95,8 @@ async function startDesiredTasks(
 async function startLocalContainers(
   count: number,
   scanType: string,
-  queueUrl: string
+  queueUrl: string,
+  shodan_api_key: string = ''
 ) {
   // Start 'count' number of local Docker containers
   for (let i = 0; i < count; i++) {
@@ -117,7 +138,7 @@ async function startLocalContainers(
           `SIXGILL_CLIENT_ID=${process.env.SIXGILL_CLIENT_ID}`,
           `SIXGILL_CLIENT_SECRET=${process.env.SIXGILL_CLIENT_SECRET}`,
           `INTELX_API_KEY=${process.env.INTELX_API_KEY}`,
-          `PE_SHODAN_API_KEYS=${process.env.PE_SHODAN_API_KEYS}`,
+          `PE_SHODAN_API_KEYS=${shodan_api_key}`,
           `WORKER_SIGNATURE_PUBLIC_KEY=${process.env.WORKER_SIGNATURE_PUBLIC_KEY}`,
           `WORKER_SIGNATURE_PRIVATE_KEY=${process.env.WORKER_SIGNATURE_PRIVATE_KEY}`,
           `ELASTICSEARCH_ENDPOINT=${process.env.ELASTICSEARCH_ENDPOINT}`,
@@ -141,13 +162,16 @@ async function startLocalContainers(
 export const handler: Handler = async (event) => {
   let desiredCount: integer;
   let scanType: string;
+
+  // Check if desired count was passed
   if (event.desiredCount) {
     desiredCount = event.desiredCount;
   } else {
     console.log('Desired count not found. Setting to 1.');
     desiredCount = 1;
   }
 
+  // Check if scan type was passed
   if (event.scanType) {
     scanType = event.scanType;
   } else {
@@ -156,39 +180,38 @@ export const handler: Handler = async (event) => {
   }
 
   try {
+    // If scanType is shodan, check if API keys were passed and split up
     if (scanType === 'shodan') {
-      await startDesiredTasks(
-        scanType,
-        desiredCount,
-        process.env.SHODAN_QUEUE_URL!
-      );
-    } else if (scanType === 'dnstwist') {
-      await startDesiredTasks(
-        scanType,
-        desiredCount,
-        process.env.DNSTWIST_QUEUE_URL!
-      );
-    } else if (scanType === 'hibp') {
-      await startDesiredTasks(
-        scanType,
-        desiredCount,
-        process.env.HIBP_QUEUE_URL!
-      );
-    } else if (scanType === 'intelx') {
-      await startDesiredTasks(
-        scanType,
-        desiredCount,
-        process.env.INTELX_QUEUE_URL!
-      );
-    } else if (scanType === 'cybersixgill') {
-      await startDesiredTasks(
-        scanType,
-        desiredCount,
-        process.env.CYBERSIXGILL_QUEUE_URL!
-      );
+      let shodanApiKeyList: string[];
+      if (event.apiKeyList) {
+        shodanApiKeyList = event.apiKeyList
+          .split(',')
+          .map((value) => value.trim());
+      } else {
+        console.error(
+          'apiKeyList must be provided for shodan and be a comma-separated string'
+        );
+        return 'Failed no apiKeyList';
+      }
+      // Check if there are enough keys for the desired number of tasks
+      if (shodanApiKeyList.length >= desiredCount) {
+        console.log(
+          'The number of API keys is greater than or equal to desiredCount.'
+        );
+      } else {
+        console.error(
+          'The number of API keys is less than desired Fargate tasks.'
+        );
+        return 'Failed no apiKeyList';
+      }
+      await startDesiredTasks(scanType, desiredCount, shodanApiKeyList);
+
+      // Run the rest of the scans normally
+    } else if (SCAN_LIST.includes(scanType)) {
+      await startDesiredTasks(scanType, desiredCount);
     } else {
       console.log(
-        'Shodan, DNSTwist, HIBP, and Cybersixgill are the only script types available right now.'
+        'Shodan, DNSTwist, HIBP, IntelX, and Cybersixgill are the only script types available right now. Must be all lowercase.'
       );
     }
   } catch (error) {

backend/src/tools/run-scanExecution.ts

@@ -1,10 +1,48 @@
 // Script to execute the scanExecution function
 import { handler as scanExecution } from '../tasks/scanExecution';
+const amqp = require('amqplib');
 
-async function localScanExecution() {
+async function localScanExecution(scan_type, desired_count, apiKeyList = '') {
   console.log('Starting...');
-  const payload = { scanType: 'dnstwist', desiredCount: 3 };
+  const payload = {
+    scanType: scan_type,
+    desiredCount: desired_count,
+    apiKeyList: apiKeyList
+  };
   scanExecution(payload, {} as any, () => null);
 }
 
-localScanExecution();
+async function sendMessageToQueue(message, queue) {
+  const connection = await amqp.connect('amqp://rabbitmq');
+  const channel = await connection.createChannel();
+
+  await channel.assertQueue(queue, { durable: true });
+
+  // Simulate sending a message to the queue
+  channel.sendToQueue(queue, Buffer.from(JSON.stringify(message)), {
+    persistent: true
+  });
+
+  console.log('Message sent:', message);
+
+  setTimeout(() => {
+    connection.close();
+  }, 500);
+}
+
+// Simulate sending a message
+const SCAN_TYPE = 'dnstwist';
+const DESIRED_COUNT = 1;
+const ORG_LIST = ['DHS', 'DOI'];
+const QUEUE = `staging-${SCAN_TYPE}-queue`;
+const API_KEY_LIST = '';
+
+for (const org of ORG_LIST) {
+  const message = {
+    scriptType: SCAN_TYPE,
+    org: org
+  };
+  sendMessageToQueue(message, QUEUE);
+}
+
+localScanExecution(SCAN_TYPE, DESIRED_COUNT, API_KEY_LIST);

backend/worker/pe-worker-entry.sh

@@ -33,13 +33,13 @@ while true; do
       MESSAGE=$(echo "$RESPONSE" | jq -r '.[0].payload')
     MESSAGE=${MESSAGE//\\\"/\"}
     echo "MESSAGE: $MESSAGE"
-
   else
     echo "Running live SQS logic..."
     MESSAGE=$(aws sqs receive-message --queue-url "$SERVICE_QUEUE_URL" --output json --max-number-of-messages 1)
     echo "MESSAGE: $MESSAGE"
   fi
 
+
   # Check if there are no more messages. If no more, then exit Fargate container
   if [ -z "$MESSAGE" ] || [ "$MESSAGE" == "null" ]; then
     echo "No more messages in the queue. Exiting."
@@ -72,7 +72,7 @@ while true; do
 
   # Run the pe-source command
   eval "$COMMAND" \
-    && cat /app/pe_reports_logging.log
+    && cat /app/pe_reports_logging.log > /app/pe_reports_logging.log
 
   # Delete the processed message from the queue
   if [ "$IS_LOCAL" = true ]; then

dev.env.example

@@ -98,12 +98,8 @@ PE_DB_USERNAME=pe
 PE_DB_PASSWORD=password
 
 
-SHODAN_QUEUE_URL =shodanQueue
+QUEUE_URL=staging-
 PE_SHODAN_API_KEYS=
-DNSTWIST_QUEUE_URL=dnstwistQueue
-HIBP_QUEUE_URL=hibpQueue
-INTELX_QUEUE_URL=intelxQueue
-CYBERSIXGILL_QUEUE_URL=cybersixgillQueue
 
 PE_FARGATE_CLUSTER_NAME=pe-staging-worker
 PE_FARGATE_TASK_DEFINITION_NAME=pe-staging-worker