precommit fixes

backend/alembic/env.py

@@ -32,8 +32,9 @@
 from app.incidents.models import Case
 from app.incidents.models import CaseAlertLink
 from app.incidents.models import Comment
+from app.incidents.models import CustomerCodeFieldName
 from app.incidents.models import FieldName
-from app.incidents.models import Notification, CustomerCodeFieldName
+from app.incidents.models import Notification
 from app.integrations.alert_creation_settings.models.alert_creation_settings import (
     AlertCreationSettings,
 )

backend/alembic/versions/7187bab836ab_add_alert_title_field_name_table.py

@@ -9,7 +9,6 @@
 from typing import Union
 
 import sqlalchemy as sa
-from sqlalchemy.dialects import mysql
 
 from alembic import op
 

backend/alembic/versions/8af27c21d10c_add_timefield_field_name_table.py

@@ -9,7 +9,6 @@
 from typing import Union
 
 import sqlalchemy as sa
-from sqlalchemy.dialects import mysql
 
 from alembic import op
 

backend/alembic/versions/aae5e3008a35_add_incident_management_notifications_.py

@@ -5,32 +5,34 @@
 Create Date: 2024-08-06 11:09:40.763589
 
 """
-from typing import Sequence, Union
+from typing import Sequence
+from typing import Union
 
-from alembic import op
 import sqlalchemy as sa
-from sqlalchemy.dialects import mysql
+
+from alembic import op
 
 # revision identifiers, used by Alembic.
-revision: str = 'aae5e3008a35'
-down_revision: Union[str, None] = '7f846a07c2b8'
+revision: str = "aae5e3008a35"
+down_revision: Union[str, None] = "7f846a07c2b8"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('incident_management_notification',
-    sa.Column('id', sa.Integer(), nullable=False),
-    sa.Column('customer_code', sa.String(length=50), nullable=False),
-    sa.Column('shuffle_workflow_id', sa.String(length=1000), nullable=False),
-    sa.Column('enabled', sa.Boolean(), nullable=False),
-    sa.PrimaryKeyConstraint('id')
+    op.create_table(
+        "incident_management_notification",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("customer_code", sa.String(length=50), nullable=False),
+        sa.Column("shuffle_workflow_id", sa.String(length=1000), nullable=False),
+        sa.Column("enabled", sa.Boolean(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
     )
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_table('incident_management_notification')
+    op.drop_table("incident_management_notification")
     # ### end Alembic commands ###

backend/alembic/versions/ba98bd0e3967_add_alert_tags.py

@@ -9,7 +9,6 @@
 from typing import Union
 
 import sqlalchemy as sa
-from sqlalchemy.dialects import mysql
 
 from alembic import op
 

backend/alembic/versions/e18dc3169b33_add_incident_management_customer_code_.py

@@ -5,31 +5,33 @@
 Create Date: 2024-08-15 08:04:02.006222
 
 """
-from typing import Sequence, Union
+from typing import Sequence
+from typing import Union
 
-from alembic import op
 import sqlalchemy as sa
-from sqlalchemy.dialects import mysql
+
+from alembic import op
 
 # revision identifiers, used by Alembic.
