Upload AES-256 key to HSM and generate hashed & encrypted Tenant Secret

from Crypto.Cipher import AES
import socket
import base64
import os
import hashlib

HSM_IP = "<HSM_IP>"
HSM_PORT = 9000

def generate_key_parts_and_aes(part_count=2, components_per_part=4, bytes_per_component=8):
    """
    Generates key parts made of components, XORs all parts to form the final AES-256 key,
    produces Salesforce SHA-256 hash, and computes KCV.
    """

    all_parts = []
    xor_accumulator = bytearray(32)   # start with all zeros; final AES key = XOR of all part bytes

    print("\n=== Generating Key Parts ===")

    for part_idx in range(part_count):
        part_components = []
        part_bytes = b''

        print(f"\nKey Part {part_idx + 1}:")

        for comp_idx in range(components_per_part):
            comp = os.urandom(bytes_per_component)
            comp_hex = comp.hex()

            part_components.append(comp_hex)
            part_bytes += comp

            print(f" Component {comp_idx + 1}: {comp_hex}")

        # Expand 16 bytes → 32 bytes (Salesforce needs 256-bit final key)
        # We duplicate the 16 bytes to reach 32 bytes
        expanded_part = part_bytes + part_bytes

        # XOR this part into the overall AES key accumulator
        xor_accumulator = bytes(a ^ b for a, b in zip(xor_accumulator, expanded_part))

        all_parts.append({
            "part_number": part_idx + 1,
            "components": part_components
        })

    # Final AES-256 key
    aes_key = bytes(xor_accumulator)
    print("\nFinal AES Key (hex):", aes_key.hex().upper())

    # Salesforce SHA-256 Tenant Secret
    sha256_digest = hashlib.sha256(aes_key).digest()
    sha256_b64 = base64.b64encode(sha256_digest).decode()

    print("SHA-256 Digest (Base64 for Salesforce):", sha256_b64)

    with open("hashed_tenant_secret.b64", "w") as f:
        f.write(sha256_b64)
        print("Hash of AES-256 key is saved to hashed_tenant_secret.b64")

    # Compute KCV (first 2 bytes of AES-ECB encrypt of zero block)
    cipher = AES.new(aes_key, AES.MODE_ECB)
    encrypted = cipher.encrypt(bytes(16))
    kcv = encrypted[:2].hex().upper()

    print("Key Check Value (KCV):", kcv)

    return aes_key, sha256_b64, kcv, all_parts

# Obtain how many key parts the user would like to user to generate the AES-256
key_parts_count = int(input("Enter number of key parts: "))
aes_key, salesforce_hash, kcv, parts = generate_key_parts_and_aes(key_parts_count)

# Perform key wrapping of the AES256 key with Salesforce public key
## Grab input from user of which slot the public Salesforce key and AES-256 key is located
print("\nPlease copy the components to Futurex Excrypt Manager to generate am AES-256 key")
print("Also, make sure that the Salesforce public key is uploaded to HSM")
salesforce_pub_key_slot = input("\nEnter the slot number where the Salesforce public key is stored in the HSM: ")
aes256_slot = input("\nEnter the slot number where the AES256 key is stored in the HSM: ")

gpwb_command = f"[AOGPWB;RC{salesforce_pub_key_slot};BGBD{aes256_slot};AS2;FS6;CE4;OP1;]"
print("\n➡️ Sending (GPWB):", gpwb_command.strip())

with socket.create_connection((HSM_IP, HSM_PORT)) as s:
    s.sendall(gpwb_command.encode())
    response = s.recv(8192).decode()  # use large buffer for RSA wrapped data

print("⬅️ Raw Response (GPWB):", repr(response))

# Parse RJ (wrapped key) and AE (check value) from the response
wrapped_key_hex, check_val = None, None
if "BG" in response:
    parts = response.strip("[]").split(";")
    wrapped_key_hex = next((p[2:] for p in parts if p.startswith("BG")), None)
    check_val = next((p[2:] for p in parts if p.startswith("AE")), None)

if wrapped_key_hex:
    wrapped_b64 = base64.b64encode(bytes.fromhex(wrapped_key_hex)).decode()
    print(f"\n Encrypted Tenant Secret (Base64):\n{wrapped_b64}")
    with open("encrypted_tenant_secret.b64", "w") as f2:
        f2.write(wrapped_b64)
    print("Tenant secret encrypted saved to encrypted_secret_encrypted.b64")
    print(f"\n Length in bytes:", len(base64.b64decode(wrapped_b64)))

if check_val:
    print(f"Check Value (AE): {check_val}")