"""Nessus scanner integration: multi-source vulns + scanner_type on schedules

Revision ID: 010
Revises: 009
Create Date: 2026-05-13 14:00:00.000000

Adds the schema needed to import Nessus findings alongside Wazuh while
deduplicating onto the existing (cve_id, asset_id) unique constraint:

- vulnerabilities.sources           — JSON array, e.g. ["wazuh","nessus","manual"]
- vulnerabilities.nessus_plugin_id  — Nessus plugin ID for the finding
- vulnerabilities.nessus_finding_uuid — stable per-finding identifier
- vulnerabilities.first_detected_by — which scanner first saw it
- vulnerabilities.cve_id widened    — Pseudo-CVE 'NESSUS-PLUGIN-{id}' for non-CVE
                                       findings (EOL/Compliance/Ciphers) exceeds 20 chars
- assets.nessus_host_uuid           — pin Nessus host after first hostname/IP match
- scan_schedules.scanner_type       — route scheduled jobs to wazuh or nessus

Backfill: every existing vulnerability is wazuh-sourced.
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


revision = '010'
down_revision = '009'
branch_labels = None
depends_on = None


def _columns(bind, table_name: str) -> set:
    """Set of existing column names for the given table."""
    return {c["name"] for c in sa.inspect(bind).get_columns(table_name)}


def _indexes(bind, table_name: str) -> set:
    """Set of existing index names for the given table."""
    return {ix["name"] for ix in sa.inspect(bind).get_indexes(table_name)}


def upgrade() -> None:
    """Idempotent — skips additions that already exist. Some columns may have
    been pre-created by Base.metadata.create_all() in development mode."""
    bind = op.get_bind()

    vuln_cols = _columns(bind, "vulnerabilities")
    asset_cols = _columns(bind, "assets")
    schedule_cols = _columns(bind, "scan_schedules")
    vuln_indexes = _indexes(bind, "vulnerabilities")
    asset_indexes = _indexes(bind, "assets")

    # --- Vulnerability multi-source columns ---
    if "sources" not in vuln_cols:
        op.add_column(
            "vulnerabilities",
            sa.Column("sources", sa.Text(), server_default="[]", nullable=False),
        )
    if "nessus_plugin_id" not in vuln_cols:
        op.add_column(
            "vulnerabilities",
            sa.Column("nessus_plugin_id", sa.String(length=20), nullable=True),
        )
    if "nessus_finding_uuid" not in vuln_cols:
        op.add_column(
            "vulnerabilities",
            sa.Column("nessus_finding_uuid", sa.String(length=64), nullable=True),
        )
    if "first_detected_by" not in vuln_cols:
        op.add_column(
            "vulnerabilities",
            sa.Column("first_detected_by", sa.String(length=20), nullable=True),
        )
    if "ix_vulnerabilities_nessus_plugin_id" not in vuln_indexes:
        op.create_index(
            "ix_vulnerabilities_nessus_plugin_id",
            "vulnerabilities",
            ["nessus_plugin_id"],
        )

    # cve_id widened to 50 chars for pseudo-CVE ids (NESSUS-PLUGIN-*).
    # ALTER COLUMN TYPE is a no-op when type already matches, but wrap in try
    # to be safe across earlier custom states.
    try:
        op.alter_column(
            "vulnerabilities",
            "cve_id",
            existing_type=sa.String(length=20),
            type_=sa.String(length=50),
            existing_nullable=False,
        )
    except Exception as e:  # noqa: BLE001
        # If a prior partial run already widened the column the above is a
        # no-op anyway; log and move on.
        import logging
        logging.getLogger("alembic.runtime.migration").info(
            "Skipping cve_id widen — already at target width or unsupported: %s", e
        )

    # --- Asset stable Nessus host pinning ---
    if "nessus_host_uuid" not in asset_cols:
        op.add_column(
            "assets",
            sa.Column("nessus_host_uuid", sa.String(length=64), nullable=True),
        )
    if "ix_assets_nessus_host_uuid" not in asset_indexes:
        op.create_index(
            "ix_assets_nessus_host_uuid",
            "assets",
            ["nessus_host_uuid"],
        )

    # --- ScanSchedule scanner_type routing ---
    if "scanner_type" not in schedule_cols:
        op.add_column(
            "scan_schedules",
            sa.Column(
                "scanner_type",
                sa.String(length=20),
                server_default="wazuh",
                nullable=False,
            ),
        )

    # --- Backfill ---
    # Mark every existing vulnerability as wazuh-sourced where sources is
    # still empty/default. Safe to run repeatedly.
    op.execute(
        """UPDATE vulnerabilities
           SET sources = '["wazuh"]', first_detected_by = 'wazuh'
           WHERE sources = '[]' OR sources IS NULL OR sources = ''"""
    )


def downgrade() -> None:
    op.drop_column('scan_schedules', 'scanner_type')

    op.drop_index('ix_assets_nessus_host_uuid', table_name='assets')
    op.drop_column('assets', 'nessus_host_uuid')

    op.alter_column(
        'vulnerabilities',
        'cve_id',
        existing_type=sa.String(length=50),
        type_=sa.String(length=20),
        existing_nullable=False,
    )

    op.drop_index('ix_vulnerabilities_nessus_plugin_id', table_name='vulnerabilities')
    op.drop_column('vulnerabilities', 'first_detected_by')
    op.drop_column('vulnerabilities', 'nessus_finding_uuid')
    op.drop_column('vulnerabilities', 'nessus_plugin_id')
    op.drop_column('vulnerabilities', 'sources')
