import sys
from struct import pack, unpack

if len(sys.argv) < 3:
    print ('encrypt.py <src file> <dst file>')
    print ('Example:')
    print ('python3 ./encrypt.py ./Firmware.bin ./EncryptedFirmware.bin')
    exit()

srcFile = sys.argv[1]
dstFile = sys.argv[2]

XTEA_KEY = [0xA56BABCD, 0x0000FFFF, 0xABCDEF12, 0x13572468]

def xtea_encrypt(block, key):
    v0, v1 = unpack("<2I", block)   # LITTLE-ENDIAN
    k = key                         # tuple из 4 uint32
    delta = 0x9E3779B9
    sum = 0

    for _ in range(32):
        v0 = (v0 + (
            (((v1 << 4) & 0xFFFFFFFF) ^ (v1 >> 5)) + v1
        ^ (sum + k[sum & 3])
        )) & 0xFFFFFFFF

        sum = (sum + delta) & 0xFFFFFFFF

        v1 = (v1 + (
            (((v0 << 4) & 0xFFFFFFFF) ^ (v0 >> 5)) + v0
        ^ (sum + k[(sum >> 11) & 3])
        )) & 0xFFFFFFFF

    return pack("<2I", v0, v1)


def xtea_decrypt(block, key):
    v0, v1 = unpack("<2I", block)   # LITTLE-ENDIAN как в STM32
    delta = 0x9E3779B9
    sum = (delta * 32) & 0xFFFFFFFF

    for _ in range(32):
        v1 = (v1 - (
            (((v0 << 4) & 0xFFFFFFFF) ^ (v0 >> 5)) + v0
        ^ (sum + key[(sum >> 11) & 3])
        )) & 0xFFFFFFFF

        sum = (sum - delta) & 0xFFFFFFFF

        v0 = (v0 - (
            (((v1 << 4) & 0xFFFFFFFF) ^ (v1 >> 5)) + v1
        ^ (sum + key[sum & 3])
        )) & 0xFFFFFFFF

    return pack("<2I", v0, v1)


dst = open(dstFile, "wb") 

with open(srcFile, "rb") as f:
    while True:
        chunk = f.read(16)
        if not chunk:
            break

        if len(chunk) < 16:
            chunk = chunk.ljust(16, b'\xFF')

        enc = b""
        for i in range(0, len(chunk), 8):
            enc += xtea_encrypt(chunk[i:i+8], XTEA_KEY)
        dst.write(enc)