migrate_db.py

#!/usr/bin/env python3
"""
Database Migration Script - Run ONCE to add encryption to existing database
"""

import sys

# Add app directory to path
sys.path.insert(0, "/var/www/transkript_app")

import json
import logging

from sqlalchemy import inspect, text

from crypto_utils import crypto
from db_models import (
    Base,
    ChatHistory,
    GeneratedImage,
    Transcription,
    User,
    VisionResult,
)
from db_ops import SessionLocal, engine

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def check_column_exists(table_name, column_name):
    """Check if a column exists in a table"""
    inspector = inspect(engine)
    columns = [col["name"] for col in inspector.get_columns(table_name)]
    return column_name in columns


def add_encryption_columns():
    """Add is_encrypted and encryption_metadata columns to existing tables"""

    tables_to_update = ["transcriptions", "chat_history", "vision_results", "generated_images"]

    import re as _re

    _IDENT_RE = _re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")

    with engine.begin() as conn:
        for table in tables_to_update:
            # Defence in depth: SQL identifiers cannot be parameterised, so
            # validate against a strict allow-list pattern before interpolation.
            if not _IDENT_RE.match(table):
                logger.error(f"❌ Invalid table identifier (refused): {table!r}")
                continue
            try:
                # Check if table exists
                result = conn.execute(
                    text("SELECT name FROM sqlite_master WHERE type='table' AND name=:n"),
                    {"n": table},
                )
                if not result.fetchone():
                    logger.warning(f"⚠️ Table {table} does not exist, skipping...")
                    continue

                # Check and add is_encrypted
                if not check_column_exists(table, "is_encrypted"):
                    conn.execute(
                        text(f"ALTER TABLE {table} ADD COLUMN is_encrypted BOOLEAN DEFAULT 0")
                    )
                    logger.info(f"✅ Added is_encrypted to {table}")
                else:
                    logger.info(f"ℹ️ Column is_encrypted already exists in {table}")

                # Check and add encryption_metadata
                if not check_column_exists(table, "encryption_metadata"):
                    conn.execute(text(f"ALTER TABLE {table} ADD COLUMN encryption_metadata TEXT"))
                    logger.info(f"✅ Added encryption_metadata to {table}")
                else:
                    logger.info(f"ℹ️ Column encryption_metadata already exists in {table}")

                # Add encrypted_path for generated_images (ADDED THIS)
                if table == "generated_images" and not check_column_exists(table, "encrypted_path"):
                    conn.execute(text(f"ALTER TABLE {table} ADD COLUMN encrypted_path TEXT"))
                    logger.info(f"✅ Added encrypted_path to {table}")

            except Exception as e:
                logger.error(f"❌ Error updating {table}: {e}")
                raise

        # Security hardening: ensure users.must_change_password exists so the
        # bootstrap-admin password rotation flag works on existing databases.
        try:
            if not check_column_exists("users", "must_change_password"):
                conn.execute(
                    text("ALTER TABLE users ADD COLUMN must_change_password BOOLEAN DEFAULT 0")
                )
                logger.info("✅ Added must_change_password to users")
        except Exception as e:
            logger.error(f"❌ Error adding must_change_password column: {e}")


def create_new_tables():
    """Create new tables that don't exist yet"""
    try:
        # This creates tables like FileUploadMetadata, DSFARecord, UserConfirmation, UserSettings, etc.
        Base.metadata.create_all(bind=engine)
        logger.info("✅ Created new tables (if any)")
    except Exception as e:
        logger.error(f"❌ Error creating new tables: {e}")
        raise


