Fix lint issues, apply ruff formatting

Author: liquidsec

blockbuster/blockbuster.py

@@ -2,9 +2,7 @@
 
 import base64
 import io
-import os
 import sys
-import asyncio
 
 import httpx
 import pytest
@@ -17,6 +15,7 @@
 # AES-CBC padding oracle simulator
 # ---------------------------------------------------------------------------
 
+
 class PaddingOracle:
     """A simulated AES-CBC padding oracle for deterministic offline testing."""
 
@@ -55,8 +54,8 @@ def check_padding(self, ciphertext_with_iv: bytes) -> bool:
 
 
 # Fixed test key / IV for reproducible tests
-AES_KEY = bytes(range(16))                     # 0x00..0x0f
-AES_IV = bytes(range(16, 32))                  # 0x10..0x1f
+AES_KEY = bytes(range(16))  # 0x00..0x0f
+AES_IV = bytes(range(16, 32))  # 0x10..0x1f
 ORACLE = PaddingOracle(AES_KEY, AES_IV, blocksize=16)
 
 
@@ -130,8 +129,14 @@ def encrypted_hello(oracle):
 # respx-based mock HTTP oracle
 # ---------------------------------------------------------------------------
 
-def _make_oracle_route(oracle_obj, encoding_mode="base64", input_mode="parameter",
-                       vuln_param="token", oracle_text="Invalid padding"):
+
+def _make_oracle_route(
+    oracle_obj,
+    encoding_mode="base64",
+    input_mode="parameter",
+    vuln_param="token",
+    oracle_text="Invalid padding",
+):
     """Return a respx side-effect function that simulates a padding oracle."""
 
     def _handler(request: httpx.Request) -> httpx.Response:
@@ -141,24 +146,27 @@ def _handler(request: httpx.Request) -> httpx.Response:
 
         if input_mode == "parameter":
             from urllib.parse import parse_qs, urlparse
+
             qs = parse_qs(urlparse(url).query)
             token = qs.get(vuln_param, [None])[0]
         elif input_mode == "querystring":
             from urllib.parse import urlparse
+
             token = urlparse(url).query
         elif input_mode == "cookie":
             cookie_header = request.headers.get("cookie", "")
             for part in cookie_header.split(";"):
                 part = part.strip()
                 if part.startswith(vuln_param + "="):
-                    token = part[len(vuln_param) + 1:]
+                    token = part[len(vuln_param) + 1 :]
                     break
 
         if token is None:
             return httpx.Response(400, text="Missing token")
 
         # Decode
         import urllib.parse as up
+
         if encoding_mode == "base64":
             raw = up.unquote_plus(token)
             raw += "=" * (len(raw) % 4)

tests/conftest.py

@@ -1,15 +1,11 @@
 """Tests for nextBlock, block retry logic, and encrypt sanity check."""
 
-import asyncio
-import base64
-import io
-import sys
 from types import SimpleNamespace
-from unittest.mock import patch, AsyncMock, MagicMock
+from unittest.mock import patch, AsyncMock
 
 import pytest
 
-from tests.conftest import make_job, PaddingOracle, AES_KEY, AES_IV, ORACLE
+from tests.conftest import make_job
 
 
 class TestNextBlockDecrypt:
@@ -24,8 +20,9 @@ async def test_successful_decrypt_returns_0(self):
         j.blocks = [[0] * 16]
         j.iv = [0] * 16
 
-        with patch.object(j, 'decryptBlock', new_callable=AsyncMock,
-                          return_value=b"decrypted_block!"):
+        with patch.object(
+            j, "decryptBlock", new_callable=AsyncMock, return_value=b"decrypted_block!"
+        ):
             result = await j.nextBlock()
 
         assert result == 0
@@ -40,8 +37,12 @@ async def test_failed_decrypt_returns_1(self):
         j.blocks = [[0] * 16]
         j.iv = [0] * 16
 
