Extension System

# src/extensions/builtin/secret_vault/base_crypto.py
from abc import ABC, abstractmethod
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

class BaseCrypto(ABC):
    """
    Base class for AES-256-GCM encryption operations.
    
    Subclasses implement key derivation strategy:
    - MachineSecretCrypto: Derive key from EK fingerprint (HKDF-SHA256)
    - SharedSecretCrypto: Use master key from environment
    """
    
    def __init__(self):
        self._key = self._derive_key()
        self.cipher = AESGCM(self._key)
    
    @abstractmethod
    def _derive_key(self) -> bytes:
        """Derive encryption key (32 bytes for AES-256)."""
        pass
    
    @abstractmethod
    def get_key_id(self) -> str:
        """Return key identifier for metadata tracking."""
        pass
    
    def encrypt(self, plaintext: str) -> Tuple[bytes, bytes, bytes]:
        """
        Encrypt plaintext secret value.
        
        Returns:
            Tuple of (ciphertext, nonce, tag)
        """
        nonce = os.urandom(12)  # 96 bits for GCM
        plaintext_bytes = plaintext.encode("utf-8")
        ciphertext_and_tag = self.cipher.encrypt(nonce, plaintext_bytes, None)
        
        # Split ciphertext and tag
        ciphertext = ciphertext_and_tag[:-16]
        tag = ciphertext_and_tag[-16:]
        
        return ciphertext, nonce, tag
    
    def decrypt(self, ciphertext: bytes, nonce: bytes, tag: bytes) -> str:
        """Decrypt encrypted secret value."""
        ciphertext_and_tag = ciphertext + tag
        plaintext_bytes = self.cipher.decrypt(nonce, ciphertext_and_tag, None)
        return plaintext_bytes.decode("utf-8")

# src/extensions/builtin/secret_vault/base_models.py
from sqlmodel import Field
from datetime import datetime

class EncryptedSecretMixin:
    """
    Mixin for AES-256-GCM encrypted secret storage.
    
    Provides common fields for all secret types:
    - encrypted_value, nonce, tag (AES-GCM encryption)
    - created_at, created_by (audit trail)
    - last_accessed_at, access_count (access tracking)
    """
    
    encrypted_value: bytes = Field(description="AES-256-GCM encrypted secret value")
    nonce: bytes = Field(description="GCM nonce (96 bits)")
    tag: bytes = Field(description="GCM authentication tag (128 bits)")
    created_at: datetime = Field(default_factory=datetime.utcnow)
    created_by: str = Field(max_length=256)
    last_accessed_at: Optional[datetime] = Field(default=None)
    access_count: int = Field(default=0)

# src/extensions/builtin/secret_vault/crypto.py
from .base_crypto import BaseCrypto
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes

class MachineSecretCrypto(BaseCrypto):
    """Machine-specific secrets — TPM-bound via EK fingerprint."""
    
    def __init__(self, ek_fingerprint: str):
        self.ek_fingerprint = ek_fingerprint
        super().__init__()
    
    def _derive_key(self) -> bytes:
        """Derive AES-256 key from EK fingerprint using HKDF-SHA256."""
        ikm = bytes.fromhex(self.ek_fingerprint)
        salt = b"ITL.ControlPlane.Attestation.SecretVault.v1"
        
        kdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            info=b"machine-secret-encryption"
        )
        return kdf.derive(ikm)
    
    def get_key_id(self) -> str:
        return f"ek-{self.ek_fingerprint[:16]}"

# src/extensions/builtin/secret_vault/shared_crypto.py
from .base_crypto import BaseCrypto
import os
import hashlib

class SharedSecretCrypto(BaseCrypto):
    """Shared secrets — master key encryption (not TPM-bound)."""
    
    def __init__(self):
        self._key_source: str = ""
        super().__init__()
    
    def _derive_key(self) -> bytes:
        """Get or generate master encryption key."""
        # Priority: ITL_SHARED_SECRET_MASTER_KEY → ITL_SHARED_SECRET_PASSPHRASE → random
        if "ITL_SHARED_SECRET_MASTER_KEY" in os.environ:
            key = bytes.fromhex(os.environ["ITL_SHARED_SECRET_MASTER_KEY"])
            self._key_source = "env:ITL_SHARED_SECRET_MASTER_KEY"
            return key
        
        if "ITL_SHARED_SECRET_PASSPHRASE" in os.environ:
            passphrase = os.environ["ITL_SHARED_SECRET_PASSPHRASE"]
            key = hashlib.sha256(passphrase.encode("utf-8")).digest()
            self._key_source = "env:ITL_SHARED_SECRET_PASSPHRASE(sha256)"
            return key
        
        # Random key (dev only — not persisted!)
        key = os.urandom(32)
        self._key_source = "random-ephemeral"
        return key
    
    def get_key_id(self) -> str:
        return self._key_source