def encrypt_existing_data():
    """Encrypt all existing plaintext data"""
    db = SessionLocal()

    try:
        logger.info("🔐 Starting data encryption...")

        # Check if crypto is initialized
        if not crypto.master_key:
            logger.error("❌ Crypto not initialized! Aborting.")
            return False

        # SQLAlchemy column comparisons require ==/!=, NOT `is` — see docs.
        trans_count = db.query(Transcription).filter(Transcription.is_encrypted == False).count()
        chat_count = db.query(ChatHistory).filter(ChatHistory.is_encrypted == False).count()
        vision_count = db.query(VisionResult).filter(VisionResult.is_encrypted == False).count()

        logger.info(
            f"Found: {trans_count} transcriptions, {chat_count} chats, {vision_count} vision results to encrypt"
        )

        # 1. Encrypt Transcriptions
        if trans_count > 0:
            logger.info(f"Encrypting {trans_count} transcriptions...")
            trans_list = db.query(Transcription).filter(Transcription.is_encrypted == False).all()

            for i, t in enumerate(trans_list, 1):
                try:
                    if t.original_text and t.original_text.strip():
                        t.original_text = crypto.encrypt_text(t.original_text)
                    if t.translated_text and t.translated_text.strip():
                        t.translated_text = crypto.encrypt_text(t.translated_text)

                    t.is_encrypted = True
                    t.encryption_metadata = json.dumps({"algorithm": "AES-256-GCM", "version": 1})

                    if i % 10 == 0:
                        logger.info(f"  Progress: {i}/{trans_count}")
                        db.commit()  # Commit in batches

                except Exception as e:
                    logger.error(f"Failed to encrypt transcription {t.id}: {e}")
                    continue

            db.commit()
            logger.info(f"✅ Encrypted {trans_count} transcriptions")

        # 2. Encrypt Chat History
        if chat_count > 0:
            logger.info(f"Encrypting {chat_count} chats...")
            chat_list = db.query(ChatHistory).filter(ChatHistory.is_encrypted == False).all()

            for i, c in enumerate(chat_list, 1):
                try:
                    if c.messages and c.messages.strip():
                        c.messages = crypto.encrypt_text(c.messages)

                    c.is_encrypted = True
                    c.encryption_metadata = json.dumps({"algorithm": "AES-256-GCM", "version": 1})

                    if i % 10 == 0:
                        logger.info(f"  Progress: {i}/{chat_count}")
                        db.commit()

                except Exception as e:
                    logger.error(f"Failed to encrypt chat {c.id}: {e}")
                    continue

            db.commit()
            logger.info(f"✅ Encrypted {chat_count} chats")

        # 3. Encrypt Vision Results
        if vision_count > 0:
            logger.info(f"Encrypting {vision_count} vision results...")
            vision_list = db.query(VisionResult).filter(VisionResult.is_encrypted == False).all()

            for i, v in enumerate(vision_list, 1):
                try:
                    if v.result and v.result.strip():
                        v.result = crypto.encrypt_text(v.result)
                    if v.prompt and v.prompt.strip():
                        v.prompt = crypto.encrypt_text(v.prompt)

                    v.is_encrypted = True
                    v.encryption_metadata = json.dumps({"algorithm": "AES-256-GCM", "version": 1})

                    if i % 10 == 0:
                        logger.info(f"  Progress: {i}/{vision_count}")
                        db.commit()

                except Exception as e:
                    logger.error(f"Failed to encrypt vision {v.id}: {e}")
                    continue

            db.commit()
            logger.info(f"✅ Encrypted {vision_count} vision results")

        # 4. Note about images (skipped - too large)
        img_count = db.query(GeneratedImage).filter(GeneratedImage.is_encrypted == False).count()
        if img_count > 0:
            logger.info(f"ℹ️ Skipping {img_count} generated images (file encryption on demand)")

        logger.info("✅ Data encryption complete!")
        return True

    except Exception as e:
        db.rollback()
        logger.exception(f"🔥 Encryption failed: {e}")
        return False
    finally:
        db.close()


def verify_encryption():
    """Verify that data is actually encrypted"""
    db = SessionLocal()
    try:
        # Check a sample transcription
        trans = db.query(Transcription).filter(Transcription.is_encrypted == True).first()
        if trans:
            # Encrypted data should be base64 (only contains A-Za-z0-9+/=)
            import re

            if trans.original_text:
                # Check if it's valid base64
                is_base64 = bool(re.match(r"^[A-Za-z0-9+/]*={0,2}$", trans.original_text))

                # Check if it's NOT readable German/English text
                has_readable_words = any(
                    word in trans.original_text.lower()
                    for word in ["der", "die", "das", "the", "and", "ist", "haben"]
                )

                if is_base64 and not has_readable_words:
                    logger.info(
                        "✅ Sample data appears encrypted (base64 format, no readable words)"
                    )
                    return True
                else:
                    logger.warning("⚠️ Sample data does not look encrypted!")
                    logger.warning(f"Preview: {trans.original_text[:100]}")
                    return False
            else:
                logger.info("ℹ️ Sample has no text content")
                return True
        else:
            logger.info("ℹ️ No encrypted transcriptions found (might be empty database)")
            return True
    finally:
        db.close()


def migrate_existing_users_to_keychains():
    """
    Mark users for migration but DON'T create keychains yet.
    Keychains will be created when they next log in.
    """
    db = SessionLocal()

    try:
        users = (
            db.query(User).filter((User.salt == None) | (User.encrypted_master_key == None)).all()
        )

        if not users:
            logger.info("ℹ️ All users already have keychains")
            return True

        logger.info(f"🔑 Marking {len(users)} users for keychain migration...")

        # Just ensure columns exist - don't fill them yet
        for user in users:
            if not user.salt:
                user.salt = None  # Explicit NULL
            if not user.encrypted_master_key:
                user.encrypted_master_key = None  # Explicit NULL

        db.commit()

        logger.info("✅ Users marked for migration. Keychains will be created on next login.")
        return True

    except Exception as e:
        db.rollback()
        logger.exception(f"🔥 Migration failed: {e}")
        return False
    finally:
        db.close()


