# Decrypt.py

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.primitives.ciphers import algorithms, modes, Cipher
from cryptography.hazmat.backends import default_backend
from os import path
from os import urandom
import base64


STUDENTNAME = "IvanusicIvan" # ne koriste se HR slova (čćžšđ)
SALT = "DAJSOLI" # pitajte profesora na vježbama

# Generate hash
def hash_me(msg, hash_function):
    digest = hashes.Hash(hash_function, backend=default_backend())
    if not isinstance(msg, bytes):
        msg = msg.encode()
    digest.update(msg)
    return digest.finalize()

def decrypt_CTR(key, iv, ciphertext, cipher=algorithms.AES):
    if not isinstance(key, bytes):
        key = key.encode()

    decryptor = Cipher(cipher(key), modes.CTR(iv), backend=default_backend()).decryptor()
    plaintext = decryptor.update(ciphertext)
    plaintext += decryptor.finalize()

    return plaintext

if __name__ =='__main__':
    #Calculating hash for filename
    filename = hash_me(str.encode(STUDENTNAME + SALT), hashes.SHA256()).hex()
    print("\nFile name is: " + filename + "\n")
    f_in = open(filename + ".enc", 'rb' )

    #Reading file
    challenge = f_in.read()
    print(challenge)
    challenge_decoded = base64.b64decode(challenge)

    #Splitting decoded text
    iv = challenge_decoded[:16]
    ciphertext = challenge_decoded[16:]
    for x in range(4096):
        x_str = str(x)
        x_encoded = x_str.encode()
        KDF = Scrypt(
            length=32,
            salt=b'salt',
            n=2**15,
            r=8,
            p=1,
            backend=default_backend()
        )
        key = KDF.derive(x_encoded)
        print(x)

        #Ciphertext decryption
        plaintext=decrypt_CTR(key, iv, ciphertext)
        if "Chuck".encode() in plaintext:
            print("Plaintext:\n")
            print(plaintext)
            break