Add Hacky Encrypt mode

Added a hacky encryption mode. Duplicated the wrapper and callback functions so they can be modified for encryption mode if needed.

Author: q3st1on

src/padding_oracle/legacy.py

@@ -33,21 +33,21 @@
     'padding_oracle',
 ]
 
-
-def padding_oracle(ciphertext: Union[bytes, str],
+def padding_oracle(payload: Union[bytes, str],
                    block_size: int,
                    oracle: OracleFunc,
                    num_threads: int = 1,
                    log_level: int = logging.INFO,
                    null_byte: bytes = b' ',
                    return_raw: bool = False,
+                   mode: Union[bool, str] = 'encrypt',
                    ) -> Union[bytes, List[int]]:
     '''
     Run padding oracle attack to decrypt ciphertext given a function to check
     wether the ciphertext can be decrypted successfully.
 
     Args:
-        ciphertext     (bytes|str) the ciphertext you want to decrypt
+        payload        (bytes|str) the payload you want to encrypt/decrypt
         block_size     (int)       block size (the ciphertext length should be
                                    multiple of this)
         oracle         (function)  a function: oracle(ciphertext: bytes) -> bool
@@ -58,33 +58,48 @@ def padding_oracle(ciphertext: 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')
+
 
     Returns:
-        plaintext (bytes|List[int]) the decrypted plaintext
+        result (bytes|List[int]) the processed payload
     '''
 
     # Check args
     if not callable(oracle):
         raise TypeError('the oracle function should be callable')
-    if not isinstance(ciphertext, (bytes, str)):
-        raise TypeError('ciphertext should have type bytes')
+    if not isinstance(payload, (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(ciphertext) % block_size == 0:
-        raise ValueError('ciphertext length should be multiple of block size')
+    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 isinstance(mode, str) and mode not in ('encrypt', 'decrypt'):
+        raise ValueError('mode must be either encrypt or decrypt')
 
     logger = get_logger()
     logger.setLevel(log_level)
 
-    ciphertext = to_bytes(ciphertext)
+    payload = to_bytes(payload)
     null_byte = to_bytes(null_byte)
 
+
+    # encryption routine
+    if mode == 'encrypt' or mode:
+        return encrypt(payload, block_size, oracle, num_threads, null_byte, return_raw, logger)
+
+    # otherwise continue with decryption as normal
+    return decrypt(payload, block_size, oracle, num_threads, null_byte, return_raw, logger):
+
+def encrypt(payload, block_size, oracle, num_threads, null_byte, return_raw, logger):
     # Wrapper to handle exceptions from the oracle function
     def wrapped_oracle(ciphertext: bytes):
         try:
@@ -105,15 +120,51 @@ def plaintext_callback(plaintext: bytes):
         plaintext = convert_to_bytes(plaintext, null_byte)
         logger.info(f'plaintext: {plaintext}')
 
-    plaintext = solve(ciphertext, block_size, wrapped_oracle, num_threads,
+    plaintext = solve(payload, 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
 
+def decrypt(payload, block_size, oracle, num_threads, null_byte, return_raw, logger):
+    # Wrapper to handle exceptions from the oracle function
+    def wrapped_oracle(ciphertext: bytes):
+        try:
+            return oracle(ciphertext)
+        except Exception as e:
+            logger.error(f'error in oracle with {ciphertext!r}, {e}')
+            logger.debug('error details: {}'.format(traceback.format_exc()))
+        return False
+
+    def result_callback(result: ResultType):
+        if isinstance(result, Fail):
+            if result.is_critical:
+                logger.critical(result.message)
+            else:
+                logger.error(result.message)
+
+    def plaintext_callback(plaintext: bytes):
+        plaintext = convert_to_bytes(plaintext, null_byte)
+        logger.info(f'plaintext: {plaintext}')
+
+    def blocks(data: bytes):
+        return [data[i:(i+block_size)] for i 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)])
+
+    plaintext_blocks = blocks(payload)
+    ciphertext_blocks = [null_byte * block_size for i in range(len(plaintext_blocks)+1)]
+    
+    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[i] = bytes_xor(plaintext_blocks[index], plaintext)
+    
+    ciphertext = b''.join(ciphertext_blocks)
+    return ciphertext
 
 def get_logger():
     logger = logging.getLogger('padding_oracle')

src/padding_oracle/solve.py

@@ -38,7 +38,6 @@
     'remove_padding',
 ]
 
-
 class Pass(NamedTuple):
     block_index: int
     solved: List[int]