sshecret/packages/sshecret-backend/src/sshecret_backend/models.py

"""Database models.

TODO:

We might want to pass on audit information from the SSH server.
This might require some changes to these schemas.

"""

import enum
import logging
import uuid
from datetime import datetime

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


LOG = logging.getLogger(__name__)


class SubSystem(enum.StrEnum):
    """Available subsystems."""

    ADMIN = enum.auto()
    SSHD = enum.auto()
    BACKEND = enum.auto()
    TEST = enum.auto()


class Operation(enum.StrEnum):
    """Various operations for the audit logging module."""

    CREATE = enum.auto()
    READ = enum.auto()
    UPDATE = enum.auto()
    DELETE = enum.auto()
    DENY = enum.auto()
    PERMIT = enum.auto()
    LOGIN = enum.auto()
    REGISTER = enum.auto()
    NONE = enum.auto()


class Base(DeclarativeBase):
    pass


class Client(Base):
    """Clients."""

    __tablename__: str = "client"

    id: Mapped[uuid.UUID] = mapped_column(
        sa.Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    name: Mapped[str] = mapped_column(sa.String, unique=True)
    description: Mapped[str | None] = mapped_column(sa.String, nullable=True)
    public_key: Mapped[str] = mapped_column(sa.Text)

    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
    )

    updated_at: Mapped[datetime | None] = mapped_column(
        sa.DateTime(timezone=True),
        server_default=sa.func.now(),
        onupdate=sa.func.now(),
    )

    secrets: Mapped[list["ClientSecret"]] = relationship(
        back_populates="client", passive_deletes=True
    )

    policies: Mapped[list["ClientAccessPolicy"]] = relationship(back_populates="client")


class ClientAccessPolicy(Base):
    """Client access policies."""

    __tablename__: str = "client_access_policy"

    id: Mapped[uuid.UUID] = mapped_column(
        sa.Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    source: Mapped[str] = mapped_column(sa.String)
    client_id: Mapped[uuid.UUID | None] = mapped_column(
        sa.Uuid(as_uuid=True), sa.ForeignKey("client.id", ondelete="CASCADE")
    )
    client: Mapped[Client] = relationship(back_populates="policies")

    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
    )

    updated_at: Mapped[datetime | None] = mapped_column(
        sa.DateTime(timezone=True),
        server_default=sa.func.now(),
        onupdate=sa.func.now(),
    )


class ClientSecret(Base):
    """A client secret."""

    __tablename__: str = "client_secret"

    id: Mapped[uuid.UUID] = mapped_column(
        sa.Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    name: Mapped[str] = mapped_column(sa.String)
    description: Mapped[str | None] = mapped_column(sa.String, nullable=True)
    secret: Mapped[str] = mapped_column(sa.String)

    client_id: Mapped[uuid.UUID | None] = mapped_column(
        sa.Uuid(as_uuid=True), sa.ForeignKey("client.id", ondelete="CASCADE")
    )
    client: Mapped[Client] = relationship(back_populates="secrets")
    invalidated: Mapped[bool] = mapped_column(default=False)

    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
    )

    updated_at: Mapped[datetime | None] = mapped_column(
        sa.DateTime(timezone=True),
        server_default=sa.func.now(),
        onupdate=sa.func.now(),
    )


class APIClient(Base):
    """A client on the API.

    This should eventually get more granular permissions.
    """

    __tablename__: str = "api_client"

    id: Mapped[uuid.UUID] = mapped_column(
        sa.Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    subsystem: Mapped[SubSystem | None] = mapped_column(sa.String, nullable=True)
    token: Mapped[str] = mapped_column(sa.String)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
    )
    updated_at: Mapped[datetime | None] = mapped_column(
        sa.DateTime(timezone=True),
        server_default=sa.func.now(),
        onupdate=sa.func.now(),
    )


class AuditLog(Base):
    """Audit log.

    This is implemented without any foreign keys to avoid losing data on
    deletions.
    """

    __tablename__: str = "audit_log"
    id: Mapped[uuid.UUID] = mapped_column(
        sa.Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4
    )

    subsystem: Mapped[SubSystem] = mapped_column(sa.String)
    message: Mapped[str] = mapped_column(sa.String)
    operation: Mapped[Operation] = mapped_column(sa.String)
    client_id: Mapped[uuid.UUID | None] = mapped_column(
        sa.Uuid(as_uuid=True), nullable=True
    )
    data: Mapped[dict[str, str] | None] = mapped_column(sa.JSON, nullable=True)
    client_name: Mapped[str | None] = mapped_column(sa.String, nullable=True)
    secret_id: Mapped[uuid.UUID | None] = mapped_column(
        sa.Uuid(as_uuid=True), nullable=True
    )
    secret_name: Mapped[str | None] = mapped_column(sa.String, nullable=True)

    origin: Mapped[str | None] = mapped_column(sa.String, nullable=True)
    timestamp: Mapped[datetime] = mapped_column(
        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
    )


def init_db(engine: sa.Engine) -> None:
    """Initialize database."""
    Base.metadata.create_all(engine)


async def init_db_async(engine: AsyncEngine) -> None:
    """Initialize database."""
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)