alfred/scripts/bootstrap.py

#!/usr/bin/env python3
"""Bootstrap script - generates .env.alfred, .env.librechat, .env.secrets and .env.make."""

import re
import secrets
import sys
from pathlib import Path

import tomllib

BASE_DIR = Path(__file__).resolve().parent.parent

def load_secrets_spec(toml_data: dict) -> dict[str, tuple[int, str]]:
    """Load secrets spec from pyproject.toml [tool.alfred.secrets]."""
    raw = toml_data.get("tool", {}).get("alfred", {}).get("secrets", {})
    result = {}
    for key, rule in raw.items():
        size_str, fmt = rule.split(":")
        result[key] = (int(size_str), fmt)
    return result


def generate_secret(size: int, fmt: str) -> str:
    match fmt:
        case "hex":
            return secrets.token_hex(size)
        case "b64":
            return secrets.token_urlsafe(size)
        case _:
            raise ValueError(f"Unknown format: {fmt}")


def load_env_file(path: Path) -> dict[str, str]:
    """Load key=value pairs from an env file, ignoring comments and blanks."""
    result = {}
    if not path.exists():
        return result
    for line in path.read_text().splitlines():
        stripped = line.strip()
        if stripped and not stripped.startswith("#") and "=" in stripped:
            key, _, value = stripped.partition("=")
            result[key.strip()] = value.strip()
    return result


def copy_example_if_missing(src: Path, dst: Path, label: str) -> None:
    """Copy src to dst only if dst doesn't exist yet."""
    if dst.exists():
        print(f"  ↻ {dst.name} already exists, skipping")
        return
    if not src.exists():
        print(f"  ⚠ {label} example not found at {src} — skipping (add it manually)")
        return
    dst.write_text(src.read_text())
    print(f"  + {dst.name} created from {src.name}")


def generate_secrets_file(path: Path, secrets_spec: dict[str, tuple[int, str]]) -> None:
    """Generate .env.secrets with missing secrets, never overwrite existing ones."""
    existing = load_env_file(path)
    lines = list(path.read_text().splitlines()) if path.exists() else [
        "# Auto-generated secrets — DO NOT COMMIT",
        "# Run 'make bootstrap' to generate missing secrets",
        "",
    ]

    added = []
    for key, (size, fmt) in secrets_spec.items():
        if key not in existing:
            value = generate_secret(size, fmt)
            lines.append(f"{key}={value}")
            added.append(key)

    path.write_text("\n".join(lines) + "\n")

    if added:
        print(f"  + Generated: {', '.join(added)}")
    else:
        print("  ↻ All secrets already exist, nothing generated")


def extract_python_version(version_string: str) -> tuple[str, str]:
    clean = re.sub(r"^[=^~><]+", "", version_string.strip())
    parts = clean.split(".")
    if len(parts) >= 2:
        return clean, f"{parts[0]}.{parts[1]}"
    raise ValueError(f"Invalid Python version: {version_string}")


def build_uris(env_alfred: Path, env_secrets: Path) -> None:
    """Build MONGO_URI and POSTGRES_URI from components and append them to .env.secrets."""
    env = {**load_env_file(env_alfred), **load_env_file(env_secrets)}
    existing = load_env_file(env_secrets)

    computed = {
        "MONGO_URI": (
            f"mongodb://{env['MONGO_USER']}:{env['MONGO_PASSWORD']}"
            f"@{env['MONGO_HOST']}:{env['MONGO_PORT']}/{env['MONGO_DB_NAME']}"
            f"?authSource=admin"
        ),
        "POSTGRES_URI": (
            f"postgresql://{env['POSTGRES_USER']}:{env['POSTGRES_PASSWORD']}"
            f"@{env['POSTGRES_HOST']}:{env['POSTGRES_PORT']}/{env['POSTGRES_DB_NAME']}"
        ),
    }

    content = env_secrets.read_text()
    added = []
    for key, value in computed.items():
        if key in existing:
            content = re.sub(rf"^{key}=.*$", f"{key}={value}", content, flags=re.MULTILINE)
        else:
            content = content.rstrip("\n") + f"\n{key}={value}\n"
            added.append(key)
    env_secrets.write_text(content)

    if added:
        print(f"  + Computed: {', '.join(added)}")
    else:
        print("  ↻ URIs updated")


def write_env_make(toml_data: dict) -> None:
    """Write .env.make from pyproject.toml."""
    project = toml_data["project"]
    alfred = toml_data["tool"]["alfred"]

    python_full, python_short = extract_python_version(project["requires-python"])

    lines = [
        "# Auto-generated from pyproject.toml — do not edit manually",
        f"ALFRED_VERSION={project['version']}",
        f"PYTHON_VERSION={python_full}",
        f"IMAGE_NAME={alfred['image_name']}",
        f"SERVICE_NAME={alfred['service_name']}",
        f"LIBRECHAT_VERSION={alfred['librechat_version']}",
        f"RAG_VERSION={alfred['rag_version']}",
        f"UV_VERSION={alfred['uv_version']}",
    ]

    env_make_path = BASE_DIR / ".env.make"
    env_make_path.write_text("\n".join(lines) + "\n")
    print(f"  + {env_make_path.name} written")


def main() -> int:
    print("🚀 Starting bootstrap...")

    toml_path = BASE_DIR / "pyproject.toml"
    if not toml_path.exists():
        print(f"❌ pyproject.toml not found: {toml_path}")
        return 1

    with open(toml_path, "rb") as f:
        toml_data = tomllib.load(f)

    print("\n📄 Env files:")
    copy_example_if_missing(
        src=BASE_DIR / ".env.example",
        dst=BASE_DIR / ".env.alfred",
        label="Alfred",
    )
    copy_example_if_missing(
        src=BASE_DIR / "librechat" / ".env.example",
        dst=BASE_DIR / ".env.librechat",
        label="LibreChat",
    )

    secrets_spec = load_secrets_spec(toml_data)
    print("\n🔐 Secrets:")
    generate_secrets_file(BASE_DIR / ".env.secrets", secrets_spec)

    print("\n🔗 URIs:")
    build_uris(BASE_DIR / ".env.alfred", BASE_DIR / ".env.secrets")

    print("\n🔧 Build config:")
    write_env_make(toml_data)

    print("\n✅ Bootstrap complete!")
    return 0


if __name__ == "__main__":
    sys.exit(main())