# zchg_sign_message.py
import hashlib
from base4096 import encode, load_frozen_alphabet
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

# -------------------------
# Config
# -------------------------
FROZEN_ALPHABET_FILE = "frozen_base4096_alphabet.txt"
PRIVATE_KEY_FILE = "zchg_private_key.pem"
MESSAGE_FILE = "message.txt"
SIG4096_FILE = "message.sig4096"
SIGNATURE_VERSION = 1
SIGNATURE_HASH = "SHA-256"
SIGNATURE_DOMAIN = "ZCHG-Base4096-Fingerprint"

# -------------------------
# Load frozen Base4096 alphabet
# -------------------------
alphabet = load_frozen_alphabet(FROZEN_ALPHABET_FILE)
CHAR_TO_INDEX = {ch: idx for idx, ch in enumerate(alphabet)}

# -------------------------
# Load message
# -------------------------
with open(MESSAGE_FILE, "rb") as f:
    message_bytes = f.read()

# SHA-256 hash of the message
message_digest = hashlib.sha256(message_bytes).digest()
print("SHA-256 digest of message:", message_digest.hex())

# -------------------------
# Load private key
# -------------------------
with open(PRIVATE_KEY_FILE, "rb") as f:
    private_key = serialization.load_pem_private_key(f.read(), password=None)

# -------------------------
# Sign message (RSA PKCS#1 v1.5 + SHA-256)
# -------------------------
signature_bytes = private_key.sign(
    message_bytes,
    padding.PKCS1v15(),
    hashes.SHA256()
)

# Encode signature in Base4096
signature_b4096 = encode(signature_bytes)

# -------------------------
# Export .sig4096 file
# -------------------------
with open(SIG4096_FILE, "w", encoding="utf-8") as f:
    f.write("---BEGIN BASE4096 SIGNATURE---\n")
    f.write(f"Version: {SIGNATURE_VERSION}\n")
    f.write(f"Hash: {SIGNATURE_HASH}\n")
    f.write(f"Domain: {SIGNATURE_DOMAIN}\n")
    f.write(f"Length: {len(signature_bytes)}\n")
    f.write("Alphabet-Fingerprint:\n")
    
    # Split into 64-character lines for readability
    for i in range(0, len(signature_b4096), 64):
        f.write(signature_b4096[i:i+64] + "\n")
    
    f.write("---END BASE4096 SIGNATURE---\n")

print(f"✅ Signature generated and saved to {SIG4096_FILE}")
