main.py

"""
Simple Key Vault API - Like Azure Key Vault but FREE
Single file, no external dependencies beyond FastAPI + SQLAlchemy
Deploy on Render via Docker Hub
"""
import os
from datetime import datetime
from typing import Optional, Dict

from fastapi import FastAPI, HTTPException, Depends, Query, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from sqlalchemy import create_engine, Column, String, Text, DateTime, ForeignKey
from sqlalchemy.orm import sessionmaker, declarative_base, Session
from dotenv import load_dotenv

load_dotenv()

# ============== API KEY CONFIGURATION ==============
# Hardcoded API key for now - change in production!
API_KEY = "Azure@123"

# ============== DATABASE CONFIGURATION ==============
DATABASE_URL = os.getenv("DATABASE_URL")
if not DATABASE_URL:
    raise RuntimeError("DATABASE_URL not set. Put it in .env or export it.")

# ============== FASTAPI APP ==============
app = FastAPI(
    title="Key Vault API",
    description="Centralized secret storage - Like Azure Key Vault but FREE",
    version="1.1.0",
)

# CORS - Allow all origins (restrict in production)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============== DATABASE SETUP ==============
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


# ============== SQLALCHEMY MODELS ==============
class Secret(Base):
    __tablename__ = "secrets"

    # Keeping current schema for compatibility (id duplicates key)
    id = Column(String, primary_key=True)
    key = Column(String, unique=True, nullable=False, index=True)
    value = Column(Text, nullable=False)
    description = Column(String, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)


class AuditLog(Base):
    """
    Stores immutable audit history for secret changes.
    old_value/new_value are stored as Text (plain) to keep it simple.
    If you want later, you can mask/encrypt values.
    """
    __tablename__ = "audit_logs"

    id = Column(String, primary_key=True)  # simple string id
    secret_key = Column(String, nullable=False, index=True)
    action = Column(String, nullable=False)  # created | updated | deleted
    old_value = Column(Text, nullable=True)
    new_value = Column(Text, nullable=True)
    old_description = Column(String, nullable=True)
    new_description = Column(String, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)


# Create tables on startup
def init_db() -> None:
    Base.metadata.create_all(bind=engine)
    print("✅ Vault database initialized")


@app.on_event("startup")
def on_startup() -> None:
    init_db()


# ============== PYDANTIC MODELS ==============
class SecretCreate(BaseModel):
    key: str
    value: str
    description: Optional[str] = None


class SecretPut(BaseModel):
    value: str
    description: Optional[str] = None


class SecretResponse(BaseModel):
    key: str
    value: str
    description: Optional[str]
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True


class AllSecretsResponse(BaseModel):
    secrets: Dict[str, str]
    count: int


class HealthResponse(BaseModel):
    status: str
    timestamp: datetime
    total_secrets: int


class AuditLogResponse(BaseModel):
    id: str
    secret_key: str
    action: str
    old_value: Optional[str]
    new_value: Optional[str]
    old_description: Optional[str]
    new_description: Optional[str]
    created_at: datetime

    class Config:
        from_attributes = True


class AuditListResponse(BaseModel):
    items: list[AuditLogResponse]
    count: int


# ============== HELPER FUNCTIONS ==============
def get_db() -> Session:
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def verify_api_key(x_api_key: str = Header(None, alias="X-API-Key")):
    """Verify the API key from the X-API-Key header"""
    if not x_api_key:
        raise HTTPException(status_code=401, detail="Missing API key. Include 'X-API-Key' header.")
    if x_api_key != API_KEY:
        raise HTTPException(status_code=403, detail="Invalid API key.")
    return x_api_key


def _audit_id() -> str:
    # No extra dependencies; unique enough for personal use
    return f"audit_{datetime.utcnow().strftime('%Y%m%d%H%M%S%f')}"


