Merge pull request #6 from djosix/allow-parallel-encryption

Allow parallel encryption

Author: djosix

README.md

@@ -64,12 +64,17 @@ plaintext = decrypt(
 
 ### Encryption
 
-Below is an example demonstrating how to encrypt arbitrary bytes. For a detailed understanding of the process, please refer to [this Pull Request](https://github.com/djosix/padding_oracle.py/pull/4). Keep in mind that, unlike the decryption process, this functionality cannot be parallelized.
+Below is an example demonstrating how to encrypt arbitrary bytes. For a detailed understanding of the process, please refer to [this Pull Request](https://github.com/djosix/padding_oracle.py/pull/4).
 
 ```python
 from padding_oracle import encrypt
 
-ciphertext = encrypt(b'YourTextHere', block_size=16, oracle=oracle)
+ciphertext = encrypt(
+    b'YourTextHere', 
+    block_size=16,
+    oracle=oracle,
+    num_threads=16,
+)
 ```
 
 ### Customized Logging
@@ -106,6 +111,10 @@ The script also includes PHP-like encoding and decoding functions:
 from padding_oracle.encoding import urlencode, urldecode, base64_encode, base64_decode
 ```
 
+### TODO
+
+- [ ] Support more padding schemes
+
 ## License
 
 This script is distributed under the MIT license.

src/padding_oracle/padding_oracle.py

@@ -65,6 +65,7 @@ def encrypt(
     plaintext: bytes | str,
     block_size: int,
     oracle: OracleFunc,
+    num_threads: int = 1,
     logger: Logger = default_logger,
 ) -> bytes:
     plaintext = to_bytes(plaintext)
@@ -81,7 +82,7 @@ def encrypt(
             iv + cipher_blocks[i],
             block_size,
             oracle,
-            1,
+            num_threads,
             block_error_logger(logger),
             encrypt_progress_logger(logger, i+1, len(plain_blocks)),
         )

src/padding_oracle/solve.py

@@ -30,11 +30,22 @@
     cast,
 )
 
+__all__ = [
+    'solve',
+]
 
-class BlockResult(NamedTuple):
-    block_index: int
-    solved: list[int | None] | None = None
-    error: None | str = None
+
+class BlockResult:
+    def __init__(
+        self,
+        block_index: int,
+        *,
+        solved: list[int | None] | None = None,
+        error: None | str = None,
+    ):
+        self.block_index = block_index
+        self.solved = solved
+        self.error = error
 
 
 OracleFunc = Callable[[bytes], bool]
@@ -60,7 +71,7 @@ def solve(
     ciphertext: bytes,
     block_size: int,
     oracle: OracleFunc,
-    parallel: int,
+    num_threads: int,
     block_callback: Callable[[BlockResult], None],
     progress_callback: Callable[[list[int | None]], None],
 ) -> list[int | None]:
@@ -70,7 +81,7 @@ def solve(
         ciphertext,
         block_size,
         oracle,
-        parallel,
+        num_threads,
         block_callback,
         progress_callback,
     )
@@ -82,7 +93,7 @@ async def solve_async(
     ciphertext: bytes,
     block_size: int,
     oracle: OracleFunc,
-    parallel: int,
+    num_threads: int,
     block_callback: Callable[[BlockResult], None],
     progess_callback: Callable[[list[int | None]], None],
 ) -> list[int | None]:
@@ -95,7 +106,7 @@ async def solve_async(
         ciphertext,
         block_size,
         oracle,
-        parallel,
+        num_threads,
         block_callback,
         progess_callback,
     )
@@ -115,10 +126,10 @@ async def solve_async(
             if result.solved is None:
                 continue
 
-            if len(result.solved) >= ctx.solved_counts[result.block_index]:
+            if len(result.solved) > ctx.solved_counts[result.block_index]:
                 update_solved(ctx, result.block_index, result.solved)
                 ctx.solved_counts[result.block_index] = len(result.solved)
-                ctx.progress_callback(list[int | None](ctx.plaintext))
+                ctx.progress_callback(ctx.plaintext)
 
         if len(ctx.tasks) == 0:
             break
@@ -140,7 +151,7 @@ def create_solve_context(
     ciphertext: bytes,
     block_size: int,
     oracle: OracleFunc,
-    parallel: int,
+    num_threads: int,
     block_callback: Callable[[BlockResult], None],
     progress_callback: Callable[[list[int | None]], None],
 ) -> Context:
@@ -154,7 +165,7 @@ def create_solve_context(
 
     plaintext = [None] * (len(cipher_blocks) - 1) * block_size
 
-    executor = ThreadPoolExecutor(parallel)
+    executor = ThreadPoolExecutor(num_threads)
     loop = asyncio.get_event_loop()
     ctx = Context(
         block_size,
@@ -198,7 +209,7 @@ async def solve_block(ctx: Context, block_index: int, C0: list[int],
         if result is not None and result.error is not None:
             return result
 
-    return BlockResult(block_index, P1_suffix)
+    return BlockResult(block_index, solved=P1_suffix)
 
 
 async def exploit_oracle(ctx: Context, block_index: int,
@@ -217,7 +228,7 @@ async def exploit_oracle(ctx: Context, block_index: int,
     invalid |= len(X1_suffix) > 0 and len(hits) != 1
     if invalid:
         message = f'invalid number of hits: {len(hits)} (block: {block_index}, byte: {index})'
-        return BlockResult(block_index, message)
+        return BlockResult(block_index, error=message)
 
     for byte in hits:
         X1_test = [byte ^ padding, *X1_suffix]

src/padding_oracle/utils.py

@@ -40,7 +40,10 @@ def to_bytes(data: str | bytes | list[int]) -> bytes:
 
 
 def to_bytes_with_default(maybe_bytes: list[int | None], default: bytes = b' ') -> bytes:
-    return [default if b is None else b for b in maybe_bytes]
+    return bytes([
+        ord(default) if (b is None or b not in range(256)) else b
+        for b in maybe_bytes
+    ])
 
 
 def to_bytes_ensure_complete(maybe_bytes: list[int | None]) -> bytes:

tests/test_padding_oracle.py

@@ -27,6 +27,7 @@ def test_padding_oracle_encrypt():
         plaintext,
         cryptor.block_size,
         cryptor.oracle,
+        num_threads=4,
     )
 
     assert cryptor.decrypt(encrypted) == plaintext

tests/test_solve.py

@@ -21,7 +21,7 @@ def nop_callback(*args, **kwargs):
         ciphertext,
         crypter.block_size,
         crypter.oracle,
-        parallel=4,
+        num_threads=4,
         block_callback=nop_callback,
         progress_callback=nop_callback,
     )