Updating nodejs code of cw logs with dlq

himsharma01 · himsharma01 · commit 6a79a615bc5f · 2023-09-11T19:27:44.000+05:30
diff --git a/cloudwatchlogs-with-dlq/DLQProcessor.js b/cloudwatchlogs-with-dlq/DLQProcessor.js
@@ -1,77 +1,82 @@
-var AWS = require("aws-sdk");
-var processLogsHandler = require('./cloudwatchlogs_lambda').processLogs;
-var getEndpointURL = require('./cloudwatchlogs_lambda').getEndpointURL;
-var DLQUtils = require("./sumo-dlq-function-utils").DLQUtils;
-var Messages = DLQUtils.Messages;
-var invokeLambdas = DLQUtils.invokeLambdas;
+const { processLogs: processLogsHandler, getEndpointURL } = require('./cloudwatchlogs_lambda');
+const { DLQUtils } = require("./sumo-dlq-function-utils");
+
+const { Messages, invokeLambdas } = DLQUtils;
 
 exports.consumeMessages = async function (env, context, callback) {
-    var sqs = new AWS.SQS({region: env.AWS_REGION});
-    var MessagesObj = new Messages(env);
-    env.SUMO_CLIENT_HEADER="dlq-aws-lambda";
+    const MessagesObj = new Messages(env);
+    env.SUMO_CLIENT_HEADER = "dlq-aws-lambda";
+
     if (!env.SUMO_ENDPOINT) {
-        let SUMO_ENDPOINT = await getEndpointURL();
-        if (SUMO_ENDPOINT instanceof Error) {
-            console.log("Error in getEndpointURL: ", SUMO_ENDPOINT);
-            callback(SUMO_ENDPOINT, null);
+        try {
+            let SUMO_ENDPOINT = await getEndpointURL();
+            env.SUMO_ENDPOINT = SUMO_ENDPOINT;
+        } catch (error) {
+            console.log("Error in getEndpointURL: ", error);
+            callback(error, null);
             return;
         }
-        env.SUMO_ENDPOINT = SUMO_ENDPOINT;
     } else {
         console.log("consumeMessages: Getting SUMO_ENDPOINT from env");
     }
-    MessagesObj.receiveMessages(10, function (err, data) {
-        var messages = (data)? data.Messages: null;
-        if (err) {
-            callback(err);
-        } else if (messages && messages.length > 0) {
-            var fail_cnt = 0, msgCount = 0;
+
+    try {
+        const messages = await MessagesObj.receiveMessages(10);
+
+
+        if (messages && messages.length > 0) {
+            let fail_cnt = 0, msgCount = 0;
             console.log("Messages Received", messages.length);
-            for (var i = 0; i < messages.length; i++) {
-                (function(idx) {
-                    var payload = JSON.parse(messages[idx].Body);
-                    var receiptHandle = messages[idx].ReceiptHandle;
+
+            for (let i = 0; i < messages.length; i++) {
+                (function (idx) {
+                    const payload = JSON.parse(messages[idx].Body);
+                    const receiptHandle = messages[idx].ReceiptHandle;
+
                     if (!(payload.awslogs && payload.awslogs.data)) {
                         console.log("Message does not contain awslogs or awslogs.data attributes", payload);
-                        //deleting msg in DLQ after injesting in sumo
-                        MessagesObj.deleteMessage(receiptHandle, function (err, data) {
-                            if (err) console.log(err, err.stack);
-                        });
+
+                        MessagesObj.deleteMessage(receiptHandle)
+                            .catch((err) => console.log(err, err.stack));
+
                         return;
                     }
-                    var logdata = payload.awslogs.data;
+
+                    const logdata = payload.awslogs.data;
+
                     processLogsHandler(env, logdata, function (err, msg) {
                         msgCount++;
+
                         if (err) {
                             console.log(err, msg);
                             fail_cnt++;
                         } else {
-                            //deleting msg in DLQ after injesting in sumo
-                            MessagesObj.deleteMessage(receiptHandle, function (err, data) {
-                                if (err) console.log(err, err.stack);
-                            });
+                            MessagesObj.deleteMessage(receiptHandle)
+                                .catch((err) => console.log(err, err.stack));
                         }
-                        if (msgCount == messages.length) {
-                            if (fail_cnt == 0 && (parseInt(env.is_worker) === 0)) {
+
+                        if (msgCount === messages.length) {
+                            if (fail_cnt === 0 && parseInt(env.is_worker) === 0) {
                                 invokeLambdas(env.AWS_REGION, parseInt(env.NUM_OF_WORKERS),
-                                              context.functionName, '{"is_worker": "1"}', context);
+                                    context.functionName, '{"is_worker": "1"}', context);
                             }
-                            callback(null, (messages.length-fail_cnt) + ' success');
+
+                            callback(null, `${messages.length - fail_cnt} success`);
                         }
                     });
                 })(i);
             }
-
         } else {
+
             callback(null, 'success');
         }
-    });
+    } catch (error) {
+        callback(error);
+    }
 };
 
 exports.handler = function (event, context, callback) {
-
-    var env = Object.assign({}, process.env);
-    env['is_worker'] = event.is_worker || 0;
+    const env = Object.assign({}, process.env);
+    env.is_worker = event.is_worker || 0;
     exports.consumeMessages(env, context, callback);
-};
-
+};
diff --git a/cloudwatchlogs-with-dlq/cloudwatchlogs_lambda.js b/cloudwatchlogs-with-dlq/cloudwatchlogs_lambda.js
@@ -19,29 +19,25 @@ var url = require('url');
 var vpcutils = require('./vpcutils');
 var SumoLogsClient = require('./sumo-dlq-function-utils').SumoLogsClient;
 var Utils = require('./sumo-dlq-function-utils').Utils;
-const AWS = require('aws-sdk');
-const ssm = new AWS.SSM();
+
+const { SSMClient, GetParameterCommand } = require("@aws-sdk/client-ssm");
 
 exports.getEndpointURL = async function() {
-  console.log('Getting SUMO_ENDPOINT from AWS SSM Parameter Store');
-  return new Promise((resolve, reject) => {
-    ssm.getParameter(
-      {
-        Name: 'SUMO_ENDPOINT',
-        WithDecryption: true
-      },
-      (err, data) => {
-        if (err) {
-          console.log(err, err.stack);
-          reject(new Error('Unable to get EndpointURL from SSM: ' + err));
-        } else {
-          // console.log(data);
-          resolve(data.Parameter.Value);
-        }
-      }
-    );
-  });
-}
+    console.log('Getting SUMO_ENDPOINT from AWS SSM Parameter Store');
+    const ssmClient = new SSMClient();
+    try {
+      const data = await ssmClient.send(
+        new GetParameterCommand({
+          Name: 'SUMO_ENDPOINT',
+          WithDecryption: true
+        })
+      );
+      return data.Parameter.Value;
+    } catch (error) {
+      console.error('Unable to get EndpointURL from SSM:', error);
+      throw new Error('Unable to get EndpointURL from SSM: ' + error);
+    }
+  }
 
 function createRecords(config, events, awslogsData) {
     var records = [];
@@ -210,4 +206,4 @@ exports.handler = function (event, context, callback) {
 
     exports.processLogs(process.env, event.awslogs.data, callback);
 
-};
+};
diff --git a/cloudwatchlogs-with-dlq/sumo-dlq-function-utils/lib/dlqutils.js b/cloudwatchlogs-with-dlq/sumo-dlq-function-utils/lib/dlqutils.js
@@ -1,46 +1,54 @@
-var AWS = require("aws-sdk");
+const { SQSClient, ReceiveMessageCommand, DeleteMessageCommand } = require("@aws-sdk/client-sqs");
+const { LambdaClient, InvokeCommand } = require("@aws-sdk/client-lambda");
 
-function Messages(env) {
-    this.sqs = new AWS.SQS({region: env.AWS_REGION});
+class Messages {
+  constructor(env) {
+    this.sqs = new SQSClient({ region: env.AWS_REGION });
     this.env = env;
-}
+  }
 
-Messages.prototype.receiveMessages = function (messageCount, callback) {
-    var params = {
-        QueueUrl: this.env.TASK_QUEUE_URL,
-        MaxNumberOfMessages: messageCount
+  async receiveMessages(messageCount) {
+    const params = {
+      QueueUrl: this.env.TASK_QUEUE_URL,
+      MaxNumberOfMessages: messageCount,
     };
-    this.sqs.receiveMessage(params, callback);
-};
-
-Messages.prototype.deleteMessage = function (receiptHandle, callback) {
-    this.sqs.deleteMessage({
-        ReceiptHandle: receiptHandle,
-        QueueUrl: this.env.TASK_QUEUE_URL
-    }, callback);
-};
-
-function invokeLambdas(awsRegion, numOfWorkers, functionName, payload, context) {
-
-    for (var i = 0; i < numOfWorkers; i++) {
-        var lambda = new AWS.Lambda({
-          region: awsRegion
-        });
-        lambda.invoke({
-            InvocationType: 'Event',
-            FunctionName: functionName,
-            Payload: payload
-        }, function(err, data) {
-           if (err) {
-               context.fail(err);
-           } else {
-               context.succeed('success');
-           }
-        });
+
+    const command = new ReceiveMessageCommand(params);
+    const response = await this.sqs.send(command);
+    return response.Messages || [];
+  }
+
+  async deleteMessage(receiptHandle) {
+    const params = {
+      ReceiptHandle: receiptHandle,
+      QueueUrl: this.env.TASK_QUEUE_URL,
+    };
+
+    const command = new DeleteMessageCommand(params);
+    await this.sqs.send(command);
+  }
+}
+
+async function invokeLambdas(awsRegion, numOfWorkers, functionName, payload, context) {
+  const lambda = new LambdaClient({ region: awsRegion });
+
+  for (let i = 0; i < numOfWorkers; i++) {
+    const command = new InvokeCommand({
+      InvocationType: 'Event',
+      FunctionName: functionName,
+      Payload: payload,
+    });
+
+    try {
+      await lambda.send(command);
+      context.succeed('success');
+    } catch (err) {
+      context.fail(err);
     }
+  }
 }
 
 module.exports = {
-    Messages: Messages,
-    invokeLambdas: invokeLambdas
-};
+  Messages,
+  invokeLambdas,
+};
diff --git a/cloudwatchlogs-with-dlq/vpcutils.js b/cloudwatchlogs-with-dlq/vpcutils.js
@@ -1,7 +1,10 @@
-var find = require('lodash').find;
-var EC2 = require('aws-sdk/clients/ec2');
+// Import the required AWS SDK modules
+const { EC2Client, DescribeNetworkInterfacesCommand } = require('@aws-sdk/client-ec2');
 var jmespath = require('jmespath');
-var ec2 = null;
+var find = require('lodash').find;
+// Create an instance of the EC2 client
+const ec2Client = new EC2Client({ region: process.env.AWS_REGION });
+
 /*
 VPC Log Format
 version         The VPC Flow Logs version.
@@ -50,19 +53,25 @@ function discardInternalTraffic(vpcCIDRPrefix, records) {
  *
  * @return `Promise` for async processing
  */
-function listNetworkInterfaces(allIPaddresses) {
-    if (!ec2) {
-        ec2 = new EC2({region: process.env.AWS_REGION});
-    }
-    var params = {
-        Filters: [
-            {
-              Name: 'private-ip-address',
-              Values: allIPaddresses
-            }
-        ]
-    }
-    return ec2.describeNetworkInterfaces(params).promise();
+async function listNetworkInterfaces(allIPaddresses) {
+  const params = {
+    Filters: [
+      {
+        Name: 'private-ip-address',
+        Values: allIPaddresses,
+      },
+    ],
+  };
+
+  const command = new DescribeNetworkInterfacesCommand(params);
+
+  try {
+    const response = await ec2Client.send(command);
+    return response;
+  } catch (err) {
+    console.log('Error in listNetworkInterfaces', err);
+    throw err;
+  }
 }
 
 /**
@@ -145,4 +154,4 @@ function includeSecurityGroupIds(records) {
 module.exports = {
     discardInternalTraffic: discardInternalTraffic,
     includeSecurityGroupIds: includeSecurityGroupIds
-};
+};
