mirror of
https://github.com/yogeshojha/rengine.git
synced 2026-10-18 14:27:17 +02:00
Yogesh Ojha
5764897fdf
An engine keeps its three intensity presets and now carries optional per-tool rate and concurrency in transport_overrides. Any value that differs from the preset shows as Custom; an empty field uses the preset. A context rate caps the engine value rather than replacing it, and rate is settable only for tools that have one. Edited through a Rates panel on the engine builder, validated and clamped, carried through create, update, duplicate, preview and launch. Claude-Session: https://claude.ai/code/session_01Bb4u8p34euie1YFzb7qCPu
214 lines
6.9 KiB
Python
214 lines
6.9 KiB
Python
import uuid
|
|
from datetime import datetime
|
|
from functools import partial
|
|
|
|
from pydantic import BaseModel, field_validator
|
|
from pydantic import Field as PydanticField
|
|
from sqlalchemy import Column
|
|
from sqlalchemy.types import JSON, Text
|
|
from sqlmodel import Field, SQLModel
|
|
|
|
from shared.definitions.intensity import clean_transport_overrides
|
|
from shared.enums.scan import Intensity
|
|
from shared.models.scan_context import ScanContextCreate
|
|
from shared.utils.datetime import utc_now
|
|
from shared.utils.validation import clean_name, clean_optional_name
|
|
|
|
|
|
class ScanEngine(SQLModel, table=True):
|
|
__tablename__ = "scan_engines"
|
|
|
|
id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
|
|
project_id: uuid.UUID = Field(foreign_key="projects.id", index=True)
|
|
created_by: uuid.UUID = Field(foreign_key="users.id")
|
|
name: str = Field(max_length=200)
|
|
description: str | None = Field(default=None, max_length=1000)
|
|
intensity: str = Field(default=Intensity.NORMAL.value)
|
|
global_headers: list = Field(
|
|
default_factory=list, sa_column=Column(JSON, nullable=False)
|
|
)
|
|
stages: dict = Field(default_factory=dict, sa_column=Column(JSON, nullable=False))
|
|
transport_overrides: dict = Field(
|
|
default_factory=dict, sa_column=Column(JSON, nullable=False)
|
|
)
|
|
yaml_source: str | None = Field(default=None, sa_column=Column(Text, nullable=True))
|
|
tool_options: dict = Field(
|
|
default_factory=dict, sa_column=Column(JSON, nullable=False)
|
|
)
|
|
builtin: bool = Field(default=False)
|
|
created_at: datetime = Field(default_factory=utc_now)
|
|
updated_at: datetime = Field(default_factory=utc_now)
|
|
last_used_at: datetime | None = Field(default=None)
|
|
|
|
|
|
class ScanEngineCreate(BaseModel):
|
|
name: str
|
|
description: str | None = None
|
|
intensity: str = Intensity.NORMAL.value
|
|
global_headers: list[str] = PydanticField(default_factory=list)
|
|
stages: dict[str, dict] = PydanticField(default_factory=dict)
|
|
transport_overrides: dict[str, dict] = PydanticField(default_factory=dict)
|
|
yaml_source: str | None = None
|
|
tool_options: dict[str, str] = PydanticField(default_factory=dict)
|
|
|
|
_validate_name = field_validator("name")(partial(clean_name, max_len=200))
|
|
_validate_transport = field_validator("transport_overrides")(
|
|
lambda cls, v: clean_transport_overrides(v) # noqa: ARG005
|
|
)
|
|
|
|
|
|
class ScanEngineUpdate(BaseModel):
|
|
name: str | None = None
|
|
description: str | None = None
|
|
intensity: str | None = None
|
|
global_headers: list[str] | None = None
|
|
stages: dict[str, dict] | None = None
|
|
transport_overrides: dict[str, dict] | None = None
|
|
yaml_source: str | None = None
|
|
tool_options: dict[str, str] | None = None
|
|
|
|
_validate_name = field_validator("name")(partial(clean_optional_name, max_len=200))
|
|
|
|
@field_validator("transport_overrides")
|
|
@classmethod
|
|
def _validate_transport(cls, value):
|
|
return None if value is None else clean_transport_overrides(value)
|
|
|
|
|
|
class EngineUsage(BaseModel):
|
|
schedules: int = 0
|
|
scans: int = 0
|
|
|
|
|
|
class ScanEngineRead(BaseModel):
|
|
id: uuid.UUID
|
|
project_id: uuid.UUID
|
|
created_by: uuid.UUID
|
|
name: str
|
|
description: str | None
|
|
intensity: str
|
|
global_headers: list[str]
|
|
stages: dict[str, dict]
|
|
transport_overrides: dict[str, dict] = PydanticField(default_factory=dict)
|
|
yaml_source: str | None
|
|
tool_options: dict[str, str] = PydanticField(default_factory=dict)
|
|
usage: EngineUsage = PydanticField(default_factory=EngineUsage)
|
|
builtin: bool = False
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
last_used_at: datetime | None
|
|
|
|
|
|
class StageField(BaseModel):
|
|
name: str
|
|
title: str
|
|
description: str | None = None
|
|
type: str
|
|
default: object = None
|
|
options: list[str] | None = None
|
|
option_labels: dict[str, str] | None = None
|
|
minimum: int | None = None
|
|
maximum: int | None = None
|
|
tier: str = "basic"
|
|
widget: str | None = None
|
|
kind: str | None = None
|
|
launch: bool = False
|
|
needs: str | None = None
|
|
unavailable_reason: str | None = None
|
|
|
|
|
|
class StageTransport(BaseModel):
|
|
tool: str
|
|
rates: dict[str, int | None]
|
|
threads: dict[str, int]
|
|
timeout: int
|
|
|
|
|
|
class StageCatalogEntry(BaseModel):
|
|
name: str
|
|
title: str
|
|
description: str
|
|
phase: str
|
|
level: int
|
|
applies_to: list[str]
|
|
tools: list[str]
|
|
api_keys: list[str]
|
|
requires_api_keys: bool = False
|
|
touches_target: bool = True
|
|
passive_capable: bool = False
|
|
always_on: bool = False
|
|
launch_fields: list[str] = PydanticField(default_factory=list)
|
|
group: str
|
|
role: str
|
|
consumes: list[str] = PydanticField(default_factory=list)
|
|
produces: list[str] = PydanticField(default_factory=list)
|
|
transport: StageTransport | None = None
|
|
defaults: dict
|
|
fields: list[StageField] = PydanticField(default_factory=list)
|
|
|
|
|
|
class StageGroupEntry(BaseModel):
|
|
key: str
|
|
label: str
|
|
|
|
|
|
class ToolOption(BaseModel):
|
|
name: str
|
|
label: str
|
|
phase: str
|
|
example: str
|
|
|
|
|
|
class EnginePreset(BaseModel):
|
|
name: str
|
|
title: str
|
|
description: str
|
|
intensity: str
|
|
stages: dict[str, dict]
|
|
|
|
|
|
class EngineCatalog(BaseModel):
|
|
phases: list[str]
|
|
stages: list[StageCatalogEntry]
|
|
rate_tools: list[str]
|
|
tool_options: list[ToolOption]
|
|
presets: list[EnginePreset]
|
|
target_types: list[str]
|
|
groups: list[StageGroupEntry] = PydanticField(default_factory=list)
|
|
seed_produces: dict[str, list[str]] = PydanticField(default_factory=dict)
|
|
|
|
|
|
class PreviewResolved(BaseModel):
|
|
header_names: list[str] = PydanticField(default_factory=list)
|
|
global_rate_limit_ceiling: int | None = None
|
|
per_tool_rate_limits: dict[str, int] = PydanticField(default_factory=dict)
|
|
preset_rates: dict[str, int] = PydanticField(default_factory=dict)
|
|
preset_threads: dict[str, int] = PydanticField(default_factory=dict)
|
|
excluded_subdomains: list[str] = PydanticField(default_factory=list)
|
|
excluded_paths: list[str] = PydanticField(default_factory=list)
|
|
excluded_ips: list[str] = PydanticField(default_factory=list)
|
|
included_subdomains: list[str] = PydanticField(default_factory=list)
|
|
follow_redirects: bool | None = None
|
|
http_protocol: str = "both"
|
|
|
|
|
|
class EnginePreviewResult(BaseModel):
|
|
phases: list = PydanticField(default_factory=list)
|
|
resolved_stages: dict[str, dict] = PydanticField(default_factory=dict)
|
|
resolved: PreviewResolved = PydanticField(default_factory=PreviewResolved)
|
|
warnings: list[str] = PydanticField(default_factory=list)
|
|
|
|
|
|
class EnginePreviewRequest(BaseModel):
|
|
target_type: str
|
|
context_id: uuid.UUID | None = None
|
|
context: ScanContextCreate | None = None
|
|
intensity: str = Intensity.NORMAL.value
|
|
stages: dict[str, dict] = PydanticField(default_factory=dict)
|
|
transport_overrides: dict[str, dict] = PydanticField(default_factory=dict)
|
|
|
|
@field_validator("transport_overrides")
|
|
@classmethod
|
|
def _validate_transport(cls, value):
|
|
return clean_transport_overrides(value)
|