step-ca-webui/core/api_server.py

from typing import List, Union

import uvicorn
from fastapi import FastAPI, Query, HTTPException, Request
from fastapi.responses import PlainTextResponse

from core.certificate_manager import CertificateManager
from core.trace_id_handler import TraceIdHandler
from shared.api_models import (
    CertificateDTO,
    CertificateGenerateRequest,
    CertificateGenerateResult,
    CertificateRenewResult,
    CertificateRevokeResult,
    CommandPreviewDTO,
    LogEntryDTO,
    LogsRequest,
    CommandInfoDTO,
)
from shared.logger import Logger, LogsFilter, Paging
from shared.models import LogSeverity

_default_response = {
    500: {
        "description": "Internal Server Error",
        "content": {"text/plain": {"example": "An unexpected error occurred"}},
    }
}


# noinspection PyPep8Naming
class APIServer:
    def __init__(
        self,
        cert_manager: CertificateManager,
        logger: Logger,
        version: str,
        port: int,
        prod_url: str = None,
    ):
        self._cert_manager = cert_manager
        self._logger = logger
        self._port = port
        self.App = FastAPI(
            title="Step-CA Management API",
            version=version,
            description="API for managing step-ca Certificate Authority",
            openapi_url="/openapi.json",
            docs_url="/swagger",
            redoc_url="/redoc",
            servers=(
                [
                    {
                        "url": f"http://localhost:{port}",
                        "description": "Local development environment",
                    },
                ]
                + [{"url": prod_url, "description": "Production environment"}]
                if prod_url
                else []
            ),
        )

        self._setup_routes()
        self._setup_handlers()

    def run(self):
        uvicorn.run(self.App, host="0.0.0.0", port=self._port)

    def _setup_routes(self):
        @self.App.get(
            "/certificates",
            response_model=Union[List[CertificateDTO], CommandPreviewDTO],
            responses=_default_response,
        )
        async def list_certificates(
            preview: bool = Query(...),
        ) -> Union[List[CertificateDTO], CommandPreviewDTO]:
            if preview:
                command = self._cert_manager.preview_list_certificates()
                return CommandPreviewDTO(command=command)
            certs = self._cert_manager.list_certificates()

            return [
                CertificateDTO(
                    id=cert.id,
                    name=cert.name,
                    status=cert.status,
                    expirationDate=cert.expiration_date,
                )
                for cert in certs
            ]

        @self.App.post(
            "/certificates/generate",
            response_model=Union[CertificateGenerateResult, CommandPreviewDTO],
            responses=_default_response,
        )
        async def generate_certificate(
            cert_request: CertificateGenerateRequest, preview: bool = Query(...)
        ):
            if preview:
                command = self._cert_manager.preview_generate_certificate(
                    cert_request.keyName, cert_request.keyType, cert_request.duration
                )
                return CommandPreviewDTO(command=command)
            cert = self._cert_manager.generate_certificate(
                cert_request.keyName, cert_request.keyType, cert_request.duration
            )

            return CertificateGenerateResult(
                success=cert.success,
                message=cert.message,
                logEntryId=cert.log_entry_id,
                certificateId=cert.certificate_id,
                certificateName=cert.certificate_name,
                expirationDate=cert.expiration_date,
            )

        @self.App.post(
            "/certificates/renew",
            response_model=Union[CertificateRenewResult, CommandPreviewDTO],
            responses=_default_response,
        )
        async def renew_certificate(
            certId: str = Query(...),
            duration: int = Query(..., description="Duration in seconds"),
            preview: bool = Query(...),
        ):
            if preview:
                command = self._cert_manager.preview_renew_certificate(certId, duration)
                return CommandPreviewDTO(command=command)
            cert = self._cert_manager.renew_certificate(certId, duration)

            return CertificateRenewResult(
                success=cert.success,
                message=cert.message,
                logEntryId=cert.log_entry_id,
                certificateId=cert.certificate_id,
                newExpirationDate=cert.new_expiration_date,
            )

        @self.App.post(
            "/certificates/revoke",
            response_model=Union[CertificateRevokeResult, CommandPreviewDTO],
            responses=_default_response,
        )
        async def revoke_certificate(
            certId: str = Query(...), preview: bool = Query(...)
        ):
            if preview:
                command = self._cert_manager.preview_revoke_certificate(certId)
                return CommandPreviewDTO(command=command)
            cert = self._cert_manager.revoke_certificate(certId)

            return CertificateRevokeResult(
                success=cert.success,
                message=cert.message,
                logEntryId=cert.log_entry_id,
                certificateId=cert.certificate_id,
                revocationDate=cert.revocation_date,
            )

        @self.App.get(
            "/logs/single", response_model=LogEntryDTO, responses=_default_response
        )
        async def get_log_entry(logId: int = Query(..., gt=0)):
            log_entry = self._logger.get_log_entry(logId)
            if not log_entry:
                raise HTTPException(status_code=404, detail="Log entry not found")

            return LogEntryDTO(
                entryId=log_entry.entry_id,
                timestamp=log_entry.timestamp,
                severity=log_entry.severity,
                message=log_entry.message,
                traceId=log_entry.trace_id,
                commandInfo=(
                    CommandInfoDTO(
                        command=log_entry.command_info.command,
                        output=log_entry.command_info.output,
                        exitCode=log_entry.command_info.exit_code,
                        action=log_entry.command_info.action,
                    )
                    if log_entry.command_info
                    else None
                ),
            )

        @self.App.post(
            "/logs", response_model=List[LogEntryDTO], responses=_default_response
        )
        async def get_logs(logs_request: LogsRequest) -> List[LogEntryDTO]:
            logs = self._logger.get_logs(
                LogsFilter(
                    trace_id=logs_request.traceId,
                    commands_only=logs_request.commandsOnly,
                    severity=logs_request.severity,
                ),
                Paging(page=logs_request.page, page_size=logs_request.pageSize),
            )

            return [
                LogEntryDTO(
                    entryId=log.entry_id,
                    timestamp=log.timestamp,
                    severity=log.severity,
                    message=log.message,
                    traceId=log.trace_id,
                    commandInfo=(
                        CommandInfoDTO(
                            command=log.command_info.command,
                            output=log.command_info.output,
                            exitCode=log.command_info.exit_code,
                            action=log.command_info.action,
                        )
                        if log.command_info
                        else None
                    ),
                )
                for log in logs
            ]

    def _setup_handlers(self):
        @self.App.exception_handler(HTTPException)
        async def custom_http_exception_handler(request: Request, exc: HTTPException):
            return PlainTextResponse(str(exc.detail), status_code=exc.status_code)

        @self.App.exception_handler(Exception)
        async def handle_all_exceptions(request: Request, exc: Exception):
            trace_id = getattr(request.state, "trace_id", "UNKNOWN")
            self._logger.log(
                LogSeverity.ERROR, f"Unhandled exception, trace_id [{trace_id}]: {exc}"
            )
            return PlainTextResponse(
                f"Internal server error, trace_id [{trace_id}]", status_code=500
            )

        @self.App.middleware("http")
        async def setup_logging_scope(request: Request, call_next):
            async with TraceIdHandler.logging_scope():
                response = await call_next(request)
            return response