<gh_stars>1-10
"""support multiple certs per instance

Revision ID: 18f343c889ec
Revises: <PASSWORD>
Create Date: 2020-07-29 23:34:18.105973

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import orm
import sqlalchemy_utils
from sqlalchemy.dialects import postgresql
from sqlalchemy_utils.types.encrypted.encrypted_type import (
    AesGcmEngine,
    StringEncryptedType,
)
from sqlalchemy.ext.declarative import declarative_base

from broker.config import config_from_env


config = config_from_env()


def db_encryption_key():
    return config.DATABASE_ENCRYPTION_KEY


# revision identifiers, used by Alembic.
revision = "18f343c889ec"
down_revision = "e1793cdca02f"
branch_labels = None
depends_on = None

Base = declarative_base()


class ServiceInstance(Base):
    __tablename__ = "service_instance"
    id = sa.Column(sa.String(36), primary_key=True)
    domain_names = sa.Column(postgresql.JSONB, default=[])
    cert_pem = sa.Column(sa.Text)
    cert_expires_at = sa.Column(sa.TIMESTAMP(timezone=True))
    private_key_pem = sa.Column(
        StringEncryptedType(sa.Text, db_encryption_key, AesGcmEngine, "pkcs5")
    )
    fullchain_pem = sa.Column(sa.Text)
    iam_server_certificate_id = sa.Column(sa.String)
    iam_server_certificate_name = sa.Column(sa.String)
    iam_server_certificate_arn = sa.Column(sa.String)
    current_certificate_id = sa.Column(sa.Integer, sa.ForeignKey("certificate.id"))
    new_certificate_id = sa.Column(sa.Integer, sa.ForeignKey("certificate.id"))
    csr_pem = sa.Column(sa.Text)


class Certificate(Base):
    __tablename__ = "certificate"
    id = sa.Column(sa.Integer, primary_key=True)
    service_instance_id = sa.Column(
        sa.String, sa.ForeignKey("service_instance.id"), nullable=False
    )
    subject_alternative_names = sa.Column(postgresql.JSONB, default=[])
    leaf_pem = sa.Column(sa.Text)
    expires_at = sa.Column(sa.TIMESTAMP(timezone=True))
    private_key_pem = sa.Column(
        StringEncryptedType(sa.Text, db_encryption_key, AesGcmEngine, "pkcs5")
    )
    fullchain_pem = sa.Column(sa.Text)
    iam_server_certificate_id = sa.Column(sa.String)
    iam_server_certificate_name = sa.Column(sa.String)
    iam_server_certificate_arn = sa.Column(sa.String)
    csr_pem = sa.Column(sa.Text)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "certificate",
        sa.Column(
            "created_at",
            sa.TIMESTAMP(timezone=True),
            server_default=sa.text("now()"),
            nullable=True,
        ),
        sa.Column("updated_at", sa.TIMESTAMP(timezone=True), nullable=True),
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("service_instance_id", sa.String(), nullable=False),
        sa.Column(
            "subject_alternative_names",
            postgresql.JSONB(astext_type=sa.Text()),
            nullable=True,
        ),
        sa.Column("leaf_pem", sa.Text(), nullable=True),
        sa.Column("expires_at", sa.TIMESTAMP(timezone=True), nullable=True),
        sa.Column(
            "private_key_pem",
            sqlalchemy_utils.types.encrypted.encrypted_type.StringEncryptedType(),
            nullable=True,
        ),
        sa.Column("csr_pem", sa.Text(), nullable=True),
        sa.Column("fullchain_pem", sa.Text(), nullable=True),
        sa.Column("iam_server_certificate_id", sa.String(), nullable=True),
        sa.Column("iam_server_certificate_name", sa.String(), nullable=True),
        sa.Column("iam_server_certificate_arn", sa.String(), nullable=True),
        sa.ForeignKeyConstraint(["service_instance_id"], ["service_instance.id"]),
        sa.PrimaryKeyConstraint("id"),
    )
    op.add_column(
        "service_instance",
        sa.Column("current_certificate_id", sa.Integer(), nullable=True),
    )
    op.add_column(
        "service_instance", sa.Column("new_certificate_id", sa.Integer(), nullable=True)
    )
    op.create_foreign_key(
        "fk__service_instance__certificate__new_certificate_id",
        "service_instance",
        "certificate",
        ["new_certificate_id"],
        ["id"],
    )
    op.create_foreign_key(
        "fk__service_instance__certificate__current_certificate_id",
        "service_instance",
        "certificate",
        ["current_certificate_id"],
        ["id"],
    )
    bind = op.get_bind()
    session = orm.Session(bind=bind)

    for service_instance in session.query(ServiceInstance):
        cert = Certificate()
        cert.service_instance_id = service_instance.id
        cert.expires_at = service_instance.cert_expires_at
        cert.leaf_pem = service_instance.cert_pem
        cert.csr_pem = service_instance.csr_pem
        cert.private_key_pem = service_instance.private_key_pem
        cert.fullchain_pem = service_instance.fullchain_pem
        cert.iam_server_certificate_arn = service_instance.iam_server_certificate_arn
        cert.iam_server_certificate_id = service_instance.iam_server_certificate_id
        cert.iam_server_certificate_name = service_instance.iam_server_certificate_name
        cert.subject_alternative_names = service_instance.domain_names
        session.add(cert)
    session.commit()
    session = orm.Session(bind=bind)
    for certificate in session.query(Certificate):
        service_instance = session.query(ServiceInstance).get(
            certificate.service_instance_id
        )
        service_instance.current_certificate_id = certificate.id
        session.add(service_instance)
    session.commit()

    op.drop_column("service_instance", "iam_server_certificate_arn")
    op.drop_column("service_instance", "cert_expires_at")
    op.drop_column("service_instance", "csr_pem")
    op.drop_column("service_instance", "iam_server_certificate_id")
    op.drop_column("service_instance", "iam_server_certificate_name")
    op.drop_column("service_instance", "private_key_pem")
    op.drop_column("service_instance", "cert_pem")
    op.drop_column("service_instance", "fullchain_pem")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "service_instance",
        sa.Column("fullchain_pem", sa.TEXT(), autoincrement=False, nullable=True),
    )
    op.add_column(
        "service_instance",
        sa.Column("cert_pem", sa.TEXT(), autoincrement=False, nullable=True),
    )
    op.add_column(
        "service_instance",
        sa.Column("private_key_pem", sa.VARCHAR(), autoincrement=False, nullable=True),
    )
    op.add_column(
        "service_instance",
        sa.Column(
            "iam_server_certificate_name",
            sa.VARCHAR(),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.add_column(
        "service_instance",
        sa.Column(
            "iam_server_certificate_id",
            sa.VARCHAR(),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.add_column(
        "service_instance",
        sa.Column("csr_pem", sa.TEXT(), autoincrement=False, nullable=True),
    )
    op.add_column(
        "service_instance",
        sa.Column(
            "cert_expires_at",
            postgresql.TIMESTAMP(timezone=True),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.add_column(
        "service_instance",
        sa.Column(
            "iam_server_certificate_arn",
            sa.VARCHAR(),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.drop_constraint(
        "fk__service_instance__certificate__current_certificate_id",
        "service_instance",
        type_="foreignkey",
    )
    op.drop_constraint(
        "fk__service_instance__certificate__new_certificate_id",
        "service_instance",
        type_="foreignkey",
    )
    op.drop_column("service_instance", "new_certificate_id")
    op.drop_column("service_instance", "current_certificate_id")
    op.drop_table("certificate")
    # ### end Alembic commands ###