-revision: str = 'e18dc3169b33'
-down_revision: Union[str, None] = 'aae5e3008a35'
+revision: str = "e18dc3169b33"
+down_revision: Union[str, None] = "aae5e3008a35"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('incident_management_customercodefieldname',
-    sa.Column('id', sa.Integer(), nullable=False),
-    sa.Column('source', sa.String(length=50), nullable=False),
-    sa.Column('field_name', sa.String(length=100), nullable=False),
-    sa.PrimaryKeyConstraint('id')
+    op.create_table(
+        "incident_management_customercodefieldname",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("source", sa.String(length=50), nullable=False),
+        sa.Column("field_name", sa.String(length=100), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
     )
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_table('incident_management_customercodefieldname')
+    op.drop_table("incident_management_customercodefieldname")
     # ### end Alembic commands ###

backend/alembic/versions/e90f3760fa40_add_assigned_to_and_case_meta.py

@@ -9,7 +9,6 @@
 from typing import Union
 
 import sqlalchemy as sa
-from sqlalchemy.dialects import mysql
 
 from alembic import op
 

backend/app/connectors/graylog/services/collector.py

@@ -14,7 +14,6 @@
 from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse
 from app.connectors.graylog.utils.universal import send_get_request
 
-
 # async def get_indices_full() -> GraylogIndicesResponse:
 #     """Get indices from Graylog.
 
@@ -47,6 +46,7 @@
 #             message="Failed to collect indices",
 #         )
 
+
 async def get_indices_full() -> GraylogIndicesResponse:
     """Get indices from Graylog.
 
@@ -84,6 +84,7 @@ async def get_indices_full() -> GraylogIndicesResponse:
             message="Failed to collect indices",
         )
 
+
 async def fetch_configured_inputs() -> Tuple[bool, List[ConfiguredInput]]:
     """
     Fetches the configured inputs from the Graylog server.

backend/app/connectors/shuffle/routes/integrations.py

@@ -3,8 +3,10 @@
 from loguru import logger
 
 from app.auth.utils import AuthHandler
-from app.connectors.shuffle.schema.integrations import IntegrationRequest, ExecuteWorkflowRequest
-from app.connectors.shuffle.services.integrations import execute_integration, execute_workflow
+from app.connectors.shuffle.schema.integrations import ExecuteWorkflowRequest
+from app.connectors.shuffle.schema.integrations import IntegrationRequest
+from app.connectors.shuffle.services.integrations import execute_integration
+from app.connectors.shuffle.services.integrations import execute_workflow
 
 shuffle_integrations_router = APIRouter()
 
@@ -45,5 +47,3 @@ async def invoke_workflow_route(request: ExecuteWorkflowRequest):
     """
     logger.info("Executing workflow")
     return await execute_workflow(request)
-
-

backend/app/connectors/shuffle/schema/integrations.py

@@ -2,9 +2,11 @@
 from typing import Dict
 from typing import List
 from typing import Optional
+
 from fastapi import HTTPException
 from pydantic import BaseModel
-from pydantic import Field, root_validator
+from pydantic import Field
+from pydantic import root_validator
 
 
 class IntegrationRequest(BaseModel):
@@ -26,6 +28,7 @@ class IntegrationRequest(BaseModel):
         example=True,
     )
 
+
 class ExecuteWorkflowRequest(BaseModel):
     workflow_id: str = Field(..., description="The ID of the workflow", example="workflow_id")
     execution_arguments: Optional[Dict[str, Any]] = Field(
@@ -37,7 +40,7 @@ class ExecuteWorkflowRequest(BaseModel):
 
     @root_validator
     def check_customer_code(cls, values):
-        execution_arguments = values.get('execution_arguments', {})
-        if 'customer_code' not in execution_arguments or not execution_arguments['customer_code']:
+        execution_arguments = values.get("execution_arguments", {})
+        if "customer_code" not in execution_arguments or not execution_arguments["customer_code"]:
             raise HTTPException(status_code=400, detail="customer_code is required")
         return values

backend/app/connectors/shuffle/services/integrations.py

@@ -1,6 +1,7 @@
 from loguru import logger
 
-from app.connectors.shuffle.schema.integrations import IntegrationRequest, ExecuteWorkflowRequest
+from app.connectors.shuffle.schema.integrations import ExecuteWorkflowRequest
+from app.connectors.shuffle.schema.integrations import IntegrationRequest
 from app.connectors.shuffle.utils.universal import send_post_request
 
 
@@ -19,6 +20,7 @@ async def execute_integration(request: IntegrationRequest) -> dict:
     logger.info(f"Response: {response}")
     return response
 
+
 async def execute_workflow(request: ExecuteWorkflowRequest) -> dict:
     """
     Execute a workflow.

backend/app/connectors/velociraptor/routes/artifacts.py

@@ -12,8 +12,9 @@
 from app.connectors.velociraptor.schema.artifacts import ArtifactReccomendationAIRequest
 from app.connectors.velociraptor.schema.artifacts import ArtifactReccomendationRequest
 from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
-from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody, CollectFileBody
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
 from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
+from app.connectors.velociraptor.schema.artifacts import CollectFileBody
 from app.connectors.velociraptor.schema.artifacts import OSPrefixEnum
 from app.connectors.velociraptor.schema.artifacts import OSPrefixModel
 from app.connectors.velociraptor.schema.artifacts import QuarantineBody
@@ -23,7 +24,8 @@
 from app.connectors.velociraptor.services.artifacts import get_artifacts
 from app.connectors.velociraptor.services.artifacts import post_to_copilot_ai_module
 from app.connectors.velociraptor.services.artifacts import quarantine_host
-from app.connectors.velociraptor.services.artifacts import run_artifact_collection, run_file_collection
+from app.connectors.velociraptor.services.artifacts import run_artifact_collection
+from app.connectors.velociraptor.services.artifacts import run_file_collection
 from app.connectors.velociraptor.services.artifacts import run_remote_command
 from app.db.db_session import get_db
 from app.db.universal_models import Agents
@@ -477,6 +479,7 @@ async def get_artifact_recommendation(request: ArtifactReccomendationAIRequest):
         ),
     )
 