# Singleton instance
def get_shared_crypto() -> SharedSecretCrypto:
    """Get singleton SharedSecretCrypto instance."""
    global _shared_crypto_instance
    if _shared_crypto_instance is None:
        _shared_crypto_instance = SharedSecretCrypto()
    return _shared_crypto_instance

# src/extensions/builtin/secret_vault/models.py
from .base_models import EncryptedSecretMixin

class SecretRow(EncryptedSecretMixin, SQLModel, table=True):
    """Machine-specific TPM-bound secret."""
    __tablename__ = "extension_secrets"
    
    secret_id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    machine_id: uuid.UUID = Field(foreign_key="machines.machine_id")
    name: str = Field(max_length=128, index=True)
    
    # Encrypted storage fields inherited from EncryptedSecretMixin:
    # - encrypted_value, nonce, tag
    # - created_at, created_by
    # - last_accessed_at, access_count

# src/extensions/builtin/secret_vault/shared_models.py
from .base_models import EncryptedSecretMixin

class SharedSecretRow(EncryptedSecretMixin, SQLModel, table=True):
    """Shared secret accessible by multiple machines."""
    __tablename__ = "extension_shared_secrets"
    
    shared_secret_id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    name: str = Field(max_length=128, unique=True, index=True)
    
    # Encrypted storage fields inherited from EncryptedSecretMixin:
    # - encrypted_value, nonce, tag
    # - created_at, created_by
    # - last_accessed_at, access_count
    
    encryption_key_id: str = Field(max_length=64)
    last_rotated_at: Optional[datetime] = Field(default=None)
    description: Optional[str] = Field(default=None, max_length=512)

import hmac
import hashlib

def verify_webhook_signature(payload: bytes, signature: str, secret: str) -> bool:
    """Verify webhook HMAC signature."""
    expected = hmac.new(
        secret.encode("utf-8"),
        payload,
        hashlib.sha256
    ).hexdigest()
    
    # Extract hex digest from "sha256=..."
    received = signature.replace("sha256=", "")
    
    return hmac.compare_digest(expected, received)

# Usage in webhook handler
@app.post("/hooks/attestation")
async def handle_attestation_webhook(request: Request):
    payload = await request.body()
    signature = request.headers.get("X-Webhook-Signature")
    
    if not verify_webhook_signature(payload, signature, WEBHOOK_SECRET):
        raise HTTPException(status_code=401, detail="Invalid signature")
    
    event = await request.json()
    # Process event...

# src/extensions/builtin/my_extension/extension.py
from extensions.base import AttestationExtension
from fastapi import APIRouter
from .models import MyRow

class MyExtension(AttestationExtension):
    @property
    def name(self) -> str:
        return "my_extension"
    
    @property
    def version(self) -> str:
        return "1.0.0"
    
    @property
    def description(self) -> str:
        return "My custom extension"
    
    def get_router(self) -> APIRouter | None:
        router = APIRouter(prefix="/api/v1/my", tags=["my"])
        
        @router.get("/hello")
        async def hello():
            return {"message": "Hello from my extension"}
        
        return router
    
    def get_models(self) -> list[type]:
        return [MyRow]
    
    def on_startup(self) -> None:
        print(f"[MyExtension] Started (v{self.version})")

# src/extensions/builtin/my_extension/__init__.py
from .extension import MyExtension

__all__ = ["MyExtension"]

# tests/extensions/test_secret_vault.py
import pytest
from extensions.builtin.secret_vault.crypto import SecretCrypto

def test_encrypt_decrypt():
    plaintext = "my-secret-value"
    ek_fp = "a" * 64
    
    encrypted, nonce, tag = SecretCrypto.encrypt(plaintext, ek_fp)
    decrypted = SecretCrypto.decrypt(encrypted, nonce, tag, ek_fp)
    
    assert decrypted == plaintext

# tests/extensions/test_secret_vault_api.py
from fastapi.testclient import TestClient
from attestation.core.app import create_app

def test_create_secret():
    client = TestClient(create_app())
    response = client.post(
        "/api/v1/secrets/machines/test-id/secrets",
        json={"name": "test", "value": "secret123"},
        headers={"Authorization": "Bearer <token>"}
    )
    assert response.status_code == 201

def get_models(self) -> list[type]:
    return [MyRow]  # Must return actual class, not string