-        with patch.object(j, 'decryptBlock', new_callable=AsyncMock,
-                          side_effect=Exception("Block failed")):
+        with patch.object(
+            j,
+            "decryptBlock",
+            new_callable=AsyncMock,
+            side_effect=Exception("Block failed"),
+        ):
             result = await j.nextBlock()
 
         assert result == 1
@@ -53,7 +54,9 @@ class TestNextBlockEncrypt:
 
     @pytest.mark.asyncio
     async def test_successful_encrypt_with_passing_sanity_check(self):
-        j = make_job(mode="encrypt", oracleMode="negative", oracleText="Invalid padding")
+        j = make_job(
+            mode="encrypt", oracleMode="negative", oracleText="Invalid padding"
+        )
         j.initialize_client()
         j.blockCount = 1
         j.currentBlock = 0
@@ -65,16 +68,20 @@ async def test_successful_encrypt_with_passing_sanity_check(self):
         def mock_request(token):
             return SimpleNamespace(text="OK", status_code=200)
 
-        with patch.object(j, 'encryptBlock', new_callable=AsyncMock, return_value=block_result):
-            with patch.object(j, 'makeRequest', side_effect=mock_request):
+        with patch.object(
+            j, "encryptBlock", new_callable=AsyncMock, return_value=block_result
+        ):
+            with patch.object(j, "makeRequest", side_effect=mock_request):
                 result = await j.nextBlock()
 
         assert result == 0
         assert j.solvedBlocks[0] == block_result
 
     @pytest.mark.asyncio
     async def test_failed_sanity_check_backs_out_block(self):
-        j = make_job(mode="encrypt", oracleMode="negative", oracleText="Invalid padding")
+        j = make_job(
+            mode="encrypt", oracleMode="negative", oracleText="Invalid padding"
+        )
         j.initialize_client()
         j.blockCount = 1
         j.currentBlock = 0
@@ -87,8 +94,10 @@ def mock_request(token):
             # Sanity check fails — the oracle text IS found (negative mode: pass = text absent)
             return SimpleNamespace(text="Invalid padding", status_code=200)
 
-        with patch.object(j, 'encryptBlock', new_callable=AsyncMock, return_value=block_result):
-            with patch.object(j, 'makeRequest', side_effect=mock_request):
+        with patch.object(
+            j, "encryptBlock", new_callable=AsyncMock, return_value=block_result
+        ):
+            with patch.object(j, "makeRequest", side_effect=mock_request):
                 result = await j.nextBlock()
 
         assert result == 1
@@ -103,8 +112,12 @@ async def test_encrypt_exception_returns_1(self):
         j.blocks = [[0] * 16]
         j.iv = [0] * 16
 
-        with patch.object(j, 'encryptBlock', new_callable=AsyncMock,
-                          side_effect=Exception("Encrypt failed")):
+        with patch.object(
+            j,
+            "encryptBlock",
+            new_callable=AsyncMock,
+            side_effect=Exception("Encrypt failed"),
+        ):
             result = await j.nextBlock()
 
         assert result == 1
@@ -136,8 +149,8 @@ async def mock_next_block():
         max_block_retries = 3
         block_failures = 0
 
-        with patch.object(j, 'nextBlock', side_effect=mock_next_block):
-            with patch('blockbuster.blockbuster.saveState'):
+        with patch.object(j, "nextBlock", side_effect=mock_next_block):
+            with patch("blockbuster.blockbuster.saveState"):
                 while j.currentBlock < j.blockCount:
                     result = await j.nextBlock()
                     if result == 0:
@@ -167,7 +180,7 @@ async def always_fail():
         max_block_retries = 3
         block_failures = 0
 
-        with patch.object(j, 'nextBlock', side_effect=always_fail):
+        with patch.object(j, "nextBlock", side_effect=always_fail):
             while j.currentBlock < j.blockCount:
                 result = await j.nextBlock()
                 if result == 0:
