#!/usr/bin/env python3 """ Database migration script for TTS Arena analytics columns and new security features. Usage: python migrate.py database.db python migrate.py instance/tts_arena.db """ import click import sqlite3 import sys import os from pathlib import Path def check_column_exists(cursor, table_name, column_name): """Check if a column exists in a table.""" cursor.execute(f"PRAGMA table_info({table_name})") columns = [row[1] for row in cursor.fetchall()] return column_name in columns def check_table_exists(cursor, table_name): """Check if a table exists in the database.""" cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,)) return cursor.fetchone() is not None def create_timeout_and_campaign_tables(cursor): """Create the new timeout and campaign tables.""" tables_created = [] # Create coordinated_voting_campaign table if not check_table_exists(cursor, "coordinated_voting_campaign"): cursor.execute(""" CREATE TABLE coordinated_voting_campaign ( id INTEGER PRIMARY KEY, model_id VARCHAR(100) NOT NULL, model_type VARCHAR(20) NOT NULL, detected_at DATETIME DEFAULT CURRENT_TIMESTAMP, time_window_hours INTEGER NOT NULL, vote_count INTEGER NOT NULL, user_count INTEGER NOT NULL, confidence_score REAL NOT NULL, status VARCHAR(20) DEFAULT 'active', admin_notes TEXT, resolved_by INTEGER, resolved_at DATETIME, FOREIGN KEY (model_id) REFERENCES model (id), FOREIGN KEY (resolved_by) REFERENCES user (id) ) """) tables_created.append("coordinated_voting_campaign") click.echo("✅ Created table 'coordinated_voting_campaign'") else: click.echo("⏭️ Table 'coordinated_voting_campaign' already exists, skipping") # Create campaign_participant table if not check_table_exists(cursor, "campaign_participant"): cursor.execute(""" CREATE TABLE campaign_participant ( id INTEGER PRIMARY KEY, campaign_id INTEGER NOT NULL, user_id INTEGER NOT NULL, votes_in_campaign INTEGER NOT NULL, first_vote_at DATETIME NOT NULL, last_vote_at DATETIME NOT NULL, suspicion_level VARCHAR(20) NOT NULL, FOREIGN KEY (campaign_id) REFERENCES coordinated_voting_campaign (id), FOREIGN KEY (user_id) REFERENCES user (id) ) """) tables_created.append("campaign_participant") click.echo("✅ Created table 'campaign_participant'") else: click.echo("⏭️ Table 'campaign_participant' already exists, skipping") # Create user_timeout table if not check_table_exists(cursor, "user_timeout"): cursor.execute(""" CREATE TABLE user_timeout ( id INTEGER PRIMARY KEY, user_id INTEGER NOT NULL, reason VARCHAR(500) NOT NULL, timeout_type VARCHAR(50) NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP, expires_at DATETIME NOT NULL, created_by INTEGER, is_active BOOLEAN DEFAULT 1, cancelled_at DATETIME, cancelled_by INTEGER, cancel_reason VARCHAR(500), related_campaign_id INTEGER, FOREIGN KEY (user_id) REFERENCES user (id), FOREIGN KEY (created_by) REFERENCES user (id), FOREIGN KEY (cancelled_by) REFERENCES user (id), FOREIGN KEY (related_campaign_id) REFERENCES coordinated_voting_campaign (id) ) """) tables_created.append("user_timeout") click.echo("✅ Created table 'user_timeout'") else: click.echo("⏭️ Table 'user_timeout' already exists, skipping") return tables_created def add_analytics_columns_and_tables(db_path): """Add analytics columns and create new security tables.""" if not os.path.exists(db_path): click.echo(f"❌ Database file not found: {db_path}", err=True) return False try: # Connect to the database conn = sqlite3.connect(db_path) cursor = conn.cursor() # Check if vote table exists cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='vote'") if not cursor.fetchone(): click.echo("❌ Vote table not found in database", err=True) return False # Define the columns to add to vote table vote_columns_to_add = [ ("session_duration_seconds", "REAL"), ("ip_address_partial", "VARCHAR(20)"), ("user_agent", "VARCHAR(500)"), ("generation_date", "DATETIME"), ("cache_hit", "BOOLEAN") ] # Define the columns to add to user table user_columns_to_add = [ ("hf_account_created", "DATETIME") ] added_columns = [] skipped_columns = [] # Add vote table columns click.echo("📊 Processing vote table columns...") for column_name, column_type in vote_columns_to_add: if check_column_exists(cursor, "vote", column_name): skipped_columns.append(f"vote.{column_name}") click.echo(f"⏭️ Column 'vote.{column_name}' already exists, skipping") else: try: cursor.execute(f"ALTER TABLE vote ADD COLUMN {column_name} {column_type}") added_columns.append(f"vote.{column_name}") click.echo(f"✅ Added column 'vote.{column_name}' ({column_type})") except sqlite3.Error as e: click.echo(f"❌ Failed to add column 'vote.{column_name}': {e}", err=True) conn.rollback() return False # Add user table columns click.echo("👤 Processing user table columns...") for column_name, column_type in user_columns_to_add: if check_column_exists(cursor, "user", column_name): skipped_columns.append(f"user.{column_name}") click.echo(f"⏭️ Column 'user.{column_name}' already exists, skipping") else: try: cursor.execute(f"ALTER TABLE user ADD COLUMN {column_name} {column_type}") added_columns.append(f"user.{column_name}") click.echo(f"✅ Added column 'user.{column_name}' ({column_type})") except sqlite3.Error as e: click.echo(f"❌ Failed to add column 'user.{column_name}': {e}", err=True) conn.rollback() return False # Create new security tables click.echo("🔒 Creating security and timeout management tables...") tables_created = create_timeout_and_campaign_tables(cursor) # Commit the changes conn.commit() conn.close() # Summary if added_columns: click.echo(f"\n🎉 Successfully added {len(added_columns)} analytics columns:") for col in added_columns: click.echo(f" • {col}") if skipped_columns: click.echo(f"\n⏭️ Skipped {len(skipped_columns)} existing columns:") for col in skipped_columns: click.echo(f" • {col}") if tables_created: click.echo(f"\n🔒 Successfully created {len(tables_created)} security tables:") for table in tables_created: click.echo(f" • {table}") if not added_columns and not skipped_columns and not tables_created: click.echo("❌ No columns or tables were processed") return False click.echo(f"\n✨ Migration completed successfully!") if tables_created: click.echo("\n🚨 New Security Features Enabled:") click.echo(" • Automatic coordinated voting campaign detection") click.echo(" • User timeout management") click.echo(" • Admin panels for security monitoring") click.echo("\nNew admin panel sections:") click.echo(" • /admin/timeouts - Manage user timeouts") click.echo(" • /admin/campaigns - View coordinated voting campaigns") return True except sqlite3.Error as e: click.echo(f"❌ Database error: {e}", err=True) return False except Exception as e: click.echo(f"❌ Unexpected error: {e}", err=True) return False @click.command() @click.argument('database_path', type=click.Path()) @click.option('--dry-run', is_flag=True, help='Show what would be done without making changes') @click.option('--backup', is_flag=True, help='Create a backup before migration') def migrate(database_path, dry_run, backup): """ Add analytics columns and security tables to the TTS Arena database. DATABASE_PATH: Path to the SQLite database file (e.g., instance/tts_arena.db) """ click.echo("🚀 TTS Arena Migration Tool") click.echo("Analytics + Security Features") click.echo("=" * 40) # Resolve the database path db_path = Path(database_path).resolve() click.echo(f"📁 Database: {db_path}") if not db_path.exists(): click.echo(f"❌ Database file not found: {db_path}", err=True) sys.exit(1) # Create backup if requested if backup: backup_path = db_path.with_suffix(f"{db_path.suffix}.backup") try: import shutil shutil.copy2(db_path, backup_path) click.echo(f"💾 Backup created: {backup_path}") except Exception as e: click.echo(f"❌ Failed to create backup: {e}", err=True) sys.exit(1) if dry_run: click.echo("\n🔍 DRY RUN MODE - No changes will be made") click.echo("\nThe following columns would be added to the 'vote' table:") click.echo(" • session_duration_seconds (REAL)") click.echo(" • ip_address_partial (VARCHAR(20))") click.echo(" • user_agent (VARCHAR(500))") click.echo(" • generation_date (DATETIME)") click.echo(" • cache_hit (BOOLEAN)") click.echo("\nThe following columns would be added to the 'user' table:") click.echo(" • hf_account_created (DATETIME)") click.echo("\nThe following security tables would be created:") click.echo(" • coordinated_voting_campaign - Track detected voting campaigns") click.echo(" • campaign_participant - Track users involved in campaigns") click.echo(" • user_timeout - Manage user timeouts/bans") click.echo("\nRun without --dry-run to apply changes.") return # Confirm before proceeding if not click.confirm(f"\n⚠️ This will modify the database at {db_path}. Continue?"): click.echo("❌ Migration cancelled") sys.exit(0) # Perform the migration click.echo("\n🔧 Starting migration...") success = add_analytics_columns_and_tables(str(db_path)) if success: click.echo("\n🎊 Migration completed successfully!") click.echo("You can now restart your TTS Arena application to use all new features.") else: click.echo("\n💥 Migration failed!") sys.exit(1) if __name__ == "__main__": migrate()