# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Add malware detection tables

Revision ID: 061ff3d24c22
Revises: b5bb5d08543d
Create Date: 2019-12-18 17:27:00.183542
"""
import sqlalchemy as sa

from alembic import op
from sqlalchemy.dialects import postgresql

revision = "061ff3d24c22"
down_revision = "b5bb5d08543d"

MalwareCheckTypes = sa.Enum("event_hook", "scheduled", name="malwarechecktypes")

MalwareCheckStates = sa.Enum(
    "enabled", "evaluation", "disabled", "wiped_out", name="malwarecheckstate"
)

MalwareCheckObjectTypes = sa.Enum(
    "File", "Release", "Project", name="malwarecheckobjecttype"
)

VerdictClassifications = sa.Enum(
    "threat", "indeterminate", "benign", name="verdictclassification"
)
VerdictConfidences = sa.Enum("low", "medium", "high", name="verdictconfidence")


def upgrade():
    op.create_table(
        "malware_checks",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            server_default=sa.text("gen_random_uuid()"),
            nullable=False,
        ),
        sa.Column("name", postgresql.CITEXT(), nullable=False),
        sa.Column("version", sa.Integer(), default=1, nullable=False),
        sa.Column("short_description", sa.String(length=128), nullable=False),
        sa.Column("long_description", sa.Text(), nullable=False),
        sa.Column("check_type", MalwareCheckTypes, nullable=False),
        sa.Column("hooked_object", MalwareCheckObjectTypes, nullable=True),
        sa.Column("schedule", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column(
            "state", MalwareCheckStates, server_default="disabled", nullable=False
        ),
        sa.Column(
            "created", sa.DateTime(), server_default=sa.text("now()"), nullable=False
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("name", "version"),
    )
    op.create_table(
        "malware_verdicts",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            server_default=sa.text("gen_random_uuid()"),
            nullable=False,
        ),
        sa.Column(
            "run_date", sa.DateTime(), server_default=sa.text("now()"), nullable=False
        ),
        sa.Column("check_id", postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column("file_id", postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column("project_id", postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column("release_id", postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column("classification", VerdictClassifications, nullable=False),
        sa.Column("confidence", VerdictConfidences, nullable=False),
        sa.Column("message", sa.Text(), nullable=True),
        sa.Column("details", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column(
            "manually_reviewed",
            sa.Boolean(),
            server_default=sa.text("false"),
            nullable=False,
        ),
        sa.Column("reviewer_verdict", VerdictClassifications, nullable=True),
        sa.Column("full_report_link", sa.String(), nullable=True),
        sa.ForeignKeyConstraint(
            ["check_id"], ["malware_checks.id"], onupdate="CASCADE", ondelete="CASCADE"
        ),
        sa.ForeignKeyConstraint(["file_id"], ["release_files.id"]),
        sa.ForeignKeyConstraint(["release_id"], ["releases.id"]),
        sa.ForeignKeyConstraint(["project_id"], ["projects.id"]),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_malware_verdicts_check_id"),
        "malware_verdicts",
        ["check_id"],
        unique=False,
    )


def downgrade():
    op.drop_index(op.f("ix_malware_verdicts_check_id"), table_name="malware_verdicts")
    op.drop_table("malware_verdicts")
    op.drop_table("malware_checks")
    MalwareCheckTypes.drop(op.get_bind())
    MalwareCheckStates.drop(op.get_bind())
    MalwareCheckObjectTypes.drop(op.get_bind())
    VerdictClassifications.drop(op.get_bind())
    VerdictConfidences.drop(op.get_bind())
