Less blocking for basic functions

b64decode/b64decode/b64decode.py

@@ -22,14 +22,33 @@
 
 """
 
-import base64
+from asyncio import sleep
+from base64 import b64decode
+from re import sub
 
 from stoq.plugins import WorkerPlugin
 from stoq import ExtractedPayload, Payload, Request, WorkerResponse
 
 
 class B64Decode(WorkerPlugin):
     async def scan(self, payload: Payload, request: Request) -> WorkerResponse:
-        decoded_content = base64.b64decode(payload.content)
+        decoded_content = b''
+        remainder = b''
+        block_size = 2 ** 24  # Pass control back to asyncio loop every 16MB
+        for block_index in range(0, len(payload.content), block_size):
+            block = remainder + sub(rb'[^A-Za-z0-9+/=]', b'', payload.content[block_index:block_index + block_size])
+            remainder_index = - (len(block) % 4)
+            if remainder_index:
+                decoded_content += b64decode(block[:remainder_index])
+                remainder = block[remainder_index:]
+            else:
+                decoded_content += b64decode(block)
+                remainder = b''
+            if b'=' in block:
+                break
+            await sleep(0)
+        if remainder:
+            remainder = remainder.ljust(len(remainder)//4 + 4, b'=')
+            decoded_content += b64decode(remainder)
         extracted = [ExtractedPayload(decoded_content)]
         return WorkerResponse(extracted=extracted)

b64decode/tests/test_b64decode.py

@@ -15,8 +15,8 @@
 #   limitations under the License.
 
 import os
-import base64
 import asynctest
+from base64 import b64encode
 
 from pathlib import Path
 
@@ -38,8 +38,17 @@ def tearDown(self) -> None:
     async def test_scan(self) -> None:
         s = Stoq(plugin_dir_list=[str(self.plugin_dir)])
         plugin = s.load_plugin(self.plugin_name)
-        payload = Payload(base64.b64encode(self.generic_data))
+        payload = Payload(b64encode(self.generic_data))
         response = await plugin.scan(payload, RequestMeta())
         self.assertIsInstance(response, WorkerResponse)
         self.assertEqual(1, len(response.extracted))
         self.assertEqual(self.generic_data, response.extracted[0].content)
+
+    async def test_scan_invalid(self) -> None:
+        s = Stoq(plugin_dir_list=[str(self.plugin_dir)])
+        plugin = s.load_plugin(self.plugin_name)
+        payload = Payload(b64encode(self.generic_data)[:-1])
+        response = await plugin.scan(payload, RequestMeta())
+        self.assertIsInstance(response, WorkerResponse)
+        self.assertEqual(1, len(response.extracted))
+        self.assertEqual(self.generic_data[:-1], response.extracted[0].content)

entropy/entropy/entropy.py

@@ -21,19 +21,24 @@
 """
 
 import math
+from asyncio import sleep
 from typing import Dict
 from collections import Counter
 
 from stoq.plugins import WorkerPlugin
 from stoq import Payload, Request, WorkerResponse
 
 
-class Hash(WorkerPlugin):
+class Entropy(WorkerPlugin):
     async def scan(self, payload: Payload, request: Request) -> WorkerResponse:
         entropy: float = 0.0
         results: Dict[str, float] = {}
         if payload.content:
-            occurences = Counter(bytearray(payload.content))
+            occurences = Counter()
+            block_size = 2**24  # Pass control back to asyncio loop every 16MB
+            for block_index in range(0, len(payload.content), block_size):
+                occurences.update(payload.content[block_index:block_index+block_size])
+                await sleep(0)
             for bc in occurences.values():
                 b = float(bc) / len(payload.content)
                 entropy -= b * math.log(b, 2)

hash/hash/hash.py

@@ -22,6 +22,7 @@
 
 """
 
+from asyncio import sleep
 from hashlib import md5, sha1, sha256
 
 from stoq.plugins import WorkerPlugin
@@ -30,10 +31,22 @@
 
 class Hash(WorkerPlugin):
     async def scan(self, payload: Payload, request: Request) -> WorkerResponse:
+        secure_hashes = {
+            'md5': md5(),
+            'sha1': sha1(),
+            'sha256': sha256(),
+        }
+        block_size = 2 ** 24  # Pass control back to asyncio loop every 16MB
+        for block_index in range(0, len(payload.content), block_size):
+            for function in secure_hashes.keys():
+                secure_hashes[function].update(payload.content[block_index:block_index + block_size])
+                await sleep(0)
+        for function in secure_hashes.keys():
+            if hasattr(secure_hashes[function], 'hexdigest'):
+                secure_hashes[function] = secure_hashes[function].hexdigest()
+            else:
+                secure_hashes[function] = secure_hashes[function].digest()
+
         return WorkerResponse(
-            results={
-                'sha256': sha256(payload.content).hexdigest(),
-                'md5': md5(payload.content).hexdigest(),
-                'sha1': sha1(payload.content).hexdigest(),
-            }
+            results=secure_hashes
         )

hash_ssdeep/hash_ssdeep/hash_ssdeep.py

@@ -23,11 +23,27 @@
 """
 
 import ssdeep
+from asyncio import sleep
 
 from stoq.plugins import WorkerPlugin
 from stoq import Payload, Request, WorkerResponse
 
 
-class HashSsdeep(WorkerPlugin):
+class HashSSDeep(WorkerPlugin):
     async def scan(self, payload: Payload, request: Request) -> WorkerResponse:
-        return WorkerResponse(results={'ssdeep': ssdeep.hash(payload.content)})
+        secure_hashes = {
+            'ssdeep': ssdeep.Hash(),
+        }
+        block_size = 2 ** 24  # Pass control back to asyncio loop every 16MB
+        for block_index in range(0, len(payload.content), block_size):
+            for function in secure_hashes.keys():
+                secure_hashes[function].update(payload.content[block_index:block_index + block_size])
+                await sleep(0)
+        for function in secure_hashes.keys():
+            if hasattr(secure_hashes[function], 'hexdigest'):
+                secure_hashes[function] = secure_hashes[function].hexdigest()
+            else:
+                secure_hashes[function] = secure_hashes[function].digest()
+        return WorkerResponse(
+            results=secure_hashes
+        )