diff --git a/b64decode/b64decode/b64decode.py b/b64decode/b64decode/b64decode.py index ce55f7b..4fb7381 100644 --- a/b64decode/b64decode/b64decode.py +++ b/b64decode/b64decode/b64decode.py @@ -22,7 +22,9 @@ """ -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 @@ -30,6 +32,23 @@ 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) diff --git a/b64decode/tests/test_b64decode.py b/b64decode/tests/test_b64decode.py index ebbe11c..edb40d2 100644 --- a/b64decode/tests/test_b64decode.py +++ b/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) diff --git a/entropy/entropy/entropy.py b/entropy/entropy/entropy.py index d92c480..94fa760 100644 --- a/entropy/entropy/entropy.py +++ b/entropy/entropy/entropy.py @@ -21,6 +21,7 @@ """ import math +from asyncio import sleep from typing import Dict from collections import Counter @@ -28,12 +29,16 @@ 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) diff --git a/hash/hash/hash.py b/hash/hash/hash.py index cd22fa6..6a55a9a 100644 --- a/hash/hash/hash.py +++ b/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 ) diff --git a/hash_ssdeep/hash_ssdeep/hash_ssdeep.py b/hash_ssdeep/hash_ssdeep/hash_ssdeep.py index 07ddb1c..9ca77a0 100644 --- a/hash_ssdeep/hash_ssdeep/hash_ssdeep.py +++ b/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 + )