import re
import sys

if len(sys.argv) < 4:
    print(f"Usage: ./{sys.argv[0]} <original minio binary> <new minio binary> <new public key>")
    exit(1)

original_minio_binary_path = sys.argv[1]
new_minio_binary_path = sys.argv[2]
new_public_key_path = sys.argv[3]

with open(original_minio_binary_path, "rb") as f:
    minio_binary = f.read()

with open(new_public_key_path, "r") as f:
    new_public_key = f.read().strip().encode("ascii")

# AIStor встраивает НЕСКОЛЬКО публичных ключей лицензии (secp384r1, SPKI/PEM):
# один — legacy-ключ subnet.min.io, второй используется для валидации лицензии.
# Заменять надо ВСЕ ключи совпадающей длины, иначе проверка подписи идёт по
# незаменённому ключу и падает с "could not verify message using any of the keys".
pem_re = re.compile(rb"-----BEGIN PUBLIC KEY-----\n.*?\n-----END PUBLIC KEY-----", re.DOTALL)

embedded_keys = pem_re.findall(minio_binary)
if not embedded_keys:
    raise ValueError("no embedded PUBLIC KEY blocks found in binary")

new_minio_binary = minio_binary
replaced = 0
for key in set(embedded_keys):
    # длина нового ключа обязана побайтово совпадать со встроенным, иначе поедут
    # смещения в бинарнике; ключи другой длины/типа пропускаем.
    if len(key) != len(new_public_key):
        continue
    new_minio_binary = new_minio_binary.replace(key, new_public_key)
    replaced += 1

if replaced == 0:
    raise ValueError(
        f"no public key of matching length ({len(new_public_key)}) found to replace"
    )
if new_minio_binary == minio_binary:
    raise ValueError("could not replace old public key(s)")

print(f"replaced {replaced} embedded public key(s)")

with open(new_minio_binary_path, "wb") as f:
    f.write(new_minio_binary)