docstrings & param checks & oracle exception handle

Author: djosix

padding_oracle/padding_oracle.py

@@ -21,6 +21,7 @@
 '''
 
 import logging
+import traceback
 from typing import Union, Callable
 from concurrent.futures import ThreadPoolExecutor
 
@@ -33,37 +34,80 @@
 
 
 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,
                    block_size: int,
                    oracle: Callable[[bytes], bool],
                    num_threads: int = 1,
                    log_level: int = logging.INFO,
                    null: bytes = b' ') -> bytes:
+    '''
+    Run padding oracle attack to decrypt cipher given a function to check wether the cipher
+    can be decrypted successfully.
+
+    Args:
+        cipher      (bytes)     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        (bytes)     the null byte if the (default: b' ')
+
+    Returns:
+        plaintext   (bytes)     the decrypted plaintext
+    '''
+
     # Check the oracle function
     assert callable(oracle), 'the oracle function should be callable'
     assert oracle.__code__.co_argcount == 1, 'expect oracle function with only 1 argument'
     assert len(cipher) % block_size == 0, 'cipher length should be multiple of block size'
+    assert isinstance(null, bytes), 'expect null with type bytes'
+    assert len(null) == 1, 'null byte should have length of 1'
 
-    logger = logging.getLogger('padding_oracle')
+    logger = _get_logger()
     logger.setLevel(log_level)
-    formatter = logging.Formatter('[%(asctime)s][%(levelname)s] %(message)s')
-    handler = logging.StreamHandler()
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
+
+    # Wrapper to handle exception from the oracle function
+    def _oracle_wrapper(i: int, j: int, cipher: bytes):
+        try:
+            return oracle(cipher)
+        except Exception as e:
+            logger.error('unhandled error at block[{}][{}]: ', i, j, e)
+            logger.debug('error details at block[{}][{}]: ', i, j, traceback.format_exc())
+        return False
 
     # The plaintext bytes list to save 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)
 
+    # 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)
@@ -80,7 +124,7 @@ def _block_decrypt_task(i, prev: bytes, block: bytes):
                 test_list = guess_list.copy()
                 test_list[-j] = k
                 oracle_futures[k] = oracle_executor.submit(
-                    oracle, bytes(test_list) + block)
+                    _oracle_wrapper, i, j, bytes(test_list) + block)
 
             for k, future in oracle_futures.items():
                 if future.result():
@@ -89,6 +133,7 @@ def _block_decrypt_task(i, prev: bytes, block: bytes):
             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)))

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='padding_oracle',
-    version='0.1.3',
+    version='0.1.4',
     author='Yuankui Lee',
     author_email='[email protected]',
     description='Threaded padding oracle automation.',