def add_user_key_columns():
    """Add salt and encrypted_master_key to users table"""
    with engine.begin() as conn:
        try:
            # Check columns
            inspector = inspect(engine)
            cols = [c["name"] for c in inspector.get_columns("users")]

            if "salt" not in cols:
                logger.info("➕ Adding 'salt' column to users...")
                conn.execute(text("ALTER TABLE users ADD COLUMN salt TEXT"))

            if "encrypted_master_key" not in cols:
                logger.info("➕ Adding 'encrypted_master_key' column to users...")
                conn.execute(text("ALTER TABLE users ADD COLUMN encrypted_master_key TEXT"))

            logger.info("✅ User table schema updated")
        except Exception as e:
            logger.error(f"❌ Failed to update user schema: {e}")


def security_audit_master_key():
    """
    Check .master_key status: is it on disk, in env, or both?
    Auto-migrate to .env file if the operator agrees.
    """
    import base64
    import os
    from pathlib import Path

    from crypto_utils import _APP_DIR

    key_file = os.path.join(_APP_DIR, ".master_key")
    env_key = os.environ.get("MASTER_ENCRYPTION_KEY")
    dotenv_path = os.path.join(_APP_DIR, ".env")

    has_file = os.path.exists(key_file)
    has_env = bool(env_key)

    if not has_file and not has_env:
        logger.info("ℹ️ No master key found — one will be generated on first app start")
        return

    if has_env and not has_file:
        logger.info("✅ Master key loaded from MASTER_ENCRYPTION_KEY env var (best practice)")
        return

    if has_file:
        with open(key_file, "rb") as f:
            raw = f.read()
        encoded = base64.b64encode(raw).decode("ascii")

        if has_env:
            # Both exist — check they match
            try:
                env_raw = base64.b64decode(env_key)
                if env_raw == raw:
                    logger.info("✅ .master_key file and MASTER_ENCRYPTION_KEY env match")
                    logger.info("   You can safely delete the .master_key file now:")
                    logger.info(f"   rm {key_file}")
                else:
                    logger.error(
                        "🔥 .master_key file and MASTER_ENCRYPTION_KEY env DIFFER! "
                        "This is dangerous — resolve immediately. The env var takes "
                        "precedence at runtime."
                    )
            except Exception:
                logger.error("🔥 MASTER_ENCRYPTION_KEY env is not valid base64!")
            return

        # File exists, env does not — offer to write it into .env
        logger.warning(
            f"⚠️ .master_key file found at {key_file} but MASTER_ENCRYPTION_KEY not in env."
        )
        logger.info(f"   Base64 value: {encoded}")

        # Check if .env already has a MASTER_ENCRYPTION_KEY line
        already_in_dotenv = False
        if os.path.exists(dotenv_path):
            with open(dotenv_path) as f:
                for line in f:
                    if line.strip().startswith("MASTER_ENCRYPTION_KEY="):
                        already_in_dotenv = True
                        break

        if already_in_dotenv:
            logger.info("   .env already contains MASTER_ENCRYPTION_KEY — load it at startup")
            logger.info(f"   Then delete: rm {key_file}")
            return

        answer = (
            input("\n   Write MASTER_ENCRYPTION_KEY into .env automatically? [y/N]: ")
            .strip()
            .lower()
        )
        if answer == "y":
            with open(dotenv_path, "a") as f:
                f.write("\n# Master encryption key (migrated from .master_key)\n")
                f.write(f"MASTER_ENCRYPTION_KEY={encoded}\n")
            logger.info(f"✅ Appended MASTER_ENCRYPTION_KEY to {dotenv_path}")
            logger.info(f"   Verify the app starts correctly, then delete: rm {key_file}")
        else:
            logger.info("   Skipped. Run `python migrate_master_key.py` manually, or")
            logger.info("   add this to your .env or systemd EnvironmentFile:")
            logger.info(f"   MASTER_ENCRYPTION_KEY={encoded}")