def add_audit_log(
    db: Session,
    *,
    secret_key: str,
    action: str,
    old_value: Optional[str],
    new_value: Optional[str],
    old_description: Optional[str],
    new_description: Optional[str],
) -> AuditLog:
    log = AuditLog(
        id=_audit_id(),
        secret_key=secret_key,
        action=action,
        old_value=old_value,
        new_value=new_value,
        old_description=old_description,
        new_description=new_description,
        created_at=datetime.utcnow(),
    )
    db.add(log)
    return log


# ============== API ENDPOINTS ==============

@app.get("/")
def root():
    """Welcome endpoint"""
    return {
        "service": "Key Vault API",
        "version": "1.1.0",
        "description": "Like Azure Key Vault but FREE",
        "endpoints": {
            "GET /": "This help message",
            "GET /health": "Health check",
            "GET /secrets": "Get all secrets (key-value pairs)",
            "GET /secrets/keys": "Get all keys only (no values)",
            "GET /secrets/{key}": "Get secret value by key",
            "GET /secrets/{key}/full": "Get full secret details",
            "PUT /secrets/{key}": "Create or update a secret (JSON body)",
            "POST /secrets": "Create a secret (fails if exists)",
            "DELETE /secrets/{key}": "Delete a secret",
            "GET /secrets/{key}/history": "Get audit history for a key",
            "GET /audit": "Get recent audit logs (all keys)",
            "GET /stats": "Vault statistics",
        },
    }


@app.get("/health", response_model=HealthResponse)
def health(db: Session = Depends(get_db)):
    """Health check endpoint"""
    total = db.query(Secret).count()
    return HealthResponse(status="healthy", timestamp=datetime.utcnow(), total_secrets=total)


