#!/usr/bin/env python3
import sys

def compute_psa_checksum(data: bytes) -> bytes:
    var_e = 0x0000  # Correct seed from firmware logic
    for i in range(len(data)):
        var_e ^= (data[i] + i) & 0xFFFF
    return var_e.to_bytes(2, 'big')  # 2-byte checksum, big-endian

def add_checksum(hexstring: str) -> str:
    try:
        data = bytes.fromhex(hexstring.strip())
    except ValueError:
        return "Invalid hex input!"
    checksum = compute_psa_checksum(data)
    return (data + checksum).hex()

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: ./hex_checksum.py <hexstring_without_checksum>")
        sys.exit(1)
    result = add_checksum(sys.argv[1])
    print(result)