77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from core.plugin_types import EmptyPluginConfig
|
|
from core.schemas.results import AuditResults
|
|
|
|
if TYPE_CHECKING:
|
|
from core.schemas.config import ResolvedPluginConfig
|
|
from reporting.models import AuditReport, RenderedReport
|
|
|
|
TReportTransport = TypeVar("TReportTransport", bound=BaseModel)
|
|
TWorkerConfig = TypeVar("TWorkerConfig", bound=BaseModel)
|
|
|
|
|
|
class BaseWorker(ABC):
|
|
config_model: type[BaseModel] = EmptyPluginConfig
|
|
|
|
def __init__(self, config: BaseModel | dict[str, Any] | None = None) -> None:
|
|
"""
|
|
Config dict will be provided on initialization, providing settings, parsed from config.yaml
|
|
"""
|
|
self.raw_config = config
|
|
self.config: BaseModel | dict[str, Any]
|
|
if isinstance(config, BaseModel):
|
|
self.config = config.model_dump(mode="python")
|
|
else:
|
|
self.config = config or {}
|
|
|
|
@abstractmethod
|
|
def run(self) -> AuditResults:
|
|
"""
|
|
Executes the scanning.
|
|
Method should return AuditResults object, otherwise the output won't be processed.
|
|
"""
|
|
...
|
|
|
|
|
|
class BaseRenderer(ABC):
|
|
content_type: str
|
|
media_type: str
|
|
file_extension: str
|
|
aliases: tuple[str, ...] = ()
|
|
config_model: type[BaseModel] = EmptyPluginConfig
|
|
|
|
@abstractmethod
|
|
def render(self, report: AuditReport, config: BaseModel) -> RenderedReport: ...
|
|
|
|
|
|
class BaseTransport(ABC):
|
|
transport_type: str
|
|
config_model: type[BaseModel] = EmptyPluginConfig
|
|
|
|
@abstractmethod
|
|
def publish(self, rendered_report: RenderedReport, transport: ResolvedPluginConfig) -> str: ...
|
|
|
|
|
|
class TypedTransport(BaseTransport, Generic[TReportTransport], ABC):
|
|
config_model: type[TReportTransport]
|
|
|
|
def publish(self, rendered_report: RenderedReport, transport: ResolvedPluginConfig) -> str:
|
|
if not isinstance(transport.config, self.config_model):
|
|
raise TypeError(
|
|
f"Transport '{self.transport_type}' expected {self.config_model.__name__}, "
|
|
f"got {type(transport.config).__name__}"
|
|
)
|
|
|
|
return self.publish_typed(rendered_report, cast(TReportTransport, transport.config))
|
|
|
|
@abstractmethod
|
|
def publish_typed(
|
|
self, rendered_report: RenderedReport, transport: TReportTransport
|
|
) -> str: ...
|