#!/usr/bin/env python

import argparse
import serial
from pathlib import Path
import json
import hmac
import hashlib

def unlock(secrets_file, port):

    with open(secrets_file, "r") as fh:
        secrets = json.load(fh)

    ser = serial.Serial(
        port=port,
        baudrate=115200,
        parity=serial.PARITY_NONE,
        stopbits=serial.STOPBITS_ONE,
        bytesize=serial.EIGHTBITS,
    )

    ser.write(b"GC\r\n")

    # Get challenge
    ser.readline()
    challenge_bytes = ser.readline().decode("UTF-8")
    challenge_bytes = bytes.fromhex(challenge_bytes.strip())
    print(f"Challenge bytes: {challenge_bytes}")

    challenge_key = bytes.fromhex(secrets["challenge_key"])
    print(f"Challenge key: {challenge_key}")

    response = f"SR {hmac.new(challenge_key, msg=challenge_bytes, digestmod=hashlib.sha256).hexdigest().upper()}\r\n"
    print(f"Response: {response}")

    ser.write(response.encode())
    ser.write(b"UNLOCK\r\n")

    ser.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--secrets-file", required=True, type=Path, help="Path to the bootloader secrets-file"
    )
    parser.add_argument("--port", required=True, help="Serial port")

    args = parser.parse_args()
    unlock(args.secrets_file, args.port)

if __name__ == "__main__":
    main()