'''
Solution for https://cryptohack.org/courses/symmetric/bean_counter/
'''

import os
import requests
from math import ceil

from Crypto.Cipher import AES

from celeste.math.util import xor_bytes

class StepUpCounter(object):
    def __init__(self, step_up=False):
        self.value = os.urandom(16).hex()
        self.step = 1
        self.stup = step_up

    def increment(self):
        if self.stup:
            self.newIV = hex(int(self.value, 16) + self.step)
        else:
            self.newIV = hex(int(self.value, 16) - self.stup)
        self.value = self.newIV[2:len(self.newIV)]
        return bytes.fromhex(self.value.zfill(32))

    def __repr__(self):
        self.increment()
        return self.value

'''
NOTE: Since step_up is used ONLY as false, we can simplify:  
class StepUpCounter(object):
    def __init__(self):
        self.value = os.urandom(16).hex()

    # SAME AS DOING NOTHING
    def increment(self):
        return self.value()

    def __repr__(self):
        return self.value
'''

URL = 'https://aes.cryptohack.org/bean_counter'

def get_encrypted() -> bytes:
    resp = requests.get(f'{URL}/encrypt/')
    encrypted = bytes.fromhex(resp.json()['encrypted'])
    return encrypted

PNG_HEADER = [
    '89', '50', '4e', '47', 
    '0d', '0a', '1a', '0a', 
    '00', '00', '00', '0d', 
    '49', '48', '44', '52'
]

def main() -> None:
    # NOTE: the counter is constant!
    ctr = StepUpCounter()
    # init = ctr.increment()
    # init_val = ctr.value
    # for i in range(2560000):
    #     if ctr.increment() != init:
    #         print('Counter Changed')
    #         print(i)
    #         break
    #     elif ctr.value != init_val:
    #         print('valued changed')
    #         print(i)
    #         break
            
    # PNGs *should* have a constant header, we will use the first 16 bytes
    # to determine what the counter value must be set to
    encrypted = get_encrypted()
    png_header = bytes.fromhex(''.join(PNG_HEADER))
    ctr_value = xor_bytes(encrypted[:16], png_header)
    # apply the counter value to the entire encrypted image (in blocks of 16 bytes)
    BLOCK_SIZE = 16
    with open('bean_flag.png', 'wb') as f:
        for i in range(ceil(len(encrypted) / BLOCK_SIZE)):
            block = encrypted[i*BLOCK_SIZE : (i+1)*BLOCK_SIZE]
            # NOTE: the block we get is not guaranteed to be block size (ie if at end)
            plaintext_block = xor_bytes(block, ctr_value[:len(block)])
            f.write(plaintext_block)

if __name__ == '__main__':
    try:
        main()
    except (KeyboardInterrupt, EOFError):
        print('\n[!] Interrupt')