@@ -188,8 +201,13 @@ class TestNextBlockEncryptKnownIV:
     @pytest.mark.asyncio
     async def test_encrypt_knowniv_appends_iv(self):
         iv = list(range(16))
-        j = make_job(mode="encrypt", oracleMode="negative", oracleText="Invalid padding",
-                     ivMode="knownIV", iv=iv)
+        j = make_job(
+            mode="encrypt",
+            oracleMode="negative",
+            oracleText="Invalid padding",
+            ivMode="knownIV",
+            iv=iv,
+        )
         j.initialize_client()
         j.blockCount = 1
         j.currentBlock = 0
@@ -202,8 +220,10 @@ def mock_request(token):
             captured_tokens.append(token)
             return SimpleNamespace(text="OK", status_code=200)
 
-        with patch.object(j, 'encryptBlock', new_callable=AsyncMock, return_value=block_result):
-            with patch.object(j, 'makeRequest', side_effect=mock_request):
+        with patch.object(
+            j, "encryptBlock", new_callable=AsyncMock, return_value=block_result
+        ):
+            with patch.object(j, "makeRequest", side_effect=mock_request):
                 result = await j.nextBlock()
 
         assert result == 0

tests/test_block_orchestration.py

@@ -1,11 +1,7 @@
 """Tests for argument parsing and config validation in async_main."""
 
 import configparser
-import io
 import json
-import os
-import sys
-import tempfile
 
 import pytest
 
@@ -101,11 +97,16 @@ def test_invalid_http_method(self, tmp_path):
     def test_invalid_post_format(self, tmp_path):
         path = self._write_config(tmp_path, {"httpMethod": "POST", "postFormat": "xml"})
         config = self._read_config(path)
-        assert config["default"]["postFormat"] not in ("form-urlencoded", "multipart", "json")
+        assert config["default"]["postFormat"] not in (
+            "form-urlencoded",
+            "multipart",
+            "json",
+        )
 
     def test_invalid_proxy_ip(self):
         """socket.inet_aton should reject invalid IPs."""
         import socket
+
         with pytest.raises(socket.error):
             socket.inet_aton("not.an.ip")
 
@@ -131,12 +132,22 @@ def test_valid_iv_modes(self, tmp_path):
         for mode in ("firstblock", "knownIV", "unknown", "anchorBlock"):
             path = self._write_config(tmp_path, {"ivMode": mode})
             config = self._read_config(path)
-            assert config["default"]["ivMode"] in ("firstblock", "knownIV", "unknown", "anchorBlock")
+            assert config["default"]["ivMode"] in (
+                "firstblock",
+                "knownIV",
+                "unknown",
+                "anchorBlock",
+            )
 
     def test_invalid_iv_mode(self, tmp_path):
         path = self._write_config(tmp_path, {"ivMode": "randomIV"})
         config = self._read_config(path)
-        assert config["default"]["ivMode"] not in ("firstblock", "knownIV", "unknown", "anchorBlock")
+        assert config["default"]["ivMode"] not in (
+            "firstblock",
+            "knownIV",
+            "unknown",
+            "anchorBlock",
+        )
 
     def test_knowniv_requires_iv(self, tmp_path):
         """In knownIV mode, iv must be provided, correct length, all ints."""
@@ -146,7 +157,9 @@ def test_knowniv_requires_iv(self, tmp_path):
         assert len(iv) == 0  # Would fail validation
 
     def test_knowniv_iv_length_mismatch(self, tmp_path):
-        path = self._write_config(tmp_path, {"ivMode": "knownIV", "iv": "[1,2,3]", "blocksize": "16"})
+        path = self._write_config(
+            tmp_path, {"ivMode": "knownIV", "iv": "[1,2,3]", "blocksize": "16"}
+        )
         config = self._read_config(path)
         iv = json.loads(config["default"]["iv"])
         blocksize = int(config["default"]["blocksize"])