kitos/Aegis
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat(refactor): PEP8, type annotations, docstrings and PyJWT security fix

Browse Source
This commit is contained in:
kitos
2026-06-11 11:09:41 +02:00
parent 98fddccd32 1f19bd8432
commit 8fea0c1ada
161 changed files with 15318 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/routers/test_templates.py
+259 -20
View File
@@ -22,35 +22,69 @@ Filters (GET /test-templates)
- offset / limit: pagination (default limit=50)
"""
# Import uuid
import uuid
# Import Optional from typing
from typing import Optional
# Import APIRouter, Depends, Query, status from fastapi
from fastapi import APIRouter, Depends, Query, status
# Import Session from sqlalchemy.orm
from sqlalchemy.orm import Session
# Import get_db from app.database
from app.database import get_db
# Import get_current_user, require_any_role from app.dependencies.auth
from app.dependencies.auth import get_current_user, require_any_role
# Import UnitOfWork from app.domain.unit_of_work
from app.domain.unit_of_work import UnitOfWork
from app.models.technique import Technique
from app.models.user import User
# Import from app.schemas.test_template
from app.schemas.test_template import (
TestTemplateCreate,
TestTemplateOut,
TestTemplateSummary,
)
# Import log_action from app.services.audit_service
from app.services.audit_service import log_action
# Import from app.services.test_template_service
from app.services.test_template_service import (
bulk_activate,
create_template as create_template_svc,
get_template_or_raise,
get_template_stats,
get_templates_by_technique as templates_by_technique,
list_templates,
soft_delete_template,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
create_template as create_template_svc,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
get_templates_by_technique as templates_by_technique,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
toggle_template_active as toggle_template_active_svc,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
update_template as update_template_svc,
)
# Assign router = APIRouter(prefix="/test-templates", tags=["test-templates"])
router = APIRouter(prefix="/test-templates", tags=["test-templates"])
@@ -60,28 +94,64 @@ router = APIRouter(prefix="/test-templates", tags=["test-templates"])
@router.get("", response_model=list[TestTemplateSummary])
# Define function _list_templates_handler
def _list_templates_handler(
# Entry: source
source: Optional[str] = Query(None, description="Filter by source (atomic_red_team, mitre, custom)"),
# Entry: platform
platform: Optional[str] = Query(None, description="Filter by platform (windows, linux, macos)"),
# Entry: severity
severity: Optional[str] = Query(None, description="Filter by severity (low, medium, high, critical)"),
# Entry: mitre_technique_id
mitre_technique_id: Optional[str] = Query(None, description="Filter by MITRE technique ID"),
# Entry: search
search: Optional[str] = Query(None, description="Search in name and description"),
# Entry: is_active
is_active: Optional[bool] = Query(None, description="Filter by active status (true/false). Omit to return all."),
# Entry: offset
offset: int = Query(0, ge=0),
# Entry: limit
limit: int = Query(50, ge=1, le=200),
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
):
"""Return a paginated, filterable list of test templates."""
) -> list:
"""Return a paginated, filterable list of test templates.
Args:
source (Optional[str]): Filter by source (``atomic_red_team``, ``mitre``, ``custom``).
platform (Optional[str]): Filter by platform (``windows``, ``linux``, ``macos``).
severity (Optional[str]): Filter by severity (``low``, ``medium``, ``high``, ``critical``).
mitre_technique_id (Optional[str]): Filter by MITRE technique ID string.
search (Optional[str]): Full-text search across name and description.
is_active (Optional[bool]): Filter by active status; omit to return all.
offset (int): Number of records to skip for pagination.
limit (int): Maximum number of records to return.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
list: Serialised list of :class:`TestTemplateSummary` objects.
"""
# Return list_templates(
return list_templates(
db,
# Keyword argument: source
source=source,
# Keyword argument: platform
platform=platform,
# Keyword argument: severity
severity=severity,
# Keyword argument: mitre_technique_id
mitre_technique_id=mitre_technique_id,
# Keyword argument: search
search=search,
# Keyword argument: is_active
is_active=is_active,
# Keyword argument: offset
offset=offset,
# Keyword argument: limit
limit=limit,
)
@@ -92,11 +162,23 @@ def _list_templates_handler(
@router.get("/stats")
# Define function template_stats
def template_stats(
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Return catalog statistics: active, by_source, by_platform."""
) -> dict:
"""Return catalog statistics: active, by_source, by_platform.
Args:
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
dict: Counts of active templates broken down by source and platform.
"""
# Return get_template_stats(db)
return get_template_stats(db)
@@ -106,27 +188,53 @@ def template_stats(
@router.patch("/bulk-activate")
# Define function bulk_activate_templates
def bulk_activate_templates(
# Entry: activate
activate: bool = Query(True, description="True to activate all, False to deactivate all"),
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Set all templates to active or inactive."""
) -> dict:
"""Set all templates to active or inactive.
Args:
activate (bool): ``True`` to activate all templates, ``False`` to deactivate all.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
dict: Confirmation message with ``affected`` count and the applied ``is_active`` flag.
"""
# Assign count = bulk_activate(db, activate=activate)
count = bulk_activate(db, activate=activate)
# Open context manager
with UnitOfWork(db) as uow:
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="bulk_activate_templates" if activate else "bulk_deactivate_templates",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=None,
# Keyword argument: details
details={"affected": count, "is_active": activate},
)
# Call uow.commit()
uow.commit()
# Return {
return {
# Literal argument value
"detail": f"{'Activated' if activate else 'Deactivated'} {count} templates",
# Literal argument value
"affected": count,
# Literal argument value
"is_active": activate,
}
@@ -137,12 +245,26 @@ def bulk_activate_templates(
@router.get("/by-technique/{mitre_id}", response_model=list[TestTemplateSummary])
# Define function _templates_by_technique_handler
def _templates_by_technique_handler(
# Entry: mitre_id
mitre_id: str,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
):
"""Return all active templates mapped to a specific MITRE technique."""
) -> list:
"""Return all active templates mapped to a specific MITRE technique.
Args:
mitre_id (str): MITRE ATT&CK technique ID (e.g. ``T1059.001``).
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
list: Serialised list of :class:`TestTemplateSummary` objects for the technique.
"""
# Return templates_by_technique(db, mitre_id)
return templates_by_technique(db, mitre_id)
@@ -152,12 +274,26 @@ def _templates_by_technique_handler(
@router.get("/{template_id}", response_model=TestTemplateOut)
# Define function get_template
def get_template(
# Entry: template_id
template_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
):
"""Return full details for a single test template."""
) -> TestTemplateOut:
"""Return full details for a single test template.
Args:
template_id (uuid.UUID): Primary key of the template to retrieve.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
TestTemplateOut: Full template detail including all fields.
"""
# Return get_template_or_raise(db, template_id)
return get_template_or_raise(db, template_id)
@@ -167,17 +303,35 @@ def get_template(
@router.post(
# Literal argument value
"",
# Keyword argument: response_model
response_model=TestTemplateOut,
# Keyword argument: status_code
status_code=status.HTTP_201_CREATED,
)
# Define function create_template
def create_template(
# Entry: payload
payload: TestTemplateCreate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Create a custom test template."""
) -> TestTemplateOut:
"""Create a custom test template.
Args:
payload (TestTemplateCreate): All fields for the new template.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead creating the template.
Returns:
TestTemplateOut: The newly created template with all fields populated.
"""
# Assign template = create_template_svc(db, **payload.model_dump())
template = create_template_svc(db, **payload.model_dump())
# Open context manager
with UnitOfWork(db) as uow:
# Flag the associated technique for review — new template available
if template.mitre_technique_id:
@@ -190,19 +344,30 @@ def create_template(
technique.review_required = True
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="create_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={
# Literal argument value
"name": template.name,
# Literal argument value
"source": template.source,
# Literal argument value
"mitre_technique_id": template.mitre_technique_id,
},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(template)
# Return template
return template
@@ -212,26 +377,52 @@ def create_template(
@router.patch("/{template_id}", response_model=TestTemplateOut)
# Define function update_template
def update_template(
# Entry: template_id
template_id: uuid.UUID,
# Entry: payload
payload: TestTemplateCreate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Update fields of an existing test template."""
) -> TestTemplateOut:
"""Update fields of an existing test template.
Args:
template_id (uuid.UUID): Primary key of the template to update.
payload (TestTemplateCreate): Fields to update; only set fields are applied.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead updating the template.
Returns:
TestTemplateOut: The updated template with refreshed field values.
"""
# Assign template = update_template_svc(db, template_id, **payload.model_dump(exclude_u...
template = update_template_svc(db, template_id, **payload.model_dump(exclude_unset=True))
# Open context manager
with UnitOfWork(db) as uow:
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="update_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={"updated_fields": list(payload.model_dump(exclude_unset=True).keys())},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(template)
# Return template
return template
@@ -241,25 +432,49 @@ def update_template(
@router.patch("/{template_id}/toggle-active", response_model=TestTemplateOut)
# Define function toggle_template_active
def toggle_template_active(
# Entry: template_id
template_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Toggle a template between active and inactive (is_active = not is_active)."""
) -> TestTemplateOut:
"""Toggle a template between active and inactive (is_active = not is_active).
Args:
template_id (uuid.UUID): Primary key of the template to toggle.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
TestTemplateOut: The template with the updated ``is_active`` flag.
"""
# Assign template = toggle_template_active_svc(db, template_id)
template = toggle_template_active_svc(db, template_id)
# Open context manager
with UnitOfWork(db) as uow:
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="toggle_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={"name": template.name, "is_active": template.is_active},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(template)
# Return template
return template
@@ -269,23 +484,47 @@ def toggle_template_active(
@router.delete("/{template_id}", status_code=status.HTTP_200_OK)
# Define function delete_template
def delete_template(
# Entry: template_id
template_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Soft-delete a test template by setting ``is_active=False``."""
) -> dict:
"""Soft-delete a test template by setting ``is_active=False``.
Args:
template_id (uuid.UUID): Primary key of the template to delete.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
dict: Confirmation message with key ``detail``.
"""
# Assign template = get_template_or_raise(db, template_id)
template = get_template_or_raise(db, template_id)
# Call soft_delete_template()
soft_delete_template(db, template_id)
# Open context manager
with UnitOfWork(db) as uow:
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="delete_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={"name": template.name},
)
# Call uow.commit()
uow.commit()
# Return {"detail": "Test template deactivated"}
return {"detail": "Test template deactivated"}