update_key.py

import base64
import logging
import random
from Crypto.PublicKey import RSA
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256

from utils import secret_encrypt, secret_decrypt

logger = logging.getLogger(__name__)

RSA_SIGN_ENCODING = 'gbk'


def generate_update_key(unique_addr: str, key_path: str = "key.pem", remote_update_key: str = "linkv3"):
    salt = base64.b64encode(
        random.randint(
            1000000000, 2147483647
        ).to_bytes(
            32, byteorder='big'
        )
    ).decode()  # 随机的一串 32 位 Base64
    logger.info(f"random salt generated: {salt}")

    array_0 = secret_encrypt(remote_update_key, salt)
    sign_source_str = unique_addr + remote_update_key + salt

    try:
        with open(key_path, "rb") as f:
            key_data = f.read()
            private_key = RSA.import_key(key_data)
    except Exception as e:
        logger.exception("failed to load RSA private key")
        raise e
    h = SHA256.new(sign_source_str.encode(RSA_SIGN_ENCODING))
    signature = pkcs1_15.new(private_key).sign(h)

    array_2 = base64.b64encode(signature).decode('ascii')
    raw_payload = f"{array_0}\\{salt}\\{array_2}"
    final_key = ""
    mapping = {
        '=': '1',
        '1': '-',
        '\\': '0',
        '0': '*'
    }
    for char in raw_payload:
        final_key += mapping.get(char, char)

    return final_key


def analyze_update_key(key: str, unique_addr: str, key_path: str = "key.pem"):
    internal_str = key.replace("1", "=") \
        .replace("-", "1") \
        .replace("0", "\\") \
        .replace("*", "0")
    # 分割数组
    array = internal_str.split("\\")

    array_0 = array[0]  # 密文 (Encrypted Payload)
    array_1 = array[1]  # Salt / Key Raw
    array_2 = array[2]  # RSA 签名 (Signature)
    remote_update_key = secret_decrypt(array_0, array_1)

    logger.info(f"Array analyze result:")
    logger.info(f"\tArray(0) [Remote update key]: {remote_update_key}")
    logger.info(f"\tArray(1) [Salt]: {array_1}")
    logger.info(f"\tArray(2) [RSA Signature]: {array_2[:30]}...\n")

    # 构造待验证的字符串
    # 逻辑: ModBase._ReponseRole + UpdateKey + array(1)
    verify_str = unique_addr + remote_update_key + array_1
    logger.info(f"Verifying RSA Signature:\n\n{verify_str}\n\n")

    try:
        with open(key_path, "rb") as f:
            key_data = f.read()
            key = RSA.import_key(key_data)

        h = SHA256.new(verify_str.encode(RSA_SIGN_ENCODING))

        signature = base64.b64decode(array_2)

        # 验证
        pkcs1_15.new(key).verify(h, signature)
        logger.info("Success. This is a valid signature.")

    except (ValueError, TypeError):
        logger.error("Invalid Signature.")


if __name__ == "__main__":
    from get_unique_address import secret_get_unique_address
    print(generate_update_key(secret_get_unique_address()))