myagent-ai 1.0.0

package/communication/crypto.py

@@ -0,0 +1,347 @@
+"""
+communication/crypto.py - Cryptographic Primitives
+===================================================
+Ed25519 identity, X25519 key exchange, AES-256-GCM encryption/decryption,
+message signing/verification, and key serialization.
+
+Uses the ``cryptography`` library.  Falls back to a pure-Python stub when it
+is not available (signing only, encryption disabled).
+"""
+from __future__ import annotations
+
+import hashlib
+import hmac
+import os
+import base64
+import json
+import struct
+from typing import Optional, Tuple
+
+# ---------------------------------------------------------------------------
+# Try importing the real crypto library
+# ---------------------------------------------------------------------------
+_HAS_CRYPTOGRAPHY = False
+try:
+    from cryptography.hazmat.primitives.asymmetric.ed25519 import (
+        Ed25519PrivateKey,
+        Ed25519PublicKey,
+    )
+    from cryptography.hazmat.primitives.asymmetric.x25519 import (
+        X25519PrivateKey,
+        X25519PublicKey,
+    )
+    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.exceptions import InvalidSignature
+
+    _HAS_CRYPTOGRAPHY = True
+except ImportError:
+    pass
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+NONCE_SIZE = 12          # 96-bit nonce for AES-256-GCM
+SHARED_SECRET_SIZE = 32  # 256-bit shared secret
+
+
+# =========================================================================
+# Key Generation
+# =========================================================================
+
+def generate_ed25519_keypair() -> Tuple[bytes, bytes]:
+    """Generate an Ed25519 key pair.
+
+    Returns:
+        (private_bytes, public_bytes) — both 32 bytes raw.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        private_key = Ed25519PrivateKey.generate()
+        priv_bytes = private_key.private_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PrivateFormat.Raw,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        pub_bytes = private_key.public_key().public_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PublicFormat.Raw,
+        )
+        return priv_bytes, pub_bytes
+    else:
+        # Fallback: generate random bytes (signing will not be available)
+        priv_bytes = os.urandom(32)
+        pub_bytes = hashlib.sha256(priv_bytes).digest()
+        return priv_bytes, pub_bytes
+
+
+def generate_x25519_keypair() -> Tuple[bytes, bytes]:
+    """Generate an X25519 key pair for Diffie-Hellman.
+
+    Returns:
+        (private_bytes, public_bytes) — both 32 bytes raw.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        private_key = X25519PrivateKey.generate()
+        priv_bytes = private_key.private_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PrivateFormat.Raw,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        pub_bytes = private_key.public_key().public_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PublicFormat.Raw,
+        )
+        return priv_bytes, pub_bytes
+    else:
+        priv_bytes = os.urandom(32)
+        pub_bytes = hashlib.sha256(priv_bytes).digest()
+        return priv_bytes, pub_bytes
+
+
+# =========================================================================
+# Ed25519 <-> X25519 Conversion
+# =========================================================================
+
+def ed25519_to_x25519_private(ed_priv: bytes) -> bytes:
+    """Convert an Ed25519 private key to X25519 private key bytes.
+
+    According to RFC 7748, we hash the 32-byte Ed25519 seed with SHA-512
+    and clamp the first 32 bytes.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        # The raw Ed25519 private key from cryptography is already the seed.
+        # Use the standard clamping procedure.
+        h = hashlib.sha512(ed_priv).digest()
+        clamped = bytearray(h[:32])
+        clamped[0] &= 248
+        clamped[31] &= 127
+        clamped[31] |= 64
+        return bytes(clamped)
+    else:
+        return hashlib.sha256(ed_priv).digest()
+
+
+def ed25519_to_x25519_public(ed_pub: bytes) -> bytes:
+    """Convert an Ed25519 public key to X25519 public key.
+
+    In practice with the cryptography library, Ed25519 and X25519 public
+    keys share the same representation for curve25519.  We return a simple
+    hash fallback when the library is not available.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        # Both use the same Montgomery u-coordinate internally.
+        return ed_pub
+    else:
+        return hashlib.sha256(ed_pub).digest()
+
+
+# =========================================================================
+# X25519 Diffie-Hellman Key Exchange
+# =========================================================================
+
+def x25519_key_exchange(
+    my_private: bytes,
+    their_public: bytes,
+) -> bytes:
+    """Perform X25519 Diffie-Hellman key exchange.
+
+    Returns:
+        32-byte shared secret.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        my_priv_key = X25519PrivateKey.from_private_bytes(my_private)
+        their_pub_key = X25519PublicKey.from_public_bytes(their_public)
+        shared = my_priv_key.exchange(their_pub_key)
+        return shared
+    else:
+        # Fallback: HMAC-based pseudo-DH (NOT secure, for testing only)
+        return hmac.new(my_private, their_public, hashlib.sha256).digest()
+
+
+# =========================================================================
+# HKDF — Derive Encryption Key from Shared Secret
+# =========================================================================
+
+def derive_shared_key(shared_secret: bytes, context: bytes = b"") -> bytes:
+    """Derive a 256-bit AES key from a raw shared secret using HKDF-SHA256.
+
+    Args:
+        shared_secret: Raw DH shared secret (32 bytes).
+        context: Optional context string for domain separation.
+
+    Returns:
+        32-byte key suitable for AES-256-GCM.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        hkdf = HKDF(
+            algorithm=hashes.SHA256(),
+            length=32,
+            salt=None,
+            info=b"myagent-comm-v1" + context,
+        )
+        return hkdf.derive(shared_secret)
+    else:
+        # Fallback: simple hash-based derivation
+        derived = hashlib.sha256(shared_secret + b"myagent-comm-v1" + context).digest()
+        return derived
+
+
+# =========================================================================
+# AES-256-GCM Encrypt / Decrypt
+# =========================================================================
+
+def encrypt(key: bytes, plaintext: bytes) -> bytes:
+    """Encrypt *plaintext* with AES-256-GCM.
+
+    Args:
+        key: 32-byte encryption key.
+        plaintext: Arbitrary data to encrypt.
+
+    Returns:
+        ``nonce (12) || ciphertext+tag`` — concatenation of nonce and sealed data.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        nonce = os.urandom(NONCE_SIZE)
+        aesgcm = AESGCM(key)
+        ct = aesgcm.encrypt(nonce, plaintext, None)
+        return nonce + ct
+    else:
+        # Fallback: XOR-based "encryption" (NOT secure)
+        nonce = os.urandom(NONCE_SIZE)
+        stream = b""
+        for i, b in enumerate(plaintext):
+            stream += bytes([b ^ key[i % len(key)]])
+        return nonce + stream
+
+
+def decrypt(key: bytes, sealed: bytes) -> bytes:
+    """Decrypt data produced by :func:`encrypt`.
+
+    Args:
+        key: 32-byte encryption key.
+        sealed: ``nonce (12) || ciphertext+tag``.
+
+    Returns:
+        Original plaintext.
+
+    Raises:
+        ValueError: If decryption fails (tampered data).
+    """
+    if _HAS_CRYPTOGRAPHY:
+        nonce = sealed[:NONCE_SIZE]
+        ct = sealed[NONCE_SIZE:]
+        aesgcm = AESGCM(key)
+        try:
+            return aesgcm.decrypt(nonce, ct, None)
+        except Exception:
+            raise ValueError("Decryption failed: data may be tampered")
+    else:
+        nonce = sealed[:NONCE_SIZE]
+        ct = sealed[NONCE_SIZE:]
+        stream = b""
+        for i, b in enumerate(ct):
+            stream += bytes([b ^ key[i % len(key)]])
+        return stream
+
+
+# =========================================================================
+# Ed25519 Sign / Verify
+# =========================================================================
+
+def sign(private_key: bytes, message: bytes) -> bytes:
+    """Sign *message* with an Ed25519 *private_key*.
+
+    Returns:
+        64-byte signature.
+    """
+    if _HAS_CRYPTOGRAPHY:
+        priv = Ed25519PrivateKey.from_private_bytes(private_key)
+        return priv.sign(message)
+    else:
+        # Fallback: HMAC-SHA512 (NOT a real Ed25519 signature)
+        return hmac.new(private_key, message, hashlib.sha512).digest()
+
+
+def verify(public_key: bytes, message: bytes, signature: bytes) -> bool:
+    """Verify an Ed25519 *signature* for *message* with *public_key*."""
+    if _HAS_CRYPTOGRAPHY:
+        try:
+            pub = Ed25519PublicKey.from_public_bytes(public_key)
+            pub.verify(signature, message)
+            return True
+        except InvalidSignature:
+            return False
+        except Exception:
+            return False
+    else:
+        expected = hmac.new(public_key, message, hashlib.sha512).digest()
+        return hmac.compare_digest(expected, signature)
+
+
+# =========================================================================
+# Key Serialization (Hex)
+# =========================================================================
+
+def serialize_public_key(pub: bytes) -> str:
+    """Serialize a 32-byte public key to a hex string."""
+    return pub.hex()
+
+
+def deserialize_public_key(hex_str: str) -> bytes:
+    """Deserialize a hex public key to 32 bytes."""
+    return bytes.fromhex(hex_str)
+
+
+def serialize_private_key(priv: bytes) -> str:
+    """Serialize a 32-byte private key to a hex string."""
+    return priv.hex()
+
+
+def deserialize_private_key(hex_str: str) -> bytes:
+    """Deserialize a hex private key to 32 bytes."""
+    return bytes.fromhex(hex_str)
+
+
+# =========================================================================
+# Agent ID Generation
+# =========================================================================
+
+def generate_agent_id(public_key: bytes) -> str:
+    """Generate a short, human-readable Agent ID from a 32-byte Ed25519 public key.
+
+    Uses the first 8 bytes (16 hex chars) of SHA-256(public_key).
+
+    Returns:
+        A 16-char lowercase hex string, e.g. ``"a1b2c3d4e5f6a7b8"``.
+    """
+    digest = hashlib.sha256(public_key).hexdigest()
+    return digest[:16]
+
+
+# =========================================================================
+# Convenience: encrypt a text message payload
+# =========================================================================
+
+def encrypt_message(
+    shared_key: bytes,
+    message_data: dict,
+) -> str:
+    """Encrypt a JSON-serializable dict to a base64 sealed string.
+
+    This is the high-level helper used by the channel layer.
+    """
+    plaintext = json.dumps(message_data, ensure_ascii=False).encode("utf-8")
+    sealed = encrypt(shared_key, plaintext)
+    return base64.b64encode(sealed).decode("ascii")
+
+
+def decrypt_message(
+    shared_key: bytes,
+    sealed_b64: str,
+) -> dict:
+    """Decrypt a base64 sealed string back to a dict."""
+    sealed = base64.b64decode(sealed_b64)
+    plaintext = decrypt(shared_key, sealed)
+    return json.loads(plaintext.decode("utf-8"))