#!/usr/bin/env python3
import sys

"""
PSA Checksum Calculator

This script calculates the 2-byte checksum used in PSA multicast protocol messages.
The checksum algorithm is based on reverse-engineered firmware behavior.

How it works:
-------------
- Initialize a 16-bit value `var_e` to 0x0000
- For each byte in the payload:
    - Add the byte's index to its value (i.e., byte + index)
    - XOR the result into `var_e`
- The result is a 2-byte checksum (big endian: high byte first, low byte second)

This checksum is appended to the payload to form a valid message.
"""

def compute_psa_checksum(data: bytes) -> bytes:
    var_e = 0x0000
    for i in range(len(data)):
        var_e ^= (data[i] + i) & 0xFFFF
    return var_e.to_bytes(2, 'big')

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)