+
 # ! WIP ! #
 @velociraptor_artifacts_router.post(
     "/collect/file",
@@ -513,4 +516,3 @@ async def collect_file(collect_artifact_body: CollectFileBody, session: AsyncSes
     #     collect_artifact_body.hostname,
     # )
     return await run_file_collection(collect_artifact_body)
-

backend/app/connectors/velociraptor/schema/artifacts.py

@@ -4,9 +4,10 @@
 from typing import List
 from typing import Optional
 
-from pydantic import BaseModel
-from pydantic import Field, validator
 from fastapi import HTTPException
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import validator
 
 
 class Artifacts(BaseModel):
@@ -83,17 +84,18 @@ class CollectArtifactBody(BaseBody):
         description="Name of the artifact for collection or command running",
     )
 
+
 class CollectFileBody(BaseBody):
     artifact_name: str = Field(
-        'Generic.Collectors.File',
+        "Generic.Collectors.File",
         description="Name of the artifact for collection or command running",
     )
-    file: str = Field('Glob\nUsers\\Administrator\\Documents\\*\n', description="File to collect")
-    root_disk: Optional[str] = Field('C:', description="Root disk to collect from")
+    file: str = Field("Glob\nUsers\\Administrator\\Documents\\*\n", description="File to collect")
+    root_disk: Optional[str] = Field("C:", description="Root disk to collect from")
 
-    @validator('artifact_name')
+    @validator("artifact_name")
     def validate_artifact_name(cls, value):
-        if value != 'Generic.Collectors.File':
+        if value != "Generic.Collectors.File":
             raise HTTPException(status_code=400, detail="Invalid artifact name. Name should be 'Generic.Collectors.File'")
         return value
 

backend/app/connectors/velociraptor/services/artifacts.py

@@ -7,7 +7,8 @@
 from app.connectors.velociraptor.schema.artifacts import Artifacts
 from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
 from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
-from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse, CollectFileBody
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
+from app.connectors.velociraptor.schema.artifacts import CollectFileBody
 from app.connectors.velociraptor.schema.artifacts import QuarantineBody
 from app.connectors.velociraptor.schema.artifacts import QuarantineResponse
 from app.connectors.velociraptor.schema.artifacts import RunCommandBody
@@ -164,6 +165,7 @@ async def run_artifact_collection(
             detail=f"Failed to run artifact collection on {collect_artifact_body}: {err}",
         )
 
+
 async def run_file_collection(
     collect_artifact_body: CollectFileBody,
 ) -> CollectArtifactResponse:
@@ -237,6 +239,7 @@ async def run_file_collection(
             detail=f"Failed to run artifact collection on {collect_artifact_body}: {err}",
         )
 
+
 async def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResponse:
     """
     Run a remote command on a client.

backend/app/connectors/wazuh_indexer/models/sigma.py

@@ -1,15 +1,7 @@
 from datetime import datetime
-from typing import Dict
-from typing import List
 from typing import Optional
-from uuid import uuid4
 
-from sqlalchemy import ForeignKey
-from sqlalchemy import PrimaryKeyConstraint
-from sqlmodel import JSON
-from sqlmodel import Column
 from sqlmodel import Field
-from sqlmodel import Relationship
 from sqlmodel import SQLModel
 from sqlmodel import Text
 

backend/app/connectors/wazuh_indexer/routes/alerts.py

@@ -7,20 +7,21 @@
 from loguru import logger
 
 from app.auth.utils import AuthHandler
-from app.incidents.schema.alert_collection import AlertsPayload
 from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
 from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHostResponse
 from app.connectors.wazuh_indexer.schema.alerts import AlertsByRuleResponse
 from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchBody
-from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchResponse, GraylogAlertsSearchBody
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchResponse
+from app.connectors.wazuh_indexer.schema.alerts import GraylogAlertsSearchBody
 from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchBody
 from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchResponse
 from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchBody
 from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchResponse
 from app.connectors.wazuh_indexer.services.alerts import get_alerts
 from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_host
-from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule, get_graylog_alerts
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule
 from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule_per_host
+from app.connectors.wazuh_indexer.services.alerts import get_graylog_alerts
 from app.connectors.wazuh_indexer.services.alerts import get_host_alerts
 from app.connectors.wazuh_indexer.services.alerts import get_index_alerts
 from app.connectors.wazuh_indexer.utils.universal import collect_indices
@@ -198,6 +199,7 @@ async def get_all_alerts_by_rule_per_host(
     logger.info("Fetching number of all alerts for all rules per host")
     return await get_alerts_by_rule_per_host(alerts_search_body)
 
+
 @wazuh_indexer_alerts_router.post(
     "/alerts/graylog",
     response_model=AlertsSearchResponse,