From 6f1e40255e9f3ac41bb93073cbeab4269150d376 Mon Sep 17 00:00:00 2001
From: Kenneth Barry <kenneth@sumologic.com>
Date: Fri, 1 Feb 2019 08:20:11 -0800
Subject: [PATCH] Added sampling functionality

---
 .../cloudwatchlogs_lambda.js                  |  7 ++++
 cloudwatchlogs-with-dlq/package.json          |  2 +-
 cloudwatchlogs-with-dlq/samplingutils.js      | 38 +++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 cloudwatchlogs-with-dlq/samplingutils.js

diff --git a/cloudwatchlogs-with-dlq/cloudwatchlogs_lambda.js b/cloudwatchlogs-with-dlq/cloudwatchlogs_lambda.js
index e405408..d9befce 100644
--- a/cloudwatchlogs-with-dlq/cloudwatchlogs_lambda.js
+++ b/cloudwatchlogs-with-dlq/cloudwatchlogs_lambda.js
@@ -17,6 +17,7 @@ var consoleFormatRegex = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}.\d{3}Z\s(\w+?-\w+
 var requestIdRegex = /(?:RequestId:|Z)\s+([\w\d\-]+)/;
 var url = require('url');
 var vpcutils = require('./vpcutils');
+var samplingutils = require('./samplingutils');
 var SumoLogsClient = require('./sumo-dlq-function-utils').SumoLogsClient;
 var Utils = require('./sumo-dlq-function-utils').Utils;
 
@@ -81,6 +82,7 @@ function getConfig(env) {
         "LogFormat": env.LOG_FORMAT || 'Others',
         "compressData": env.COMPRESS_DATA || true,
         "vpcCIDRPrefix": env.VPC_CIDR_PREFIX || '',
+        "samplingPercent": env.SAMPLING_PERCENT || '',
         "includeLogInfo": ("INCLUDE_LOG_INFO" in env) ? env.INCLUDE_LOG_INFO === "true" : false,
         "includeSecurityGroupInfo": ("INCLUDE_SECURITY_GROUP_INFO" in env) ? env.INCLUDE_SECURITY_GROUP_INFO === "true" : false
     };
@@ -112,10 +114,15 @@ function transformRecords(config, records) {
 
 function filterRecords(config, records) {
     var filteredRecords = records;
+    console.log(config)
     if (config.LogFormat.startsWith("VPC") && config.vpcCIDRPrefix) {
         filteredRecords = vpcutils.discardInternalTraffic(config.vpcCIDRPrefix, records);
         console.log(records.length - filteredRecords.length + " records discarded as InternalTraffic");
     }
+    if (config.samplingPercent) {
+        filteredRecords = samplingutils.sampleTraffic(config.samplingPercent, records);
+        console.log(records.length - filteredRecords.length + " records discarded via Sampling");
+    }
     return filteredRecords;
 }
 
diff --git a/cloudwatchlogs-with-dlq/package.json b/cloudwatchlogs-with-dlq/package.json
index 69dec72..ca42b01 100644
--- a/cloudwatchlogs-with-dlq/package.json
+++ b/cloudwatchlogs-with-dlq/package.json
@@ -10,7 +10,7 @@
   "devDependencies": {},
   "scripts": {
     "test": "node -e 'require('./test').test()'",
-    "build": "rm -f cloudwatchlogs-with-dlq.zip && zip -r cloudwatchlogs-with-dlq.zip DLQProcessor.js cloudwatchlogs_lambda.js vpcutils.js package.json sumo-dlq-function-utils/ node_modules/",
+    "build": "rm -f cloudwatchlogs-with-dlq.zip && zip -r cloudwatchlogs-with-dlq.zip DLQProcessor.js cloudwatchlogs_lambda.js vpcutils.js samplingutils.js package.json sumo-dlq-function-utils/ node_modules/",
     "prod_deploy": "python -c 'from test_cwl_lambda import prod_deploy;prod_deploy()'"
   },
   "author": "Himanshu Pal",
diff --git a/cloudwatchlogs-with-dlq/samplingutils.js b/cloudwatchlogs-with-dlq/samplingutils.js
new file mode 100644
index 0000000..2bc8c6f
--- /dev/null
+++ b/cloudwatchlogs-with-dlq/samplingutils.js
@@ -0,0 +1,38 @@
+function sampleTraffic(samplingPercent, records) {
+    if (!samplingPercent) {
+        return records;
+    }
+    var filteredRecords = [];
+    var toBeSamplesRecords = [];
+    records.forEach(function (log) {
+        if ("logLevel" in log.message.message) {
+            if(Number.isInteger(log.message.message.logLevel)){
+                if (log.message.message.logLevel >= 4) {
+                    filteredRecords.push(log);
+                } else {
+                    toBeSamplesRecords.push(log);
+                }
+            } else {
+                filteredRecords.push(log);
+            }
+        } else {
+            filteredRecords.push(log);
+        }
+    });
+    if(samplingPercent == 0) {
+        var increment = 0;
+    } else {
+        var increment = 100/samplingPercent;
+    }
+    
+    if(increment > 0) {
+        for (var i = 0; i < toBeSamplesRecords.length; i += increment) {
+            filteredRecords.push(toBeSamplesRecords[i]);
+        }
+    }
+    return filteredRecords;
+}
+
+module.exports = {
+    sampleTraffic: sampleTraffic
+};
\ No newline at end of file