# verify_sig4096.py
import hashlib
from base4096 import decode, load_frozen_alphabet
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, ec
from cryptography.exceptions import InvalidSignature

# -------------------------
# Config
# -------------------------
FROZEN_ALPHABET_FILE = "frozen_base4096_alphabet.txt"
SIG4096_FILE = "base4096.sig4096"   # Your real signature file
MESSAGE_FILE = "message.txt"        # Signed message
PUBLIC_KEY_FILE = "public_key.pem"  # PEM public key

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

# -------------------------
# Parse .sig4096 file
# -------------------------
def extract_alphabet_fingerprint(sig_file):
    fingerprint_lines = []
    recording = False
    with open(sig_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line.startswith("Alphabet-Fingerprint:"):
                recording = True
                line = line[len("Alphabet-Fingerprint:"):].strip()
            elif line.startswith("---END BASE4096 SIGNATURE---"):
                recording = False
            if recording and line:
                fingerprint_lines.append(line)
    return "".join(fingerprint_lines)

alphabet_fingerprint = extract_alphabet_fingerprint(SIG4096_FILE)
print(f"Decoded Base4096 fingerprint length: {len(alphabet_fingerprint)} chars")

# -------------------------
# Decode signature
# -------------------------
signature_bytes = decode(alphabet_fingerprint)
print("Signature length (bytes):", len(signature_bytes))

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

# Hash for info (optional)
digest = hashlib.sha256(message_bytes).digest()
print("SHA-256 digest of message:", digest.hex())

# -------------------------
# Load public key
# -------------------------
with open(PUBLIC_KEY_FILE, "rb") as f:
    pubkey = serialization.load_pem_public_key(f.read())

# -------------------------
# Verify signature
# -------------------------
try:
    if isinstance(pubkey, ec.EllipticCurvePublicKey):
        pubkey.verify(signature_bytes, message_bytes, ec.ECDSA(hashes.SHA256()))
    else:
        # Assume RSA
        pubkey.verify(
            signature_bytes,
            message_bytes,
            padding.PKCS1v15(),
            hashes.SHA256()
        )
    print("✅ Signature is valid!")
except InvalidSignature:
    print("❌ Signature is INVALID!")
