SneakyScan/app/migrations/versions/009_remove_cidrs.py

"""Remove CIDR table - make sites IP-only

Revision ID: 009
Revises: 008
Create Date: 2025-11-19

This migration removes the SiteCIDR table entirely, making sites purely
IP-based. CIDRs are now only used as a convenience for bulk IP addition,
not stored as permanent entities.

Changes:
- Set all site_ips.site_cidr_id to NULL (preserve all IPs)
- Drop foreign key from site_ips to site_cidrs
- Drop site_cidrs table
- Remove site_cidr_id column from site_ips

All existing IPs are preserved. They become "standalone" IPs without
a CIDR parent.
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import text


# revision identifiers, used by Alembic
revision = '009'
down_revision = '008'
branch_labels = None
depends_on = None


def upgrade():
    """
    Remove CIDR table and make all IPs standalone.
    """

    connection = op.get_bind()
    inspector = sa.inspect(connection)

    print("\n=== Migration 009: Remove CIDR Table ===\n")

    # Get counts before migration
    try:
        total_cidrs = connection.execute(text('SELECT COUNT(*) FROM site_cidrs')).scalar()
        total_ips = connection.execute(text('SELECT COUNT(*) FROM site_ips')).scalar()
        ips_with_cidr = connection.execute(text(
            'SELECT COUNT(*) FROM site_ips WHERE site_cidr_id IS NOT NULL'
        )).scalar()

        print(f"Before migration:")
        print(f"  - Total CIDRs: {total_cidrs}")
        print(f"  - Total IPs: {total_ips}")
        print(f"  - IPs linked to CIDRs: {ips_with_cidr}")
        print(f"  - Standalone IPs: {total_ips - ips_with_cidr}\n")
    except Exception as e:
        print(f"Could not get pre-migration stats: {e}\n")

    # Step 1: Set all site_cidr_id to NULL (preserve all IPs as standalone)
    print("Step 1: Converting all IPs to standalone (nulling CIDR associations)...")
    try:
        result = connection.execute(text("""
            UPDATE site_ips
            SET site_cidr_id = NULL
            WHERE site_cidr_id IS NOT NULL
        """))
        print(f"  ✓ Converted {result.rowcount} IPs to standalone\n")
    except Exception as e:
        print(f"  ⚠ Error or already done: {e}\n")

    # Step 2: Drop foreign key constraint from site_ips to site_cidrs
    print("Step 2: Dropping foreign key constraint from site_ips to site_cidrs...")
    foreign_keys = inspector.get_foreign_keys('site_ips')
    fk_to_drop = None

    for fk in foreign_keys:
        if fk['referred_table'] == 'site_cidrs':
            fk_to_drop = fk['name']
            break

    if fk_to_drop:
        try:
            op.drop_constraint(fk_to_drop, 'site_ips', type_='foreignkey')
            print(f"  ✓ Dropped foreign key constraint: {fk_to_drop}\n")
        except Exception as e:
            print(f"  ⚠ Could not drop foreign key: {e}\n")
    else:
        print("  ⚠ Foreign key constraint not found or already dropped\n")

    # Step 3: Drop index on site_cidr_id (if exists)
    print("Step 3: Dropping index on site_cidr_id...")
    indexes = inspector.get_indexes('site_ips')
    index_to_drop = None

    for idx in indexes:
        if 'site_cidr_id' in idx['column_names']:
            index_to_drop = idx['name']
            break

    if index_to_drop:
        try:
            op.drop_index(index_to_drop, table_name='site_ips')
            print(f"  ✓ Dropped index: {index_to_drop}\n")
        except Exception as e:
            print(f"  ⚠ Could not drop index: {e}\n")
    else:
        print("  ⚠ Index not found or already dropped\n")

    # Step 4: Drop site_cidrs table
    print("Step 4: Dropping site_cidrs table...")
    tables = inspector.get_table_names()

    if 'site_cidrs' in tables:
        try:
            op.drop_table('site_cidrs')
            print("  ✓ Dropped site_cidrs table\n")
        except Exception as e:
            print(f"  ⚠ Could not drop table: {e}\n")
    else:
        print("  ⚠ Table site_cidrs not found or already dropped\n")

    # Step 5: Drop site_cidr_id column from site_ips
    print("Step 5: Dropping site_cidr_id column from site_ips...")
    site_ips_columns = [col['name'] for col in inspector.get_columns('site_ips')]

    if 'site_cidr_id' in site_ips_columns:
        try:
            op.drop_column('site_ips', 'site_cidr_id')
            print("  ✓ Dropped site_cidr_id column from site_ips\n")
        except Exception as e:
            print(f"  ⚠ Could not drop column: {e}\n")
    else:
        print("  ⚠ Column site_cidr_id not found or already dropped\n")

    # Get counts after migration
    try:
        final_ips = connection.execute(text('SELECT COUNT(*) FROM site_ips')).scalar()
        total_sites = connection.execute(text('SELECT COUNT(*) FROM sites')).scalar()

        print("After migration:")
        print(f"  - Total sites: {total_sites}")
        print(f"  - Total IPs (all standalone): {final_ips}")
        print(f"  - CIDRs: N/A (table removed)")
    except Exception as e:
        print(f"Could not get post-migration stats: {e}")

    print("\n✓ Migration 009 complete: Sites are now IP-only")
    print("  All IPs preserved as standalone. CIDRs can still be used")
    print("  via the API/UI for bulk IP creation, but are not stored.\n")


def downgrade():
    """
    Recreate site_cidrs table (CANNOT restore original CIDR associations).

    WARNING: This downgrade creates an empty site_cidrs table structure but
    cannot restore the original CIDR-to-IP associations since that data was
    deleted. All IPs will remain standalone.
    """

    connection = op.get_bind()

    print("\n=== Downgrade 009: Recreate CIDR Table Structure ===\n")
    print("⚠ WARNING: Cannot restore original CIDR associations!")
    print("  The site_cidrs table structure will be recreated but will be empty.")
    print("  All IPs will remain standalone. This is a PARTIAL downgrade.\n")

    # Step 1: Recreate site_cidrs table (empty)
    print("Step 1: Recreating site_cidrs table structure...")
    try:
        op.create_table(
            'site_cidrs',
            sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
            sa.Column('site_id', sa.Integer(), nullable=False),
            sa.Column('cidr', sa.String(length=45), nullable=False, comment='CIDR notation (e.g., 10.0.0.0/24)'),
            sa.Column('created_at', sa.DateTime(), nullable=False),
            sa.PrimaryKeyConstraint('id'),
            sa.ForeignKeyConstraint(['site_id'], ['sites.id'], ),
            sa.UniqueConstraint('site_id', 'cidr', name='uix_site_cidr')
        )
        print("  ✓ Recreated site_cidrs table (empty)\n")
    except Exception as e:
        print(f"  ⚠ Could not create table: {e}\n")

    # Step 2: Add site_cidr_id column back to site_ips (nullable)
    print("Step 2: Adding site_cidr_id column back to site_ips...")
    try:
        op.add_column('site_ips', sa.Column('site_cidr_id', sa.Integer(), nullable=True, comment='FK to site_cidrs (optional, for grouping)'))
        print("  ✓ Added site_cidr_id column (nullable)\n")
    except Exception as e:
        print(f"  ⚠ Could not add column: {e}\n")

    # Step 3: Add foreign key constraint
    print("Step 3: Adding foreign key constraint...")
    try:
        op.create_foreign_key('fk_site_ips_site_cidr_id', 'site_ips', 'site_cidrs', ['site_cidr_id'], ['id'])
        print("  ✓ Created foreign key constraint\n")
    except Exception as e:
        print(f"  ⚠ Could not create foreign key: {e}\n")

    # Step 4: Add index on site_cidr_id
    print("Step 4: Adding index on site_cidr_id...")
    try:
        op.create_index('ix_site_ips_site_cidr_id', 'site_ips', ['site_cidr_id'], unique=False)
        print("  ✓ Created index on site_cidr_id\n")
    except Exception as e:
        print(f"  ⚠ Could not create index: {e}\n")

    print("✓ Downgrade complete: CIDR table structure restored (but empty)")
    print("  All IPs remain standalone. You would need to manually recreate")
    print("  CIDR records and associate IPs with them.\n")