#!/usr/bin/env python3

"""
Recover plaintext from encrypted AES/ECB ciphertext using an oracle.

Scenario:
You can send whatever you want to the server and you get it back
encrypted with an attached secret. (The secret was attached before the
encryption process). Since ECB is *stateless*, meaning each plaintext
block that is the same will result in the same ciphertext block (as long
as the same key is used), we can recover the plaintext by sending a
block-sized - 1 string of known gibberish and bruteforcing the last
character, until it is what we initially received from the server.

E.g.:
    oracle(x) -> cbc_enc(x + 'mysecret')
    oracle('aaaaaaa') -> cbc_enc('aaaaaaa' + 'mysecret')                    | = b00143de4d11c36f8c632341895783ff
    oracle('aaaaaaaaaaaaaaaa') -> cbc_enc('aaaaaaaaaaaaaaaa' + 'mysecret')  | = da030c17f299c08f5f63527b29ed520f 5f38d2da07d2354b5e215c91857d9303
    oracle('aaaaaaaaaaaaaaa') -> cbc_enc('aaaaaaaaaaaaaaa' + 'mysecret')    | = bcd0a89e64f9a900670c75c6e7730d17 9b252036f625dcada2d7082df6be3e41

    oracle('aaaaaaaaaaaaaaa0') -> cbc_enc('aaaaaaaaaaaaaaa0' + 'mysecret')  | = 727bcc09edf00a740885b39ef7e36f35 5f38d2da07d2354b5e215c91857d9303
    oracle('aaaaaaaaaaaaaaa1') -> cbc_enc('aaaaaaaaaaaaaaa1' + 'mysecret')  | = 25fdb4acee3191a64ed4818c02f5815a 5f38d2da07d2354b5e215c91857d9303
    ...
    oracle('aaaaaaaaaaaaaaam') -> cbc_enc('aaaaaaaaaaaaaaam' + 'mysecret')  | = bcd0a89e64f9a900670c75c6e7730d17 5f38d2da07d2354b5e215c91857d9303

    --> as you can see, the first block of 'aaaaaaaaaaaaaaam' and 'aaaaaaaaaaaaaaa' are
        are the same, suggesting that the secret coming after our plaintext starts with an 'm'

E.g.:
    --> to use the example from above with the code below, just do the following:
    
             ---- oracle function
            |        ------ "target" block; block which secret will shift into
            |       |   ------ block size
            v       v  v
    recover(oracle, 0, 16)
    recover(oracle, 0, 16, align=0, only_printable=True, anim=True)
                                 ^                 ^          ^
                                 |                 |           ---- show fancy recovery status while recovering?
                                 |                  ----- try to restore only printable characters?
                                  ---- in case your plaintext isn't the first thing in the ciphertext
"""

import os
from math import ceil
from base64 import b64decode
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from string import printable
from typing import *

KEY: bytes = os.urandom(16)
SECRET: bytes = b'mysecret'
OFF: int = 0

def oracle(txt: bytes) -> bytes:
    cipher: AES = AES.new(KEY, AES.MODE_ECB)
    return cipher.encrypt(pad(b' '*OFF + txt + SECRET, len(KEY)))

def recover(f: Callable[[bytes], bytes], b: int, k: int, align: int = 0, only_printable: bool = True, anim: bool = False) -> bytes:
    secret: bytes = b''
    l: int = len(f(b'x'*align+b'a'*k)[(b+1)*k:])
    _b: int = b
    for i in range(l):
        msg: bytes = b'x'*align + b'a'*(k-i%k-1)
        goal: bytes = f(msg)[k*_b:k*(_b+1)]
        msg = b'x'*align + (msg + secret)[-k+1:]
        for x in (printable.encode() if only_printable else range(0x100)):
            if anim:
                print(f'[*] {secret+bytes([x,])} ... ', end='\r')
            cu: bytes = f(msg+bytes([x,]))[k*b:k*(b+1)]
            if cu == goal:
                secret += bytes([x,])
                break
        if i%k == k-1:
            _b += 1
    return secret

def main():
    rec: bytes = recover(oracle, max(ceil(OFF/len(KEY)), 1), len(KEY), align=len(KEY)-OFF)      # align to fix cross-block known plaintext ... 
    print(f'[*] Recovered secret: {rec} (len={len(rec)})')

if __name__ == '__main__':
    main()