import subprocess
import sys
from typing import Union
import binascii
import base64

# Check if the pycryptodome library is installed and install it if not - not best practice but works for this demo
try:
    from Crypto.Cipher import AES
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "pycryptodome"])
    from Crypto.Cipher import AES

# Little helper function to pad the text with zeros to make it a multiple of the block size
def zero_pad(text, block_size):
    num_zeros = (-len(text)) % block_size
    return text + b"\0" * num_zeros

# Main function to encrypt/decrypt a string using AES ECB mode compatible with MQL4/5
def CryptMQL(
    text: Union[str, bytes],
    key: Union[str, bytes],
    encrypt=True,
    keyIsHex=False,
    encryptedIsHex=False,
):
    # convert the key to bytes
    key = binascii.unhexlify(key)[:32] if keyIsHex else key.encode()[:32]

    # create a new AES cipher object with the key
    cipher = AES.new(key, AES.MODE_ECB)

    if encrypt:
        # convert the string to bytes, pad it with zeros and then encrypt
        encrypted_string = cipher.encrypt(zero_pad(text.encode(), AES.block_size))

        # return the encrypted string in desired format
        return (
            binascii.hexlify(encrypted_string).upper().decode()
            if encryptedIsHex
            else base64.b64encode(encrypted_string).decode()
        )
    else:
        # convert the text to bytes
        text = binascii.unhexlify(text) if encryptedIsHex else base64.b64decode(text)

        # decrypt the string and return
        return cipher.decrypt(text).decode()

# 2 wrapper functions to encrypt/decrypt compatible with Bcrypt library https://www.mql5.com/en/code/16378
def EncryptMQL(plaintext, keystring):
    return CryptMQL(text=plaintext,key=keystring,encrypt=True,keyIsHex=False,encryptedIsHex=True)

def DecryptMQL(ciphertext, keystring):
    return CryptMQL(text=ciphertext,key=keystring,encrypt=False,keyIsHex=False,encryptedIsHex=True)

# Demo

keystring = "thisisyoursecretkeydonttellnoone"
plaintext = "Mother Tell Your Children Not To Walk My Way"
keyhex = binascii.hexlify(keystring.encode()).upper().decode()

print("\nplaintext:  ", plaintext)
print("keystring:  ", keystring)
print("keyhex:     ", keyhex, "\n")

print("1. Encrypt string and output the result as HEX")
encrypted = EncryptMQL(plaintext, keystring)
print(f"encrypted: {encrypted}\n")

print("2. Decrypt HEX input")
decrypted = DecryptMQL(encrypted, keystring)
print(f"decrypted: {decrypted}\n")

print("3. Encrypt string and output the result as Base64")
encrypted = CryptMQL(plaintext, keystring, encrypt=True, encryptedIsHex=False)
print(f"encrypted: {encrypted}\n")

print("4. Decrypt Base64 input")
decrypted = CryptMQL(encrypted, keystring, encrypt=False, encryptedIsHex=False)
print(f"decrypted: {decrypted}\n")

print("5. Encrypt using HEX key")
encrypted = CryptMQL(plaintext, keyhex, encrypt=True, keyIsHex=True)
print(f"encrypted: {encrypted}\n")

print("6. Decrypt using HEX key")
decrypted = CryptMQL(encrypted, keyhex, encrypt=False, keyIsHex=True)
print(f"decrypted: {decrypted}\n")