wallet.py

"""
wallet.py — Cryptographic wallet for the blockchain.

A wallet:
  1. Generates an RSA key-pair (private + public key).
  2. Derives a wallet address from the public key (SHA-256 → first 40 hex chars).
  3. Signs transactions with the private key so they cannot be forged.
  4. Verifies signatures using the corresponding public key.

Why RSA signatures?
  - Only the true owner of the private key can sign a transaction.
  - Anyone can verify the signature using the public key.
  - If you don't hold the private key, you cannot spend from that address.
"""

import hashlib
import json
import os
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives           import hashes, serialization
from cryptography.hazmat.backends             import default_backend
from cryptography.exceptions                  import InvalidSignature


class Wallet:
    """
    Represents a user's cryptographic identity on the blockchain.

    Attributes:
        private_key : RSA private key object (keep this SECRET).
        public_key  : RSA public key object  (share freely).
        address     : 40-character hex address derived from the public key.
    """

    def __init__(self):
        """
        Generate a fresh 2048-bit RSA key-pair and derive the wallet address.

        Key size 2048 bits is the current industry standard — long enough to
        be secure, short enough to be practical for demonstration purposes.
        """
        # ── Generate RSA private key ──────────────────────────────────────────
        # public_exponent=65537 is the conventional, mathematically safe choice.
        self.private_key = rsa.generate_private_key(
            public_exponent = 65537,
            key_size        = 2048,
            backend         = default_backend(),
        )

        # ── Derive the matching public key ────────────────────────────────────
        self.public_key = self.private_key.public_key()

        # ── Derive the wallet address from the public key ─────────────────────
        # (Mirrors Bitcoin's address derivation, simplified for clarity)
        self.address = self._derive_address()

    # ──────────────────────────────────────────────────────────────────────────
    # ADDRESS DERIVATION
    # ──────────────────────────────────────────────────────────────────────────

    def _derive_address(self) -> str:
        """
        Derive a short wallet address from the public key.

        Steps:
          1. Serialise the public key to DER bytes (a compact binary format).
          2. Hash those bytes with SHA-256.
          3. Take the first 40 hex characters as the address.

        The result looks like: "a3f9b2c1d0e4f5a6b7c8d9e0f1a2b3c4d5e6f7a8"

        Returns:
            A 40-character lowercase hex string.
        """
        # Serialise public key to DER (Distinguished Encoding Rules) bytes
        pub_bytes = self.public_key.public_bytes(
            encoding = serialization.Encoding.DER,
            format   = serialization.PublicFormat.SubjectPublicKeyInfo,
        )

        # Hash and truncate to 40 hex chars (160 bits — same as Bitcoin's RIPEMD)
        return hashlib.sha256(pub_bytes).hexdigest()[:40]

    # ──────────────────────────────────────────────────────────────────────────
    # SIGNING & VERIFICATION
    # ──────────────────────────────────────────────────────────────────────────

    def sign_transaction(self, transaction: dict) -> bytes:
        """
        Cryptographically sign a transaction with this wallet's private key.

        The signature proves that:
          (a) the owner of this address authorised the transaction, and
          (b) the transaction data has not been altered since signing.

        Args:
            transaction: A dict with keys sender, recipient, amount, fee, note.

        Returns:
            A raw bytes signature (typically ~256 bytes for RSA-2048).
        """
        # Serialise the transaction to a canonical JSON string, then encode to bytes
        tx_bytes = json.dumps(transaction, sort_keys=True).encode("utf-8")

        # Sign using RSA-PSS with SHA-256 (PSS = Probabilistic Signature Scheme,
        # more secure than the older PKCS#1 v1.5 padding)
        signature = self.private_key.sign(
            tx_bytes,
            padding.PSS(
                mgf        = padding.MGF1(hashes.SHA256()),
                salt_length= padding.PSS.MAX_LENGTH,
            ),
            hashes.SHA256(),
        )
        return signature

    @staticmethod
    def verify_signature(public_key_pem: str, transaction: dict, signature: bytes) -> bool:
        """
        Verify a transaction's signature using the sender's public key.

        Returns True only if:
          - The signature was created by the private key matching public_key_pem.
          - The transaction data is exactly as it was when signed.

        Args:
            public_key_pem: PEM-encoded public key string.
            transaction   : The transaction dict (must match what was signed).
            signature     : The raw bytes signature to verify.

        Returns:
            True if the signature is valid; False otherwise.
        """
        # Load the public key from its PEM representation
        public_key = serialization.load_pem_public_key(
            public_key_pem.encode("utf-8"),
            backend = default_backend(),
        )

        # Re-serialise the transaction the same way as during signing
        tx_bytes = json.dumps(transaction, sort_keys=True).encode("utf-8")

        try:
            # verify() raises InvalidSignature if the check fails
            public_key.verify(
                signature,
                tx_bytes,
                padding.PSS(
                    mgf        = padding.MGF1(hashes.SHA256()),
                    salt_length= padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
            return True          # No exception → signature is valid
        except InvalidSignature:
            return False         # Signature doesn't match

    # ──────────────────────────────────────────────────────────────────────────
    # KEY EXPORT / IMPORT (PEM format)
    # ──────────────────────────────────────────────────────────────────────────

    def get_public_key_pem(self) -> str:
        """
        Export the public key as a PEM-encoded string.
        PEM is the standard ASCII armor format (base64-wrapped DER with headers).
        Safe to share publicly.
        """
        return self.public_key.public_bytes(
            encoding = serialization.Encoding.PEM,
            format   = serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode("utf-8")

    def get_private_key_pem(self, password: bytes | None = None) -> str:
        """
        Export the private key as a PEM-encoded string.

        Args:
            password: Optional bytes password to encrypt the exported key.
                      If None, the key is exported unencrypted (less safe).

        Returns:
            PEM string of the private key. KEEP THIS SECRET.
        """
        encryption = (
            serialization.BestAvailableEncryption(password)
            if password
            else serialization.NoEncryption()
        )
        return self.private_key.private_bytes(
            encoding           = serialization.Encoding.PEM,
            format             = serialization.PrivateFormat.PKCS8,
            encryption_algorithm = encryption,
        ).decode("utf-8")

    def save_to_file(self, filepath: str, password: bytes | None = None):
        """
        Save this wallet's keys and address to a JSON file.

        Args:
            filepath: Path to write the wallet file (e.g., "my_wallet.json").
            password: Optional password to encrypt the private key in the file.
        """
        wallet_data = {
            "address"    : self.address,
            "public_key" : self.get_public_key_pem(),
            "private_key": self.get_private_key_pem(password),
        }
        with open(filepath, "w") as f:
            json.dump(wallet_data, f, indent=2)
        print(f"💾  Wallet saved to '{filepath}'. Guard the private key!")

    @classmethod
    def load_from_file(cls, filepath: str, password: bytes | None = None) -> "Wallet":
        """
        Reconstruct a Wallet from a JSON file saved by save_to_file().

        Args:
            filepath: Path to the wallet JSON file.
            password: Password used when saving, if the private key was encrypted.

        Returns:
            A fully initialised Wallet object.
        """
        with open(filepath, "r") as f:
            data = json.load(f)

        wallet = cls.__new__(cls)   # Create instance WITHOUT calling __init__

        # Reload private key from PEM bytes
        wallet.private_key = serialization.load_pem_private_key(
            data["private_key"].encode("utf-8"),
            password = password,
            backend  = default_backend(),
        )
        wallet.public_key = wallet.private_key.public_key()
        wallet.address    = data["address"]

        print(f"🔑  Wallet loaded from '{filepath}'. Address: {wallet.address}")
        return wallet

    def __repr__(self) -> str:
        return f"Wallet(address={self.address})"