@app.get("/secrets", response_model=AllSecretsResponse)
def get_all_secrets(db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Get ALL secrets as key-value pairs."""
    secrets = db.query(Secret).all()
    return AllSecretsResponse(secrets={s.key: s.value for s in secrets}, count=len(secrets))


@app.get("/secrets/keys")
def get_all_keys(db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Get all secret keys (without values)"""
    secrets = db.query(Secret.key).all()
    return {"keys": [s[0] for s in secrets], "count": len(secrets)}


@app.get("/secrets/{key}")
def get_secret(key: str, db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Get a secret value by key."""
    secret = db.query(Secret).filter(Secret.key == key).first()
    if not secret:
        raise HTTPException(status_code=404, detail=f"Secret '{key}' not found")
    return {"key": secret.key, "value": secret.value}


@app.get("/secrets/{key}/full", response_model=SecretResponse)
def get_secret_full(key: str, db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Get full secret details including metadata"""
    secret = db.query(Secret).filter(Secret.key == key).first()
    if not secret:
        raise HTTPException(status_code=404, detail=f"Secret '{key}' not found")
    return secret


@app.put("/secrets/{key}")
def create_or_update_secret(key: str, body: SecretPut, db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """
    Create a new secret or update existing one.
    Body: {"value": "...", "description": "..."}
    """
    secret = db.query(Secret).filter(Secret.key == key).first()
    now = datetime.utcnow()

    try:
        if secret:
            old_value = secret.value
            old_desc = secret.description

            secret.value = body.value
            secret.description = body.description
            secret.updated_at = now

            add_audit_log(
                db,
                secret_key=key,
                action="updated",
                old_value=old_value,
                new_value=body.value,
                old_description=old_desc,
                new_description=body.description,
            )
            action = "updated"
        else:
            secret = Secret(
                id=key,
                key=key,
                value=body.value,
                description=body.description,
                created_at=now,
                updated_at=now,
            )
            db.add(secret)

            add_audit_log(
                db,
                secret_key=key,
                action="created",
                old_value=None,
                new_value=body.value,
                old_description=None,
                new_description=body.description,
            )
            action = "created"

        db.commit()
        db.refresh(secret)
        return {
            "key": secret.key,
            "value": secret.value,
            "description": secret.description,
            "created_at": secret.created_at,
            "updated_at": secret.updated_at,
            "action": action,
        }
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=f"Failed to write secret: {e}")


@app.post("/secrets", response_model=SecretResponse)
def create_secret(secret: SecretCreate, db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Create a new secret (fails if key already exists)."""
    existing = db.query(Secret).filter(Secret.key == secret.key).first()
    if existing:
        raise HTTPException(
            status_code=400,
            detail=f"Secret '{secret.key}' already exists. Use PUT to update.",
        )

    now = datetime.utcnow()
    new_secret = Secret(
        id=secret.key,
        key=secret.key,
        value=secret.value,
        description=secret.description,
        created_at=now,
        updated_at=now,
    )

    try:
        db.add(new_secret)
        add_audit_log(
            db,
            secret_key=secret.key,
            action="created",
            old_value=None,
            new_value=secret.value,
            old_description=None,
            new_description=secret.description,
        )
        db.commit()
        db.refresh(new_secret)
        return new_secret
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=f"Failed to create secret: {e}")


@app.delete("/secrets/{key}")
def delete_secret(key: str, db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Delete a secret"""
    secret = db.query(Secret).filter(Secret.key == key).first()
    if not secret:
        raise HTTPException(status_code=404, detail=f"Secret '{key}' not found")

    try:
        add_audit_log(
            db,
            secret_key=key,
            action="deleted",
            old_value=secret.value,
            new_value=None,
            old_description=secret.description,
            new_description=None,
        )
        db.delete(secret)
        db.commit()
        return {"status": "deleted", "key": key}
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=f"Failed to delete secret: {e}")


@app.get("/secrets/{key}/history", response_model=AuditListResponse)
def get_secret_history(
    key: str,
    limit: int = Query(100, ge=1, le=1000),
    db: Session = Depends(get_db),
    api_key: str = Depends(verify_api_key),
):
    """Get audit history for a specific key (recent first)."""
    items = (
        db.query(AuditLog)
        .filter(AuditLog.secret_key == key)
        .order_by(AuditLog.created_at.desc())
        .limit(limit)
        .all()
    )
    return AuditListResponse(items=items, count=len(items))


@app.get("/audit", response_model=AuditListResponse)
def get_audit(
    limit: int = Query(200, ge=1, le=2000),
    key: Optional[str] = None,
    db: Session = Depends(get_db),
    api_key: str = Depends(verify_api_key),
):
    """Get recent audit logs across all keys (optionally filter by key)."""
    q = db.query(AuditLog)
    if key:
        q = q.filter(AuditLog.secret_key == key)

    items = q.order_by(AuditLog.created_at.desc()).limit(limit).all()
    return AuditListResponse(items=items, count=len(items))


@app.get("/stats")
def get_stats(db: Session = Depends(get_db), api_key: str = Depends(verify_api_key)):
    """Get vault statistics"""
    total = db.query(Secret).count()
    return {
        "total_secrets": total,
        "service": "Key Vault API",
        "version": "1.1.0",
        "status": "running",
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)


# ============== MOUNT GRADIO FOR HF SPACES ==============
def mount_gradio_ui():
    """Mount Gradio UI under /ui path for HF Spaces deployment"""
    try:
        import gradio as gr
        # Lazy import to avoid circular dependency
        import importlib.util
        spec = importlib.util.spec_from_file_location("app_gradio", "app_gradio.py")
        app_gradio = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(app_gradio)
        
        # Mount Gradio under /ui
        app.mount("/ui", gr.mount_gradio_app(app, app_gradio.blocks, path="/"))
        print("✅ Gradio UI mounted at /ui")
    except Exception as e:
        print(f"⚠️  Could not mount Gradio: {e}")


# Auto-mount when running in HF Spaces (PORT env var set)
if os.getenv("PORT"):
    mount_gradio_ui()