#!/usr/bin/env python3

import os
from math import ceil
from base64 import b64decode
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from random import randbytes, randint
from string import printable
from typing import Callable, List, Optional

DBG: bool = False
KEY: bytes = os.urandom(16)
SECRET: bytes = b'whatever'
OFF: int = randint(1, 32)  # use this to generate some stuff before str ...
RND: bytes = randbytes(OFF)

def oracle(txt: bytes) -> bytes:
    cipher: AES = AES.new(KEY, AES.MODE_ECB)
    p: bytes = pad(RND + txt + SECRET, len(KEY))
    if DBG:
        chnks: Callable[[bytes], List[bytes]] = lambda s: \
               [ s[i:i+len(KEY)] for i in range(0, len(s), len(KEY)) ]
        print(f'[*] ORC: ')
        print('\n'.join(f'    {c}' for c in chnks(p)))
    return cipher.encrypt(p)
    
def recover(f: Callable[[bytes], bytes], k_sz: int,
            block: Optional[int] = None,
            align: Optional[int] = None,
            only_printable: bool = True) -> bytes:
    secret: bytes = b''
    chnks: Callable[[bytes], List[bytes]] = lambda s: \
           [ s[i:i+k_sz] for i in range(0, len(s), k_sz) ]
    mtchc: Callable[[bytes, bytes], int] = lambda a, b: \
           [ x==y for x, y in zip(chnks(a), chnks(b)) ].index(False)
    # first, determine the first block we can use for leaking ...
    if block is None:
        # eliminate possibility of random match ...
        ds: List[int] = []
        js: List[int] = []
        for x in (b'a', b'b',):
            a: bytes = f(x*(2*k_sz))
            c: int = mtchc(a, f(x*(2*k_sz-1)))
            d: int = 0
            j: int = 1
            while (d := mtchc(a, f(x*(2*k_sz-j)))) == c:
                j += 1
            ds.append(d)
            js.append(k_sz-j+1)
        block = max(ds)
        align = max(js)
    # now, start recovery ...
    l: int = len(f(b'x'*align+b'a'*k_sz)[(block+1)*k_sz:])
    _b: int = block
    for i in range(l):
        msg: bytes = b'x'*align + b'a'*(k_sz-i%k_sz-1)
        goal: bytes = f(msg)[k_sz*_b:k_sz*(_b+1)]
        msg = b'x'*align + (msg + secret)[-k_sz+1:]
        for x in (printable.encode() if only_printable else range(0x100)):
            cu: bytes = f(msg+bytes([x,]))[k_sz*block:k_sz*(block+1)]
            if cu == goal:
                secret += bytes([x,])
                break
        if i%k_sz == k_sz-1:
            _b += 1
    return secret      

def main():
    rec: bytes = recover(oracle, len(KEY))
    print(f'[*] Recovered secret: {rec} (len={len(rec)})')

if __name__ == '__main__':
    main()