Project is now named Celeste

celeste/attacks/paddingoracle.py

@@ -0,0 +1,106 @@
+from math import inf, ceil
+from typing import Callable
+
+from celeste.math.util import clamp_max
+
+SPACER = b'\xff' # arbitrary spacing character
+
+class DecryptFailed(Exception):
+    pass
+    
+def round_to_blocks(n: int, block_size: int) -> int:
+    return ceil(n / block_size)
+
+def crack_secret_len(cipher: Callable[[str], str], 
+                     max_iters: int = inf) -> int | None:
+    # calculate length for i = 1
+    init_len = len(cipher(SPACER))
+    secret_len = None
+    i = 2
+    while True:
+        if i - 2 > max_iters:
+            break
+        elif len(cipher(SPACER*i)) != init_len:
+            secret_len = init_len - i
+            break
+        i += 1
+    return secret_len
+
+'''
+
+NOTE: pad_if_perfect exists for PKCS#7 which will add a full block
+NOTE: of padding if the input is perfectly alligned to the blocks already. 
+'''
+def crack(cipher: Callable[[str], str],
+          padfn: Callable[[bytes, int], bytes],
+          charset: list,
+          block_size: int,
+          max_secret_iters: int = inf,
+          batch_size: int = 1,
+          pad_if_perfect: bool = True,
+          debug: bool = False) -> str | None:
+    if len(charset) % batch_size:
+        raise ValueError(f'batch_size={batch_size} does not divide len(charset)={len(charset)}')
+    # calculate the secret length
+    secret_len = crack_secret_len(cipher, max_iters=max_secret_iters)
+    if debug:
+        print(f'[+] Found secret length: {secret_len}')
+
+    # calculate how many blocks the secret stretches over
+    # NOTE: secret_block_len - secret_len represents the number of
+    # NOTE: bytes required to make the secret fill the blocks with no padding
+    secret_block_len = round_to_blocks(secret_len, block_size) * block_size
+    default_push = (secret_block_len - secret_len)
+    known = b''
+    while True:
+        # the "full tail" is all characters we know + the 1 we're cracking (target)
+        # the "current tail" is the characters in the same block as the target
+        full_tail_bytes = len(known) + 1
+        tail_bytes = clamp_max(full_tail_bytes, block_size - 1)
+        # generate ALL possible tails (avoid padding if no padding required)
+        tails = [c + known[:tail_bytes] for c in charset]
+        if len(tails[0]) != block_size:
+            tails = [padfn(tail, 16) for tail in tails]
+        # calculate the "push" applied to the secret
+        push_size = (default_push + full_tail_bytes) % block_size
+
+        matched = False
+        NUM_BATCHES = len(tails) // batch_size
+        for i in range(NUM_BATCHES):
+            if debug:
+                print(f'{int(i/NUM_BATCHES*100)}%', end='\r')
+            batch = tails[i*batch_size : (i+1)*batch_size]
+            batch = b''.join(batch) + (SPACER * push_size) # apply spacing
+
+            # encrypt batch and split the ciphertext into blocks
+            ciphertext = cipher(batch)
+            num_blocks = len(ciphertext)//block_size
+            blocks = [ciphertext[i*block_size : (i+1)*block_size] for i in range(num_blocks)]
+
+            oracle_pos = round_to_blocks(full_tail_bytes, block_size)
+            if pad_if_perfect and (push_size + secret_len) % block_size == 0:
+                oracle_pos += 1
+            
+            for j, cipher_block in enumerate(blocks[:batch_size]):
+                if cipher_block == blocks[-oracle_pos]:
+                    char = charset[i*batch_size + j]
+                    known = char + known
+                    if debug:
+                        print(f'[*] Found Tail: {known}')
+                    matched = True
+                    break
+            if matched:
+                break
+        if not matched:
+            break
+        elif len(known) == secret_len:
+            if debug:
+                print('[+] SUCCESS')
+            return known
+    # if we reached the end (no return)
+    # then the attack failed
+    err_msg = 'Padding oracle attack failed'
+    if not debug:
+        raise DecryptFailed(err_msg)
+    print(f'\n[!] {err_msg}')
+