def security_audit_legacy_users():
    """
    Detect admin123/user123 accounts and flag them for password rotation.
    Optionally rename or delete them interactively.
    """
    db = SessionLocal()
    try:
        LEGACY_NAMES = ("admin123", "user123")
        legacy = []
        for name in LEGACY_NAMES:
            u = db.query(User).filter(User.username == name).first()
            if u:
                legacy.append(u)

        if not legacy:
            logger.info("✅ No legacy default users (admin123/user123) found")
            return

        logger.warning(
            f"⚠️ Found {len(legacy)} legacy default user(s) with publicly known passwords:"
        )
        for u in legacy:
            admin_flag = " [ADMIN]" if u.is_admin else ""
            logger.warning(f"   - {u.username} (id={u.id}){admin_flag}")

        # Always set must_change_password
        for u in legacy:
            if hasattr(u, "must_change_password"):
                u.must_change_password = True
        db.commit()
        logger.info("   Flagged must_change_password=True on all legacy users")

        # Offer interactive remediation
        print()
        print("   Options:")
        print("   [1] Keep users but force password change on next login (already done)")
        print("   [2] Delete legacy users entirely (destructive — loses their data)")
        print("   [3] Skip (handle manually later)")
        choice = input("   Choose [1/2/3] (default 1): ").strip()

        if choice == "2":
            for u in legacy:
                uname = u.username
                db.delete(u)
                logger.info(f"   Deleted user '{uname}'")
            db.commit()
            logger.info("✅ Legacy users deleted. Create new admin via BOOTSTRAP_ADMIN_* env vars.")
        else:
            logger.info("   Keeping legacy users. They MUST change passwords on next login.")
    except Exception as e:
        db.rollback()
        logger.error(f"❌ Error auditing legacy users: {e}")
    finally:
        db.close()


def security_audit_pickle_tokens():
    """Warn about and optionally delete legacy pickle token files."""
    import glob
    from pathlib import Path

    here = Path(__file__).parent
    stale = sorted(glob.glob(str(here / "yt_token_*.pickle")))
    if not stale:
        logger.info("✅ No legacy yt_token_*.pickle files found")
        return

    logger.warning(f"⚠️ Found {len(stale)} legacy pickle token file(s) (RCE risk):")
    for p in stale:
        logger.warning(f"   - {p}")

    answer = input("\n   Delete these pickle files? [y/N]: ").strip().lower()
    if answer == "y":
        import os

        for p in stale:
            os.remove(p)
            logger.info(f"   Deleted {p}")
        logger.info("✅ Legacy pickle tokens removed. Re-authenticate YouTube channels.")
    else:
        logger.info("   Skipped. Delete manually or re-authenticate to migrate to JSON tokens.")


def main():
    """Unified migration + security hardening workflow."""

    print()
    print("=" * 64)
    print("  DATABASE MIGRATION & SECURITY HARDENING")
    print("=" * 64)
    print()
    print("  This script is safe to run multiple times (idempotent).")
    print("  It will:")
    print("    1. Add missing columns to existing tables")
    print("    2. Create any new tables from the ORM models")
    print("    3. Encrypt any remaining plaintext data")
    print("    4. Migrate users to per-user encryption keychains")
    print("    5. Verify encryption is working")
    print("    6. Audit .master_key → env var migration")
    print("    7. Detect and remediate legacy default users")
    print("    8. Detect and clean up legacy pickle token files")
    print()

    response = input("  Continue? (type 'YES'): ")
    if response != "YES":
        print("  Cancelled.")
        return

    print()

    # ── PHASE 1: Schema migrations ──────────────────────────────────────
    try:
        logger.info("Phase 1/3: Schema migrations")
        logger.info("  Adding encryption columns...")
        add_encryption_columns()
        logger.info("  Adding user key columns...")
        add_user_key_columns()
        logger.info("  Creating new tables...")
        create_new_tables()
        print()
    except Exception as e:
        logger.exception(f"🔥 Schema migration failed: {e}")
        return

    # ── PHASE 2: Data migrations ────────────────────────────────────────
    try:
        logger.info("Phase 2/3: Data migrations")
        logger.info("  Encrypting existing plaintext data...")
        success = encrypt_existing_data()
        if not success:
            logger.error("❌ Data encryption failed!")
            return
        logger.info("  Migrating users to per-user keychains...")
        if not migrate_existing_users_to_keychains():
            logger.error("❌ User keychain migration failed!")
            return
        logger.info("  Verifying encryption...")
        verify_encryption()
        print()
    except Exception as e:
        logger.exception(f"🔥 Data migration failed: {e}")
        return

    # ── PHASE 3: Security hardening ─────────────────────────────────────
    logger.info("Phase 3/3: Security hardening")
    print()
    security_audit_master_key()
    print()
    security_audit_legacy_users()
    print()
    security_audit_pickle_tokens()

    print()
    print("=" * 64)
    print("  MIGRATION COMPLETE")
    print("=" * 64)
    print()
    print("  Next steps:")
    print("  1. Restart the app:  sudo systemctl restart transkript.service")
    print("  2. Verify existing data decrypts:  open the app, load a chat/transcription")
    print("  3. If you migrated .master_key → .env, delete the .master_key file")
    print("  4. Rotate any remaining legacy user passwords")
    print()