#!/usr/bin/python3

# Run me like this:
# $ python3 bleichenbacher.py "coach+username+100.00"
# or select "Bleichenbacher" from the VS Code debugger

from roots import *

import hashlib
import sys

def main():
    if len(sys.argv) < 2:
        print(f"usage: {sys.argv[0]} MESSAGE", file=sys.stderr)
        sys.exit(-1)
    message = sys.argv[1]

    # Create a bytearray for working with byte values
    forgedSignature = bytearray()
    
    # Add the starting bytes
    forgedSignature.append(0x00)
    forgedSignature.append(0x01)
    
    # Add an 0xFF byte, since only one is needed to fool the broken parsing 
    # method
    forgedSignature.append(0xFF)
    
    # Add an 0x00 byte, which is required between the 0xFF byte(s) and the 
    # ASN.1 bytes for the hash algorithm
    forgedSignature.append(0x00)
    
    # Add the ASN.1 bytes for SHA-256
    forgedSignature.extend(bytes().fromhex("3031300D060960864801650304020105000420"))
    
    # Get the SHA-256 hash of the provided message
    messageHash = hashlib.sha256()
    messageHash.update(message.encode())
    
    # Add the message hash to the forged signature
    forgedSignature.extend(messageHash.digest())
    
    # Add 201 arbitrary bytes (these end up being 0x00, but that's okay)
    forgedSignature.extend(bytes(201))
    
    # Take the cube root of the forged signature
    cubeRoot = integer_nthroot(bytes_to_integer(forgedSignature), 3)
    
    # Since the integer_nthroot() function may not be able to calculate an 
    # exact cube root, rounding may need to be performed (adding 1 should 
    # do the trick)
    if cubeRoot[1]:
        
        forgedSignature = cubeRoot[0]
        
    else:
        
        forgedSignature = cubeRoot[0] + 1
        
    # Print the forged signature in base64
    print(bytes_to_base64(integer_to_bytes(forgedSignature, 256)))

if __name__ == '__main__':
    main()