Implement chunking algorithm

This commit adds a chunking algorithm which will split payloads based on
the maximum allowed size and length of data for the Nightfall API.

Author: levlaz

CHANGELOG

@@ -3,6 +3,13 @@ Changelog
 
 Here you can see the full list of changes between each Nightfall release.
 
+Version 0.3.0
+-------------
+
+Released on June 13, 2021
+
+- Implement basic chunking algorithm to split payloads per the API limts.
+
 Version 0.2.0
 -------------
 

docs/api.rst

@@ -1,9 +1,6 @@
 API
 ===
 
-.. note::
-    Unless otherwise noted all arguments are of the :class:`str` type.
-
 API Object
 ----------
 

nightfall/api.py

@@ -11,7 +11,18 @@
 import logging
 
 class Api():
-    """A python interface into the Nightfall API"""
+    """A python interface for the Nightfall API.
+
+    .. data:: MAX_PAYLOAD_SIZE 
+
+        Maximum payload size that the Nightfall API will accept 
+
+    .. data:: MAX_NUM_ITEMS
+
+        Maximum number of items that the Nightfall API will accept
+    """
+    MAX_PAYLOAD_SIZE = 450_000
+    MAX_NUM_ITEMS = 50_000
 
     def __init__(self, token, condition_set):
         """Instantiate a new nightfall.Api object.
@@ -27,26 +38,73 @@ def __init__(self, token, condition_set):
             'x-api-key': self.token
         }
 
+    def make_payloads(self, data):
+        """Turn a list of strings into a list of acceptable payloads.
+
+        Creates chunks based on the MAX_PAYLOAD_SIZE and MAX_NUM_ITEMS 
+        constants.
+
+        :param data: list of string
+        :returns: list of list of strings
+        """
+        cur_chunk_bytes = 0
+        cur_chunk = []
+        chunks = []
+
+        for i in data:
+            if cur_chunk_bytes + len(i) >= self.MAX_PAYLOAD_SIZE or \
+                len(cur_chunk) >= self.MAX_NUM_ITEMS:
+                chunks.append(cur_chunk)
+
+                cur_chunk_bytes = len(i)
+
+                if len(i) < self.MAX_PAYLOAD_SIZE:
+                    cur_chunk = [i]
+                else:
+                    cur_chunk = []
+                    for i in range(0, len(i), self.MAX_PAYLOAD_SIZE):
+                        chunks.append([i[i:i+self.MAX_PAYLOAD_SIZE]])
+
+            else:
+                cur_chunk.append(i)
+                cur_chunk_bytes += len(i)
+        if cur_chunk:
+            chunks.append(cur_chunk)
+
+        return chunks
+
+
     def scan(self, data):
-        """Scan a piece of data with Nightfall.
+        """Scan lists of data with Nightfall.
+
+        This method will convert the list of strings into chunks if necessary 
+        and then makes one or more requests to the Nightfall API to scan the 
+        data.
 
-        :param data: Data to scan.
+        :param data: list of strings to scan.
         :type data: list
+
+        :returns: list of list of resposes for items in payload.
         """
-        payload = {
-            'payload': data,
-            'config': {
-                'conditionSetUUID': self.condition_set
+        responses = []
+
+        for i in self.make_payloads(data):
+            payload = {
+                'payload': i,
+                'config': {
+                    'conditionSetUUID': self.condition_set
+                }
             }
-        }
 
-        response = requests.post(
-            url=self._url,
-            headers=self._headers,
-            data=json.dumps(payload)
-        )
-        response.raise_for_status()
+            response = requests.post(
+                url=self._url,
+                headers=self._headers,
+                data=json.dumps(payload)
+            )
+            response.raise_for_status()
+            responses += response.json()
 
-        return response
+        return responses
+
 
 

setup.py

@@ -16,7 +16,7 @@ def readme():
 
 setup(
     name="nightfall",
-    version="0.2.0",
+    version="0.3.0",
     description="Python SDK for Nightfall",
     long_description=readme(),
     long_description_content_type="text/markdown",

tests/nightfall/test_api.py

@@ -13,12 +13,32 @@ def setUp(self):
             os.getenv('NIGHTFALL_CONDITION_SET')
             )
 
-    def testBadScan(self):
-        with self.assertRaises(requests.exceptions.HTTPError) as e:
-            resp = self.client.scan('')
-    
     def testScan(self):
-        """Test basics of API, can submit a request and receive a response."""
+        """Test basics of API, can submit a request and receive a response.
+        
+        This test assumes that you don't have a condition set matching the 'testing' string.
+        """
         resp = self.client.scan(['testing'])
-        self.assertEqual(resp.status_code, 200)
+        self.assertEqual(resp[0], None)
+
+    def testChunking(self):
+        """Test chunking algorithm."""
+
+        # a list of 10 strings that are 100k bytes each should turn into a 
+        # list of three lists
+        large_list = []
+        large_string = "x" * 100000
+
+        for i in range(0,10):
+            large_list.append(large_string)
 
+        chunks = self.client.make_payloads(large_list)
+        self.assertEqual(len(chunks), 3)
+
+        # a list of 100,000 single byte items should turn into two lists
+        many_list = []
+        for i in range(0, 100000):
+            many_list.append('x')
+
+        chunks = self.client.make_payloads(many_list)
+        self.assertEqual(len(chunks), 2)