From b9c5a5bf3e705e3806d5bf32da995d1261133174 Mon Sep 17 00:00:00 2001
From: Emile Clark-Boman <eclarkboman@gmail.com>
Date: Thu, 26 Jun 2025 01:26:27 +1000
Subject: [PATCH] bug fixes for paddingoracle attack

---
 crack.py                     | 12 ++++++++----
 imp/attacks/paddingoracle.py | 28 ++++++++++++++++++++--------
 test.py                      |  6 ++++--
 3 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/crack.py b/crack.py
index 5422cd4..c749cdf 100644
--- a/crack.py
+++ b/crack.py
@@ -1,7 +1,11 @@
 import requests
 
 from imp.constants import PRINTABLE
-from imp.attacks.paddingoracle import paddingoracle
+from imp.attacks import paddingoracle
+
+from Crypto.Util.Padding import pad
+
+import string
 
 NIBBLESET = [n.to_bytes() for n in range(256)]
 
@@ -9,8 +13,8 @@ HOST = 'https://aes.cryptohack.org'
 ENDPOINT = '/ecb_oracle/encrypt/{0}/'
 URI = HOST + ENDPOINT
 
-# CHARSET = [c.encode() for c in string.printable]
-CHARSET = NIBBLESET # OVERRIDE
+CHARSET = [c.encode() for c in string.printable]
+# CHARSET = NIBBLESET # OVERRIDE
 
 FLAG_LEN = 26  # flag length, calculated by hand
 SPACER = b'\x0f' # arbitrary spacing character
@@ -30,7 +34,7 @@ def encrypt(b: bytes) -> bytes:
     return bytes.fromhex(resp['ciphertext'])
 
 def main() -> None:
-    paddingoracle(encrypt, NIBBLESET, 64, batch_size=16, debug=True)
+    paddingoracle.crack(encrypt, pad, CHARSET, 64, batch_size=2, debug=True)
 
                
 
diff --git a/imp/attacks/paddingoracle.py b/imp/attacks/paddingoracle.py
index e419535..d9fd1a5 100644
--- a/imp/attacks/paddingoracle.py
+++ b/imp/attacks/paddingoracle.py
@@ -26,12 +26,18 @@ def crack_secret_len(cipher: Callable[[str], str],
         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)}')
@@ -51,10 +57,12 @@ def crack(cipher: Callable[[str], str],
         # 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
-        tails = [padfn(c+known[:tail_bytes], 16) for c in charset]
+        # 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 + tail_bytes
+        push_size = (default_push + full_tail_bytes) % block_size
 
         matched = False
         NUM_BATCHES = len(tails) // batch_size
@@ -66,24 +74,28 @@ def crack(cipher: Callable[[str], str],
 
             # encrypt batch and split the ciphertext into blocks
             ciphertext = cipher(batch)
-            index = i*block_size
             num_blocks = len(ciphertext)//block_size
-            blocks = [ciphertext[index:index+block_size] for i in range(num_blocks)]
+            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)
-            print(oracle_pos)
+            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
-                    matched = True
                     if debug:
                         print(f'[*] Found Tail: {known}')
-                        input()
+                    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
diff --git a/test.py b/test.py
index 74ba07b..6902349 100644
--- a/test.py
+++ b/test.py
@@ -11,13 +11,15 @@ KEY = b'you wont get me!'
 FLAG = b'imbaud{omg_you_catched_me}'
 CIPHER = AES.new(KEY, AES.MODE_ECB)
 
-def encrypt(b: bytes) -> bytes:
+def encrypt(b: bytes, debug=False) -> bytes:
     padded = pad(b + FLAG, 16)
+    if debug:
+        print(padded)
     # print(padded)
     return CIPHER.encrypt(padded)
 
 def main() -> None:
-    paddingoracle.crack(encrypt, pad, CHARSET, 16, batch_size=1, debug=True)
+    paddingoracle.crack(encrypt, pad, CHARSET, 16, batch_size=50, debug=True)
 
 if __name__ == '__main__':
     try: