From c599b707aee65b5c49f5ea028cb8daa1f26eba74 Mon Sep 17 00:00:00 2001 From: taylor_socfortress <111797488+taylorwalton@users.noreply.github.com> Date: Fri, 9 Feb 2024 10:18:03 -0600 Subject: [PATCH] Precommit fixes (#137) * precommits * precommit adjustments --- .pre-commit-config.yaml | 2 +- .../app/agents/dfir_iris/services/cases.py | 5 +- backend/app/agents/routes/agents.py | 33 ++--- backend/app/agents/schema/agents.py | 4 +- backend/app/agents/services/modify.py | 3 +- backend/app/agents/services/status.py | 19 +-- backend/app/agents/services/sync.py | 13 +- .../app/agents/velociraptor/schema/agents.py | 3 +- .../agents/velociraptor/services/agents.py | 9 +- backend/app/agents/wazuh/schema/agents.py | 6 +- backend/app/agents/wazuh/services/agents.py | 18 +-- .../agents/wazuh/services/vulnerabilities.py | 9 +- backend/app/auth/models/users.py | 8 +- backend/app/auth/routes/auth.py | 32 +++-- backend/app/auth/schema/user.py | 3 +- backend/app/auth/services/universal.py | 7 +- backend/app/auth/utils.py | 13 +- .../app/connectors/cortex/routes/analyzers.py | 19 +-- .../app/connectors/cortex/schema/analyzers.py | 10 +- .../connectors/cortex/services/analyzers.py | 23 ++-- .../app/connectors/cortex/utils/universal.py | 10 +- .../app/connectors/dfir_iris/routes/alerts.py | 40 +++--- .../app/connectors/dfir_iris/routes/assets.py | 8 +- .../app/connectors/dfir_iris/routes/cases.py | 46 +++---- .../app/connectors/dfir_iris/routes/notes.py | 19 +-- .../app/connectors/dfir_iris/routes/users.py | 25 ++-- .../app/connectors/dfir_iris/schema/admin.py | 7 +- .../app/connectors/dfir_iris/schema/alerts.py | 8 +- .../app/connectors/dfir_iris/schema/assets.py | 6 +- .../app/connectors/dfir_iris/schema/cases.py | 11 +- .../app/connectors/dfir_iris/schema/notes.py | 7 +- .../connectors/dfir_iris/services/alerts.py | 31 ++--- .../connectors/dfir_iris/services/assets.py | 11 +- .../connectors/dfir_iris/services/cases.py | 36 +++-- .../connectors/dfir_iris/services/notes.py | 25 ++-- .../connectors/dfir_iris/services/users.py | 8 +- .../connectors/dfir_iris/utils/universal.py | 15 ++- .../event_shipper/utils/universal.py | 1 + .../connectors/grafana/routes/dashboards.py | 13 +- .../connectors/grafana/schema/dashboards.py | 8 +- .../connectors/grafana/services/dashboards.py | 34 ++--- .../app/connectors/grafana/utils/universal.py | 10 +- .../connectors/graylog/routes/collector.py | 26 ++-- .../app/connectors/graylog/routes/events.py | 17 +-- .../connectors/graylog/routes/management.py | 52 ++++---- .../connectors/graylog/routes/monitoring.py | 22 ++- .../connectors/graylog/routes/pipelines.py | 44 +++--- .../app/connectors/graylog/routes/streams.py | 6 +- .../connectors/graylog/schema/collector.py | 7 +- .../app/connectors/graylog/schema/events.py | 5 +- .../connectors/graylog/schema/monitoring.py | 6 +- .../connectors/graylog/schema/pipelines.py | 3 +- .../app/connectors/graylog/schema/streams.py | 3 +- .../connectors/graylog/services/collector.py | 41 +++--- .../app/connectors/graylog/services/events.py | 35 ++--- .../connectors/graylog/services/management.py | 33 +++-- .../connectors/graylog/services/monitoring.py | 33 ++--- .../connectors/graylog/services/pipelines.py | 34 ++--- .../connectors/graylog/services/streams.py | 16 +-- .../app/connectors/graylog/utils/universal.py | 9 +- .../app/connectors/influxdb/routes/alerts.py | 6 +- .../connectors/influxdb/services/alerts.py | 14 +- .../connectors/influxdb/utils/universal.py | 8 +- backend/app/connectors/models.py | 7 +- backend/app/connectors/routes.py | 27 ++-- backend/app/connectors/schema.py | 3 +- backend/app/connectors/services.py | 28 ++-- .../connectors/shuffle/routes/workflows.py | 21 ++- .../connectors/shuffle/schema/workflows.py | 8 +- .../connectors/shuffle/services/workflows.py | 11 +- .../app/connectors/shuffle/utils/universal.py | 13 +- .../app/connectors/sublime/models/alerts.py | 7 +- .../app/connectors/sublime/routes/alerts.py | 20 +-- .../app/connectors/sublime/schema/alerts.py | 6 +- .../app/connectors/sublime/services/alerts.py | 27 ++-- .../app/connectors/sublime/utils/universal.py | 7 +- backend/app/connectors/utils.py | 9 +- .../velociraptor/routes/artifacts.py | 50 ++++--- .../connectors/velociraptor/routes/flows.py | 18 ++- .../velociraptor/schema/artifacts.py | 8 +- .../connectors/velociraptor/schema/flows.py | 7 +- .../velociraptor/services/artifacts.py | 30 ++--- .../connectors/velociraptor/services/flows.py | 13 +- .../velociraptor/utils/universal.py | 24 ++-- .../connectors/wazuh_indexer/routes/alerts.py | 42 +++--- .../wazuh_indexer/routes/monitoring.py | 25 ++-- .../connectors/wazuh_indexer/schema/alerts.py | 13 +- .../wazuh_indexer/schema/indices.py | 3 +- .../wazuh_indexer/schema/monitoring.py | 7 +- .../wazuh_indexer/services/alerts.py | 66 ++++----- .../wazuh_indexer/services/monitoring.py | 42 +++--- .../wazuh_indexer/utils/universal.py | 28 ++-- .../connectors/wazuh_manager/models/rules.py | 3 +- .../connectors/wazuh_manager/routes/rules.py | 27 ++-- .../connectors/wazuh_manager/schema/rules.py | 6 +- .../wazuh_manager/services/rules.py | 29 ++-- .../wazuh_manager/utils/universal.py | 10 +- .../routes/decommission.py | 12 +- .../customer_provisioning/routes/provision.py | 33 +++-- .../schema/decommission.py | 3 +- .../customer_provisioning/schema/grafana.py | 3 +- .../customer_provisioning/schema/graylog.py | 6 +- .../customer_provisioning/schema/provision.py | 8 +- .../schema/wazuh_worker.py | 3 +- .../services/decommission.py | 22 ++- .../services/dfir_iris.py | 12 +- .../customer_provisioning/services/grafana.py | 19 ++- .../customer_provisioning/services/graylog.py | 25 ++-- .../services/provision.py | 57 +++----- .../services/wazuh_manager.py | 3 +- backend/app/customers/routes/customers.py | 49 +++---- backend/app/customers/schema/customers.py | 6 +- backend/app/db/db_populate.py | 33 ++--- backend/app/db/db_session.py | 13 +- backend/app/db/db_setup.py | 19 ++- backend/app/db/universal_models.py | 41 ++---- .../app/healthchecks/agents/routes/agents.py | 35 ++--- .../app/healthchecks/agents/schema/agents.py | 11 +- .../healthchecks/agents/services/agents.py | 47 +++---- .../alert_creation/general/routes/alert.py | 27 ++-- .../alert_creation/general/schema/alert.py | 11 +- .../alert_creation/general/services/alert.py | 49 +++---- .../general/services/alert_multi_exclude.py | 20 +-- .../alert_creation/office365/routes/alert.py | 30 +++-- .../office365/schema/exchange.py | 11 +- .../office365/schema/threat_intel.py | 9 +- .../office365/services/exchange.py | 27 ++-- .../office365/services/threat_intel.py | 27 ++-- .../models/alert_creation_settings.py | 7 +- .../routes/alert_creation_settings.py | 38 +++--- .../schema/alert_creation_settings.py | 3 +- .../alert_escalation/routes/general_alert.py | 15 ++- .../alert_escalation/schema/general_alert.py | 11 +- .../services/general_alert.py | 46 +++---- .../ask_socfortress/routes/ask_socfortress.py | 12 +- .../ask_socfortress/schema/ask_socfortress.py | 3 +- .../services/ask_socfortress.py | 23 ++-- .../integrations/dnstwist/routes/analyze.py | 13 +- .../integrations/dnstwist/schema/analyze.py | 6 +- .../integrations/dnstwist/services/analyze.py | 7 +- .../integrations/mimecast/routes/mimecast.py | 28 ++-- .../integrations/mimecast/routes/provision.py | 12 +- .../integrations/mimecast/schema/mimecast.py | 24 ++-- .../integrations/mimecast/schema/provision.py | 9 +- .../mimecast/services/mimecast.py | 23 ++-- .../mimecast/services/provision.py | 49 +++---- .../models/customer_integration_settings.py | 7 +- .../models/monitoring_alert.py | 3 +- .../routes/monitoring_alert.py | 26 ++-- .../monitoring_alert/routes/provision.py | 24 ++-- .../schema/monitoring_alert.py | 13 +- .../monitoring_alert/schema/provision.py | 8 +- .../monitoring_alert/services/provision.py | 42 ++++-- .../monitoring_alert/services/suricata.py | 56 ++++---- .../monitoring_alert/services/wazuh.py | 52 ++++---- .../monitoring_alert/utils/db_operations.py | 3 +- .../office365/routes/provision.py | 25 ++-- .../office365/schema/provision.py | 7 +- .../office365/services/provision.py | 126 +++++++----------- backend/app/integrations/routes.py | 93 ++++++------- backend/app/integrations/schema.py | 6 +- backend/app/integrations/utils/alerts.py | 9 +- backend/app/integrations/utils/collection.py | 4 +- .../app/integrations/utils/event_shipper.py | 14 +- backend/app/integrations/utils/schema.py | 7 +- backend/app/integrations/utils/utils.py | 6 +- backend/app/middleware/exception_handlers.py | 13 +- backend/app/middleware/logger.py | 18 +-- backend/app/routers/agents.py | 3 +- backend/app/routers/alert_creation.py | 3 +- .../app/routers/alert_creation_settings.py | 3 +- backend/app/routers/ask_socfortress.py | 3 +- backend/app/routers/auth.py | 3 +- backend/app/routers/connectors.py | 3 +- backend/app/routers/cortex.py | 3 +- backend/app/routers/customer_provisioning.py | 3 +- backend/app/routers/customers.py | 3 +- backend/app/routers/dfir_iris.py | 3 +- backend/app/routers/dnstwist.py | 3 +- backend/app/routers/grafana.py | 3 +- backend/app/routers/graylog.py | 3 +- backend/app/routers/healthcheck.py | 3 +- backend/app/routers/influxdb.py | 3 +- backend/app/routers/integrations.py | 3 +- backend/app/routers/logs.py | 3 +- backend/app/routers/mimecast.py | 3 +- backend/app/routers/monitoring_alert.py | 3 +- backend/app/routers/office365.py | 3 +- backend/app/routers/scheduler.py | 3 +- backend/app/routers/shuffle.py | 3 +- backend/app/routers/smtp.py | 3 +- backend/app/routers/sublime.py | 3 +- backend/app/routers/threat_intel.py | 3 +- backend/app/routers/velociraptor.py | 3 +- backend/app/routers/wazuh_indexer.py | 3 +- backend/app/routers/wazuh_manager.py | 3 +- backend/app/schedulers/models/scheduler.py | 3 +- backend/app/schedulers/routes/scheduler.py | 10 +- backend/app/schedulers/scheduler.py | 31 ++--- backend/app/schedulers/services/agent_sync.py | 7 +- .../schedulers/services/invoke_mimecast.py | 29 ++-- .../schedulers/services/monitoring_alert.py | 26 ++-- backend/app/schedulers/utils/universal.py | 3 +- backend/app/smtp/routes/configure.py | 9 +- backend/app/smtp/routes/reports.py | 9 +- backend/app/smtp/services/create_report.py | 6 +- backend/app/smtp/services/reports.py | 3 +- .../app/threat_intel/routes/socfortress.py | 16 ++- .../app/threat_intel/schema/socfortress.py | 3 +- .../app/threat_intel/services/socfortress.py | 18 +-- backend/app/utils.py | 63 ++++----- backend/copilot.py | 93 +++++++------ git_tasks/Vagrantfile | 3 - git_tasks/pyproject.toml => pyproject.toml | 0 214 files changed, 1875 insertions(+), 1786 deletions(-) delete mode 100644 git_tasks/Vagrantfile rename git_tasks/pyproject.toml => pyproject.toml (100%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a109ba45..c614869e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,7 +10,7 @@ repos: - id: check-json - id: trailing-whitespace - id: check-added-large-files - - id: detect-private-key + #- id: detect-private-key # Not using since uploading nginx key - id: requirements-txt-fixer args: ["backend/requirements.txt", "backend/requirements.in"] diff --git a/backend/app/agents/dfir_iris/services/cases.py b/backend/app/agents/dfir_iris/services/cases.py index 814cb8e0..9cb060a6 100644 --- a/backend/app/agents/dfir_iris/services/cases.py +++ b/backend/app/agents/dfir_iris/services/cases.py @@ -1,10 +1,11 @@ from typing import List +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.agents.dfir_iris.schema.cases import AssetCaseIDResponse from app.connectors.dfir_iris.services.assets import get_case_assets from app.connectors.dfir_iris.services.cases import get_all_cases -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession async def collect_agent_soc_cases( diff --git a/backend/app/agents/routes/agents.py b/backend/app/agents/routes/agents.py index 4ff8dba0..6cf12022 100644 --- a/backend/app/agents/routes/agents.py +++ b/backend/app/agents/routes/agents.py @@ -1,15 +1,21 @@ +from fastapi import APIRouter +from fastapi import BackgroundTasks +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy import delete +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.agents.dfir_iris.services.cases import collect_agent_soc_cases -from app.agents.schema.agents import ( - AgentModifyResponse, - AgentsResponse, - OutdatedVelociraptorAgentsResponse, - OutdatedWazuhAgentsResponse, - SyncedAgentsResponse, -) -from app.agents.services.status import ( - get_outdated_agents_velociraptor, - get_outdated_agents_wazuh, -) +from app.agents.schema.agents import AgentModifyResponse +from app.agents.schema.agents import AgentsResponse +from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse +from app.agents.schema.agents import OutdatedWazuhAgentsResponse +from app.agents.schema.agents import SyncedAgentsResponse +from app.agents.services.status import get_outdated_agents_velociraptor +from app.agents.services.status import get_outdated_agents_wazuh from app.agents.services.sync import sync_agents from app.agents.velociraptor.services.agents import delete_agent_velociraptor from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse @@ -23,11 +29,6 @@ # App specific imports # from app.db.db_session import session from app.db.universal_models import Agents -from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy import delete -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select agents_router = APIRouter() diff --git a/backend/app/agents/schema/agents.py b/backend/app/agents/schema/agents.py index c4769ed2..def6dc44 100644 --- a/backend/app/agents/schema/agents.py +++ b/backend/app/agents/schema/agents.py @@ -1,9 +1,11 @@ from typing import List +from pydantic import BaseModel +from pydantic import Field + from app.agents.velociraptor.schema.agents import VelociraptorAgent from app.agents.wazuh.schema.agents import WazuhAgent from app.db.universal_models import Agents -from pydantic import BaseModel, Field class AgentsResponse(BaseModel): diff --git a/backend/app/agents/services/modify.py b/backend/app/agents/services/modify.py index fc544f1b..38bb4c29 100644 --- a/backend/app/agents/services/modify.py +++ b/backend/app/agents/services/modify.py @@ -1,7 +1,8 @@ +from fastapi import HTTPException + import app.agents.wazuh.services.agents as wazuh_services from app.db.db_session import session from app.db.universal_models import Agents -from fastapi import HTTPException def delete_agent_db(agent_id: str): diff --git a/backend/app/agents/services/status.py b/backend/app/agents/services/status.py index b213d2f5..d4b278c6 100644 --- a/backend/app/agents/services/status.py +++ b/backend/app/agents/services/status.py @@ -1,17 +1,16 @@ from typing import List -from app.agents.schema.agents import ( - OutdatedVelociraptorAgentsResponse, - OutdatedWazuhAgentsResponse, -) -from app.connectors.velociraptor.utils.universal import UniversalService -from app.db.db_session import session -from app.db.universal_models import Agents from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse +from app.agents.schema.agents import OutdatedWazuhAgentsResponse +from app.connectors.velociraptor.utils.universal import UniversalService +from app.db.db_session import session +from app.db.universal_models import Agents + def get_agent(agent_id: str) -> List[Agents]: """ @@ -100,11 +99,7 @@ async def get_outdated_agents_velociraptor( ) agents_result = await session.execute(select(Agents)) agents = agents_result.scalars().all() - outdated_velociraptor_agents = [ - agent - for agent in agents - if agent.velociraptor_agent_version != server_version - ] + outdated_velociraptor_agents = [agent for agent in agents if agent.velociraptor_agent_version != server_version] return OutdatedVelociraptorAgentsResponse( message="Outdated Velociraptor agents fetched successfully.", diff --git a/backend/app/agents/services/sync.py b/backend/app/agents/services/sync.py index a2164eb0..9cf15eba 100644 --- a/backend/app/agents/services/sync.py +++ b/backend/app/agents/services/sync.py @@ -1,15 +1,18 @@ from typing import List +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + import app.agents.velociraptor.services.agents as velociraptor_services import app.agents.wazuh.services.agents as wazuh_services -from app.agents.schema.agents import SyncedAgent, SyncedAgentsResponse +from app.agents.schema.agents import SyncedAgent +from app.agents.schema.agents import SyncedAgentsResponse from app.agents.velociraptor.schema.agents import VelociraptorAgent -from app.agents.wazuh.schema.agents import WazuhAgent, WazuhAgentsList +from app.agents.wazuh.schema.agents import WazuhAgent +from app.agents.wazuh.schema.agents import WazuhAgentsList from app.connectors.models import Connectors from app.db.universal_models import Agents -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select async def fetch_wazuh_agents() -> WazuhAgentsList: diff --git a/backend/app/agents/velociraptor/schema/agents.py b/backend/app/agents/velociraptor/schema/agents.py index 0362fd82..e3446df7 100644 --- a/backend/app/agents/velociraptor/schema/agents.py +++ b/backend/app/agents/velociraptor/schema/agents.py @@ -1,7 +1,8 @@ from datetime import datetime from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class VelociraptorAgent(BaseModel): diff --git a/backend/app/agents/velociraptor/services/agents.py b/backend/app/agents/velociraptor/services/agents.py index 5c9cb571..492c19df 100644 --- a/backend/app/agents/velociraptor/services/agents.py +++ b/backend/app/agents/velociraptor/services/agents.py @@ -1,10 +1,11 @@ from datetime import datetime +from fastapi import HTTPException +from loguru import logger + from app.agents.schema.agents import AgentModifyResponse from app.agents.velociraptor.schema.agents import VelociraptorAgent from app.connectors.velociraptor.utils.universal import UniversalService -from fastapi import HTTPException -from loguru import logger def create_query(query: str) -> str: @@ -45,9 +46,7 @@ async def collect_velociraptor_agent(agent_name: str) -> VelociraptorAgent: ) try: - vql_last_seen_at = ( - f"select last_seen_at from clients(search='host:{agent_name}')" - ) + vql_last_seen_at = f"select last_seen_at from clients(search='host:{agent_name}')" last_seen_at = await velociraptor_service._get_last_seen_timestamp( vql_last_seen_at, ) diff --git a/backend/app/agents/wazuh/schema/agents.py b/backend/app/agents/wazuh/schema/agents.py index b534a69e..00f482c8 100644 --- a/backend/app/agents/wazuh/schema/agents.py +++ b/backend/app/agents/wazuh/schema/agents.py @@ -1,7 +1,9 @@ from datetime import datetime -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class WazuhAgent(BaseModel): diff --git a/backend/app/agents/wazuh/services/agents.py b/backend/app/agents/wazuh/services/agents.py index 3d0bb9e7..3a82c83b 100644 --- a/backend/app/agents/wazuh/services/agents.py +++ b/backend/app/agents/wazuh/services/agents.py @@ -1,12 +1,12 @@ -from app.agents.schema.agents import AgentModifyResponse -from app.agents.wazuh.schema.agents import WazuhAgent, WazuhAgentsList -from app.connectors.wazuh_manager.utils.universal import ( - send_delete_request, - send_get_request, -) from fastapi import HTTPException from loguru import logger +from app.agents.schema.agents import AgentModifyResponse +from app.agents.wazuh.schema.agents import WazuhAgent +from app.agents.wazuh.schema.agents import WazuhAgentsList +from app.connectors.wazuh_manager.utils.universal import send_delete_request +from app.connectors.wazuh_manager.utils.universal import send_get_request + async def collect_wazuh_agents() -> WazuhAgentsList: """ @@ -29,11 +29,7 @@ async def collect_wazuh_agents() -> WazuhAgentsList: try: if agents_collected.get("success"): wazuh_agents_list = [] - for agent in ( - agents_collected.get("data", {}) - .get("data", {}) - .get("affected_items", []) - ): + for agent in agents_collected.get("data", {}).get("data", {}).get("affected_items", []): os_name = agent.get("os", {}).get("name", "Unknown") last_keep_alive = agent.get("lastKeepAlive", "Unknown") agent_group_list = agent.get("group", []) diff --git a/backend/app/agents/wazuh/services/vulnerabilities.py b/backend/app/agents/wazuh/services/vulnerabilities.py index a12afda5..004f544f 100644 --- a/backend/app/agents/wazuh/services/vulnerabilities.py +++ b/backend/app/agents/wazuh/services/vulnerabilities.py @@ -1,13 +1,12 @@ from typing import List -from app.agents.wazuh.schema.agents import ( - WazuhAgentVulnerabilities, - WazuhAgentVulnerabilitiesResponse, -) -from app.connectors.wazuh_manager.utils.universal import send_get_request from fastapi import HTTPException from loguru import logger +from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilities +from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse +from app.connectors.wazuh_manager.utils.universal import send_get_request + async def collect_agent_vulnerabilities(agent_id: str): """ diff --git a/backend/app/auth/models/users.py b/backend/app/auth/models/users.py index 8403d35f..f2cecb17 100644 --- a/backend/app/auth/models/users.py +++ b/backend/app/auth/models/users.py @@ -6,8 +6,12 @@ from typing import Optional import bcrypt -from pydantic import BaseModel, EmailStr, validator -from sqlmodel import Field, Relationship, SQLModel +from pydantic import BaseModel +from pydantic import EmailStr +from pydantic import validator +from sqlmodel import Field +from sqlmodel import Relationship +from sqlmodel import SQLModel class Role(SQLModel, table=True): diff --git a/backend/app/auth/routes/auth.py b/backend/app/auth/routes/auth.py index 89e0c137..061f797e 100644 --- a/backend/app/auth/routes/auth.py +++ b/backend/app/auth/routes/auth.py @@ -1,22 +1,28 @@ from datetime import timedelta -from app.auth.models.users import ( - PasswordReset, - PasswordResetToken, - User, - UserInput, - UserLogin, -) -from app.auth.schema.auth import Token, UserLoginResponse, UserResponse -from app.auth.schema.user import UserBaseResponse -from app.auth.services.universal import find_user, select_all_users -from app.auth.utils import AuthHandler -from app.db.db_session import get_db -from fastapi import APIRouter, Depends, HTTPException, Security, status +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from fastapi import status from fastapi.security import OAuth2PasswordRequestForm from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from app.auth.models.users import PasswordReset +from app.auth.models.users import PasswordResetToken +from app.auth.models.users import User +from app.auth.models.users import UserInput +from app.auth.models.users import UserLogin +from app.auth.schema.auth import Token +from app.auth.schema.auth import UserLoginResponse +from app.auth.schema.auth import UserResponse +from app.auth.schema.user import UserBaseResponse +from app.auth.services.universal import find_user +from app.auth.services.universal import select_all_users +from app.auth.utils import AuthHandler +from app.db.db_session import get_db + ACCESS_TOKEN_EXPIRE_MINUTES = 1440 auth_router = APIRouter() diff --git a/backend/app/auth/schema/user.py b/backend/app/auth/schema/user.py index 078a3dc4..29633e74 100644 --- a/backend/app/auth/schema/user.py +++ b/backend/app/auth/schema/user.py @@ -1,6 +1,7 @@ from typing import List -from pydantic import BaseModel, EmailStr +from pydantic import BaseModel +from pydantic import EmailStr class UserBase(BaseModel): diff --git a/backend/app/auth/services/universal.py b/backend/app/auth/services/universal.py index 18b059cd..9142f3a9 100644 --- a/backend/app/auth/services/universal.py +++ b/backend/app/auth/services/universal.py @@ -1,11 +1,14 @@ -from app.auth.models.users import Password, Role, User -from app.db.db_session import async_engine from loguru import logger # ! New with Async from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import select +from app.auth.models.users import Password +from app.auth.models.users import Role +from app.auth.models.users import User +from app.db.db_session import async_engine + passwords_in_memory = {} diff --git a/backend/app/auth/utils.py b/backend/app/auth/utils.py index 9234d98d..15d71c5a 100644 --- a/backend/app/auth/utils.py +++ b/backend/app/auth/utils.py @@ -1,12 +1,17 @@ -from datetime import datetime, timedelta +from datetime import datetime +from datetime import timedelta import jwt -from app.auth.services.universal import find_user, get_role -from fastapi import Depends, HTTPException -from fastapi.security import OAuth2PasswordBearer, SecurityScopes +from fastapi import Depends +from fastapi import HTTPException +from fastapi.security import OAuth2PasswordBearer +from fastapi.security import SecurityScopes from loguru import logger from passlib.context import CryptContext +from app.auth.services.universal import find_user +from app.auth.services.universal import get_role + class AuthHandler: security = OAuth2PasswordBearer( diff --git a/backend/app/connectors/cortex/routes/analyzers.py b/backend/app/connectors/cortex/routes/analyzers.py index 4ae93b99..ce13ebfc 100644 --- a/backend/app/connectors/cortex/routes/analyzers.py +++ b/backend/app/connectors/cortex/routes/analyzers.py @@ -1,15 +1,18 @@ from typing import List -from app.auth.utils import AuthHandler -from app.connectors.cortex.schema.analyzers import ( - AnalyzersResponse, - RunAnalyzerBody, - RunAnalyzerResponse, -) -from app.connectors.cortex.services.analyzers import get_analyzers, run_analyzer -from fastapi import APIRouter, Depends, HTTPException, Security +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security from loguru import logger +from app.auth.utils import AuthHandler +from app.connectors.cortex.schema.analyzers import AnalyzersResponse +from app.connectors.cortex.schema.analyzers import RunAnalyzerBody +from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse +from app.connectors.cortex.services.analyzers import get_analyzers +from app.connectors.cortex.services.analyzers import run_analyzer + # App specific imports diff --git a/backend/app/connectors/cortex/schema/analyzers.py b/backend/app/connectors/cortex/schema/analyzers.py index 8e1e78ba..db8cf30d 100644 --- a/backend/app/connectors/cortex/schema/analyzers.py +++ b/backend/app/connectors/cortex/schema/analyzers.py @@ -1,8 +1,14 @@ import ipaddress import re -from typing import Any, Dict, List, Optional, Tuple +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple -from pydantic import BaseModel, Field, validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator HASH_REGEX = re.compile( r"[a-fA-F\d]{32}|[a-fA-F\d]{64}", diff --git a/backend/app/connectors/cortex/services/analyzers.py b/backend/app/connectors/cortex/services/analyzers.py index 64972d2a..baa13ae4 100644 --- a/backend/app/connectors/cortex/services/analyzers.py +++ b/backend/app/connectors/cortex/services/analyzers.py @@ -1,19 +1,20 @@ -from typing import Dict, List, Union +from typing import Dict +from typing import List +from typing import Union -from app.connectors.cortex.schema.analyzers import ( - AnalyzerJobData, - AnalyzersResponse, - RunAnalyzerBody, - RunAnalyzerResponse, -) -from app.connectors.cortex.utils.universal import ( # Importing create_cortex_client; Importing from universal.py - create_cortex_client, - run_and_wait_for_analyzer, -) from cortex4py.api import Api from fastapi import HTTPException from loguru import logger +from app.connectors.cortex.schema.analyzers import AnalyzerJobData +from app.connectors.cortex.schema.analyzers import AnalyzersResponse +from app.connectors.cortex.schema.analyzers import RunAnalyzerBody +from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse +from app.connectors.cortex.utils.universal import ( + create_cortex_client, # Importing create_cortex_client; Importing from universal.py +) +from app.connectors.cortex.utils.universal import run_and_wait_for_analyzer + async def fetch_analyzers(api: Api) -> List[Dict]: """ diff --git a/backend/app/connectors/cortex/utils/universal.py b/backend/app/connectors/cortex/utils/universal.py index 0ece853a..3d946476 100644 --- a/backend/app/connectors/cortex/utils/universal.py +++ b/backend/app/connectors/cortex/utils/universal.py @@ -1,13 +1,15 @@ import time -from typing import Any, Dict +from typing import Any +from typing import Dict -from app.connectors.cortex.schema.analyzers import AnalyzerJobData -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session from cortex4py.api import Api from fastapi import HTTPException from loguru import logger +from app.connectors.cortex.schema.analyzers import AnalyzerJobData +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session + async def verify_cortex_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: """ diff --git a/backend/app/connectors/dfir_iris/routes/alerts.py b/backend/app/connectors/dfir_iris/routes/alerts.py index 145de060..e4dcde1a 100644 --- a/backend/app/connectors/dfir_iris/routes/alerts.py +++ b/backend/app/connectors/dfir_iris/routes/alerts.py @@ -1,26 +1,26 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler -from app.connectors.dfir_iris.schema.alerts import ( - AlertResponse, - AlertsResponse, - BookmarkedAlertsResponse, - CaseCreationResponse, - DeleteAlertResponse, - DeleteMultipleAlertsRequest, - FilterAlertsRequest, -) -from app.connectors.dfir_iris.services.alerts import ( - bookmark_alert, - create_case, - delete_alert, - get_alert, - get_alerts, - get_bookmarked_alerts, -) +from app.connectors.dfir_iris.schema.alerts import AlertResponse +from app.connectors.dfir_iris.schema.alerts import AlertsResponse +from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse +from app.connectors.dfir_iris.schema.alerts import CaseCreationResponse +from app.connectors.dfir_iris.schema.alerts import DeleteAlertResponse +from app.connectors.dfir_iris.schema.alerts import DeleteMultipleAlertsRequest +from app.connectors.dfir_iris.schema.alerts import FilterAlertsRequest +from app.connectors.dfir_iris.services.alerts import bookmark_alert +from app.connectors.dfir_iris.services.alerts import create_case +from app.connectors.dfir_iris.services.alerts import delete_alert +from app.connectors.dfir_iris.services.alerts import get_alert +from app.connectors.dfir_iris.services.alerts import get_alerts +from app.connectors.dfir_iris.services.alerts import get_bookmarked_alerts from app.connectors.dfir_iris.utils.universal import check_alert_exists from app.db.db_session import get_db -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession # App specific imports diff --git a/backend/app/connectors/dfir_iris/routes/assets.py b/backend/app/connectors/dfir_iris/routes/assets.py index 12fc203b..cbf9d7a6 100644 --- a/backend/app/connectors/dfir_iris/routes/assets.py +++ b/backend/app/connectors/dfir_iris/routes/assets.py @@ -1,9 +1,13 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler from app.connectors.dfir_iris.schema.assets import AssetResponse from app.connectors.dfir_iris.services.assets import get_case_assets from app.connectors.dfir_iris.utils.universal import check_case_exists -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger # App specific imports diff --git a/backend/app/connectors/dfir_iris/routes/cases.py b/backend/app/connectors/dfir_iris/routes/cases.py index 4963b7de..83e00ff7 100644 --- a/backend/app/connectors/dfir_iris/routes/cases.py +++ b/backend/app/connectors/dfir_iris/routes/cases.py @@ -1,31 +1,31 @@ from datetime import timedelta +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler -from app.connectors.dfir_iris.schema.cases import ( - CaseOlderThanBody, - CaseResponse, - CasesBreachedResponse, - ClosedCaseResponse, - PurgeCaseResponse, - ReopenedCaseResponse, - SingleCaseBody, - SingleCaseResponse, - TimeUnit, -) -from app.connectors.dfir_iris.services.cases import ( - close_case, - delete_single_case, - get_all_cases, - get_cases_older_than, - get_single_case, - purge_cases, - reopen_case, -) +from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody +from app.connectors.dfir_iris.schema.cases import CaseResponse +from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse +from app.connectors.dfir_iris.schema.cases import ClosedCaseResponse +from app.connectors.dfir_iris.schema.cases import PurgeCaseResponse +from app.connectors.dfir_iris.schema.cases import ReopenedCaseResponse +from app.connectors.dfir_iris.schema.cases import SingleCaseBody +from app.connectors.dfir_iris.schema.cases import SingleCaseResponse +from app.connectors.dfir_iris.schema.cases import TimeUnit +from app.connectors.dfir_iris.services.cases import close_case +from app.connectors.dfir_iris.services.cases import delete_single_case +from app.connectors.dfir_iris.services.cases import get_all_cases +from app.connectors.dfir_iris.services.cases import get_cases_older_than +from app.connectors.dfir_iris.services.cases import get_single_case +from app.connectors.dfir_iris.services.cases import purge_cases +from app.connectors.dfir_iris.services.cases import reopen_case from app.connectors.dfir_iris.utils.universal import check_case_exists from app.db.db_session import get_db -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession async def verify_case_exists(case_id: int) -> int: diff --git a/backend/app/connectors/dfir_iris/routes/notes.py b/backend/app/connectors/dfir_iris/routes/notes.py index 66fac8a9..d186b917 100644 --- a/backend/app/connectors/dfir_iris/routes/notes.py +++ b/backend/app/connectors/dfir_iris/routes/notes.py @@ -1,15 +1,18 @@ from typing import Optional +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler -from app.connectors.dfir_iris.schema.notes import ( - NoteCreationBody, - NoteCreationResponse, - NotesResponse, -) -from app.connectors.dfir_iris.services.notes import create_case_note, get_case_notes +from app.connectors.dfir_iris.schema.notes import NoteCreationBody +from app.connectors.dfir_iris.schema.notes import NoteCreationResponse +from app.connectors.dfir_iris.schema.notes import NotesResponse +from app.connectors.dfir_iris.services.notes import create_case_note +from app.connectors.dfir_iris.services.notes import get_case_notes from app.connectors.dfir_iris.utils.universal import check_case_exists -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger async def verify_case_exists(case_id: int) -> int: diff --git a/backend/app/connectors/dfir_iris/routes/users.py b/backend/app/connectors/dfir_iris/routes/users.py index 67acb8e9..a5e82cbe 100644 --- a/backend/app/connectors/dfir_iris/routes/users.py +++ b/backend/app/connectors/dfir_iris/routes/users.py @@ -1,17 +1,18 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler from app.connectors.dfir_iris.schema.alerts import AlertResponse -from app.connectors.dfir_iris.schema.users import User, UsersResponse -from app.connectors.dfir_iris.services.users import ( - assign_user_to_alert, - delete_user_from_alert, - get_users, -) -from app.connectors.dfir_iris.utils.universal import ( - check_alert_exists, - check_user_exists, -) -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger +from app.connectors.dfir_iris.schema.users import User +from app.connectors.dfir_iris.schema.users import UsersResponse +from app.connectors.dfir_iris.services.users import assign_user_to_alert +from app.connectors.dfir_iris.services.users import delete_user_from_alert +from app.connectors.dfir_iris.services.users import get_users +from app.connectors.dfir_iris.utils.universal import check_alert_exists +from app.connectors.dfir_iris.utils.universal import check_user_exists def verify_user_exists(user_id: int) -> int: diff --git a/backend/app/connectors/dfir_iris/schema/admin.py b/backend/app/connectors/dfir_iris/schema/admin.py index 1625b774..033e29d9 100644 --- a/backend/app/connectors/dfir_iris/schema/admin.py +++ b/backend/app/connectors/dfir_iris/schema/admin.py @@ -1,8 +1,11 @@ import uuid from datetime import datetime -from typing import Dict, List, Optional +from typing import Dict +from typing import List +from typing import Optional -from pydantic import UUID4, BaseModel +from pydantic import UUID4 +from pydantic import BaseModel class CreateCustomerData(BaseModel): diff --git a/backend/app/connectors/dfir_iris/schema/alerts.py b/backend/app/connectors/dfir_iris/schema/alerts.py index c9f7fcb9..b284a101 100644 --- a/backend/app/connectors/dfir_iris/schema/alerts.py +++ b/backend/app/connectors/dfir_iris/schema/alerts.py @@ -1,7 +1,11 @@ from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class AlertsResponse(BaseModel): diff --git a/backend/app/connectors/dfir_iris/schema/assets.py b/backend/app/connectors/dfir_iris/schema/assets.py index 2b4f27e6..c7318376 100644 --- a/backend/app/connectors/dfir_iris/schema/assets.py +++ b/backend/app/connectors/dfir_iris/schema/assets.py @@ -1,6 +1,8 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class AssetState(BaseModel): diff --git a/backend/app/connectors/dfir_iris/schema/cases.py b/backend/app/connectors/dfir_iris/schema/cases.py index 19a6f73b..78c3cb94 100644 --- a/backend/app/connectors/dfir_iris/schema/cases.py +++ b/backend/app/connectors/dfir_iris/schema/cases.py @@ -1,8 +1,13 @@ -from datetime import date, timedelta +from datetime import date +from datetime import timedelta from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict +from typing import List +from typing import Optional +from typing import Union -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class CaseModel(BaseModel): diff --git a/backend/app/connectors/dfir_iris/schema/notes.py b/backend/app/connectors/dfir_iris/schema/notes.py index 1218d196..e66d4785 100644 --- a/backend/app/connectors/dfir_iris/schema/notes.py +++ b/backend/app/connectors/dfir_iris/schema/notes.py @@ -1,6 +1,9 @@ -from typing import Dict, List, Optional +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class CustomAttributes(BaseModel): diff --git a/backend/app/connectors/dfir_iris/services/alerts.py b/backend/app/connectors/dfir_iris/services/alerts.py index dd126a14..6b2d367c 100644 --- a/backend/app/connectors/dfir_iris/services/alerts.py +++ b/backend/app/connectors/dfir_iris/services/alerts.py @@ -1,23 +1,20 @@ -from app.connectors.dfir_iris.schema.alerts import ( - AlertResponse, - AlertsResponse, - BookmarkedAlertsResponse, - CaseCreationResponse, - DeleteAlertResponse, - FilterAlertsRequest, -) -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) -from app.integrations.alert_creation_settings.models.alert_creation_settings import ( - AlertCreationSettings, -) from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.connectors.dfir_iris.schema.alerts import AlertResponse +from app.connectors.dfir_iris.schema.alerts import AlertsResponse +from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse +from app.connectors.dfir_iris.schema.alerts import CaseCreationResponse +from app.connectors.dfir_iris.schema.alerts import DeleteAlertResponse +from app.connectors.dfir_iris.schema.alerts import FilterAlertsRequest +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert +from app.integrations.alert_creation_settings.models.alert_creation_settings import ( + AlertCreationSettings, +) + async def get_customer_code(session: AsyncSession, customer_id: int) -> str: """ @@ -180,9 +177,7 @@ def construct_case_creation_params(alert_details: dict) -> dict: "case_tags": alert_details["alert_tags"], "escalation_note": "Case created from CoPilot", "iocs_import_list": [ioc["ioc_uuid"] for ioc in alert_details["iocs"]], - "assets_import_list": [ - asset["asset_uuid"] for asset in alert_details["assets"] - ], + "assets_import_list": [asset["asset_uuid"] for asset in alert_details["assets"]], } # Replace None values with the string "None" diff --git a/backend/app/connectors/dfir_iris/services/assets.py b/backend/app/connectors/dfir_iris/services/assets.py index f8a1da5a..18236d7f 100644 --- a/backend/app/connectors/dfir_iris/services/assets.py +++ b/backend/app/connectors/dfir_iris/services/assets.py @@ -1,10 +1,11 @@ -from app.connectors.dfir_iris.schema.assets import Asset, AssetResponse, AssetState -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_case, -) from fastapi import HTTPException +from app.connectors.dfir_iris.schema.assets import Asset +from app.connectors.dfir_iris.schema.assets import AssetResponse +from app.connectors.dfir_iris.schema.assets import AssetState +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_case + async def get_case_assets(case_id: int) -> AssetResponse: """ diff --git a/backend/app/connectors/dfir_iris/services/cases.py b/backend/app/connectors/dfir_iris/services/cases.py index 48ef92a1..165176ca 100644 --- a/backend/app/connectors/dfir_iris/services/cases.py +++ b/backend/app/connectors/dfir_iris/services/cases.py @@ -1,29 +1,27 @@ from datetime import datetime -from typing import Dict, List - -from app.connectors.dfir_iris.schema.cases import ( - CaseOlderThanBody, - CaseResponse, - CasesBreachedResponse, - ClosedCaseResponse, - PurgeCaseResponse, - ReopenedCaseResponse, - SingleCaseBody, - SingleCaseResponse, -) -from app.connectors.dfir_iris.utils.universal import ( - create_dfir_iris_client, - fetch_and_parse_data, -) -from app.integrations.alert_creation_settings.models.alert_creation_settings import ( - AlertCreationSettings, -) +from typing import Dict +from typing import List + from dfir_iris_client.case import Case from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody +from app.connectors.dfir_iris.schema.cases import CaseResponse +from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse +from app.connectors.dfir_iris.schema.cases import ClosedCaseResponse +from app.connectors.dfir_iris.schema.cases import PurgeCaseResponse +from app.connectors.dfir_iris.schema.cases import ReopenedCaseResponse +from app.connectors.dfir_iris.schema.cases import SingleCaseBody +from app.connectors.dfir_iris.schema.cases import SingleCaseResponse +from app.connectors.dfir_iris.utils.universal import create_dfir_iris_client +from app.connectors.dfir_iris.utils.universal import fetch_and_parse_data +from app.integrations.alert_creation_settings.models.alert_creation_settings import ( + AlertCreationSettings, +) + async def get_client_and_cases() -> Dict: """ diff --git a/backend/app/connectors/dfir_iris/services/notes.py b/backend/app/connectors/dfir_iris/services/notes.py index 4eb8bde4..28e35ed6 100644 --- a/backend/app/connectors/dfir_iris/services/notes.py +++ b/backend/app/connectors/dfir_iris/services/notes.py @@ -1,19 +1,18 @@ -from typing import Any, Dict, List - -from app.connectors.dfir_iris.schema.notes import ( - NoteCreationBody, - NoteCreationResponse, - NoteDetails, - NoteDetailsResponse, - NotesResponse, -) -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_case, -) +from typing import Any +from typing import Dict +from typing import List + from dfir_iris_client.case import Case from loguru import logger +from app.connectors.dfir_iris.schema.notes import NoteCreationBody +from app.connectors.dfir_iris.schema.notes import NoteCreationResponse +from app.connectors.dfir_iris.schema.notes import NoteDetails +from app.connectors.dfir_iris.schema.notes import NoteDetailsResponse +from app.connectors.dfir_iris.schema.notes import NotesResponse +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_case + async def process_notes(notes: List[Dict], case_id: int) -> List[Dict]: """ diff --git a/backend/app/connectors/dfir_iris/services/users.py b/backend/app/connectors/dfir_iris/services/users.py index 2cdec2cf..50a2edbf 100644 --- a/backend/app/connectors/dfir_iris/services/users.py +++ b/backend/app/connectors/dfir_iris/services/users.py @@ -1,10 +1,8 @@ from app.connectors.dfir_iris.schema.alerts import AlertResponse from app.connectors.dfir_iris.schema.users import UsersResponse -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, - initialize_client_and_user, -) +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert +from app.connectors.dfir_iris.utils.universal import initialize_client_and_user async def get_users() -> UsersResponse: diff --git a/backend/app/connectors/dfir_iris/utils/universal.py b/backend/app/connectors/dfir_iris/utils/universal.py index 4ce91b72..9d4167aa 100644 --- a/backend/app/connectors/dfir_iris/utils/universal.py +++ b/backend/app/connectors/dfir_iris/utils/universal.py @@ -1,18 +1,25 @@ -from typing import Any, Callable, Dict, Optional, Tuple, Union +from typing import Any +from typing import Callable +from typing import Dict +from typing import Optional +from typing import Tuple +from typing import Union import requests -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session from dfir_iris_client.admin import AdminHelper from dfir_iris_client.alert import Alert from dfir_iris_client.case import Case from dfir_iris_client.customer import Customer -from dfir_iris_client.helper.utils import assert_api_resp, get_data_from_resp +from dfir_iris_client.helper.utils import assert_api_resp +from dfir_iris_client.helper.utils import get_data_from_resp from dfir_iris_client.session import ClientSession from dfir_iris_client.users import User from fastapi import HTTPException from loguru import logger +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session + async def verify_dfir_iris_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: """ diff --git a/backend/app/connectors/event_shipper/utils/universal.py b/backend/app/connectors/event_shipper/utils/universal.py index 59b12266..918d3bef 100644 --- a/backend/app/connectors/event_shipper/utils/universal.py +++ b/backend/app/connectors/event_shipper/utils/universal.py @@ -1,6 +1,7 @@ from typing import Optional import asyncgelf + from app.connectors.utils import get_connector_info_from_db from app.db.db_session import get_db_session diff --git a/backend/app/connectors/grafana/routes/dashboards.py b/backend/app/connectors/grafana/routes/dashboards.py index 27043fb8..5d023615 100644 --- a/backend/app/connectors/grafana/routes/dashboards.py +++ b/backend/app/connectors/grafana/routes/dashboards.py @@ -1,11 +1,12 @@ +from fastapi import APIRouter +from fastapi import Body +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler -from app.connectors.grafana.schema.dashboards import ( - DashboardProvisionRequest, - GrafanaDashboardResponse, -) +from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest +from app.connectors.grafana.schema.dashboards import GrafanaDashboardResponse from app.connectors.grafana.services.dashboards import provision_dashboards -from fastapi import APIRouter, Body, Security -from loguru import logger # App specific imports diff --git a/backend/app/connectors/grafana/schema/dashboards.py b/backend/app/connectors/grafana/schema/dashboards.py index 5906e9b6..b477949e 100644 --- a/backend/app/connectors/grafana/schema/dashboards.py +++ b/backend/app/connectors/grafana/schema/dashboards.py @@ -1,7 +1,9 @@ from enum import Enum from typing import List -from pydantic import BaseModel, Field, validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator class GrafanaDashboard(BaseModel): @@ -81,9 +83,7 @@ class DashboardProvisionRequest(BaseModel): @validator("dashboards", each_item=True) def check_dashboard_exists(cls, e): - valid_dashboards = { - item.name: item for item in list(WazuhDashboard) + list(Office365Dashboard) - } + valid_dashboards = {item.name: item for item in list(WazuhDashboard) + list(Office365Dashboard)} if e not in valid_dashboards: raise ValueError(f'Dashboard identifier "{e}" is not recognized.') return e diff --git a/backend/app/connectors/grafana/services/dashboards.py b/backend/app/connectors/grafana/services/dashboards.py index 1ecd801b..91138ee1 100644 --- a/backend/app/connectors/grafana/services/dashboards.py +++ b/backend/app/connectors/grafana/services/dashboards.py @@ -1,18 +1,17 @@ import json from pathlib import Path -from app.connectors.grafana.schema.dashboards import ( - DashboardProvisionRequest, - GrafanaDashboard, - GrafanaDashboardResponse, - MimecastDashboard, - Office365Dashboard, - WazuhDashboard, -) -from app.connectors.grafana.utils.universal import create_grafana_client from fastapi import HTTPException from loguru import logger +from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest +from app.connectors.grafana.schema.dashboards import GrafanaDashboard +from app.connectors.grafana.schema.dashboards import GrafanaDashboardResponse +from app.connectors.grafana.schema.dashboards import MimecastDashboard +from app.connectors.grafana.schema.dashboards import Office365Dashboard +from app.connectors.grafana.schema.dashboards import WazuhDashboard +from app.connectors.grafana.utils.universal import create_grafana_client + def get_dashboard_path(dashboard_info: tuple) -> Path: """ @@ -26,9 +25,7 @@ def get_dashboard_path(dashboard_info: tuple) -> Path: """ folder_name, file_name = dashboard_info current_file = Path(__file__) # Path to the current file - base_dir = ( - current_file.parent.parent - ) # Move up two levels to the 'grafana' directory + base_dir = current_file.parent.parent # Move up two levels to the 'grafana' directory return base_dir / "dashboards" / folder_name / file_name @@ -147,12 +144,7 @@ async def provision_dashboards( provisioned_dashboards = [] errors = [] - valid_dashboards = { - item.name: item - for item in list(WazuhDashboard) - + list(Office365Dashboard) - + list(MimecastDashboard) - } + valid_dashboards = {item.name: item for item in list(WazuhDashboard) + list(Office365Dashboard) + list(MimecastDashboard)} for dashboard_name in dashboard_request.dashboards: dashboard_enum = valid_dashboards[dashboard_name] @@ -175,11 +167,7 @@ async def provision_dashboards( ) success = len(errors) == 0 - message = ( - "All dashboards provisioned successfully" - if success - else "Some dashboards failed to provision" - ) + message = "All dashboards provisioned successfully" if success else "Some dashboards failed to provision" return GrafanaDashboardResponse( provisioned_dashboards=provisioned_dashboards, success=success, diff --git a/backend/app/connectors/grafana/utils/universal.py b/backend/app/connectors/grafana/utils/universal.py index 4b36b16c..36505e63 100644 --- a/backend/app/connectors/grafana/utils/universal.py +++ b/backend/app/connectors/grafana/utils/universal.py @@ -1,12 +1,14 @@ -from typing import Any, Dict +from typing import Any +from typing import Dict -from app.connectors.grafana.schema.organization import GrafanaCreateOrganizationResponse -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session from fastapi import HTTPException from grafana_client import GrafanaApi from loguru import logger +from app.connectors.grafana.schema.organization import GrafanaCreateOrganizationResponse +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session + async def construct_grafana_url( connector_url: str, diff --git a/backend/app/connectors/graylog/routes/collector.py b/backend/app/connectors/graylog/routes/collector.py index 71107b85..5c17c8eb 100644 --- a/backend/app/connectors/graylog/routes/collector.py +++ b/backend/app/connectors/graylog/routes/collector.py @@ -1,19 +1,17 @@ -from app.auth.utils import AuthHandler -from app.connectors.graylog.schema.collector import ( - ConfiguredInputsResponse, - GraylogIndicesResponse, - GraylogInputsResponse, - RunningInputsResponse, -) -from app.connectors.graylog.services.collector import ( - get_indices_full, - get_inputs, - get_inputs_configured, - get_inputs_running, -) -from fastapi import APIRouter, Security +from fastapi import APIRouter +from fastapi import Security from loguru import logger +from app.auth.utils import AuthHandler +from app.connectors.graylog.schema.collector import ConfiguredInputsResponse +from app.connectors.graylog.schema.collector import GraylogIndicesResponse +from app.connectors.graylog.schema.collector import GraylogInputsResponse +from app.connectors.graylog.schema.collector import RunningInputsResponse +from app.connectors.graylog.services.collector import get_indices_full +from app.connectors.graylog.services.collector import get_inputs +from app.connectors.graylog.services.collector import get_inputs_configured +from app.connectors.graylog.services.collector import get_inputs_running + # App specific imports diff --git a/backend/app/connectors/graylog/routes/events.py b/backend/app/connectors/graylog/routes/events.py index 517841b4..59a3a7c1 100644 --- a/backend/app/connectors/graylog/routes/events.py +++ b/backend/app/connectors/graylog/routes/events.py @@ -1,13 +1,14 @@ -from app.auth.utils import AuthHandler -from app.connectors.graylog.schema.events import ( - AlertQuery, - GraylogAlertsResponse, - GraylogEventDefinitionsResponse, -) -from app.connectors.graylog.services.events import get_alerts, get_event_definitions -from fastapi import APIRouter, Security +from fastapi import APIRouter +from fastapi import Security from loguru import logger +from app.auth.utils import AuthHandler +from app.connectors.graylog.schema.events import AlertQuery +from app.connectors.graylog.schema.events import GraylogAlertsResponse +from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse +from app.connectors.graylog.services.events import get_alerts +from app.connectors.graylog.services.events import get_event_definitions + # App specific imports diff --git a/backend/app/connectors/graylog/routes/management.py b/backend/app/connectors/graylog/routes/management.py index 33ff7184..5ed6e91b 100644 --- a/backend/app/connectors/graylog/routes/management.py +++ b/backend/app/connectors/graylog/routes/management.py @@ -1,34 +1,32 @@ from typing import List +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler -from app.connectors.graylog.schema.management import ( - DeletedIndexBody, - DeletedIndexResponse, - StartInputBody, - StartInputResponse, - StartStreamBody, - StartStreamResponse, - StopInputBody, - StopInputResponse, - StopStreamBody, - StopStreamResponse, - UrlWhitelistEntryResponse, -) -from app.connectors.graylog.services.collector import ( - get_index_names, - get_input_ids, - get_url_whitelist_entries, -) -from app.connectors.graylog.services.management import ( - delete_index, - start_input, - start_stream, - stop_input, - stop_stream, -) +from app.connectors.graylog.schema.management import DeletedIndexBody +from app.connectors.graylog.schema.management import DeletedIndexResponse +from app.connectors.graylog.schema.management import StartInputBody +from app.connectors.graylog.schema.management import StartInputResponse +from app.connectors.graylog.schema.management import StartStreamBody +from app.connectors.graylog.schema.management import StartStreamResponse +from app.connectors.graylog.schema.management import StopInputBody +from app.connectors.graylog.schema.management import StopInputResponse +from app.connectors.graylog.schema.management import StopStreamBody +from app.connectors.graylog.schema.management import StopStreamResponse +from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse +from app.connectors.graylog.services.collector import get_index_names +from app.connectors.graylog.services.collector import get_input_ids +from app.connectors.graylog.services.collector import get_url_whitelist_entries +from app.connectors.graylog.services.management import delete_index +from app.connectors.graylog.services.management import start_input +from app.connectors.graylog.services.management import start_stream +from app.connectors.graylog.services.management import stop_input +from app.connectors.graylog.services.management import stop_stream from app.connectors.graylog.services.streams import get_stream_ids -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger graylog_management_router = APIRouter() diff --git a/backend/app/connectors/graylog/routes/monitoring.py b/backend/app/connectors/graylog/routes/monitoring.py index f2b4cb11..e019c2e6 100644 --- a/backend/app/connectors/graylog/routes/monitoring.py +++ b/backend/app/connectors/graylog/routes/monitoring.py @@ -1,17 +1,15 @@ -from app.auth.utils import AuthHandler -from app.connectors.graylog.schema.monitoring import ( - GraylogEventNotificationsResponse, - GraylogMessagesResponse, - GraylogMetricsResponse, -) -from app.connectors.graylog.services.monitoring import ( - get_event_notifications, - get_messages, - get_metrics, -) -from fastapi import APIRouter, Security +from fastapi import APIRouter +from fastapi import Security from loguru import logger +from app.auth.utils import AuthHandler +from app.connectors.graylog.schema.monitoring import GraylogEventNotificationsResponse +from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse +from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse +from app.connectors.graylog.services.monitoring import get_event_notifications +from app.connectors.graylog.services.monitoring import get_messages +from app.connectors.graylog.services.monitoring import get_metrics + # App specific imports diff --git a/backend/app/connectors/graylog/routes/pipelines.py b/backend/app/connectors/graylog/routes/pipelines.py index 04931024..82e67df9 100644 --- a/backend/app/connectors/graylog/routes/pipelines.py +++ b/backend/app/connectors/graylog/routes/pipelines.py @@ -1,24 +1,23 @@ -from typing import Dict, List +from typing import Dict +from typing import List -from app.auth.utils import AuthHandler -from app.connectors.graylog.schema.pipelines import ( - GraylogPipelinesResponse, - GraylogPipelinesResponseWithRuleID, - Pipeline, - PipelineRule, - PipelineRulesResponse, - PipelineWithRuleID, - Stage, - StageWithRuleID, -) -from app.connectors.graylog.services.pipelines import ( - get_pipeline_rule_by_id, - get_pipeline_rules, - get_pipelines, -) -from fastapi import APIRouter, Security +from fastapi import APIRouter +from fastapi import Security from loguru import logger +from app.auth.utils import AuthHandler +from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse +from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponseWithRuleID +from app.connectors.graylog.schema.pipelines import Pipeline +from app.connectors.graylog.schema.pipelines import PipelineRule +from app.connectors.graylog.schema.pipelines import PipelineRulesResponse +from app.connectors.graylog.schema.pipelines import PipelineWithRuleID +from app.connectors.graylog.schema.pipelines import Stage +from app.connectors.graylog.schema.pipelines import StageWithRuleID +from app.connectors.graylog.services.pipelines import get_pipeline_rule_by_id +from app.connectors.graylog.services.pipelines import get_pipeline_rules +from app.connectors.graylog.services.pipelines import get_pipelines + # App specific imports @@ -57,9 +56,7 @@ def transform_stages_with_rule_ids( """ new_stages = [] for stage in stages: - rule_ids = [ - rule_title_to_id.get(rule_title, None) for rule_title in stage.rules - ] + rule_ids = [rule_title_to_id.get(rule_title, None) for rule_title in stage.rules] new_stage = StageWithRuleID(**stage.dict(), rule_ids=rule_ids) new_stages.append(new_stage) return new_stages @@ -123,10 +120,7 @@ async def get_all_pipelines_with_rule_ids() -> GraylogPipelinesResponseWithRuleI pipeline_rules_response.pipeline_rules, ) - new_pipelines = [ - transform_pipeline_with_rule_ids(pipeline, rule_title_to_id) - for pipeline in pipelines_response.pipelines - ] + new_pipelines = [transform_pipeline_with_rule_ids(pipeline, rule_title_to_id) for pipeline in pipelines_response.pipelines] return GraylogPipelinesResponseWithRuleID( pipelines=new_pipelines, diff --git a/backend/app/connectors/graylog/routes/streams.py b/backend/app/connectors/graylog/routes/streams.py index b14b3c36..baecd554 100644 --- a/backend/app/connectors/graylog/routes/streams.py +++ b/backend/app/connectors/graylog/routes/streams.py @@ -1,8 +1,10 @@ +from fastapi import APIRouter +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler from app.connectors.graylog.schema.streams import GraylogStreamsResponse from app.connectors.graylog.services.streams import get_streams -from fastapi import APIRouter, Security -from loguru import logger # App specific imports diff --git a/backend/app/connectors/graylog/schema/collector.py b/backend/app/connectors/graylog/schema/collector.py index d28a5fa0..c5761075 100644 --- a/backend/app/connectors/graylog/schema/collector.py +++ b/backend/app/connectors/graylog/schema/collector.py @@ -1,6 +1,9 @@ -from typing import Dict, List, Optional +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class Document(BaseModel): diff --git a/backend/app/connectors/graylog/schema/events.py b/backend/app/connectors/graylog/schema/events.py index cfbe92ba..42cdb778 100644 --- a/backend/app/connectors/graylog/schema/events.py +++ b/backend/app/connectors/graylog/schema/events.py @@ -1,4 +1,7 @@ -from typing import Dict, List, Optional, Union +from typing import Dict +from typing import List +from typing import Optional +from typing import Union from pydantic import BaseModel diff --git a/backend/app/connectors/graylog/schema/monitoring.py b/backend/app/connectors/graylog/schema/monitoring.py index 6e2a0995..afc4932a 100644 --- a/backend/app/connectors/graylog/schema/monitoring.py +++ b/backend/app/connectors/graylog/schema/monitoring.py @@ -1,6 +1,8 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class GraylogMessages(BaseModel): diff --git a/backend/app/connectors/graylog/schema/pipelines.py b/backend/app/connectors/graylog/schema/pipelines.py index 18b0e600..cd53418f 100644 --- a/backend/app/connectors/graylog/schema/pipelines.py +++ b/backend/app/connectors/graylog/schema/pipelines.py @@ -1,4 +1,5 @@ -from typing import List, Optional +from typing import List +from typing import Optional from pydantic import BaseModel diff --git a/backend/app/connectors/graylog/schema/streams.py b/backend/app/connectors/graylog/schema/streams.py index f63e75c0..f3a77870 100644 --- a/backend/app/connectors/graylog/schema/streams.py +++ b/backend/app/connectors/graylog/schema/streams.py @@ -1,4 +1,5 @@ -from typing import List, Optional +from typing import List +from typing import Optional from pydantic import BaseModel diff --git a/backend/app/connectors/graylog/services/collector.py b/backend/app/connectors/graylog/services/collector.py index 9661c280..56c381f3 100644 --- a/backend/app/connectors/graylog/services/collector.py +++ b/backend/app/connectors/graylog/services/collector.py @@ -1,19 +1,19 @@ -from typing import List, Tuple - -from app.connectors.graylog.schema.collector import ( - ConfiguredInput, - ConfiguredInputsResponse, - GraylogIndexItem, - GraylogIndicesResponse, - GraylogInputsResponse, - RunningInput, - RunningInputsResponse, -) -from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse -from app.connectors.graylog.utils.universal import send_get_request +from typing import List +from typing import Tuple + from fastapi import HTTPException from loguru import logger +from app.connectors.graylog.schema.collector import ConfiguredInput +from app.connectors.graylog.schema.collector import ConfiguredInputsResponse +from app.connectors.graylog.schema.collector import GraylogIndexItem +from app.connectors.graylog.schema.collector import GraylogIndicesResponse +from app.connectors.graylog.schema.collector import GraylogInputsResponse +from app.connectors.graylog.schema.collector import RunningInput +from app.connectors.graylog.schema.collector import RunningInputsResponse +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. @@ -33,10 +33,7 @@ async def get_indices_full() -> GraylogIndicesResponse: raise HTTPException(status_code=500, detail="Failed to collect indices key") # Convert the dictionary to a list of GraylogIndexItem - indices_list = [ - GraylogIndexItem(index_name=name, index_info=info) - for name, info in indices_data.items() - ] + indices_list = [GraylogIndexItem(index_name=name, index_info=info) for name, info in indices_data.items()] return GraylogIndicesResponse( indices=indices_list, @@ -62,10 +59,7 @@ async def fetch_configured_inputs() -> Tuple[bool, List[ConfiguredInput]]: success = configured_inputs_collected.get("success", False) if success: - return True, [ - ConfiguredInput(**input_data) - for input_data in configured_inputs_collected["data"]["inputs"] - ] + return True, [ConfiguredInput(**input_data) for input_data in configured_inputs_collected["data"]["inputs"]] else: logger.error("Failed to fetch configured inputs") return False, [] @@ -84,10 +78,7 @@ async def fetch_running_inputs() -> Tuple[bool, List[RunningInput]]: success = running_inputs_collected.get("success", False) if success: - return True, [ - RunningInput(**input_data) - for input_data in running_inputs_collected["data"]["states"] - ] + return True, [RunningInput(**input_data) for input_data in running_inputs_collected["data"]["states"]] else: logger.error("Failed to fetch running inputs") return False, [] diff --git a/backend/app/connectors/graylog/services/events.py b/backend/app/connectors/graylog/services/events.py index 277e3ff5..ef1b530a 100644 --- a/backend/app/connectors/graylog/services/events.py +++ b/backend/app/connectors/graylog/services/events.py @@ -1,17 +1,17 @@ -from app.connectors.graylog.schema.events import ( - AlertEvent, - AlertQuery, - Alerts, - Context, - EventDefinition, - GraylogAlertsResponse, - GraylogEventDefinitionsResponse, - Parameters, -) -from app.connectors.graylog.utils.universal import send_get_request, send_post_request from fastapi import HTTPException from loguru import logger +from app.connectors.graylog.schema.events import AlertEvent +from app.connectors.graylog.schema.events import AlertQuery +from app.connectors.graylog.schema.events import Alerts +from app.connectors.graylog.schema.events import Context +from app.connectors.graylog.schema.events import EventDefinition +from app.connectors.graylog.schema.events import GraylogAlertsResponse +from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse +from app.connectors.graylog.schema.events import Parameters +from app.connectors.graylog.utils.universal import send_get_request +from app.connectors.graylog.utils.universal import send_post_request + async def get_event_definitions() -> GraylogEventDefinitionsResponse: """Get event definitions from Graylog. @@ -25,9 +25,7 @@ async def get_event_definitions() -> GraylogEventDefinitionsResponse: ) if event_definitions_collected["success"]: try: - event_definitions_data = event_definitions_collected["data"][ - "event_definitions" - ] + event_definitions_data = event_definitions_collected["data"]["event_definitions"] except KeyError: raise HTTPException( status_code=500, @@ -35,10 +33,7 @@ async def get_event_definitions() -> GraylogEventDefinitionsResponse: ) # Convert the dictionary to a list of GraylogIndexItem - event_definitions_list = [ - EventDefinition(**event_definition_data) - for event_definition_data in event_definitions_data - ] + event_definitions_list = [EventDefinition(**event_definition_data) for event_definition_data in event_definitions_data] return GraylogEventDefinitionsResponse( event_definitions=event_definitions_list, @@ -78,9 +73,7 @@ async def get_alerts(alert_query: AlertQuery) -> GraylogAlertsResponse: except KeyError: raise HTTPException(status_code=500, detail="Failed to collect data key") # Convert raw event data to Event objects - event_objects = [ - AlertEvent(**event_data) for event_data in raw_alerts_data["events"] - ] + event_objects = [AlertEvent(**event_data) for event_data in raw_alerts_data["events"]] # Build the Alerts object alerts = Alerts( diff --git a/backend/app/connectors/graylog/services/management.py b/backend/app/connectors/graylog/services/management.py index 64bcbe86..0bba931f 100644 --- a/backend/app/connectors/graylog/services/management.py +++ b/backend/app/connectors/graylog/services/management.py @@ -1,23 +1,20 @@ -from app.connectors.graylog.schema.management import ( - DeletedIndexBody, - DeletedIndexResponse, - StartInputBody, - StartInputResponse, - StartStreamBody, - StartStreamResponse, - StopInputBody, - StopInputResponse, - StopStreamBody, - StopStreamResponse, -) -from app.connectors.graylog.services.collector import get_index_names -from app.connectors.graylog.utils.universal import ( - send_delete_request, - send_post_request, - send_put_request, -) from loguru import logger +from app.connectors.graylog.schema.management import DeletedIndexBody +from app.connectors.graylog.schema.management import DeletedIndexResponse +from app.connectors.graylog.schema.management import StartInputBody +from app.connectors.graylog.schema.management import StartInputResponse +from app.connectors.graylog.schema.management import StartStreamBody +from app.connectors.graylog.schema.management import StartStreamResponse +from app.connectors.graylog.schema.management import StopInputBody +from app.connectors.graylog.schema.management import StopInputResponse +from app.connectors.graylog.schema.management import StopStreamBody +from app.connectors.graylog.schema.management import StopStreamResponse +from app.connectors.graylog.services.collector import get_index_names +from app.connectors.graylog.utils.universal import send_delete_request +from app.connectors.graylog.utils.universal import send_post_request +from app.connectors.graylog.utils.universal import send_put_request + async def delete_index(index_name: DeletedIndexBody) -> DeletedIndexResponse: """Delete an index from Graylog. diff --git a/backend/app/connectors/graylog/services/monitoring.py b/backend/app/connectors/graylog/services/monitoring.py index 69d02ac6..54184f4e 100644 --- a/backend/app/connectors/graylog/services/monitoring.py +++ b/backend/app/connectors/graylog/services/monitoring.py @@ -1,16 +1,15 @@ -from app.connectors.graylog.schema.monitoring import ( - GraylogEventNotificationsResponse, - GraylogMessages, - GraylogMessagesResponse, - GraylogMetricsResponse, - GraylogThroughputMetrics, - GraylogThroughputMetricsCollection, - GraylogUncommittedJournalEntries, -) -from app.connectors.graylog.utils.universal import send_get_request from fastapi import HTTPException from loguru import logger +from app.connectors.graylog.schema.monitoring import GraylogEventNotificationsResponse +from app.connectors.graylog.schema.monitoring import GraylogMessages +from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse +from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse +from app.connectors.graylog.schema.monitoring import GraylogThroughputMetrics +from app.connectors.graylog.schema.monitoring import GraylogThroughputMetricsCollection +from app.connectors.graylog.schema.monitoring import GraylogUncommittedJournalEntries +from app.connectors.graylog.utils.universal import send_get_request + async def get_messages(page_number: int) -> GraylogMessagesResponse: """Get messages from Graylog. @@ -109,10 +108,7 @@ def filter_and_create_throughput_metrics(merged_metrics: dict) -> list: Returns: list: A list of GraylogThroughputMetrics objects. """ - model_fields = [ - field_info.alias - for field_info in GraylogThroughputMetricsCollection.__fields__.values() - ] + model_fields = [field_info.alias for field_info in GraylogThroughputMetricsCollection.__fields__.values()] throughput_metrics_list = [ GraylogThroughputMetrics(metric=metric_name, value=metric_data.get("value", 0)) for metric_name, metric_data in merged_metrics.items() @@ -132,19 +128,14 @@ async def get_metrics() -> GraylogMetricsResponse: throughput_metrics_collected = await fetch_metrics_from_graylog() uncommitted_journal_entries_collected = await fetch_uncommitted_journal_entries() try: - if ( - throughput_metrics_collected["success"] - and uncommitted_journal_entries_collected["success"] - ): + if throughput_metrics_collected["success"] and uncommitted_journal_entries_collected["success"]: merged_metrics = merge_metrics_data(throughput_metrics_collected) throughput_metrics_list = filter_and_create_throughput_metrics( merged_metrics, ) uncommitted_journal_entries = GraylogUncommittedJournalEntries( - uncommitted_journal_entries=uncommitted_journal_entries_collected[ - "data" - ]["uncommitted_journal_entries"], + uncommitted_journal_entries=uncommitted_journal_entries_collected["data"]["uncommitted_journal_entries"], ) return GraylogMetricsResponse( diff --git a/backend/app/connectors/graylog/services/pipelines.py b/backend/app/connectors/graylog/services/pipelines.py index d1ecca36..51c6cd36 100644 --- a/backend/app/connectors/graylog/services/pipelines.py +++ b/backend/app/connectors/graylog/services/pipelines.py @@ -1,19 +1,17 @@ -from app.connectors.graylog.schema.pipelines import ( - CreatePipeline, - CreatePipelineRule, - GraylogPipelinesResponse, - Pipeline, - PipelineRule, - PipelineRulesResponse, -) -from app.connectors.graylog.utils.universal import send_get_request, send_post_request -from app.customer_provisioning.schema.graylog import ( - StreamConnectionToPipelineRequest, - StreamConnectionToPipelineResponse, -) from fastapi import HTTPException from loguru import logger +from app.connectors.graylog.schema.pipelines import CreatePipeline +from app.connectors.graylog.schema.pipelines import CreatePipelineRule +from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse +from app.connectors.graylog.schema.pipelines import Pipeline +from app.connectors.graylog.schema.pipelines import PipelineRule +from app.connectors.graylog.schema.pipelines import PipelineRulesResponse +from app.connectors.graylog.utils.universal import send_get_request +from app.connectors.graylog.utils.universal import send_post_request +from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest +from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineResponse + async def get_pipelines() -> GraylogPipelinesResponse: """Get pipelines from Graylog. @@ -30,10 +28,7 @@ async def get_pipelines() -> GraylogPipelinesResponse: ) try: if pipelines_collected["success"]: - pipelines_list = [ - Pipeline(**pipeline_data) - for pipeline_data in pipelines_collected["data"] - ] + pipelines_list = [Pipeline(**pipeline_data) for pipeline_data in pipelines_collected["data"]] return GraylogPipelinesResponse( pipelines=pipelines_list, success=True, @@ -63,10 +58,7 @@ async def get_pipeline_rules() -> PipelineRulesResponse: ) try: if pipeline_rules_collected["success"]: - pipeline_rules_list = [ - PipelineRule(**pipeline_rule_data) - for pipeline_rule_data in pipeline_rules_collected["data"] - ] + pipeline_rules_list = [PipelineRule(**pipeline_rule_data) for pipeline_rule_data in pipeline_rules_collected["data"]] return PipelineRulesResponse( pipeline_rules=pipeline_rules_list, success=True, diff --git a/backend/app/connectors/graylog/services/streams.py b/backend/app/connectors/graylog/services/streams.py index 30565bd2..8c1fd797 100644 --- a/backend/app/connectors/graylog/services/streams.py +++ b/backend/app/connectors/graylog/services/streams.py @@ -1,10 +1,12 @@ from typing import List -from app.connectors.graylog.schema.streams import GraylogStreamsResponse, Stream -from app.connectors.graylog.utils.universal import send_get_request from fastapi import HTTPException from loguru import logger +from app.connectors.graylog.schema.streams import GraylogStreamsResponse +from app.connectors.graylog.schema.streams import Stream +from app.connectors.graylog.utils.universal import send_get_request + async def get_streams() -> GraylogStreamsResponse: """Get streams from Graylog. @@ -19,10 +21,7 @@ async def get_streams() -> GraylogStreamsResponse: streams_collected = await send_get_request(endpoint="/api/streams") try: if streams_collected["success"]: - streams_list = [ - Stream(**stream_data) - for stream_data in streams_collected["data"]["streams"] - ] + streams_list = [Stream(**stream_data) for stream_data in streams_collected["data"]["streams"]] return GraylogStreamsResponse( streams=streams_list, success=True, @@ -60,10 +59,7 @@ async def get_stream_ids() -> List[str]: streams_collected = await send_get_request(endpoint="/api/streams") try: if streams_collected["success"]: - return [ - stream_data["id"] - for stream_data in streams_collected["data"]["streams"] - ] + return [stream_data["id"] for stream_data in streams_collected["data"]["streams"]] else: return [] except KeyError as e: diff --git a/backend/app/connectors/graylog/utils/universal.py b/backend/app/connectors/graylog/utils/universal.py index 68fe873c..a4b6a037 100644 --- a/backend/app/connectors/graylog/utils/universal.py +++ b/backend/app/connectors/graylog/utils/universal.py @@ -1,11 +1,14 @@ -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional import requests -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session from fastapi import HTTPException from loguru import logger +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session + HEADERS = {"X-Requested-By": "CoPilot"} diff --git a/backend/app/connectors/influxdb/routes/alerts.py b/backend/app/connectors/influxdb/routes/alerts.py index f358f283..fb255f82 100644 --- a/backend/app/connectors/influxdb/routes/alerts.py +++ b/backend/app/connectors/influxdb/routes/alerts.py @@ -1,8 +1,10 @@ +from fastapi import APIRouter +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler from app.connectors.influxdb.schema.alerts import InfluxDBAlertsResponse from app.connectors.influxdb.services.alerts import get_alerts -from fastapi import APIRouter, Security -from loguru import logger # App specific imports diff --git a/backend/app/connectors/influxdb/services/alerts.py b/backend/app/connectors/influxdb/services/alerts.py index 66a815b7..9b90cdfd 100644 --- a/backend/app/connectors/influxdb/services/alerts.py +++ b/backend/app/connectors/influxdb/services/alerts.py @@ -1,13 +1,13 @@ from typing import List -from app.connectors.influxdb.schema.alerts import InfluxDBAlert, InfluxDBAlertsResponse -from app.connectors.influxdb.utils.universal import ( - create_influxdb_client, - get_influxdb_organization, -) from fastapi import HTTPException from loguru import logger +from app.connectors.influxdb.schema.alerts import InfluxDBAlert +from app.connectors.influxdb.schema.alerts import InfluxDBAlertsResponse +from app.connectors.influxdb.utils.universal import create_influxdb_client +from app.connectors.influxdb.utils.universal import get_influxdb_organization + # Constants BUCKET_NAME = "_monitoring" @@ -46,9 +46,7 @@ async def process_alert_records(result) -> List[InfluxDBAlert]: for table in result: for record in table.records: alert = InfluxDBAlert( - time=record.values.get("time").isoformat() - if record.values.get("time") - else None, + time=record.values.get("time").isoformat() if record.values.get("time") else None, message=record.values.get("message"), checkID=record.values.get("checkID"), checkName=record.values.get("checkName"), diff --git a/backend/app/connectors/influxdb/utils/universal.py b/backend/app/connectors/influxdb/utils/universal.py index 252e82a0..e6a68cf4 100644 --- a/backend/app/connectors/influxdb/utils/universal.py +++ b/backend/app/connectors/influxdb/utils/universal.py @@ -1,11 +1,13 @@ -from typing import Any, Dict +from typing import Any +from typing import Dict -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session from fastapi import HTTPException from influxdb_client.client.influxdb_client_async import InfluxDBClientAsync from loguru import logger +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session + async def verify_influxdb_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: """ diff --git a/backend/app/connectors/models.py b/backend/app/connectors/models.py index 780681c1..0dfe876f 100644 --- a/backend/app/connectors/models.py +++ b/backend/app/connectors/models.py @@ -1,7 +1,10 @@ from datetime import datetime -from typing import List, Optional +from typing import List +from typing import Optional -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field +from sqlmodel import Relationship +from sqlmodel import SQLModel class ConnectorHistory(SQLModel, table=True): diff --git a/backend/app/connectors/routes.py b/backend/app/connectors/routes.py index a80ebbc0..5cbf3ddc 100644 --- a/backend/app/connectors/routes.py +++ b/backend/app/connectors/routes.py @@ -1,21 +1,24 @@ from typing import Union -from app.auth.utils import AuthHandler -from app.connectors.schema import ( - ConnectorListResponse, - ConnectorResponse, - ConnectorsListResponse, - UpdateConnector, - VerifyConnectorResponse, -) -from app.connectors.services import ConnectorServices -from app.db.db_session import get_db - ## Auth Things -from fastapi import APIRouter, Depends, File, HTTPException, Security, UploadFile +from fastapi import APIRouter +from fastapi import Depends +from fastapi import File +from fastapi import HTTPException +from fastapi import Security +from fastapi import UploadFile from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from app.auth.utils import AuthHandler +from app.connectors.schema import ConnectorListResponse +from app.connectors.schema import ConnectorResponse +from app.connectors.schema import ConnectorsListResponse +from app.connectors.schema import UpdateConnector +from app.connectors.schema import VerifyConnectorResponse +from app.connectors.services import ConnectorServices +from app.db.db_session import get_db + connector_router = APIRouter() diff --git a/backend/app/connectors/schema.py b/backend/app/connectors/schema.py index e728f8e1..d96bd03b 100644 --- a/backend/app/connectors/schema.py +++ b/backend/app/connectors/schema.py @@ -1,5 +1,6 @@ from datetime import datetime -from typing import List, Optional +from typing import List +from typing import Optional from pydantic import BaseModel diff --git a/backend/app/connectors/services.py b/backend/app/connectors/services.py index b0220c1a..b27d3551 100644 --- a/backend/app/connectors/services.py +++ b/backend/app/connectors/services.py @@ -1,8 +1,18 @@ import os from datetime import datetime -from typing import List, Optional, Type, Union +from typing import List +from typing import Optional +from typing import Type +from typing import Union import aiofiles +from fastapi import UploadFile +from loguru import logger +from pydantic import BaseModel +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select +from werkzeug.utils import secure_filename + from app.connectors.cortex.utils.universal import verify_cortex_connection from app.connectors.dfir_iris.utils.universal import verify_dfir_iris_connection from app.connectors.grafana.utils.universal import verify_grafana_connection @@ -22,16 +32,8 @@ from app.threat_intel.services.socfortress import ( verifiy_socfortress_threat_intel_connector, ) -from app.utils import ( - verify_alert_creation_provisioning_connection, - verify_wazuh_worker_provisioning_connection, -) -from fastapi import UploadFile -from loguru import logger -from pydantic import BaseModel -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select -from werkzeug.utils import secure_filename +from app.utils import verify_alert_creation_provisioning_connection +from app.utils import verify_wazuh_worker_provisioning_connection UPLOAD_FOLDER = "file-store" UPLOAD_FOLDER = os.path.join( @@ -396,9 +398,7 @@ def allowed_file(filename): Returns: bool: True if the file is allowed, False otherwise. """ - return ( - "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS - ) + return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS @classmethod async def save_file( diff --git a/backend/app/connectors/shuffle/routes/workflows.py b/backend/app/connectors/shuffle/routes/workflows.py index 10351bb5..eac7a8a5 100644 --- a/backend/app/connectors/shuffle/routes/workflows.py +++ b/backend/app/connectors/shuffle/routes/workflows.py @@ -1,16 +1,15 @@ -from app.auth.utils import AuthHandler -from app.connectors.shuffle.schema.workflows import ( - WorkflowExecutionBodyModel, - WorkflowExecutionResponseModel, - WorkflowsResponse, -) -from app.connectors.shuffle.services.workflows import ( - get_workflow_executions, - get_workflows, -) -from fastapi import APIRouter, HTTPException, Security +from fastapi import APIRouter +from fastapi import HTTPException +from fastapi import Security from loguru import logger +from app.auth.utils import AuthHandler +from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel +from app.connectors.shuffle.schema.workflows import WorkflowExecutionResponseModel +from app.connectors.shuffle.schema.workflows import WorkflowsResponse +from app.connectors.shuffle.services.workflows import get_workflow_executions +from app.connectors.shuffle.services.workflows import get_workflows + shuffle_workflows_router = APIRouter() diff --git a/backend/app/connectors/shuffle/schema/workflows.py b/backend/app/connectors/shuffle/schema/workflows.py index b341ae15..e14d660a 100644 --- a/backend/app/connectors/shuffle/schema/workflows.py +++ b/backend/app/connectors/shuffle/schema/workflows.py @@ -1,6 +1,10 @@ -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class WorkflowsResponse(BaseModel): diff --git a/backend/app/connectors/shuffle/services/workflows.py b/backend/app/connectors/shuffle/services/workflows.py index 1ac6df8c..7a07d2dd 100644 --- a/backend/app/connectors/shuffle/services/workflows.py +++ b/backend/app/connectors/shuffle/services/workflows.py @@ -1,14 +1,13 @@ from typing import List -from app.connectors.shuffle.schema.workflows import ( - WorkflowExecutionBodyModel, - WorkflowExecutionStatusResponseModel, - WorkflowsResponse, -) -from app.connectors.shuffle.utils.universal import send_get_request from fastapi import HTTPException from loguru import logger +from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel +from app.connectors.shuffle.schema.workflows import WorkflowExecutionStatusResponseModel +from app.connectors.shuffle.schema.workflows import WorkflowsResponse +from app.connectors.shuffle.utils.universal import send_get_request + def remove_large_images_from_actions(workflows: List) -> List: """ diff --git a/backend/app/connectors/shuffle/utils/universal.py b/backend/app/connectors/shuffle/utils/universal.py index 8b25356f..eb365dc1 100644 --- a/backend/app/connectors/shuffle/utils/universal.py +++ b/backend/app/connectors/shuffle/utils/universal.py @@ -1,11 +1,14 @@ -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional import requests -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session from fastapi import HTTPException from loguru import logger +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session + async def verify_shuffle_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: """ @@ -164,9 +167,7 @@ def send_post_request( return { "data": response.json(), "success": False if response.status_code >= 400 else True, - "message": "Successfully retrieved data" - if response.status_code < 400 - else "Failed to retrieve data", + "message": "Successfully retrieved data" if response.status_code < 400 else "Failed to retrieve data", } except Exception as e: logger.debug(f"Response: {response}") diff --git a/backend/app/connectors/sublime/models/alerts.py b/backend/app/connectors/sublime/models/alerts.py index 7e093968..6b681396 100644 --- a/backend/app/connectors/sublime/models/alerts.py +++ b/backend/app/connectors/sublime/models/alerts.py @@ -1,7 +1,10 @@ import datetime -from typing import List, Optional +from typing import List +from typing import Optional -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field +from sqlmodel import Relationship +from sqlmodel import SQLModel class FlaggedRule(SQLModel, table=True): diff --git a/backend/app/connectors/sublime/routes/alerts.py b/backend/app/connectors/sublime/routes/alerts.py index 3124f0b0..2cd8d4ab 100644 --- a/backend/app/connectors/sublime/routes/alerts.py +++ b/backend/app/connectors/sublime/routes/alerts.py @@ -1,15 +1,17 @@ -from app.auth.utils import AuthHandler -from app.connectors.sublime.schema.alerts import ( - AlertRequestBody, - AlertResponseBody, - SublimeAlertsResponse, -) -from app.connectors.sublime.services.alerts import collect_alerts, store_sublime_alert -from app.db.db_session import get_db -from fastapi import APIRouter, Depends, Security +from fastapi import APIRouter +from fastapi import Depends +from fastapi import Security from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from app.auth.utils import AuthHandler +from app.connectors.sublime.schema.alerts import AlertRequestBody +from app.connectors.sublime.schema.alerts import AlertResponseBody +from app.connectors.sublime.schema.alerts import SublimeAlertsResponse +from app.connectors.sublime.services.alerts import collect_alerts +from app.connectors.sublime.services.alerts import store_sublime_alert +from app.db.db_session import get_db + sublime_alerts_router = APIRouter() diff --git a/backend/app/connectors/sublime/schema/alerts.py b/backend/app/connectors/sublime/schema/alerts.py index 8ec8c194..c68994e1 100644 --- a/backend/app/connectors/sublime/schema/alerts.py +++ b/backend/app/connectors/sublime/schema/alerts.py @@ -1,7 +1,9 @@ import datetime -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class FlaggedRule(BaseModel): diff --git a/backend/app/connectors/sublime/services/alerts.py b/backend/app/connectors/sublime/services/alerts.py index 9f57fe44..bd6a778e 100644 --- a/backend/app/connectors/sublime/services/alerts.py +++ b/backend/app/connectors/sublime/services/alerts.py @@ -1,27 +1,24 @@ import json from typing import List -from app.connectors.sublime.models.alerts import ( - FlaggedRule, - Mailbox, - Recipient, - Sender, - SublimeAlerts, - TriggeredAction, -) -from app.connectors.sublime.schema.alerts import ( - AlertRequestBody, - AlertResponseBody, - SublimeAlertsResponse, - SublimeAlertsSchema, -) -from app.connectors.sublime.utils.universal import send_get_request from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy.orm import selectinload +from app.connectors.sublime.models.alerts import FlaggedRule +from app.connectors.sublime.models.alerts import Mailbox +from app.connectors.sublime.models.alerts import Recipient +from app.connectors.sublime.models.alerts import Sender +from app.connectors.sublime.models.alerts import SublimeAlerts +from app.connectors.sublime.models.alerts import TriggeredAction +from app.connectors.sublime.schema.alerts import AlertRequestBody +from app.connectors.sublime.schema.alerts import AlertResponseBody +from app.connectors.sublime.schema.alerts import SublimeAlertsResponse +from app.connectors.sublime.schema.alerts import SublimeAlertsSchema +from app.connectors.sublime.utils.universal import send_get_request + def create_sublime_alert(alert_request_body: AlertRequestBody) -> SublimeAlerts: """ diff --git a/backend/app/connectors/sublime/utils/universal.py b/backend/app/connectors/sublime/utils/universal.py index 7d8cd361..ac65564d 100644 --- a/backend/app/connectors/sublime/utils/universal.py +++ b/backend/app/connectors/sublime/utils/universal.py @@ -1,9 +1,12 @@ -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional import requests +from loguru import logger + from app.connectors.utils import get_connector_info_from_db from app.db.db_session import get_db_session -from loguru import logger async def verify_sublime_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: diff --git a/backend/app/connectors/utils.py b/backend/app/connectors/utils.py index be159410..b0c64da2 100644 --- a/backend/app/connectors/utils.py +++ b/backend/app/connectors/utils.py @@ -1,11 +1,14 @@ -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional -from app.connectors.models import Connectors -from app.connectors.schema import ConnectorResponse from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.connectors.models import Connectors +from app.connectors.schema import ConnectorResponse + # ! New with Async async def get_connector_info_from_db( diff --git a/backend/app/connectors/velociraptor/routes/artifacts.py b/backend/app/connectors/velociraptor/routes/artifacts.py index 2ad03e32..48c99591 100644 --- a/backend/app/connectors/velociraptor/routes/artifacts.py +++ b/backend/app/connectors/velociraptor/routes/artifacts.py @@ -1,30 +1,30 @@ from typing import List -from app.auth.utils import AuthHandler -from app.connectors.velociraptor.schema.artifacts import ( - ArtifactsResponse, - CollectArtifactBody, - CollectArtifactResponse, - OSPrefixEnum, - OSPrefixModel, - QuarantineBody, - QuarantineResponse, - RunCommandBody, - RunCommandResponse, -) -from app.connectors.velociraptor.services.artifacts import ( - get_artifacts, - quarantine_host, - run_artifact_collection, - run_remote_command, -) -from app.db.db_session import get_db -from app.db.universal_models import Agents -from fastapi import APIRouter, Depends, HTTPException, Security +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.auth.utils import AuthHandler +from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse +from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody +from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse +from app.connectors.velociraptor.schema.artifacts import OSPrefixEnum +from app.connectors.velociraptor.schema.artifacts import OSPrefixModel +from app.connectors.velociraptor.schema.artifacts import QuarantineBody +from app.connectors.velociraptor.schema.artifacts import QuarantineResponse +from app.connectors.velociraptor.schema.artifacts import RunCommandBody +from app.connectors.velociraptor.schema.artifacts import RunCommandResponse +from app.connectors.velociraptor.services.artifacts import get_artifacts +from app.connectors.velociraptor.services.artifacts import quarantine_host +from app.connectors.velociraptor.services.artifacts import run_artifact_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 + # App specific imports @@ -68,9 +68,7 @@ def verify_os_prefix_exists(os_prefix: str) -> str: detail=f"OS prefix {os_prefix} does not exist.", ) - return OSPrefixEnum[ - os_prefix_upper - ].value # Use the uppercase version for Enum matching + return OSPrefixEnum[os_prefix_upper].value # Use the uppercase version for Enum matching def get_os_prefix_from_os_name(os_name: str) -> str: @@ -223,9 +221,7 @@ async def get_all_artifacts_for_os_prefix( # Get all the artifacts names that begin with the OS prefix artifacts = await get_artifacts() artifacts = artifacts.artifacts - artifacts_for_os_prefix = [ - artifact for artifact in artifacts if artifact.name.startswith(os_prefix) - ] + artifacts_for_os_prefix = [artifact for artifact in artifacts if artifact.name.startswith(os_prefix)] return ArtifactsResponse( success=True, message=f"All artifacts for OS prefix {os_prefix} retrieved", diff --git a/backend/app/connectors/velociraptor/routes/flows.py b/backend/app/connectors/velociraptor/routes/flows.py index 0ee05bfe..c824f60a 100644 --- a/backend/app/connectors/velociraptor/routes/flows.py +++ b/backend/app/connectors/velociraptor/routes/flows.py @@ -1,13 +1,19 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.auth.utils import AuthHandler from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse -from app.connectors.velociraptor.schema.flows import FlowResponse, RetrieveFlowRequest -from app.connectors.velociraptor.services.flows import get_flow, get_flows +from app.connectors.velociraptor.schema.flows import FlowResponse +from app.connectors.velociraptor.schema.flows import RetrieveFlowRequest +from app.connectors.velociraptor.services.flows import get_flow +from app.connectors.velociraptor.services.flows import get_flows from app.db.db_session import get_db from app.db.universal_models import Agents -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select velociraptor_flows_router = APIRouter() diff --git a/backend/app/connectors/velociraptor/schema/artifacts.py b/backend/app/connectors/velociraptor/schema/artifacts.py index dce07789..25ccc6e0 100644 --- a/backend/app/connectors/velociraptor/schema/artifacts.py +++ b/backend/app/connectors/velociraptor/schema/artifacts.py @@ -1,7 +1,11 @@ from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class Artifacts(BaseModel): diff --git a/backend/app/connectors/velociraptor/schema/flows.py b/backend/app/connectors/velociraptor/schema/flows.py index 22c9e484..27b63d85 100644 --- a/backend/app/connectors/velociraptor/schema/flows.py +++ b/backend/app/connectors/velociraptor/schema/flows.py @@ -1,8 +1,11 @@ -from typing import List, Optional +from typing import List +from typing import Optional from fastapi import HTTPException from loguru import logger -from pydantic import BaseModel, Field, root_validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import root_validator class FlowSpecParameter(BaseModel): diff --git a/backend/app/connectors/velociraptor/services/artifacts.py b/backend/app/connectors/velociraptor/services/artifacts.py index 7c02871b..2b4262de 100644 --- a/backend/app/connectors/velociraptor/services/artifacts.py +++ b/backend/app/connectors/velociraptor/services/artifacts.py @@ -1,17 +1,16 @@ -from app.connectors.velociraptor.schema.artifacts import ( - Artifacts, - ArtifactsResponse, - CollectArtifactBody, - CollectArtifactResponse, - QuarantineBody, - QuarantineResponse, - RunCommandBody, - RunCommandResponse, -) -from app.connectors.velociraptor.utils.universal import UniversalService from fastapi import HTTPException from loguru import logger +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 +from app.connectors.velociraptor.schema.artifacts import QuarantineBody +from app.connectors.velociraptor.schema.artifacts import QuarantineResponse +from app.connectors.velociraptor.schema.artifacts import RunCommandBody +from app.connectors.velociraptor.schema.artifacts import RunCommandResponse +from app.connectors.velociraptor.utils.universal import UniversalService + def create_query(query: str) -> str: """ @@ -61,10 +60,7 @@ def get_artifact_key(analyzer_body: CollectArtifactBody) -> str: f"env=dict(Command='{analyzer_body.command}'))" ) else: - return ( - f"collect_client(client_id='{analyzer_body.velociraptor_id}', " - f"artifacts=['{analyzer_body.artifact_name}'])" - ) + return f"collect_client(client_id='{analyzer_body.velociraptor_id}', " f"artifacts=['{analyzer_body.artifact_name}'])" async def get_artifacts() -> ArtifactsResponse: @@ -140,9 +136,7 @@ async def run_artifact_collection( message=results["message"], results=results["results"], ) - except ( - HTTPException - ) as he: # Catch HTTPException separately to propagate the original message + except HTTPException as he: # Catch HTTPException separately to propagate the original message logger.error( f"HTTPException while running artifact collection on {collect_artifact_body}: {he.detail}", ) diff --git a/backend/app/connectors/velociraptor/services/flows.py b/backend/app/connectors/velociraptor/services/flows.py index 9a0e4b5a..53d423be 100644 --- a/backend/app/connectors/velociraptor/services/flows.py +++ b/backend/app/connectors/velociraptor/services/flows.py @@ -1,13 +1,12 @@ -from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse -from app.connectors.velociraptor.schema.flows import ( - FlowClientSession, - FlowResponse, - RetrieveFlowRequest, -) -from app.connectors.velociraptor.utils.universal import UniversalService from fastapi import HTTPException from loguru import logger +from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse +from app.connectors.velociraptor.schema.flows import FlowClientSession +from app.connectors.velociraptor.schema.flows import FlowResponse +from app.connectors.velociraptor.schema.flows import RetrieveFlowRequest +from app.connectors.velociraptor.utils.universal import UniversalService + def create_query(query: str) -> str: """ diff --git a/backend/app/connectors/velociraptor/utils/universal.py b/backend/app/connectors/velociraptor/utils/universal.py index 88d202e4..5c4b5dac 100644 --- a/backend/app/connectors/velociraptor/utils/universal.py +++ b/backend/app/connectors/velociraptor/utils/universal.py @@ -1,14 +1,18 @@ import json from datetime import datetime -from typing import Any, Dict +from typing import Any +from typing import Dict import grpc import pyvelociraptor -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import AsyncSessionLocal, get_db_session from fastapi import HTTPException from loguru import logger -from pyvelociraptor import api_pb2, api_pb2_grpc +from pyvelociraptor import api_pb2 +from pyvelociraptor import api_pb2_grpc + +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import AsyncSessionLocal +from app.db.db_session import get_db_session async def verify_velociraptor_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: @@ -244,12 +248,8 @@ async def get_client_id(self, client_name: str): """ # Formulate queries try: - vql_client_id = ( - f"select client_id,os_info from clients(search='host:{client_name}')" - ) - vql_last_seen_at = ( - f"select last_seen_at from clients(search='host:{client_name}')" - ) + vql_client_id = f"select client_id,os_info from clients(search='host:{client_name}')" + vql_last_seen_at = f"select last_seen_at from clients(search='host:{client_name}')" # Get the last seen timestamp logger.info(f"Getting last seen at timestamp for {client_name}") @@ -322,6 +322,4 @@ async def _is_offline(self, last_seen_at: float): Returns: bool: True if the client is offline, False otherwise. """ - return ( - datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000) - ).total_seconds() > 30 + return (datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000)).total_seconds() > 30 diff --git a/backend/app/connectors/wazuh_indexer/routes/alerts.py b/backend/app/connectors/wazuh_indexer/routes/alerts.py index b54e5a32..e512c3ba 100644 --- a/backend/app/connectors/wazuh_indexer/routes/alerts.py +++ b/backend/app/connectors/wazuh_indexer/routes/alerts.py @@ -1,28 +1,28 @@ from typing import List +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger + from app.auth.utils import AuthHandler -from app.connectors.wazuh_indexer.schema.alerts import ( - AlertsByHostResponse, - AlertsByRulePerHostResponse, - AlertsByRuleResponse, - AlertsSearchBody, - AlertsSearchResponse, - HostAlertsSearchBody, - HostAlertsSearchResponse, - IndexAlertsSearchBody, - IndexAlertsSearchResponse, -) -from app.connectors.wazuh_indexer.services.alerts import ( - get_alerts, - get_alerts_by_host, - get_alerts_by_rule, - get_alerts_by_rule_per_host, - get_host_alerts, - get_index_alerts, -) +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 +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 +from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule_per_host +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 -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger # App specific imports diff --git a/backend/app/connectors/wazuh_indexer/routes/monitoring.py b/backend/app/connectors/wazuh_indexer/routes/monitoring.py index 0c7d6d69..2a41a874 100644 --- a/backend/app/connectors/wazuh_indexer/routes/monitoring.py +++ b/backend/app/connectors/wazuh_indexer/routes/monitoring.py @@ -1,21 +1,20 @@ from typing import Union +from fastapi import APIRouter +from fastapi import HTTPException +from fastapi import Security + from app.auth.utils import AuthHandler -from app.connectors.wazuh_indexer.schema.monitoring import ( - ClusterHealthResponse, - IndicesStatsResponse, - NodeAllocationResponse, - ShardsResponse, -) +from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse +from app.connectors.wazuh_indexer.schema.monitoring import IndicesStatsResponse +from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocationResponse +from app.connectors.wazuh_indexer.schema.monitoring import ShardsResponse # from app.connectors.wazuh_indexer.schema import WazuhIndexerResponse, WazuhIndexerListResponse -from app.connectors.wazuh_indexer.services.monitoring import ( - cluster_healthcheck, - indices_stats, - node_allocation, - shards, -) -from fastapi import APIRouter, HTTPException, Security +from app.connectors.wazuh_indexer.services.monitoring import cluster_healthcheck +from app.connectors.wazuh_indexer.services.monitoring import indices_stats +from app.connectors.wazuh_indexer.services.monitoring import node_allocation +from app.connectors.wazuh_indexer.services.monitoring import shards wazuh_indexer_router = APIRouter() diff --git a/backend/app/connectors/wazuh_indexer/schema/alerts.py b/backend/app/connectors/wazuh_indexer/schema/alerts.py index 5246d0d0..99a614ac 100644 --- a/backend/app/connectors/wazuh_indexer/schema/alerts.py +++ b/backend/app/connectors/wazuh_indexer/schema/alerts.py @@ -1,7 +1,12 @@ from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field, validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator class Alert(BaseModel): @@ -116,8 +121,6 @@ class AlertsByRulePerHostResponse(BaseModel): ############# ! PASSABLE MESSAGES FROM ES CLIENT ! ############# class SkippableWazuhIndexerClientErrors(Enum): - NO_MAPPING_FOR_TIMESTAMP = ( - "No mapping found for [timestamp_utc] in order to sort on" - ) + NO_MAPPING_FOR_TIMESTAMP = "No mapping found for [timestamp_utc] in order to sort on" # Add other error messages here, for example: # ANOTHER_ERROR = "Another specific error message" diff --git a/backend/app/connectors/wazuh_indexer/schema/indices.py b/backend/app/connectors/wazuh_indexer/schema/indices.py index 2e12c007..57716753 100644 --- a/backend/app/connectors/wazuh_indexer/schema/indices.py +++ b/backend/app/connectors/wazuh_indexer/schema/indices.py @@ -1,6 +1,7 @@ from typing import Dict -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class Indices(BaseModel): diff --git a/backend/app/connectors/wazuh_indexer/schema/monitoring.py b/backend/app/connectors/wazuh_indexer/schema/monitoring.py index 3327db77..a4d6477e 100644 --- a/backend/app/connectors/wazuh_indexer/schema/monitoring.py +++ b/backend/app/connectors/wazuh_indexer/schema/monitoring.py @@ -1,6 +1,9 @@ -from typing import List, Optional, Union +from typing import List +from typing import Optional +from typing import Union -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class ClusterHealth(BaseModel): diff --git a/backend/app/connectors/wazuh_indexer/services/alerts.py b/backend/app/connectors/wazuh_indexer/services/alerts.py index 229019ee..170d2d1f 100644 --- a/backend/app/connectors/wazuh_indexer/services/alerts.py +++ b/backend/app/connectors/wazuh_indexer/services/alerts.py @@ -1,29 +1,29 @@ -from typing import Dict, List, Optional, Type - -from app.connectors.wazuh_indexer.schema.alerts import ( - AlertsByHost, - AlertsByHostResponse, - AlertsByRule, - AlertsByRulePerHost, - AlertsByRulePerHostResponse, - AlertsByRuleResponse, - AlertsSearchBody, - AlertsSearchResponse, - CollectAlertsResponse, - HostAlertsSearchBody, - HostAlertsSearchResponse, - IndexAlertsSearchBody, - IndexAlertsSearchResponse, - SkippableWazuhIndexerClientErrors, -) -from app.connectors.wazuh_indexer.utils.universal import ( - AlertsQueryBuilder, - collect_indices, - create_wazuh_indexer_client, -) +from typing import Dict +from typing import List +from typing import Optional +from typing import Type + from fastapi import HTTPException from loguru import logger +from app.connectors.wazuh_indexer.schema.alerts import AlertsByHost +from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse +from app.connectors.wazuh_indexer.schema.alerts import AlertsByRule +from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHost +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 +from app.connectors.wazuh_indexer.schema.alerts import CollectAlertsResponse +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.schema.alerts import SkippableWazuhIndexerClientErrors +from app.connectors.wazuh_indexer.utils.universal import AlertsQueryBuilder +from app.connectors.wazuh_indexer.utils.universal import collect_indices +from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client + async def collect_and_aggregate_alerts( field_names: List[str], @@ -48,17 +48,11 @@ async def collect_and_aggregate_alerts( alerts_response = await collect_alerts_generic(index_name, body=search_body) if alerts_response.success: for alert in alerts_response.alerts: - composite_key = tuple( - alert["_source"][field] for field in field_names - ) - aggregated_alerts_dict[composite_key] = ( - aggregated_alerts_dict.get(composite_key, 0) + 1 - ) + composite_key = tuple(alert["_source"][field] for field in field_names) + aggregated_alerts_dict[composite_key] = aggregated_alerts_dict.get(composite_key, 0) + 1 except HTTPException as e: detail_str = str(e.detail) # Convert to string to make sure it's comparable - if any( - err.value in detail_str for err in SkippableWazuhIndexerClientErrors - ): + if any(err.value in detail_str for err in SkippableWazuhIndexerClientErrors): logger.warning( f"Skipping index {index_name} due to specific error: {e.detail}", ) @@ -149,9 +143,7 @@ async def get_alerts_generic( ) alerts_summary = [] indices = await collect_indices() - index_list = ( - [index_name] if index_name else indices.indices_list - ) # Use the provided index_name or get all indices + index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices for index_name in index_list: try: @@ -170,9 +162,7 @@ async def get_alerts_generic( ) except HTTPException as e: detail_str = str(e.detail) # Convert to string to make sure it's comparable - if any( - err.value in detail_str for err in SkippableWazuhIndexerClientErrors - ): + if any(err.value in detail_str for err in SkippableWazuhIndexerClientErrors): logger.warning( f"Skipping index {index_name} due to specific error: {e.detail}", ) diff --git a/backend/app/connectors/wazuh_indexer/services/monitoring.py b/backend/app/connectors/wazuh_indexer/services/monitoring.py index b23bfeeb..823e3835 100644 --- a/backend/app/connectors/wazuh_indexer/services/monitoring.py +++ b/backend/app/connectors/wazuh_indexer/services/monitoring.py @@ -1,23 +1,21 @@ -from typing import Dict, Union - -from app.connectors.wazuh_indexer.schema.monitoring import ( - ClusterHealth, - ClusterHealthResponse, - IndicesStats, - IndicesStatsResponse, - NodeAllocation, - NodeAllocationResponse, - Shards, - ShardsResponse, -) -from app.connectors.wazuh_indexer.utils.universal import ( - create_wazuh_indexer_client, - format_indices_stats, - format_node_allocation, - format_shards, -) +from typing import Dict +from typing import Union + from loguru import logger +from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealth +from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse +from app.connectors.wazuh_indexer.schema.monitoring import IndicesStats +from app.connectors.wazuh_indexer.schema.monitoring import IndicesStatsResponse +from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocation +from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocationResponse +from app.connectors.wazuh_indexer.schema.monitoring import Shards +from app.connectors.wazuh_indexer.schema.monitoring import ShardsResponse +from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client +from app.connectors.wazuh_indexer.utils.universal import format_indices_stats +from app.connectors.wazuh_indexer.utils.universal import format_node_allocation +from app.connectors.wazuh_indexer.utils.universal import format_shards + async def cluster_healthcheck() -> Union[ClusterHealthResponse, Dict[str, str]]: """ @@ -64,9 +62,7 @@ async def node_allocation() -> Union[NodeAllocationResponse, Dict[str, bool]]: raw_node_allocation_data, ) - node_allocation_models = [ - NodeAllocation(**node) for node in formatted_node_allocation_data - ] + node_allocation_models = [NodeAllocation(**node) for node in formatted_node_allocation_data] return NodeAllocationResponse( node_allocation=node_allocation_models, @@ -97,9 +93,7 @@ async def indices_stats() -> Union[IndicesStatsResponse, Dict[str, str]]: raw_indices_stats_data, ) - indices_stats_models = [ - IndicesStats(**index) for index in formatted_indices_stats_data - ] + indices_stats_models = [IndicesStats(**index) for index in formatted_indices_stats_data] return IndicesStatsResponse( indices_stats=indices_stats_models, diff --git a/backend/app/connectors/wazuh_indexer/utils/universal.py b/backend/app/connectors/wazuh_indexer/utils/universal.py index c97dadfe..a09e5a1c 100644 --- a/backend/app/connectors/wazuh_indexer/utils/universal.py +++ b/backend/app/connectors/wazuh_indexer/utils/universal.py @@ -1,13 +1,19 @@ -from datetime import datetime, timedelta -from typing import Any, Dict, Iterable, Tuple +from datetime import datetime +from datetime import timedelta +from typing import Any +from typing import Dict +from typing import Iterable +from typing import Tuple -from app.connectors.utils import get_connector_info_from_db -from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel, Indices -from app.db.db_session import get_db_session from elasticsearch7 import Elasticsearch from fastapi import HTTPException from loguru import logger +from app.connectors.utils import get_connector_info_from_db +from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel +from app.connectors.wazuh_indexer.schema.indices import Indices +from app.db.db_session import get_db_session + async def verify_wazuh_indexer_credentials( attributes: Dict[str, Any], @@ -189,9 +195,7 @@ async def collect_indices() -> Indices: indices_list = list(indices_dict.keys()) # Check if the index is valid index_config = IndexConfigModel() - indices_list = [ - index for index in indices_list if index_config.is_valid_index(index) - ] + indices_list = [index for index in indices_list if index_config.is_valid_index(index)] return Indices( indices_list=indices_list, success=True, @@ -226,9 +230,7 @@ def _get_time_range_start(timerange: str) -> str: ) start = datetime.utcnow() - delta - return ( - start.isoformat() + "Z" - ) # Elasticsearch expects the time in ISO format with a Z at the end + return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end def __init__(self): self.query = { @@ -350,9 +352,7 @@ def _get_time_range_start(timerange: str) -> str: ) start = datetime.utcnow() - delta - return ( - start.isoformat() + "Z" - ) # Elasticsearch expects the time in ISO format with a Z at the end + return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end def __init__(self): self.query = { diff --git a/backend/app/connectors/wazuh_manager/models/rules.py b/backend/app/connectors/wazuh_manager/models/rules.py index 42902474..fabcf94f 100644 --- a/backend/app/connectors/wazuh_manager/models/rules.py +++ b/backend/app/connectors/wazuh_manager/models/rules.py @@ -1,7 +1,8 @@ import datetime from typing import Optional -from sqlmodel import Field, SQLModel +from sqlmodel import Field +from sqlmodel import SQLModel class DisabledRule(SQLModel, table=True): diff --git a/backend/app/connectors/wazuh_manager/routes/rules.py b/backend/app/connectors/wazuh_manager/routes/rules.py index 320c6e26..b332d62c 100644 --- a/backend/app/connectors/wazuh_manager/routes/rules.py +++ b/backend/app/connectors/wazuh_manager/routes/rules.py @@ -1,24 +1,27 @@ # App specific imports +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.auth.routes.auth import AuthHandler from app.connectors.wazuh_manager.models.rules import DisabledRule -from app.connectors.wazuh_manager.schema.rules import ( - AllDisabledRuleResponse, - RuleDisable, - RuleDisableResponse, - RuleEnable, - RuleEnableResponse, -) +from app.connectors.wazuh_manager.schema.rules import AllDisabledRuleResponse +from app.connectors.wazuh_manager.schema.rules import RuleDisable +from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse +from app.connectors.wazuh_manager.schema.rules import RuleEnable +from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse # from app.connectors.wazuh_manager.schema.rules import RuleExclude # from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse -from app.connectors.wazuh_manager.services.rules import disable_rule, enable_rule +from app.connectors.wazuh_manager.services.rules import disable_rule +from app.connectors.wazuh_manager.services.rules import enable_rule # from app.connectors.wazuh_manager.services.rules import exclude_rule from app.db.db_session import get_db -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select # from app.connectors.wazuh_manager.schema.rules import RuleExclude # from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse diff --git a/backend/app/connectors/wazuh_manager/schema/rules.py b/backend/app/connectors/wazuh_manager/schema/rules.py index 12d3c73e..6498bc1a 100644 --- a/backend/app/connectors/wazuh_manager/schema/rules.py +++ b/backend/app/connectors/wazuh_manager/schema/rules.py @@ -1,6 +1,8 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class RuleDisable(BaseModel): diff --git a/backend/app/connectors/wazuh_manager/services/rules.py b/backend/app/connectors/wazuh_manager/services/rules.py index a0f3fa29..dfd57bd9 100644 --- a/backend/app/connectors/wazuh_manager/services/rules.py +++ b/backend/app/connectors/wazuh_manager/services/rules.py @@ -1,25 +1,26 @@ import re from enum import Enum -from typing import Any, Dict, List, Tuple, Union +from typing import Any +from typing import Dict +from typing import List +from typing import Tuple +from typing import Union # import pcre2 import xmltodict -from app.connectors.wazuh_manager.schema.rules import ( - RuleDisable, - RuleDisableResponse, - RuleEnable, - RuleEnableResponse, - RuleExclude, - RuleExcludeResponse, -) -from app.connectors.wazuh_manager.utils.universal import ( - restart_service, - send_get_request, - send_put_request, -) from fastapi import HTTPException from loguru import logger +from app.connectors.wazuh_manager.schema.rules import RuleDisable +from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse +from app.connectors.wazuh_manager.schema.rules import RuleEnable +from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse +from app.connectors.wazuh_manager.schema.rules import RuleExclude +from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse +from app.connectors.wazuh_manager.utils.universal import restart_service +from app.connectors.wazuh_manager.utils.universal import send_get_request +from app.connectors.wazuh_manager.utils.universal import send_put_request + async def fetch_filename(rule_id: str) -> str: """ diff --git a/backend/app/connectors/wazuh_manager/utils/universal.py b/backend/app/connectors/wazuh_manager/utils/universal.py index e2d331fa..0ef82d59 100644 --- a/backend/app/connectors/wazuh_manager/utils/universal.py +++ b/backend/app/connectors/wazuh_manager/utils/universal.py @@ -1,10 +1,14 @@ -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional import requests -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import AsyncSessionLocal, get_db_session from loguru import logger +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import AsyncSessionLocal +from app.db.db_session import get_db_session + async def verify_wazuh_manager_credentials( attributes: Dict[str, Any], diff --git a/backend/app/customer_provisioning/routes/decommission.py b/backend/app/customer_provisioning/routes/decommission.py index 3f95ae0b..a68e6a38 100644 --- a/backend/app/customer_provisioning/routes/decommission.py +++ b/backend/app/customer_provisioning/routes/decommission.py @@ -1,12 +1,16 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.auth.utils import AuthHandler from app.customer_provisioning.schema.decommission import DecommissionCustomerResponse from app.customer_provisioning.services.decommission import decomission_wazuh_customer from app.db.db_session import get_db from app.db.universal_models import CustomersMeta -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select # App specific imports diff --git a/backend/app/customer_provisioning/routes/provision.py b/backend/app/customer_provisioning/routes/provision.py index d69a1818..c8341c49 100644 --- a/backend/app/customer_provisioning/routes/provision.py +++ b/backend/app/customer_provisioning/routes/provision.py @@ -1,21 +1,26 @@ -from app.auth.utils import AuthHandler -from app.connectors.grafana.schema.dashboards import Office365Dashboard, WazuhDashboard -from app.customer_provisioning.schema.provision import ( - CustomerProvisionResponse, - CustomersMetaResponse, - CustomerSubsctipion, - GetDashboardsResponse, - GetSubscriptionsResponse, - ProvisionNewCustomer, -) -from app.customer_provisioning.services.provision import provision_wazuh_customer -from app.db.db_session import get_db -from app.db.universal_models import Customers, CustomersMeta -from fastapi import APIRouter, Body, Depends, HTTPException, Security +from fastapi import APIRouter +from fastapi import Body +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.auth.utils import AuthHandler +from app.connectors.grafana.schema.dashboards import Office365Dashboard +from app.connectors.grafana.schema.dashboards import WazuhDashboard +from app.customer_provisioning.schema.provision import CustomerProvisionResponse +from app.customer_provisioning.schema.provision import CustomersMetaResponse +from app.customer_provisioning.schema.provision import CustomerSubsctipion +from app.customer_provisioning.schema.provision import GetDashboardsResponse +from app.customer_provisioning.schema.provision import GetSubscriptionsResponse +from app.customer_provisioning.schema.provision import ProvisionNewCustomer +from app.customer_provisioning.services.provision import provision_wazuh_customer +from app.db.db_session import get_db +from app.db.universal_models import Customers +from app.db.universal_models import CustomersMeta + customer_provisioning_router = APIRouter() diff --git a/backend/app/customer_provisioning/schema/decommission.py b/backend/app/customer_provisioning/schema/decommission.py index dcf5275b..19afb3fa 100644 --- a/backend/app/customer_provisioning/schema/decommission.py +++ b/backend/app/customer_provisioning/schema/decommission.py @@ -1,6 +1,7 @@ from typing import List -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class DecommissionedData(BaseModel): diff --git a/backend/app/customer_provisioning/schema/grafana.py b/backend/app/customer_provisioning/schema/grafana.py index 5551cbef..ac922cc7 100644 --- a/backend/app/customer_provisioning/schema/grafana.py +++ b/backend/app/customer_provisioning/schema/grafana.py @@ -1,7 +1,8 @@ from datetime import datetime from typing import Dict -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field # ! Organization ! # diff --git a/backend/app/customer_provisioning/schema/graylog.py b/backend/app/customer_provisioning/schema/graylog.py index 723eb2c7..02e84c23 100644 --- a/backend/app/customer_provisioning/schema/graylog.py +++ b/backend/app/customer_provisioning/schema/graylog.py @@ -1,6 +1,8 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field # ! INDEX SETS ! # diff --git a/backend/app/customer_provisioning/schema/provision.py b/backend/app/customer_provisioning/schema/provision.py index 57def4a0..b028a596 100644 --- a/backend/app/customer_provisioning/schema/provision.py +++ b/backend/app/customer_provisioning/schema/provision.py @@ -1,10 +1,14 @@ import re from enum import Enum -from typing import List, Optional +from typing import List +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest from app.db.universal_models import CustomersMeta -from pydantic import BaseModel, Field, validator class CustomerSubsctipion(Enum): diff --git a/backend/app/customer_provisioning/schema/wazuh_worker.py b/backend/app/customer_provisioning/schema/wazuh_worker.py index ede33500..295a39c8 100644 --- a/backend/app/customer_provisioning/schema/wazuh_worker.py +++ b/backend/app/customer_provisioning/schema/wazuh_worker.py @@ -1,4 +1,5 @@ -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class ProvisionWorkerRequest(BaseModel): diff --git a/backend/app/customer_provisioning/services/decommission.py b/backend/app/customer_provisioning/services/decommission.py index 3f3af7d6..b23b0794 100644 --- a/backend/app/customer_provisioning/services/decommission.py +++ b/backend/app/customer_provisioning/services/decommission.py @@ -1,21 +1,19 @@ import requests +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.customer_provisioning.schema.decommission import DecommissionCustomerResponse -from app.customer_provisioning.schema.wazuh_worker import ( - DecommissionWorkerRequest, - DecommissionWorkerResponse, -) +from app.customer_provisioning.schema.wazuh_worker import DecommissionWorkerRequest +from app.customer_provisioning.schema.wazuh_worker import DecommissionWorkerResponse from app.customer_provisioning.services.dfir_iris import delete_customer from app.customer_provisioning.services.grafana import delete_grafana_organization -from app.customer_provisioning.services.graylog import delete_index_set, delete_stream -from app.customer_provisioning.services.wazuh_manager import ( - delete_wazuh_agents, - delete_wazuh_groups, - gather_wazuh_agents, -) +from app.customer_provisioning.services.graylog import delete_index_set +from app.customer_provisioning.services.graylog import delete_stream +from app.customer_provisioning.services.wazuh_manager import delete_wazuh_agents +from app.customer_provisioning.services.wazuh_manager import delete_wazuh_groups +from app.customer_provisioning.services.wazuh_manager import gather_wazuh_agents from app.db.universal_models import CustomersMeta from app.utils import get_connector_attribute -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession async def decomission_wazuh_customer( diff --git a/backend/app/customer_provisioning/services/dfir_iris.py b/backend/app/customer_provisioning/services/dfir_iris.py index 8d1c6501..a1b16625 100644 --- a/backend/app/customer_provisioning/services/dfir_iris.py +++ b/backend/app/customer_provisioning/services/dfir_iris.py @@ -1,12 +1,12 @@ -from app.connectors.dfir_iris.schema.admin import CreateCustomerResponse, ListCustomers -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_admin, - initialize_client_and_customer, -) from fastapi import HTTPException from loguru import logger +from app.connectors.dfir_iris.schema.admin import CreateCustomerResponse +from app.connectors.dfir_iris.schema.admin import ListCustomers +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_admin +from app.connectors.dfir_iris.utils.universal import initialize_client_and_customer + async def check_customer_exists(customer_name: str) -> bool: """ diff --git a/backend/app/customer_provisioning/services/grafana.py b/backend/app/customer_provisioning/services/grafana.py index e64d84d6..a9602555 100644 --- a/backend/app/customer_provisioning/services/grafana.py +++ b/backend/app/customer_provisioning/services/grafana.py @@ -1,17 +1,16 @@ +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.connectors.grafana.utils.universal import create_grafana_client from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client -from app.customer_provisioning.schema.grafana import ( - GrafanaDatasource, - GrafanaDataSourceCreationResponse, - GrafanaFolderCreationResponse, - GrafanaOrganizationCreation, - NodesVersionResponse, -) +from app.customer_provisioning.schema.grafana import GrafanaDatasource +from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse +from app.customer_provisioning.schema.grafana import GrafanaFolderCreationResponse +from app.customer_provisioning.schema.grafana import GrafanaOrganizationCreation +from app.customer_provisioning.schema.grafana import NodesVersionResponse from app.customer_provisioning.schema.provision import ProvisionNewCustomer from app.utils import get_connector_attribute -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession ################# ! GRAFANA PROVISIONING ! ################# diff --git a/backend/app/customer_provisioning/services/graylog.py b/backend/app/customer_provisioning/services/graylog.py index ad7e5b28..e83d8e1a 100644 --- a/backend/app/customer_provisioning/services/graylog.py +++ b/backend/app/customer_provisioning/services/graylog.py @@ -1,23 +1,20 @@ import json from datetime import datetime -from app.connectors.graylog.services.pipelines import get_pipelines -from app.connectors.graylog.utils.universal import ( - send_delete_request, - send_post_request, -) -from app.customer_provisioning.schema.graylog import ( - GraylogIndexSetCreationResponse, - StreamConnectionToPipelineRequest, - StreamConnectionToPipelineResponse, - StreamCreationResponse, - TimeBasedIndexSet, - WazuhEventStream, -) -from app.customer_provisioning.schema.provision import ProvisionNewCustomer from fastapi import HTTPException from loguru import logger +from app.connectors.graylog.services.pipelines import get_pipelines +from app.connectors.graylog.utils.universal import send_delete_request +from app.connectors.graylog.utils.universal import send_post_request +from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse +from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest +from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineResponse +from app.customer_provisioning.schema.graylog import StreamCreationResponse +from app.customer_provisioning.schema.graylog import TimeBasedIndexSet +from app.customer_provisioning.schema.graylog import WazuhEventStream +from app.customer_provisioning.schema.provision import ProvisionNewCustomer + ######### ! GRAYLOG PROVISIONING ! ############ # ! INDEX SETS ! # diff --git a/backend/app/customer_provisioning/services/provision.py b/backend/app/customer_provisioning/services/provision.py index 33e26657..717d187f 100644 --- a/backend/app/customer_provisioning/services/provision.py +++ b/backend/app/customer_provisioning/services/provision.py @@ -1,41 +1,32 @@ import requests +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest from app.connectors.grafana.services.dashboards import provision_dashboards from app.connectors.graylog.services.management import start_stream from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest -from app.customer_provisioning.schema.provision import ( - CustomerProvisionMeta, - CustomerProvisionResponse, - ProvisionNewCustomer, -) -from app.customer_provisioning.schema.wazuh_worker import ( - ProvisionWorkerRequest, - ProvisionWorkerResponse, -) +from app.customer_provisioning.schema.provision import CustomerProvisionMeta +from app.customer_provisioning.schema.provision import CustomerProvisionResponse +from app.customer_provisioning.schema.provision import ProvisionNewCustomer +from app.customer_provisioning.schema.wazuh_worker import ProvisionWorkerRequest +from app.customer_provisioning.schema.wazuh_worker import ProvisionWorkerResponse from app.customer_provisioning.services.dfir_iris import create_customer -from app.customer_provisioning.services.grafana import ( - create_grafana_datasource, - create_grafana_folder, - create_grafana_organization, -) -from app.customer_provisioning.services.graylog import ( - connect_stream_to_pipeline, - create_event_stream, - create_index_set, - get_pipeline_id, -) -from app.customer_provisioning.services.wazuh_manager import ( - apply_group_configurations, - create_wazuh_groups, -) +from app.customer_provisioning.services.grafana import create_grafana_datasource +from app.customer_provisioning.services.grafana import create_grafana_folder +from app.customer_provisioning.services.grafana import create_grafana_organization +from app.customer_provisioning.services.graylog import connect_stream_to_pipeline +from app.customer_provisioning.services.graylog import create_event_stream +from app.customer_provisioning.services.graylog import create_index_set +from app.customer_provisioning.services.graylog import get_pipeline_id +from app.customer_provisioning.services.wazuh_manager import apply_group_configurations +from app.customer_provisioning.services.wazuh_manager import create_wazuh_groups from app.db.universal_models import CustomersMeta from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, ) from app.utils import get_connector_attribute -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession # ! MAIN FUNCTION ! # @@ -61,9 +52,7 @@ async def provision_wazuh_customer( # Initialize an empty dictionary to store the meta data provision_meta_data = {} provision_meta_data["index_set_id"] = (await create_index_set(request)).data.id - provision_meta_data["stream_id"] = ( - await create_event_stream(request, provision_meta_data["index_set_id"]) - ).data.stream_id + provision_meta_data["stream_id"] = (await create_event_stream(request, provision_meta_data["index_set_id"])).data.stream_id provision_meta_data["pipeline_ids"] = await get_pipeline_id(subscription="Wazuh") stream_and_pipeline = StreamConnectionToPipelineRequest( stream_id=provision_meta_data["stream_id"], @@ -77,9 +66,7 @@ async def provision_wazuh_customer( ) await create_wazuh_groups(request) await apply_group_configurations(request) - provision_meta_data["grafana_organization_id"] = ( - await create_grafana_organization(request) - ).orgId + provision_meta_data["grafana_organization_id"] = (await create_grafana_organization(request)).orgId provision_meta_data["wazuh_datasource_uid"] = ( await create_grafana_datasource( request=request, @@ -102,9 +89,7 @@ async def provision_wazuh_customer( ), ) - provision_meta_data["iris_customer_id"] = ( - await create_customer(request.customer_name) - ).data.customer_id + provision_meta_data["iris_customer_id"] = (await create_customer(request.customer_name)).data.customer_id customer_provision_meta = CustomerProvisionMeta(**provision_meta_data) customer_meta = await update_customer_meta_table( diff --git a/backend/app/customer_provisioning/services/wazuh_manager.py b/backend/app/customer_provisioning/services/wazuh_manager.py index 1fc462ab..6efc4436 100644 --- a/backend/app/customer_provisioning/services/wazuh_manager.py +++ b/backend/app/customer_provisioning/services/wazuh_manager.py @@ -1,6 +1,8 @@ from pathlib import Path from typing import List +from loguru import logger + from app.connectors.wazuh_manager.utils.universal import ( send_delete_request as send_wazuh_delete_request, ) @@ -15,7 +17,6 @@ ) from app.customer_provisioning.schema.provision import ProvisionNewCustomer from app.customer_provisioning.schema.wazuh_manager import WazuhAgentsTemplatePaths -from loguru import logger ######### ! WAZUH MANAGER PROVISIONING ! ############ diff --git a/backend/app/customers/routes/customers.py b/backend/app/customers/routes/customers.py index f9522fea..1c067080 100644 --- a/backend/app/customers/routes/customers.py +++ b/backend/app/customers/routes/customers.py @@ -1,32 +1,33 @@ -from app.auth.utils import AuthHandler - -# App specific imports -from app.customers.schema.customers import ( - AgentModel, - AgentsResponse, - CustomerFullResponse, - CustomerMetaRequestBody, - CustomerMetaResponse, - CustomerRequestBody, - CustomerResponse, - CustomersResponse, -) -from app.db.db_session import get_db -from app.db.universal_models import Agents, Customers, CustomersMeta -from app.healthchecks.agents.schema.agents import ( - AgentHealthCheckResponse, - TimeCriteriaModel, -) -from app.healthchecks.agents.services.agents import ( - velociraptor_agents_healthcheck, - wazuh_agents_healthcheck, -) -from fastapi import APIRouter, Depends, HTTPException, Query, Security +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Query +from fastapi import Security from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from starlette.status import HTTP_401_UNAUTHORIZED +from app.auth.utils import AuthHandler + +# App specific imports +from app.customers.schema.customers import AgentModel +from app.customers.schema.customers import AgentsResponse +from app.customers.schema.customers import CustomerFullResponse +from app.customers.schema.customers import CustomerMetaRequestBody +from app.customers.schema.customers import CustomerMetaResponse +from app.customers.schema.customers import CustomerRequestBody +from app.customers.schema.customers import CustomerResponse +from app.customers.schema.customers import CustomersResponse +from app.db.db_session import get_db +from app.db.universal_models import Agents +from app.db.universal_models import Customers +from app.db.universal_models import CustomersMeta +from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse +from app.healthchecks.agents.schema.agents import TimeCriteriaModel +from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck +from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck + customers_router = APIRouter() diff --git a/backend/app/customers/schema/customers.py b/backend/app/customers/schema/customers.py index 6056cd45..e064ab76 100644 --- a/backend/app/customers/schema/customers.py +++ b/backend/app/customers/schema/customers.py @@ -1,7 +1,9 @@ from datetime import datetime -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class CustomerRequestBody(BaseModel): diff --git a/backend/app/db/db_populate.py b/backend/app/db/db_populate.py index afce8e74..8b62fe85 100644 --- a/backend/app/db/db_populate.py +++ b/backend/app/db/db_populate.py @@ -1,17 +1,18 @@ import os -from app.auth.models.users import Role -from app.connectors.models import Connectors -from app.integrations.models.customer_integration_settings import ( - AvailableIntegrations, - AvailableIntegrationsAuthKeys, -) from dotenv import load_dotenv from loguru import logger from sqlalchemy import and_ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.auth.models.users import Role +from app.connectors.models import Connectors +from app.integrations.models.customer_integration_settings import AvailableIntegrations +from app.integrations.models.customer_integration_settings import ( + AvailableIntegrationsAuthKeys, +) + load_dotenv() @@ -287,8 +288,7 @@ async def add_available_integrations_if_not_exist(session: AsyncSession): for available_integration_data in available_integrations_list: query = select(AvailableIntegrations).where( - AvailableIntegrations.integration_name - == available_integration_data["integration_name"], + AvailableIntegrations.integration_name == available_integration_data["integration_name"], ) result = await session.execute(query) existing_available_integration = result.scalars().first() @@ -382,28 +382,21 @@ async def add_available_integrations_auth_keys_if_not_exist(session: AsyncSessio Returns: None """ - available_integrations_auth_keys_list = ( - await get_available_integrations_auth_keys_list(session=session) - ) + available_integrations_auth_keys_list = await get_available_integrations_auth_keys_list(session=session) for available_integration_auth_keys_data in available_integrations_auth_keys_list: query = select(AvailableIntegrations).where( - AvailableIntegrations.integration_name - == available_integration_auth_keys_data["integration_name"], + AvailableIntegrations.integration_name == available_integration_auth_keys_data["integration_name"], ) result = await session.execute(query) existing_integration = result.scalars().first() if existing_integration: - available_integration_auth_keys_data[ - "integration_id" - ] = existing_integration.id + available_integration_auth_keys_data["integration_id"] = existing_integration.id auth_key_query = select(AvailableIntegrationsAuthKeys).where( and_( - AvailableIntegrationsAuthKeys.integration_id - == existing_integration.id, - AvailableIntegrationsAuthKeys.auth_key_name - == available_integration_auth_keys_data["auth_key_name"], + AvailableIntegrationsAuthKeys.integration_id == existing_integration.id, + AvailableIntegrationsAuthKeys.auth_key_name == available_integration_auth_keys_data["auth_key_name"], ), ) auth_key_result = await session.execute(auth_key_query) diff --git a/backend/app/db/db_session.py b/backend/app/db/db_session.py index 66059e00..bb0f8ffc 100644 --- a/backend/app/db/db_session.py +++ b/backend/app/db/db_session.py @@ -1,6 +1,8 @@ # ! Old Testing without Async +from sqlmodel import Session +from sqlmodel import create_engine + from settings import SQLALCHEMY_DATABASE_URI -from sqlmodel import Session, create_engine engine = create_engine( SQLALCHEMY_DATABASE_URI, @@ -8,14 +10,17 @@ ) session = "placeholder" -from contextlib import asynccontextmanager, contextmanager +from contextlib import asynccontextmanager +from contextlib import contextmanager from loguru import logger -from settings import SQLALCHEMY_DATABASE_URI from sqlalchemy import create_engine -from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.orm import sessionmaker +from settings import SQLALCHEMY_DATABASE_URI + # create async engine for SQLite using aiosqlite async_engine = create_async_engine(SQLALCHEMY_DATABASE_URI, echo=False) sync_engine = create_engine( diff --git a/backend/app/db/db_setup.py b/backend/app/db/db_setup.py index f6b3b1c2..48fde787 100644 --- a/backend/app/db/db_setup.py +++ b/backend/app/db/db_setup.py @@ -1,20 +1,17 @@ -from app.auth.services.universal import ( - create_admin_user, - create_scheduler_user, - remove_scheduler_user, -) -from app.db.db_populate import ( - add_available_integrations_auth_keys_if_not_exist, - add_available_integrations_if_not_exist, - add_connectors_if_not_exist, - add_roles_if_not_exist, -) from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession # ! New with Async from sqlmodel import SQLModel +from app.auth.services.universal import create_admin_user +from app.auth.services.universal import create_scheduler_user +from app.auth.services.universal import remove_scheduler_user +from app.db.db_populate import add_available_integrations_auth_keys_if_not_exist +from app.db.db_populate import add_available_integrations_if_not_exist +from app.db.db_populate import add_connectors_if_not_exist +from app.db.db_populate import add_roles_if_not_exist + async def create_tables(async_engine): """ diff --git a/backend/app/db/universal_models.py b/backend/app/db/universal_models.py index d372eeb8..25246ece 100644 --- a/backend/app/db/universal_models.py +++ b/backend/app/db/universal_models.py @@ -1,7 +1,9 @@ from datetime import datetime from typing import Optional -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field +from sqlmodel import Relationship +from sqlmodel import SQLModel class Customers(SQLModel, table=True): @@ -71,22 +73,12 @@ def update_from_model(self, customer_meta): self.customer_meta_grafana_org_id = customer_meta.customer_meta_grafana_org_id self.customer_meta_wazuh_group = customer_meta.customer_meta_wazuh_group self.customer_meta_index_retention = customer_meta.customer_meta_index_retention - self.customer_meta_wazuh_registration_port = ( - customer_meta.customer_meta_wazuh_registration_port - ) - self.customer_meta_wazuh_log_ingestion_port = ( - customer_meta.customer_meta_wazuh_log_ingestion_port - ) + self.customer_meta_wazuh_registration_port = customer_meta.customer_meta_wazuh_registration_port + self.customer_meta_wazuh_log_ingestion_port = customer_meta.customer_meta_wazuh_log_ingestion_port self.customer_meta_wazuh_api_port = customer_meta.customer_meta_wazuh_api_port - self.customer_meta_wazuh_auth_password = ( - customer_meta.customer_meta_wazuh_auth_password - ) - self.customer_meta_iris_customer_id = ( - customer_meta.customer_meta_iris_customer_id - ) - self.customer_meta_office365_organization_id = ( - customer_meta.customer_meta_office365_organization_id - ) + self.customer_meta_wazuh_auth_password = customer_meta.customer_meta_wazuh_auth_password + self.customer_meta_iris_customer_id = customer_meta.customer_meta_iris_customer_id + self.customer_meta_office365_organization_id = customer_meta.customer_meta_office365_organization_id class Agents(SQLModel, table=True): @@ -111,9 +103,7 @@ class Agents(SQLModel, table=True): def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code): # Check if agent_last_seen is 'Unknown' and set wazuh_last_seen accordingly if wazuh_agent.agent_last_seen == "Unknown": - wazuh_last_seen_value = ( - "1970-01-01T00:00:00+00:00" # default datetime value - ) + wazuh_last_seen_value = "1970-01-01T00:00:00+00:00" # default datetime value else: wazuh_last_seen_value = wazuh_agent.agent_last_seen_as_datetime @@ -125,19 +115,14 @@ def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code): label=wazuh_agent.agent_label, wazuh_last_seen=wazuh_last_seen_value, wazuh_agent_version=wazuh_agent.wazuh_agent_version, - velociraptor_id=velociraptor_agent.client_id - if velociraptor_agent.client_id - else "n/a", + velociraptor_id=velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a", velociraptor_last_seen=velociraptor_agent.client_last_seen_as_datetime, velociraptor_agent_version=velociraptor_agent.client_version, customer_code=customer_code, ) def update_from_model(self, wazuh_agent, velociraptor_agent, customer_code): - if ( - wazuh_agent.agent_last_seen == "Unknown" - or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00" - ): + if wazuh_agent.agent_last_seen == "Unknown" or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00": wazuh_last_seen_value = datetime.strptime( "1970-01-01T00:00:00+00:00", "%Y-%m-%dT%H:%M:%S%z", @@ -152,9 +137,7 @@ def update_from_model(self, wazuh_agent, velociraptor_agent, customer_code): self.label = wazuh_agent.agent_label self.wazuh_last_seen = wazuh_last_seen_value self.wazuh_agent_version = wazuh_agent.wazuh_agent_version - self.velociraptor_id = ( - velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a" - ) + self.velociraptor_id = velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a" self.velociraptor_last_seen = velociraptor_agent.client_last_seen_as_datetime self.velociraptor_agent_version = velociraptor_agent.client_version self.customer_code = customer_code diff --git a/backend/app/healthchecks/agents/routes/agents.py b/backend/app/healthchecks/agents/routes/agents.py index a9b810f3..0d74087e 100644 --- a/backend/app/healthchecks/agents/routes/agents.py +++ b/backend/app/healthchecks/agents/routes/agents.py @@ -1,24 +1,25 @@ -from app.auth.utils import AuthHandler -from app.db.db_session import get_db -from app.db.universal_models import Agents -from app.healthchecks.agents.schema.agents import ( - AgentHealthCheckResponse, - HostLogsSearchBody, - HostLogsSearchResponse, - TimeCriteriaModel, -) -from app.healthchecks.agents.services.agents import ( - host_logs, - velociraptor_agent_healthcheck, - velociraptor_agents_healthcheck, - wazuh_agent_healthcheck, - wazuh_agents_healthcheck, -) -from fastapi import APIRouter, Depends, HTTPException, Query, Security +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Query +from fastapi import Security from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.auth.utils import AuthHandler +from app.db.db_session import get_db +from app.db.universal_models import Agents +from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse +from app.healthchecks.agents.schema.agents import HostLogsSearchBody +from app.healthchecks.agents.schema.agents import HostLogsSearchResponse +from app.healthchecks.agents.schema.agents import TimeCriteriaModel +from app.healthchecks.agents.services.agents import host_logs +from app.healthchecks.agents.services.agents import velociraptor_agent_healthcheck +from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck +from app.healthchecks.agents.services.agents import wazuh_agent_healthcheck +from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck + healtcheck_agents_router = APIRouter() diff --git a/backend/app/healthchecks/agents/schema/agents.py b/backend/app/healthchecks/agents/schema/agents.py index 40f7e3be..a85c6816 100644 --- a/backend/app/healthchecks/agents/schema/agents.py +++ b/backend/app/healthchecks/agents/schema/agents.py @@ -1,7 +1,12 @@ from datetime import datetime -from typing import Any, Dict, List, Optional - -from pydantic import BaseModel, Field, validator +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator class AgentModel(BaseModel): diff --git a/backend/app/healthchecks/agents/services/agents.py b/backend/app/healthchecks/agents/services/agents.py index c91a235e..691418ab 100644 --- a/backend/app/healthchecks/agents/services/agents.py +++ b/backend/app/healthchecks/agents/services/agents.py @@ -1,24 +1,23 @@ -from datetime import datetime, timedelta -from typing import Optional, Type - -from app.connectors.wazuh_indexer.utils.universal import ( - LogsQueryBuilder, - collect_indices, - create_wazuh_indexer_client, -) -from app.healthchecks.agents.schema.agents import ( - AgentHealthCheckResponse, - AgentModel, - CollectLogsResponse, - ExtendedAgentModel, - HostLogsSearchBody, - HostLogsSearchResponse, - LogsSearchBody, - TimeCriteriaModel, -) +from datetime import datetime +from datetime import timedelta +from typing import Optional +from typing import Type + from fastapi import HTTPException from loguru import logger +from app.connectors.wazuh_indexer.utils.universal import LogsQueryBuilder +from app.connectors.wazuh_indexer.utils.universal import collect_indices +from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client +from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse +from app.healthchecks.agents.schema.agents import AgentModel +from app.healthchecks.agents.schema.agents import CollectLogsResponse +from app.healthchecks.agents.schema.agents import ExtendedAgentModel +from app.healthchecks.agents.schema.agents import HostLogsSearchBody +from app.healthchecks.agents.schema.agents import HostLogsSearchResponse +from app.healthchecks.agents.schema.agents import LogsSearchBody +from app.healthchecks.agents.schema.agents import TimeCriteriaModel + def is_wazuh_agent_unhealthy( agent: AgentModel, @@ -44,9 +43,7 @@ def is_wazuh_agent_unhealthy( return ExtendedAgentModel(**agent.dict(), unhealthy_wazuh_agent=True) # Calculate the total time delta based on the criteria - total_minutes = ( - time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60 - ) + total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60 time_delta = timedelta(minutes=total_minutes) is_unhealthy = (current_time - wazuh_last_seen) > time_delta @@ -77,9 +74,7 @@ def is_velociraptor_agent_unhealthy( return ExtendedAgentModel(**agent.dict(), unhealthy_velociraptor_agent=True) # Calculate the total time delta based on the criteria - total_minutes = ( - time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60 - ) + total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60 time_delta = timedelta(minutes=total_minutes) is_unhealthy = (current_time - velociraptor_last_seen) > time_delta @@ -277,9 +272,7 @@ async def get_logs_generic( ) logs_summary = [] indices = await collect_indices() - index_list = ( - [index_name] if index_name else indices.indices_list - ) # Use the provided index_name or get all indices + index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices for index_name in index_list: try: diff --git a/backend/app/integrations/alert_creation/general/routes/alert.py b/backend/app/integrations/alert_creation/general/routes/alert.py index eac0155b..1ae3b9e2 100644 --- a/backend/app/integrations/alert_creation/general/routes/alert.py +++ b/backend/app/integrations/alert_creation/general/routes/alert.py @@ -1,16 +1,17 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.db.db_session import get_db -from app.integrations.alert_creation.general.schema.alert import ( - CreateAlertRequest, - CreateAlertResponse, -) +from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest +from app.integrations.alert_creation.general.schema.alert import CreateAlertResponse from app.integrations.alert_creation.general.services.alert import create_alert from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, ) -from fastapi import APIRouter, Depends, HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select general_alerts_router = APIRouter() @@ -35,15 +36,12 @@ async def is_rule_id_valid( result = await session.execute( select(AlertCreationSettings).where( - AlertCreationSettings.customer_code - == create_alert_request.agent_labels_customer, + AlertCreationSettings.customer_code == create_alert_request.agent_labels_customer, ), ) settings = result.scalars().first() - if settings and str(create_alert_request.rule_id) in ( - settings.excluded_wazuh_rules or "" - ).split(","): + if settings and str(create_alert_request.rule_id) in (settings.excluded_wazuh_rules or "").split(","): return False return True @@ -69,8 +67,7 @@ async def is_customer_code_valid( result = await session.execute( select(AlertCreationSettings).where( - AlertCreationSettings.customer_code - == create_alert_request.agent_labels_customer, + AlertCreationSettings.customer_code == create_alert_request.agent_labels_customer, ), ) settings = result.scalars().first() diff --git a/backend/app/integrations/alert_creation/general/schema/alert.py b/backend/app/integrations/alert_creation/general/schema/alert.py index 9edf76d7..e8e99035 100644 --- a/backend/app/integrations/alert_creation/general/schema/alert.py +++ b/backend/app/integrations/alert_creation/general/schema/alert.py @@ -1,7 +1,12 @@ from enum import Enum -from typing import Any, Dict, List, Optional - -from pydantic import BaseModel, Extra, Field +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from pydantic import BaseModel +from pydantic import Extra +from pydantic import Field class ValidIocFields(Enum): diff --git a/backend/app/integrations/alert_creation/general/services/alert.py b/backend/app/integrations/alert_creation/general/services/alert.py index 00c2008e..aee1f9a4 100644 --- a/backend/app/integrations/alert_creation/general/services/alert.py +++ b/backend/app/integrations/alert_creation/general/services/alert.py @@ -1,33 +1,29 @@ -from typing import Optional, Set +from typing import Optional +from typing import Set + +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession from app.agents.routes.agents import get_agent from app.agents.schema.agents import AgentsResponse -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) -from app.integrations.alert_creation.general.schema.alert import ( - CreateAlertRequest, - CreateAlertResponse, - IrisAlertContext, - IrisAlertPayload, - IrisAsset, - IrisIoc, - ValidIocFields, -) +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert +from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest +from app.integrations.alert_creation.general.schema.alert import CreateAlertResponse +from app.integrations.alert_creation.general.schema.alert import IrisAlertContext +from app.integrations.alert_creation.general.schema.alert import IrisAlertPayload +from app.integrations.alert_creation.general.schema.alert import IrisAsset +from app.integrations.alert_creation.general.schema.alert import IrisIoc +from app.integrations.alert_creation.general.schema.alert import ValidIocFields from app.integrations.alert_creation.general.services.alert_multi_exclude import ( AlertDetailsService, ) -from app.integrations.utils.alerts import ( - get_asset_type_id, - send_to_shuffle, - validate_ioc_type, -) +from app.integrations.utils.alerts import get_asset_type_id +from app.integrations.utils.alerts import send_to_shuffle +from app.integrations.utils.alerts import validate_ioc_type from app.integrations.utils.schema import ShufflePayload from app.utils import get_customer_alert_settings -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession def valid_ioc_fields() -> Set[str]: @@ -57,13 +53,8 @@ async def construct_alert_source_link( The alert source link. """ # Check if the alert has a process id and that it is not "No process ID found" - if ( - hasattr(alert_details, "process_id") - and alert_details.process_id != "No process ID found" - ): - query_string = ( - f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20" - ) + if hasattr(alert_details, "process_id") and alert_details.process_id != "No process ID found": + query_string = f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20" else: query_string = f"%22query%22:%22_id:%5C%22{alert_details.id}%5C%22%20AND%20" diff --git a/backend/app/integrations/alert_creation/general/services/alert_multi_exclude.py b/backend/app/integrations/alert_creation/general/services/alert_multi_exclude.py index 85d15fa4..a2ecaff7 100644 --- a/backend/app/integrations/alert_creation/general/services/alert_multi_exclude.py +++ b/backend/app/integrations/alert_creation/general/services/alert_multi_exclude.py @@ -1,13 +1,17 @@ -from typing import Any, Dict, List, Tuple +from typing import Any +from typing import Dict +from typing import List +from typing import Tuple + +from elasticsearch7 import NotFoundError +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationEventConfig, ) from app.utils import get_customer_alert_event_configs -from elasticsearch7 import NotFoundError -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession class AlertDetailsService: @@ -65,9 +69,7 @@ def search_alerts_with_syslog_level(self) -> List[Tuple[str, str]]: result = self.es.search(index="_all", body=query) # Extract (index, id) pairs from the result - index_id_pairs = [ - (hit["_index"], hit["_id"]) for hit in result["hits"]["hits"] - ] + index_id_pairs = [(hit["_index"], hit["_id"]) for hit in result["hits"]["hits"]] logger.info( f"Found {len(index_id_pairs)} alerts with syslog_level of 'ALERT' within the last 1 hour.", ) @@ -237,9 +239,7 @@ async def collect_alert_timeline_process_id( logger.info(f"Total alert timeline hits: {total_hits}") # Build and sort the list of events - events = [ - event["_source"] for event in alert_timeline_events["hits"]["hits"] - ] + events = [event["_source"] for event in alert_timeline_events["hits"]["hits"]] events.sort(key=lambda x: x["timestamp_utc"]) # return self.process_events(events) diff --git a/backend/app/integrations/alert_creation/office365/routes/alert.py b/backend/app/integrations/alert_creation/office365/routes/alert.py index 178fbb5e..59105467 100644 --- a/backend/app/integrations/alert_creation/office365/routes/alert.py +++ b/backend/app/integrations/alert_creation/office365/routes/alert.py @@ -1,13 +1,28 @@ # from app.alerts.office365.services.threat_intel import create_threat_intel_alert +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from loguru import logger +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + from app.db.db_session import get_db from app.integrations.alert_creation.office365.schema.exchange import ( Office365ExchangeAlertBase, +) +from app.integrations.alert_creation.office365.schema.exchange import ( Office365ExchangeAlertRequest, +) +from app.integrations.alert_creation.office365.schema.exchange import ( Office365ExchangeAlertResponse, +) +from app.integrations.alert_creation.office365.schema.exchange import ( ValidOffice365Workloads, ) from app.integrations.alert_creation.office365.schema.threat_intel import ( Office365ThreatIntelAlertRequest, +) +from app.integrations.alert_creation.office365.schema.threat_intel import ( Office365ThreatIntelAlertResponse, ) from app.integrations.alert_creation.office365.services.exchange import ( @@ -19,10 +34,6 @@ from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, ) -from fastapi import APIRouter, Depends, HTTPException -from loguru import logger -from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession office365_alerts_router = APIRouter() @@ -47,8 +58,7 @@ async def is_office365_organization_id_valid( result = await session.execute( select(AlertCreationSettings).where( - AlertCreationSettings.office365_organization_id - == create_alert_request.data_office365_OrganizationId, + AlertCreationSettings.office365_organization_id == create_alert_request.data_office365_OrganizationId, ), ) settings = result.scalars().first() @@ -71,9 +81,7 @@ async def create_office365_exchange_alert( session: AsyncSession = Depends(get_db), ): logger.info(f"create_alert_request: {create_alert_request}") - if create_alert_request.data_office365_Workload not in [ - workload.value for workload in ValidOffice365Workloads - ]: + if create_alert_request.data_office365_Workload not in [workload.value for workload in ValidOffice365Workloads]: logger.info(f"Invalid workload: {create_alert_request.data_office365_Workload}") raise HTTPException(status_code=400, detail="Invalid workload") logger.info(f"Workload is valid: {create_alert_request.data_office365_Workload}") @@ -91,9 +99,7 @@ async def create_office365_threat_intel_alert( session: AsyncSession = Depends(get_db), ): logger.info(f"create_alert_request: {create_alert_request}") - if create_alert_request.data_office365_Workload not in [ - workload.value for workload in ValidOffice365Workloads - ]: + if create_alert_request.data_office365_Workload not in [workload.value for workload in ValidOffice365Workloads]: logger.info(f"Invalid workload: {create_alert_request.data_office365_Workload}") raise HTTPException(status_code=400, detail="Invalid workload") logger.info(f"Workload is valid: {create_alert_request.data_office365_Workload}") diff --git a/backend/app/integrations/alert_creation/office365/schema/exchange.py b/backend/app/integrations/alert_creation/office365/schema/exchange.py index 41400dd6..725e9daf 100644 --- a/backend/app/integrations/alert_creation/office365/schema/exchange.py +++ b/backend/app/integrations/alert_creation/office365/schema/exchange.py @@ -1,7 +1,12 @@ from enum import Enum -from typing import Any, Dict, List, Optional - -from pydantic import BaseModel, Extra, Field +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from pydantic import BaseModel +from pydantic import Extra +from pydantic import Field class ValidOffice365Workloads(Enum): diff --git a/backend/app/integrations/alert_creation/office365/schema/threat_intel.py b/backend/app/integrations/alert_creation/office365/schema/threat_intel.py index 7d382376..957b92be 100644 --- a/backend/app/integrations/alert_creation/office365/schema/threat_intel.py +++ b/backend/app/integrations/alert_creation/office365/schema/threat_intel.py @@ -1,7 +1,12 @@ from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Extra, Field +from pydantic import BaseModel +from pydantic import Extra +from pydantic import Field class ValidOffice365Workloads(Enum): diff --git a/backend/app/integrations/alert_creation/office365/services/exchange.py b/backend/app/integrations/alert_creation/office365/services/exchange.py index 469bbe86..cb2953dd 100644 --- a/backend/app/integrations/alert_creation/office365/services/exchange.py +++ b/backend/app/integrations/alert_creation/office365/services/exchange.py @@ -1,23 +1,26 @@ -from typing import Optional, Set +from typing import Optional +from typing import Set -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert from app.integrations.alert_creation.general.schema.alert import ValidIocFields +from app.integrations.alert_creation.office365.schema.exchange import IrisAlertContext +from app.integrations.alert_creation.office365.schema.exchange import IrisAlertPayload +from app.integrations.alert_creation.office365.schema.exchange import IrisAsset +from app.integrations.alert_creation.office365.schema.exchange import IrisIoc from app.integrations.alert_creation.office365.schema.exchange import ( - IrisAlertContext, - IrisAlertPayload, - IrisAsset, - IrisIoc, Office365ExchangeAlertRequest, +) +from app.integrations.alert_creation.office365.schema.exchange import ( Office365ExchangeAlertResponse, ) -from app.integrations.utils.alerts import send_to_shuffle, validate_ioc_type +from app.integrations.utils.alerts import send_to_shuffle +from app.integrations.utils.alerts import validate_ioc_type from app.integrations.utils.schema import ShufflePayload from app.utils import get_customer_alert_settings_office365 -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession def valid_ioc_fields() -> Set[str]: diff --git a/backend/app/integrations/alert_creation/office365/services/threat_intel.py b/backend/app/integrations/alert_creation/office365/services/threat_intel.py index 4e351288..da9fe334 100644 --- a/backend/app/integrations/alert_creation/office365/services/threat_intel.py +++ b/backend/app/integrations/alert_creation/office365/services/threat_intel.py @@ -1,23 +1,30 @@ -from typing import Optional, Set +from typing import Optional +from typing import Set -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert from app.integrations.alert_creation.general.schema.alert import ValidIocFields from app.integrations.alert_creation.office365.schema.threat_intel import ( IrisAlertContext, +) +from app.integrations.alert_creation.office365.schema.threat_intel import ( IrisAlertPayload, - IrisAsset, - IrisIoc, +) +from app.integrations.alert_creation.office365.schema.threat_intel import IrisAsset +from app.integrations.alert_creation.office365.schema.threat_intel import IrisIoc +from app.integrations.alert_creation.office365.schema.threat_intel import ( Office365ThreatIntelAlertRequest, +) +from app.integrations.alert_creation.office365.schema.threat_intel import ( Office365ThreatIntelAlertResponse, ) -from app.integrations.utils.alerts import send_to_shuffle, validate_ioc_type +from app.integrations.utils.alerts import send_to_shuffle +from app.integrations.utils.alerts import validate_ioc_type from app.integrations.utils.schema import ShufflePayload from app.utils import get_customer_alert_settings_office365 -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession def valid_ioc_fields() -> Set[str]: diff --git a/backend/app/integrations/alert_creation_settings/models/alert_creation_settings.py b/backend/app/integrations/alert_creation_settings/models/alert_creation_settings.py index 98e7001b..a66a5548 100644 --- a/backend/app/integrations/alert_creation_settings/models/alert_creation_settings.py +++ b/backend/app/integrations/alert_creation_settings/models/alert_creation_settings.py @@ -1,6 +1,9 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field +from sqlmodel import Relationship +from sqlmodel import SQLModel class Condition(SQLModel, table=True): diff --git a/backend/app/integrations/alert_creation_settings/routes/alert_creation_settings.py b/backend/app/integrations/alert_creation_settings/routes/alert_creation_settings.py index 7cdb4d62..0e5065e9 100644 --- a/backend/app/integrations/alert_creation_settings/routes/alert_creation_settings.py +++ b/backend/app/integrations/alert_creation_settings/routes/alert_creation_settings.py @@ -1,24 +1,39 @@ from typing import List +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select +from sqlalchemy.orm import joinedload + from app.db.db_session import get_db from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationEventConfig, +) +from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, +) +from app.integrations.alert_creation_settings.models.alert_creation_settings import ( EventOrder, ) from app.integrations.alert_creation_settings.schema.alert_creation_settings import ( AlertCreationEventConfigResponse, +) +from app.integrations.alert_creation_settings.schema.alert_creation_settings import ( AlertCreationSettingsCreate, +) +from app.integrations.alert_creation_settings.schema.alert_creation_settings import ( AlertCreationSettingsResponse, +) +from app.integrations.alert_creation_settings.schema.alert_creation_settings import ( EventOrderCreate, +) +from app.integrations.alert_creation_settings.schema.alert_creation_settings import ( EventOrderResponse, ) from app.utils import get_customer_alert_event_configs -from fastapi import APIRouter, Depends, HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select -from sqlalchemy.orm import joinedload alert_creation_settings_router = APIRouter() @@ -79,8 +94,7 @@ async def create_alert_creation_settings( result = await session.execute( select(AlertCreationSettings).where( - AlertCreationSettings.customer_code - == alert_creation_settings.customer_code, + AlertCreationSettings.customer_code == alert_creation_settings.customer_code, ), ) settings = result.scalars().first() @@ -215,9 +229,7 @@ async def add_event_order( # Query the EventOrder instance again to ensure event_configs are loaded result = await session.execute( - select(EventOrder) - .options(joinedload(EventOrder.event_configs)) - .where(EventOrder.id == event_order_db.id), + select(EventOrder).options(joinedload(EventOrder.event_configs)).where(EventOrder.id == event_order_db.id), ) event_order_db = result.scalars().first() @@ -266,11 +278,7 @@ async def update_event_orders( for event_order in event_orders: # Check if an EventOrder with the given order_label already exists existing_order = next( - ( - order - for order in settings.event_orders - if order.order_label == event_order.order_label - ), + (order for order in settings.event_orders if order.order_label == event_order.order_label), None, ) diff --git a/backend/app/integrations/alert_creation_settings/schema/alert_creation_settings.py b/backend/app/integrations/alert_creation_settings/schema/alert_creation_settings.py index 28c6ba85..a08fda54 100644 --- a/backend/app/integrations/alert_creation_settings/schema/alert_creation_settings.py +++ b/backend/app/integrations/alert_creation_settings/schema/alert_creation_settings.py @@ -1,4 +1,5 @@ -from typing import List, Optional +from typing import List +from typing import Optional from pydantic import BaseModel diff --git a/backend/app/integrations/alert_escalation/routes/general_alert.py b/backend/app/integrations/alert_escalation/routes/general_alert.py index 6795e61c..89f6c445 100644 --- a/backend/app/integrations/alert_escalation/routes/general_alert.py +++ b/backend/app/integrations/alert_escalation/routes/general_alert.py @@ -1,13 +1,14 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler from app.db.db_session import get_db -from app.integrations.alert_escalation.schema.general_alert import ( - CreateAlertRequest, - CreateAlertResponse, -) +from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest +from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse from app.integrations.alert_escalation.services.general_alert import create_alert -from fastapi import APIRouter, Depends, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession integration_general_alerts_router = APIRouter() diff --git a/backend/app/integrations/alert_escalation/schema/general_alert.py b/backend/app/integrations/alert_escalation/schema/general_alert.py index d1c005d6..e0bac64e 100644 --- a/backend/app/integrations/alert_escalation/schema/general_alert.py +++ b/backend/app/integrations/alert_escalation/schema/general_alert.py @@ -1,7 +1,12 @@ from enum import Enum -from typing import Any, Dict, List, Optional - -from pydantic import BaseModel, Extra, Field +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from pydantic import BaseModel +from pydantic import Extra +from pydantic import Field class ValidIocFields(Enum): diff --git a/backend/app/integrations/alert_escalation/services/general_alert.py b/backend/app/integrations/alert_escalation/services/general_alert.py index 1a1c3aea..bbb9c36d 100644 --- a/backend/app/integrations/alert_escalation/services/general_alert.py +++ b/backend/app/integrations/alert_escalation/services/general_alert.py @@ -1,34 +1,33 @@ -from typing import Optional, Set +from typing import Optional +from typing import Set + +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select # from app.integrations.alert_escalation.utils.universal import get_agent_data from app.agents.routes.agents import get_agent from app.agents.schema.agents import AgentsResponse -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert from app.connectors.utils import get_connector_info_from_db from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, ) -from app.integrations.alert_escalation.schema.general_alert import ( - CreateAlertRequest, - CreateAlertResponse, - GenericAlertModel, - GenericSourceModel, - IrisAlertContext, - IrisAlertPayload, - IrisAsset, - IrisIoc, - ValidIocFields, -) -from app.integrations.utils.alerts import get_asset_type_id, validate_ioc_type +from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest +from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse +from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel +from app.integrations.alert_escalation.schema.general_alert import GenericSourceModel +from app.integrations.alert_escalation.schema.general_alert import IrisAlertContext +from app.integrations.alert_escalation.schema.general_alert import IrisAlertPayload +from app.integrations.alert_escalation.schema.general_alert import IrisAsset +from app.integrations.alert_escalation.schema.general_alert import IrisIoc +from app.integrations.alert_escalation.schema.general_alert import ValidIocFields +from app.integrations.utils.alerts import get_asset_type_id +from app.integrations.utils.alerts import validate_ioc_type from app.utils import get_customer_alert_settings -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select async def is_customer_code_valid(customer_code: str, session: AsyncSession) -> bool: @@ -85,10 +84,7 @@ async def construct_alert_source_link( The alert source link. """ # Check if the alert has a process id and that it is not "No process ID found" - if ( - hasattr(alert_details, "process_id") - and alert_details._source.process_id != "No process ID found" - ): + if hasattr(alert_details, "process_id") and alert_details._source.process_id != "No process ID found": query_string = f"%22query%22:%22process_id:%5C%22{alert_details._source.process_id}%5C%22%20AND%20" else: query_string = f"%22query%22:%22_id:%5C%22{alert_details._id}%5C%22%20AND%20" diff --git a/backend/app/integrations/ask_socfortress/routes/ask_socfortress.py b/backend/app/integrations/ask_socfortress/routes/ask_socfortress.py index ae2d3af0..48b34930 100644 --- a/backend/app/integrations/ask_socfortress/routes/ask_socfortress.py +++ b/backend/app/integrations/ask_socfortress/routes/ask_socfortress.py @@ -1,16 +1,22 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler from app.db.db_session import get_db from app.integrations.ask_socfortress.schema.ask_socfortress import ( AskSocfortressRequest, +) +from app.integrations.ask_socfortress.schema.ask_socfortress import ( AskSocfortressSigmaResponse, ) from app.integrations.ask_socfortress.services.ask_socfortress import ( ask_socfortress_lookup, ) from app.utils import get_connector_attribute -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession # App specific imports diff --git a/backend/app/integrations/ask_socfortress/schema/ask_socfortress.py b/backend/app/integrations/ask_socfortress/schema/ask_socfortress.py index 4378b81f..7876aa53 100644 --- a/backend/app/integrations/ask_socfortress/schema/ask_socfortress.py +++ b/backend/app/integrations/ask_socfortress/schema/ask_socfortress.py @@ -1,4 +1,5 @@ -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class AskSocfortressRequest(BaseModel): diff --git a/backend/app/integrations/ask_socfortress/services/ask_socfortress.py b/backend/app/integrations/ask_socfortress/services/ask_socfortress.py index aa3ecbdc..4c85e6fa 100644 --- a/backend/app/integrations/ask_socfortress/services/ask_socfortress.py +++ b/backend/app/integrations/ask_socfortress/services/ask_socfortress.py @@ -1,23 +1,28 @@ -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional import httpx +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.connectors.utils import get_connector_info_from_db from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client from app.db.db_session import get_db_session -from app.integrations.alert_escalation.schema.general_alert import ( - CreateAlertRequest, - GenericAlertModel, - GenericSourceModel, -) +from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest +from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel +from app.integrations.alert_escalation.schema.general_alert import GenericSourceModel from app.integrations.ask_socfortress.schema.ask_socfortress import ( AskSocfortressRequest, +) +from app.integrations.ask_socfortress.schema.ask_socfortress import ( AskSocfortressSigmaRequest, +) +from app.integrations.ask_socfortress.schema.ask_socfortress import ( AskSocfortressSigmaResponse, ) from app.utils import get_connector_attribute -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession async def get_single_alert_details( diff --git a/backend/app/integrations/dnstwist/routes/analyze.py b/backend/app/integrations/dnstwist/routes/analyze.py index 4341e052..c63c0fb9 100644 --- a/backend/app/integrations/dnstwist/routes/analyze.py +++ b/backend/app/integrations/dnstwist/routes/analyze.py @@ -1,12 +1,13 @@ import regex -from app.integrations.dnstwist.schema.analyze import ( - DomainAnalysisResponse, - DomainRequestBody, -) -from app.integrations.dnstwist.services.analyze import analyze_domain -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException from loguru import logger +from app.integrations.dnstwist.schema.analyze import DomainAnalysisResponse +from app.integrations.dnstwist.schema.analyze import DomainRequestBody +from app.integrations.dnstwist.services.analyze import analyze_domain + dnstwist_router = APIRouter() diff --git a/backend/app/integrations/dnstwist/schema/analyze.py b/backend/app/integrations/dnstwist/schema/analyze.py index dec7f41c..0fd72df9 100644 --- a/backend/app/integrations/dnstwist/schema/analyze.py +++ b/backend/app/integrations/dnstwist/schema/analyze.py @@ -1,6 +1,8 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class DomainData(BaseModel): diff --git a/backend/app/integrations/dnstwist/services/analyze.py b/backend/app/integrations/dnstwist/services/analyze.py index 4679fa77..64a1fbc6 100644 --- a/backend/app/integrations/dnstwist/services/analyze.py +++ b/backend/app/integrations/dnstwist/services/analyze.py @@ -1,10 +1,9 @@ import dnstwist -from app.integrations.dnstwist.schema.analyze import ( - DomainAnalysisResponse, - DomainRequestBody, -) from loguru import logger +from app.integrations.dnstwist.schema.analyze import DomainAnalysisResponse +from app.integrations.dnstwist.schema.analyze import DomainRequestBody + def analyze_domain(domain: DomainRequestBody) -> DomainAnalysisResponse: """ diff --git a/backend/app/integrations/mimecast/routes/mimecast.py b/backend/app/integrations/mimecast/routes/mimecast.py index f7c71612..3f926a44 100644 --- a/backend/app/integrations/mimecast/routes/mimecast.py +++ b/backend/app/integrations/mimecast/routes/mimecast.py @@ -1,20 +1,20 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler from app.db.db_session import get_db -from app.integrations.mimecast.schema.mimecast import ( - MimecastAuthKeys, - MimecastRequest, - MimecastResponse, - MimecastTTPURLSRequest, -) -from app.integrations.mimecast.services.mimecast import get_ttp_urls, invoke_mimecast +from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys +from app.integrations.mimecast.schema.mimecast import MimecastRequest +from app.integrations.mimecast.schema.mimecast import MimecastResponse +from app.integrations.mimecast.schema.mimecast import MimecastTTPURLSRequest +from app.integrations.mimecast.services.mimecast import get_ttp_urls +from app.integrations.mimecast.services.mimecast import invoke_mimecast from app.integrations.routes import find_customer_integration -from app.integrations.utils.utils import ( - extract_mimecast_auth_keys, - get_customer_integration_response, -) -from fastapi import APIRouter, Depends, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession +from app.integrations.utils.utils import extract_mimecast_auth_keys +from app.integrations.utils.utils import get_customer_integration_response integration_mimecast_router = APIRouter() diff --git a/backend/app/integrations/mimecast/routes/provision.py b/backend/app/integrations/mimecast/routes/provision.py index 2e360945..46feeb67 100644 --- a/backend/app/integrations/mimecast/routes/provision.py +++ b/backend/app/integrations/mimecast/routes/provision.py @@ -1,15 +1,15 @@ +from fastapi import APIRouter +from fastapi import Depends +from sqlalchemy.ext.asyncio import AsyncSession + from app.db.db_session import get_db from app.integrations.mimecast.schema.mimecast import MimecastScheduledResponse -from app.integrations.mimecast.schema.provision import ( - ProvisionMimecastRequest, - ProvisionMimecastResponse, -) +from app.integrations.mimecast.schema.provision import ProvisionMimecastRequest +from app.integrations.mimecast.schema.provision import ProvisionMimecastResponse from app.integrations.mimecast.services.provision import provision_mimecast from app.integrations.utils.utils import get_customer_integration_response from app.schedulers.models.scheduler import CreateSchedulerRequest from app.schedulers.scheduler import add_scheduler_jobs -from fastapi import APIRouter, Depends -from sqlalchemy.ext.asyncio import AsyncSession integration_mimecast_scheduler_router = APIRouter() diff --git a/backend/app/integrations/mimecast/schema/mimecast.py b/backend/app/integrations/mimecast/schema/mimecast.py index dc51bb33..82799636 100644 --- a/backend/app/integrations/mimecast/schema/mimecast.py +++ b/backend/app/integrations/mimecast/schema/mimecast.py @@ -2,11 +2,17 @@ import hashlib import hmac import uuid -from datetime import datetime, timedelta +from datetime import datetime +from datetime import timedelta from enum import Enum -from typing import Dict, List, Optional +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field, HttpUrl, root_validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import HttpUrl +from pydantic import root_validator class PipelineRuleTitles(Enum): @@ -141,9 +147,7 @@ class MimecastHeaders(BaseModel): ) class Config: - allow_population_by_field_name = ( - True # This allows field population by both alias and field name - ) + allow_population_by_field_name = True # This allows field population by both alias and field name class MimecastTTPURLSRequest(BaseModel): @@ -199,9 +203,7 @@ def set_time_bounds(cls, values): elif unit == "w": lower_bound = now - timedelta(weeks=amount) - values["lower_bound"] = ( - lower_bound.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" - ) + values["lower_bound"] = lower_bound.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" values["upper_bound"] = now.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z" return values @@ -248,9 +250,7 @@ class DataItem(BaseModel): scanResult: str = Field(..., description="Scan result.") class Config: - allow_population_by_field_name = ( - True # This allows field population by both alias and field name - ) + allow_population_by_field_name = True # This allows field population by both alias and field name class RequestBody(BaseModel): diff --git a/backend/app/integrations/mimecast/schema/provision.py b/backend/app/integrations/mimecast/schema/provision.py index f8599e00..b38ea34f 100644 --- a/backend/app/integrations/mimecast/schema/provision.py +++ b/backend/app/integrations/mimecast/schema/provision.py @@ -1,6 +1,11 @@ -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from pydantic import BaseModel, Field, root_validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import root_validator class ProvisionMimecastRequest(BaseModel): diff --git a/backend/app/integrations/mimecast/services/mimecast.py b/backend/app/integrations/mimecast/services/mimecast.py index 07344e92..eeb366bf 100644 --- a/backend/app/integrations/mimecast/services/mimecast.py +++ b/backend/app/integrations/mimecast/services/mimecast.py @@ -12,21 +12,20 @@ import aiofiles import requests -from app.integrations.mimecast.schema.mimecast import ( - DataItem, - MimecastAPIEndpointResponse, - MimecastAuthKeys, - MimecastRequest, - MimecastResponse, - MimecastTTPURLSRequest, - RequestBody, - TtpURLResponseBody, -) +from fastapi import HTTPException +from loguru import logger + +from app.integrations.mimecast.schema.mimecast import DataItem +from app.integrations.mimecast.schema.mimecast import MimecastAPIEndpointResponse +from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys +from app.integrations.mimecast.schema.mimecast import MimecastRequest +from app.integrations.mimecast.schema.mimecast import MimecastResponse +from app.integrations.mimecast.schema.mimecast import MimecastTTPURLSRequest +from app.integrations.mimecast.schema.mimecast import RequestBody +from app.integrations.mimecast.schema.mimecast import TtpURLResponseBody from app.integrations.utils.collection import send_post_request from app.integrations.utils.event_shipper import event_shipper from app.integrations.utils.schema import EventShipperPayload -from fastapi import HTTPException -from loguru import logger async def get_checkpoint_filename(customer_code: str): diff --git a/backend/app/integrations/mimecast/services/provision.py b/backend/app/integrations/mimecast/services/provision.py index 3046f1a3..abf741db 100644 --- a/backend/app/integrations/mimecast/services/provision.py +++ b/backend/app/integrations/mimecast/services/provision.py @@ -1,40 +1,33 @@ import json from datetime import datetime -from app.connectors.grafana.schema.dashboards import ( - DashboardProvisionRequest, - MimecastDashboard, -) +from loguru import logger +from sqlalchemy import and_ +from sqlalchemy import update +from sqlalchemy.ext.asyncio import AsyncSession + +from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest +from app.connectors.grafana.schema.dashboards import MimecastDashboard from app.connectors.grafana.services.dashboards import provision_dashboards from app.connectors.grafana.utils.universal import create_grafana_client from app.connectors.graylog.services.management import start_stream from app.connectors.graylog.utils.universal import send_post_request -from app.customer_provisioning.schema.grafana import ( - GrafanaDatasource, - GrafanaDataSourceCreationResponse, -) -from app.customer_provisioning.schema.graylog import ( - GraylogIndexSetCreationResponse, - StreamCreationResponse, - TimeBasedIndexSet, -) -from app.customer_provisioning.services.grafana import ( - create_grafana_folder, - get_opensearch_version, -) -from app.customers.routes.customers import get_customer, get_customer_meta -from app.integrations.mimecast.schema.provision import ( - MimecastEventStream, - ProvisionMimecastRequest, - ProvisionMimecastResponse, -) +from app.customer_provisioning.schema.grafana import GrafanaDatasource +from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse +from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse +from app.customer_provisioning.schema.graylog import StreamCreationResponse +from app.customer_provisioning.schema.graylog import TimeBasedIndexSet +from app.customer_provisioning.services.grafana import create_grafana_folder +from app.customer_provisioning.services.grafana import get_opensearch_version +from app.customers.routes.customers import get_customer +from app.customers.routes.customers import get_customer_meta +from app.integrations.mimecast.schema.provision import MimecastEventStream +from app.integrations.mimecast.schema.provision import ProvisionMimecastRequest +from app.integrations.mimecast.schema.provision import ProvisionMimecastResponse from app.integrations.models.customer_integration_settings import CustomerIntegrations from app.integrations.routes import create_integration_meta from app.integrations.schema import CustomerIntegrationsMetaSchema from app.utils import get_connector_attribute -from loguru import logger -from sqlalchemy import and_, update -from sqlalchemy.ext.asyncio import AsyncSession ################## ! GRAYLOG ! ################## @@ -223,9 +216,7 @@ async def create_grafana_datasource( grafana_client = await create_grafana_client("Grafana") # Switch to the newly created organization grafana_client.user.switch_actual_user_organisation( - ( - await get_customer_meta(customer_code, session) - ).customer_meta.customer_meta_grafana_org_id, + (await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, ) datasource_payload = GrafanaDatasource( name="MIMECAST", diff --git a/backend/app/integrations/models/customer_integration_settings.py b/backend/app/integrations/models/customer_integration_settings.py index 12e5b831..5f9982b3 100644 --- a/backend/app/integrations/models/customer_integration_settings.py +++ b/backend/app/integrations/models/customer_integration_settings.py @@ -1,7 +1,10 @@ -from typing import List, Optional +from typing import List +from typing import Optional from sqlalchemy import Text -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field +from sqlmodel import Relationship +from sqlmodel import SQLModel class AvailableIntegrations(SQLModel, table=True): diff --git a/backend/app/integrations/monitoring_alert/models/monitoring_alert.py b/backend/app/integrations/monitoring_alert/models/monitoring_alert.py index e8ebd0b8..8aef05af 100644 --- a/backend/app/integrations/monitoring_alert/models/monitoring_alert.py +++ b/backend/app/integrations/monitoring_alert/models/monitoring_alert.py @@ -1,6 +1,7 @@ from typing import Optional -from sqlmodel import Field, SQLModel +from sqlmodel import Field +from sqlmodel import SQLModel class MonitoringAlerts(SQLModel, table=True): diff --git a/backend/app/integrations/monitoring_alert/routes/monitoring_alert.py b/backend/app/integrations/monitoring_alert/routes/monitoring_alert.py index 8b8056cb..f04887b0 100644 --- a/backend/app/integrations/monitoring_alert/routes/monitoring_alert.py +++ b/backend/app/integrations/monitoring_alert/routes/monitoring_alert.py @@ -1,22 +1,32 @@ from typing import List +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.auth.utils import AuthHandler from app.db.db_session import get_db from app.db.universal_models import CustomersMeta from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts +from app.integrations.monitoring_alert.schema.monitoring_alert import GraylogPostRequest from app.integrations.monitoring_alert.schema.monitoring_alert import ( - GraylogPostRequest, GraylogPostResponse, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( MonitoringAlertsRequestModel, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( MonitoringWazuhAlertsRequestModel, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhAnalysisResponse, ) from app.integrations.monitoring_alert.services.suricata import analyze_suricata_alerts from app.integrations.monitoring_alert.services.wazuh import analyze_wazuh_alerts -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select monitoring_alerts_router = APIRouter() @@ -147,8 +157,7 @@ async def run_wazuh_analysis( monitoring_alerts = await session.execute( select(MonitoringAlerts).where( - (MonitoringAlerts.customer_code == request.customer_code) - & (MonitoringAlerts.alert_source == "WAZUH"), + (MonitoringAlerts.customer_code == request.customer_code) & (MonitoringAlerts.alert_source == "WAZUH"), ), ) monitoring_alerts = monitoring_alerts.scalars().all() @@ -196,8 +205,7 @@ async def run_suricata_analysis( monitoring_alerts = await session.execute( select(MonitoringAlerts).where( - (MonitoringAlerts.customer_code == request.customer_code) - & (MonitoringAlerts.alert_source == "SURICATA"), + (MonitoringAlerts.customer_code == request.customer_code) & (MonitoringAlerts.alert_source == "SURICATA"), ), ) monitoring_alerts = monitoring_alerts.scalars().all() diff --git a/backend/app/integrations/monitoring_alert/routes/provision.py b/backend/app/integrations/monitoring_alert/routes/provision.py index ebf5fdd3..d21421cf 100644 --- a/backend/app/integrations/monitoring_alert/routes/provision.py +++ b/backend/app/integrations/monitoring_alert/routes/provision.py @@ -1,21 +1,29 @@ +from fastapi import APIRouter +from fastapi import HTTPException +from loguru import logger + from app.connectors.graylog.routes.events import get_all_event_definitions from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse +from app.integrations.monitoring_alert.schema.provision import AvailableMonitoringAlerts from app.integrations.monitoring_alert.schema.provision import ( - AvailableMonitoringAlerts, AvailableMonitoringAlertsResponse, +) +from app.integrations.monitoring_alert.schema.provision import ( ProvisionMonitoringAlertRequest, +) +from app.integrations.monitoring_alert.schema.provision import ( ProvisionWazuhMonitoringAlertResponse, ) from app.integrations.monitoring_alert.services.provision import ( provision_suricata_monitoring_alert, +) +from app.integrations.monitoring_alert.services.provision import ( provision_wazuh_monitoring_alert, ) from app.integrations.utils.event_shipper import event_shipper from app.integrations.utils.schema import EventShipperPayload from app.schedulers.models.scheduler import CreateSchedulerRequest from app.schedulers.scheduler import add_scheduler_jobs -from fastapi import APIRouter, HTTPException -from loguru import logger monitoring_alerts_provision_router = APIRouter() @@ -79,10 +87,7 @@ async def check_if_event_definition_exists(event_definition: str) -> bool: logger.info( f"Event definitions collected: {event_definitions_response.event_definitions}", ) - if event_definition in [ - event_definition.title - for event_definition in event_definitions_response.event_definitions - ]: + if event_definition in [event_definition.title for event_definition in event_definitions_response.event_definitions]: raise HTTPException( status_code=400, detail=f"Event definition {event_definition} already exists", @@ -99,10 +104,7 @@ async def get_available_monitoring_alerts_route() -> AvailableMonitoringAlertsRe """ Get the available monitoring alerts. """ - alerts = [ - {"name": alert.name.replace("_", " "), "value": alert.value} - for alert in AvailableMonitoringAlerts - ] + alerts = [{"name": alert.name.replace("_", " "), "value": alert.value} for alert in AvailableMonitoringAlerts] return AvailableMonitoringAlertsResponse( success=True, message="Alerts retrieved successfully", diff --git a/backend/app/integrations/monitoring_alert/schema/monitoring_alert.py b/backend/app/integrations/monitoring_alert/schema/monitoring_alert.py index 765b8a63..71c41c69 100644 --- a/backend/app/integrations/monitoring_alert/schema/monitoring_alert.py +++ b/backend/app/integrations/monitoring_alert/schema/monitoring_alert.py @@ -1,8 +1,15 @@ from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any +from typing import Dict +from typing import List +from typing import Optional -from app.integrations.alert_creation.general.schema.alert import IrisAsset, IrisIoc -from pydantic import BaseModel, Extra, Field +from pydantic import BaseModel +from pydantic import Extra +from pydantic import Field + +from app.integrations.alert_creation.general.schema.alert import IrisAsset +from app.integrations.alert_creation.general.schema.alert import IrisIoc class MonitoringAlertsRequestModel(BaseModel): diff --git a/backend/app/integrations/monitoring_alert/schema/provision.py b/backend/app/integrations/monitoring_alert/schema/provision.py index ee0c1416..34d8d833 100644 --- a/backend/app/integrations/monitoring_alert/schema/provision.py +++ b/backend/app/integrations/monitoring_alert/schema/provision.py @@ -1,8 +1,12 @@ from enum import Enum -from typing import Dict, List, Optional +from typing import Dict +from typing import List +from typing import Optional from fastapi import HTTPException -from pydantic import BaseModel, Field, validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator class AvailableMonitoringAlerts(str, Enum): diff --git a/backend/app/integrations/monitoring_alert/services/provision.py b/backend/app/integrations/monitoring_alert/services/provision.py index c761b77a..1d70ad66 100644 --- a/backend/app/integrations/monitoring_alert/services/provision.py +++ b/backend/app/integrations/monitoring_alert/services/provision.py @@ -1,27 +1,49 @@ import os from typing import Optional +from dotenv import load_dotenv +from fastapi import HTTPException +from loguru import logger + from app.connectors.graylog.routes.monitoring import get_all_event_notifications from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse from app.connectors.graylog.schema.monitoring import GraylogEventNotificationsResponse from app.connectors.graylog.services.collector import get_url_whitelist_entries -from app.connectors.graylog.utils.universal import send_post_request, send_put_request +from app.connectors.graylog.utils.universal import send_post_request +from app.connectors.graylog.utils.universal import send_put_request from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertProvisionConfig, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertProvisionFieldSpecItem, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertProvisionModel, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertProvisionNotification, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertProvisionNotificationSettings, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertProvisionProvider, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogAlertWebhookNotificationModel, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogUrlWhitelistEntries, +) +from app.integrations.monitoring_alert.schema.provision import ( GraylogUrlWhitelistEntryConfig, +) +from app.integrations.monitoring_alert.schema.provision import ( ProvisionMonitoringAlertRequest, +) +from app.integrations.monitoring_alert.schema.provision import ( ProvisionWazuhMonitoringAlertResponse, ) -from dotenv import load_dotenv -from fastapi import HTTPException -from loguru import logger load_dotenv() import uuid @@ -72,10 +94,7 @@ async def check_if_url_whitelist_entry_exists(url: str) -> bool: logger.info( f"Url whitelist entries collected: {url_whitelist_entries_response.url_whitelist_entries}", ) - if url in [ - url_whitelist_entry.value - for url_whitelist_entry in url_whitelist_entries_response.url_whitelist_entries.entries - ]: + if url in [url_whitelist_entry.value for url_whitelist_entry in url_whitelist_entries_response.url_whitelist_entries.entries]: logger.info(f"Url whitelist entry {url} already exists") return True return False @@ -103,9 +122,7 @@ async def get_notification_id(notification_title: str) -> Optional[str]: logger.info( f"Event notifications collected: {event_notifications_response.event_notifications}", ) - for ( - event_notification - ) in event_notifications_response.event_notifications.notifications: + for event_notification in event_notifications_response.event_notifications.notifications: if event_notification.title == notification_title: return event_notification.id return None @@ -184,8 +201,7 @@ async def check_if_event_notification_exists(event_notification: str) -> bool: f"Event notifications collected: {event_notifications_response.event_notifications}", ) if event_notification in [ - event_notification.title - for event_notification in event_notifications_response.event_notifications.notifications + event_notification.title for event_notification in event_notifications_response.event_notifications.notifications ]: return True return False diff --git a/backend/app/integrations/monitoring_alert/services/suricata.py b/backend/app/integrations/monitoring_alert/services/suricata.py index d263e6cb..2807f077 100644 --- a/backend/app/integrations/monitoring_alert/services/suricata.py +++ b/backend/app/integrations/monitoring_alert/services/suricata.py @@ -1,20 +1,21 @@ import json -from typing import Optional, Set +from typing import Optional +from typing import Set + +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession from app.agents.routes.agents import get_agent from app.agents.schema.agents import AgentsResponse -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client from app.db.universal_models import CustomersMeta -from app.integrations.alert_creation.general.schema.alert import ( - CreateAlertRequest, - IrisAsset, - IrisIoc, - ValidIocFields, -) +from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest +from app.integrations.alert_creation.general.schema.alert import IrisAsset +from app.integrations.alert_creation.general.schema.alert import IrisIoc +from app.integrations.alert_creation.general.schema.alert import ValidIocFields from app.integrations.alert_creation.general.services.alert_multi_exclude import ( AlertDetailsService, ) @@ -27,19 +28,25 @@ from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts from app.integrations.monitoring_alert.schema.monitoring_alert import ( FilterAlertsRequest, - SuricataAlertModel, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import SuricataAlertModel +from app.integrations.monitoring_alert.schema.monitoring_alert import ( SuricataIrisAlertContext, - SuricataIrisAsset, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import SuricataIrisAsset +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhAnalysisResponse, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhIrisAlertContext, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhIrisAlertPayload, ) from app.integrations.monitoring_alert.utils.db_operations import remove_alert_id -from app.integrations.utils.alerts import get_asset_type_id, validate_ioc_type +from app.integrations.utils.alerts import get_asset_type_id +from app.integrations.utils.alerts import validate_ioc_type from app.utils import get_customer_alert_settings -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession def valid_ioc_fields() -> Set[str]: @@ -69,13 +76,8 @@ async def construct_alert_source_link( The alert source link. """ # Check if the alert has a process id and that it is not "No process ID found" - if ( - hasattr(alert_details, "process_id") - and alert_details.process_id != "No process ID found" - ): - query_string = ( - f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20" - ) + if hasattr(alert_details, "process_id") and alert_details.process_id != "No process ID found": + query_string = f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20" else: query_string = f"%22query%22:%22_id:%5C%22{alert_details.id}%5C%22%20AND%20" @@ -218,11 +220,7 @@ async def check_if_open_alert_exists_in_iris(alert_details: SuricataAlertModel) lambda: alert_client.filter_alerts(**params), ) logger.info(f"Alert exists: {alert_exists['data']['alerts']}") - return ( - alert_exists["data"]["alerts"][0]["alert_id"] - if alert_exists["data"]["alerts"] - else [] - ) + return alert_exists["data"]["alerts"][0]["alert_id"] if alert_exists["data"]["alerts"] else [] def construct_params(request: FilterAlertsRequest) -> dict: diff --git a/backend/app/integrations/monitoring_alert/services/wazuh.py b/backend/app/integrations/monitoring_alert/services/wazuh.py index d3894edf..a2b72883 100644 --- a/backend/app/integrations/monitoring_alert/services/wazuh.py +++ b/backend/app/integrations/monitoring_alert/services/wazuh.py @@ -1,20 +1,21 @@ import json -from typing import Optional, Set +from typing import Optional +from typing import Set + +from fastapi import HTTPException +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession from app.agents.routes.agents import get_agent from app.agents.schema.agents import AgentsResponse -from app.connectors.dfir_iris.utils.universal import ( - fetch_and_validate_data, - initialize_client_and_alert, -) +from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data +from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client from app.db.universal_models import CustomersMeta -from app.integrations.alert_creation.general.schema.alert import ( - CreateAlertRequest, - IrisAsset, - IrisIoc, - ValidIocFields, -) +from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest +from app.integrations.alert_creation.general.schema.alert import IrisAsset +from app.integrations.alert_creation.general.schema.alert import IrisIoc +from app.integrations.alert_creation.general.schema.alert import ValidIocFields from app.integrations.alert_creation.general.services.alert_multi_exclude import ( AlertDetailsService, ) @@ -27,17 +28,21 @@ from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts from app.integrations.monitoring_alert.schema.monitoring_alert import ( FilterAlertsRequest, - WazuhAlertModel, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import WazuhAlertModel +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhAnalysisResponse, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhIrisAlertContext, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhIrisAlertPayload, ) from app.integrations.monitoring_alert.utils.db_operations import remove_alert_id -from app.integrations.utils.alerts import get_asset_type_id, validate_ioc_type +from app.integrations.utils.alerts import get_asset_type_id +from app.integrations.utils.alerts import validate_ioc_type from app.utils import get_customer_alert_settings -from fastapi import HTTPException -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession def valid_ioc_fields() -> Set[str]: @@ -67,13 +72,8 @@ async def construct_alert_source_link( The alert source link. """ # Check if the alert has a process id and that it is not "No process ID found" - if ( - hasattr(alert_details, "process_id") - and alert_details.process_id != "No process ID found" - ): - query_string = ( - f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20" - ) + if hasattr(alert_details, "process_id") and alert_details.process_id != "No process ID found": + query_string = f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20" else: query_string = f"%22query%22:%22_id:%5C%22{alert_details.id}%5C%22%20AND%20" @@ -214,11 +214,7 @@ async def check_if_open_alert_exists_in_iris(alert_details: WazuhAlertModel) -> lambda: alert_client.filter_alerts(**params), ) logger.info(f"Alert exists: {alert_exists['data']['alerts']}") - return ( - alert_exists["data"]["alerts"][0]["alert_id"] - if alert_exists["data"]["alerts"] - else [] - ) + return alert_exists["data"]["alerts"][0]["alert_id"] if alert_exists["data"]["alerts"] else [] def construct_params(request: FilterAlertsRequest) -> dict: diff --git a/backend/app/integrations/monitoring_alert/utils/db_operations.py b/backend/app/integrations/monitoring_alert/utils/db_operations.py index 41d1e82f..c4faccc4 100644 --- a/backend/app/integrations/monitoring_alert/utils/db_operations.py +++ b/backend/app/integrations/monitoring_alert/utils/db_operations.py @@ -1,8 +1,9 @@ -from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select +from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts + async def remove_alert_id(alert_id: str, session: AsyncSession) -> None: """ diff --git a/backend/app/integrations/office365/routes/provision.py b/backend/app/integrations/office365/routes/provision.py index 618d09a1..31fd49e1 100644 --- a/backend/app/integrations/office365/routes/provision.py +++ b/backend/app/integrations/office365/routes/provision.py @@ -1,20 +1,21 @@ from typing import Dict +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler from app.db.db_session import get_db -from app.integrations.office365.schema.provision import ( - ProvisionOffice365AuthKeys, - ProvisionOffice365Request, - ProvisionOffice365Response, -) +from app.integrations.office365.schema.provision import ProvisionOffice365AuthKeys +from app.integrations.office365.schema.provision import ProvisionOffice365Request +from app.integrations.office365.schema.provision import ProvisionOffice365Response from app.integrations.office365.services.provision import provision_office365 -from app.integrations.routes import ( - find_customer_integration, - get_customer_integrations_by_customer_code, -) -from app.integrations.schema import CustomerIntegrations, CustomerIntegrationsResponse -from fastapi import APIRouter, Depends, HTTPException, Security -from sqlalchemy.ext.asyncio import AsyncSession +from app.integrations.routes import find_customer_integration +from app.integrations.routes import get_customer_integrations_by_customer_code +from app.integrations.schema import CustomerIntegrations +from app.integrations.schema import CustomerIntegrationsResponse integration_office365_router = APIRouter() diff --git a/backend/app/integrations/office365/schema/provision.py b/backend/app/integrations/office365/schema/provision.py index f8034789..6d5ad582 100644 --- a/backend/app/integrations/office365/schema/provision.py +++ b/backend/app/integrations/office365/schema/provision.py @@ -1,7 +1,10 @@ from enum import Enum -from typing import Any, Dict +from typing import Any +from typing import Dict -from pydantic import BaseModel, Field, root_validator +from pydantic import BaseModel +from pydantic import Field +from pydantic import root_validator class PipelineRuleTitles(Enum): diff --git a/backend/app/integrations/office365/services/provision.py b/backend/app/integrations/office365/services/provision.py index 0da37559..73d77f47 100644 --- a/backend/app/integrations/office365/services/provision.py +++ b/backend/app/integrations/office365/services/provision.py @@ -4,65 +4,50 @@ from typing import List import requests -from app.connectors.grafana.schema.dashboards import ( - DashboardProvisionRequest, - Office365Dashboard, -) +from dotenv import load_dotenv +from fastapi import HTTPException +from loguru import logger +from sqlalchemy import and_ +from sqlalchemy import update +from sqlalchemy.ext.asyncio import AsyncSession + +from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest +from app.connectors.grafana.schema.dashboards import Office365Dashboard from app.connectors.grafana.services.dashboards import provision_dashboards from app.connectors.grafana.utils.universal import create_grafana_client -from app.connectors.graylog.schema.pipelines import ( - CreatePipeline, - CreatePipelineRule, - GraylogPipelinesResponse, - PipelineRulesResponse, -) +from app.connectors.graylog.schema.pipelines import CreatePipeline +from app.connectors.graylog.schema.pipelines import CreatePipelineRule +from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse +from app.connectors.graylog.schema.pipelines import PipelineRulesResponse from app.connectors.graylog.services.management import start_stream -from app.connectors.graylog.services.pipelines import ( - connect_stream_to_pipeline, - create_pipeline_graylog, - create_pipeline_rule, - get_pipeline_id, - get_pipeline_rules, - get_pipelines, -) +from app.connectors.graylog.services.pipelines import connect_stream_to_pipeline +from app.connectors.graylog.services.pipelines import create_pipeline_graylog +from app.connectors.graylog.services.pipelines import create_pipeline_rule +from app.connectors.graylog.services.pipelines import get_pipeline_id +from app.connectors.graylog.services.pipelines import get_pipeline_rules +from app.connectors.graylog.services.pipelines import get_pipelines from app.connectors.graylog.utils.universal import send_post_request -from app.connectors.wazuh_manager.utils.universal import ( - send_get_request, - send_put_request, -) -from app.customer_provisioning.schema.grafana import ( - GrafanaDatasource, - GrafanaDataSourceCreationResponse, -) -from app.customer_provisioning.schema.graylog import ( - GraylogIndexSetCreationResponse, - Office365EventStream, - StreamConnectionToPipelineRequest, - StreamCreationResponse, - TimeBasedIndexSet, -) -from app.customer_provisioning.services.grafana import ( - create_grafana_folder, - get_opensearch_version, -) -from app.customers.routes.customers import get_customer, get_customer_meta +from app.connectors.wazuh_manager.utils.universal import send_get_request +from app.connectors.wazuh_manager.utils.universal import send_put_request +from app.customer_provisioning.schema.grafana import GrafanaDatasource +from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse +from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse +from app.customer_provisioning.schema.graylog import Office365EventStream +from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest +from app.customer_provisioning.schema.graylog import StreamCreationResponse +from app.customer_provisioning.schema.graylog import TimeBasedIndexSet +from app.customer_provisioning.services.grafana import create_grafana_folder +from app.customer_provisioning.services.grafana import get_opensearch_version +from app.customers.routes.customers import get_customer +from app.customers.routes.customers import get_customer_meta from app.integrations.models.customer_integration_settings import CustomerIntegrations -from app.integrations.office365.schema.provision import ( - PipelineRuleTitles, - PipelineTitles, - ProvisionOffice365AuthKeys, - ProvisionOffice365Response, -) -from app.integrations.utils.schema import ( - PraecoAlertConfig, - PraecoProvisionAlertResponse, -) +from app.integrations.office365.schema.provision import PipelineRuleTitles +from app.integrations.office365.schema.provision import PipelineTitles +from app.integrations.office365.schema.provision import ProvisionOffice365AuthKeys +from app.integrations.office365.schema.provision import ProvisionOffice365Response +from app.integrations.utils.schema import PraecoAlertConfig +from app.integrations.utils.schema import PraecoProvisionAlertResponse from app.utils import get_connector_attribute -from dotenv import load_dotenv -from fastapi import HTTPException -from loguru import logger -from sqlalchemy import and_, update -from sqlalchemy.ext.asyncio import AsyncSession load_dotenv() @@ -488,9 +473,7 @@ async def pipeline_rules_exists(pipeline_rules: PipelineRulesResponse) -> List[s return [ rule_title.value for rule_title in PipelineRuleTitles - if not any( - rule.title == rule_title.value for rule in pipeline_rules.pipeline_rules - ) + if not any(rule.title == rule_title.value for rule in pipeline_rules.pipeline_rules) ] @@ -601,12 +584,7 @@ async def create_wazuh_alert_rule(rule_title: str) -> None: Creates the 'WAZUH CREATE FIELD SYSLOG LEVEL - ALERT' pipeline rule. """ rule_source = ( - f'rule "{rule_title}"\n' - "when\n" - " to_long($message.rule_level) > 11\n" - "then\n" - ' set_field("syslog_level", "ALERT");\n' - "end" + f'rule "{rule_title}"\n' "when\n" " to_long($message.rule_level) > 11\n" "then\n" ' set_field("syslog_level", "ALERT");\n' "end" ) await create_pipeline_rule( CreatePipelineRule( @@ -636,9 +614,7 @@ async def pipeline_exists(pipelines: GraylogPipelinesResponse) -> List[str]: return [ pipeline_title.value for pipeline_title in PipelineTitles - if not any( - pipeline.title == pipeline_title.value for pipeline in pipelines.pipelines - ) + if not any(pipeline.title == pipeline_title.value for pipeline in pipelines.pipelines) ] @@ -699,9 +675,7 @@ async def create_grafana_datasource( grafana_client = await create_grafana_client("Grafana") # Switch to the newly created organization grafana_client.user.switch_actual_user_organisation( - ( - await get_customer_meta(customer_code, session) - ).customer_meta.customer_meta_grafana_org_id, + (await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, ) datasource_payload = GrafanaDatasource( name="O365", @@ -784,9 +758,7 @@ async def provision_office365( await check_pipeline() # Create Index Set - index_set_id = ( - await create_index_set(customer_code=customer_code, session=session) - ).data.id + index_set_id = (await create_index_set(customer_code=customer_code, session=session)).data.id logger.info(f"Index set: {index_set_id}") # Create event stream stream_id = ( @@ -810,23 +782,17 @@ async def provision_office365( await start_stream(stream_id=stream_id) # Grafana Deployment - office365_datasource_uid = ( - await create_grafana_datasource(customer_code=customer_code, session=session) - ).datasource.uid + office365_datasource_uid = (await create_grafana_datasource(customer_code=customer_code, session=session)).datasource.uid grafana_o365_folder_id = ( await create_grafana_folder( - organization_id=( - await get_customer_meta(customer_code, session) - ).customer_meta.customer_meta_grafana_org_id, + organization_id=(await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, folder_title="OFFICE 365", ) ).id await provision_dashboards( DashboardProvisionRequest( dashboards=[dashboard.name for dashboard in Office365Dashboard], - organizationId=( - await get_customer_meta(customer_code, session) - ).customer_meta.customer_meta_grafana_org_id, + organizationId=(await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, folderId=grafana_o365_folder_id, datasourceUid=office365_datasource_uid, ), diff --git a/backend/app/integrations/routes.py b/backend/app/integrations/routes.py index 4034ce1b..4905852d 100644 --- a/backend/app/integrations/routes.py +++ b/backend/app/integrations/routes.py @@ -1,42 +1,49 @@ -from typing import List, Optional +from typing import List +from typing import Optional + +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy import delete +from sqlalchemy import update +from sqlalchemy.exc import NoResultFound +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select +from sqlalchemy.orm import joinedload from app.auth.utils import AuthHandler from app.db.db_session import get_db -from app.db.universal_models import Customers, CustomersMeta +from app.db.universal_models import Customers +from app.db.universal_models import CustomersMeta from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, ) +from app.integrations.models.customer_integration_settings import AvailableIntegrations +from app.integrations.models.customer_integration_settings import CustomerIntegrations from app.integrations.models.customer_integration_settings import ( - AvailableIntegrations, - CustomerIntegrations, CustomerIntegrationsMeta, - IntegrationAuthKeys, - IntegrationConfig, - IntegrationService, - IntegrationSubscription, ) -from app.integrations.schema import ( - AuthKey, - AvailableIntegrationsResponse, - CreateIntegrationAuthKeys, - CreateIntegrationService, - CustomerIntegrationCreate, - CustomerIntegrationCreateResponse, - CustomerIntegrationDeleteResponse, - CustomerIntegrationsMetaResponse, - CustomerIntegrationsMetaSchema, - CustomerIntegrationsResponse, - DeleteCustomerIntegration, - IntegrationWithAuthKeys, - UpdateCustomerIntegration, +from app.integrations.models.customer_integration_settings import IntegrationAuthKeys +from app.integrations.models.customer_integration_settings import IntegrationConfig +from app.integrations.models.customer_integration_settings import IntegrationService +from app.integrations.models.customer_integration_settings import ( + IntegrationSubscription, ) -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy import delete, update -from sqlalchemy.exc import NoResultFound -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select -from sqlalchemy.orm import joinedload +from app.integrations.schema import AuthKey +from app.integrations.schema import AvailableIntegrationsResponse +from app.integrations.schema import CreateIntegrationAuthKeys +from app.integrations.schema import CreateIntegrationService +from app.integrations.schema import CustomerIntegrationCreate +from app.integrations.schema import CustomerIntegrationCreateResponse +from app.integrations.schema import CustomerIntegrationDeleteResponse +from app.integrations.schema import CustomerIntegrationsMetaResponse +from app.integrations.schema import CustomerIntegrationsMetaSchema +from app.integrations.schema import CustomerIntegrationsResponse +from app.integrations.schema import DeleteCustomerIntegration +from app.integrations.schema import IntegrationWithAuthKeys +from app.integrations.schema import UpdateCustomerIntegration integration_settings_router = APIRouter() @@ -61,9 +68,7 @@ async def fetch_available_integrations(session: AsyncSession): integrations_with_auth_keys = [] for integration in unique_integrations: - auth_keys = [ - AuthKey(auth_key_name=key.auth_key_name) for key in integration.auth_keys - ] + auth_keys = [AuthKey(auth_key_name=key.auth_key_name) for key in integration.auth_keys] integration_data = IntegrationWithAuthKeys( id=integration.id, integration_name=integration.integration_name, @@ -97,9 +102,7 @@ async def validate_integration_auth_keys( Validate if the integration auth keys are valid. """ available_integrations = await fetch_available_integrations(session) - integration = [ - ai for ai in available_integrations if ai.integration_name == integration_name - ][0] + integration = [ai for ai in available_integrations if ai.integration_name == integration_name][0] available_auth_keys = [ak.auth_key_name for ak in integration.auth_keys] # loop through the `available_auth_keys` and check if the `integration_auth_keys` contains the `auth_key_name` for auth_key in available_auth_keys: @@ -120,9 +123,7 @@ async def validate_integration_auth_key_update( """ logger.info(f"integration_auth_key: {integration_auth_key}") available_integrations = await fetch_available_integrations(session) - integration = [ - ai for ai in available_integrations if ai.integration_name == integration_name - ][0] + integration = [ai for ai in available_integrations if ai.integration_name == integration_name][0] available_auth_keys = [ak.auth_key_name for ak in integration.auth_keys] for auth_key in integration_auth_key: if auth_key.auth_key_name not in available_auth_keys: @@ -500,20 +501,14 @@ def process_customer_integrations(customer_integrations_data): """ processed_customer_integrations = [] for ci in customer_integrations_data: - first_service_id = ( - ci.integration_subscriptions[0].integration_service_id - if ci.integration_subscriptions - else None - ) + first_service_id = ci.integration_subscriptions[0].integration_service_id if ci.integration_subscriptions else None customer_integration_obj = CustomerIntegrations( id=ci.id, customer_code=ci.customer_code, customer_name=ci.customer_name, integration_subscriptions=ci.integration_subscriptions, integration_service_id=first_service_id, - integration_service_name=ci.integration_subscriptions[ - 0 - ].integration_service.service_name + integration_service_name=ci.integration_subscriptions[0].integration_service.service_name if ci.integration_subscriptions else None, deployed=ci.deployed, @@ -936,10 +931,8 @@ async def delete_integration_meta( """ try: stmt = delete(CustomerIntegrationsMeta).where( - CustomerIntegrationsMeta.customer_code - == customer_integration_meta.customer_code, - CustomerIntegrationsMeta.integration_name - == customer_integration_meta.integration_name, + CustomerIntegrationsMeta.customer_code == customer_integration_meta.customer_code, + CustomerIntegrationsMeta.integration_name == customer_integration_meta.integration_name, ) await session.execute(stmt) await session.commit() diff --git a/backend/app/integrations/schema.py b/backend/app/integrations/schema.py index 74df1da2..130471a0 100644 --- a/backend/app/integrations/schema.py +++ b/backend/app/integrations/schema.py @@ -1,6 +1,8 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class AuthKey(BaseModel): diff --git a/backend/app/integrations/utils/alerts.py b/backend/app/integrations/utils/alerts.py index 186df0af..495efd88 100644 --- a/backend/app/integrations/utils/alerts.py +++ b/backend/app/integrations/utils/alerts.py @@ -1,16 +1,19 @@ import ipaddress import re from abc import ABC -from typing import Dict, Optional, Union +from typing import Dict +from typing import Optional +from typing import Union import httpx import regex -from app.integrations.utils.schema import ShufflePayload -from app.utils import get_customer_alert_settings from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from app.integrations.utils.schema import ShufflePayload +from app.utils import get_customer_alert_settings + #################### ! DFIR IRIS ASSET VALIDATOR ! #################### class AssetValidator(ABC): diff --git a/backend/app/integrations/utils/collection.py b/backend/app/integrations/utils/collection.py index 3c121593..d4746113 100644 --- a/backend/app/integrations/utils/collection.py +++ b/backend/app/integrations/utils/collection.py @@ -1,5 +1,7 @@ import asyncio -from typing import Any, Dict, Optional +from typing import Any +from typing import Dict +from typing import Optional import httpx from loguru import logger diff --git a/backend/app/integrations/utils/event_shipper.py b/backend/app/integrations/utils/event_shipper.py index 0fd74825..6d3c1e83 100644 --- a/backend/app/integrations/utils/event_shipper.py +++ b/backend/app/integrations/utils/event_shipper.py @@ -1,15 +1,15 @@ import asyncio -from typing import Any, Dict +from typing import Any +from typing import Dict + +from fastapi import HTTPException +from loguru import logger from app.connectors.event_shipper.utils.universal import create_gelf_logger from app.connectors.utils import get_connector_info_from_db from app.db.db_session import get_db_session -from app.integrations.utils.schema import ( - EventShipperPayload, - EventShipperPayloadResponse, -) -from fastapi import HTTPException -from loguru import logger +from app.integrations.utils.schema import EventShipperPayload +from app.integrations.utils.schema import EventShipperPayloadResponse async def get_gelf_logger(): diff --git a/backend/app/integrations/utils/schema.py b/backend/app/integrations/utils/schema.py index 4f81b25c..1a88120b 100644 --- a/backend/app/integrations/utils/schema.py +++ b/backend/app/integrations/utils/schema.py @@ -1,6 +1,9 @@ -from typing import List, Optional +from typing import List +from typing import Optional -from pydantic import BaseModel, Extra, Field +from pydantic import BaseModel +from pydantic import Extra +from pydantic import Field class WazuhOSInfo(BaseModel): diff --git a/backend/app/integrations/utils/utils.py b/backend/app/integrations/utils/utils.py index a62a3932..bcf2d7d3 100644 --- a/backend/app/integrations/utils/utils.py +++ b/backend/app/integrations/utils/utils.py @@ -1,11 +1,13 @@ from typing import Dict -from app.integrations.routes import get_customer_integrations_by_customer_code -from app.integrations.schema import CustomerIntegrations, CustomerIntegrationsResponse from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from app.integrations.routes import get_customer_integrations_by_customer_code +from app.integrations.schema import CustomerIntegrations +from app.integrations.schema import CustomerIntegrationsResponse + async def get_customer_integration_response( customer_code: str, diff --git a/backend/app/middleware/exception_handlers.py b/backend/app/middleware/exception_handlers.py index 6b4addda..b3bd2247 100644 --- a/backend/app/middleware/exception_handlers.py +++ b/backend/app/middleware/exception_handlers.py @@ -1,11 +1,16 @@ -from app.auth.utils import AuthHandler -from app.db.db_session import async_engine # Make sure to import the async engine -from app.utils import ErrorType, Logger, ValidationErrorItem, ValidationErrorResponse -from fastapi import HTTPException, Request +from fastapi import HTTPException +from fastapi import Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from sqlalchemy.ext.asyncio import AsyncSession +from app.auth.utils import AuthHandler +from app.db.db_session import async_engine # Make sure to import the async engine +from app.utils import ErrorType +from app.utils import Logger +from app.utils import ValidationErrorItem +from app.utils import ValidationErrorResponse + # Utility function to get user_id from request async def get_user_id_from_request(request: Request, logger_instance): diff --git a/backend/app/middleware/logger.py b/backend/app/middleware/logger.py index 862ae28d..8514140c 100644 --- a/backend/app/middleware/logger.py +++ b/backend/app/middleware/logger.py @@ -1,9 +1,11 @@ +from fastapi import HTTPException +from fastapi import Request +from fastapi.responses import JSONResponse +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler from app.db.db_session import async_engine from app.utils import Logger -from fastapi import HTTPException, Request -from fastapi.responses import JSONResponse -from sqlalchemy.ext.asyncio import AsyncSession EXCLUDED_PATHS = ["/auth/token", "/auth/register"] INTERNAL_SERVER_ERROR = 500 @@ -54,20 +56,14 @@ async def handle_exception(e, user_id, request, logger_instance): JSONResponse: The response containing the error message. """ try: - user_id = ( - await logger_instance.get_user_id_from_request(request) - if user_id is None - else user_id - ) + user_id = await logger_instance.get_user_id_from_request(request) if user_id is None else user_id except HTTPException as http_exc: return JSONResponse( status_code=http_exc.status_code, content={"message": str(http_exc), "success": False}, ) await logger_instance.log_error(user_id, request, e) - status_code = ( - e.status_code if isinstance(e, HTTPException) else INTERNAL_SERVER_ERROR - ) + status_code = e.status_code if isinstance(e, HTTPException) else INTERNAL_SERVER_ERROR return JSONResponse( status_code=status_code, content={"message": str(e), "success": False}, diff --git a/backend/app/routers/agents.py b/backend/app/routers/agents.py index 5f04f696..e9fd910e 100644 --- a/backend/app/routers/agents.py +++ b/backend/app/routers/agents.py @@ -1,6 +1,7 @@ -from app.agents.routes.agents import agents_router from fastapi import APIRouter +from app.agents.routes.agents import agents_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/alert_creation.py b/backend/app/routers/alert_creation.py index 36291851..dff27306 100644 --- a/backend/app/routers/alert_creation.py +++ b/backend/app/routers/alert_creation.py @@ -1,8 +1,9 @@ +from fastapi import APIRouter + from app.integrations.alert_creation.general.routes.alert import general_alerts_router from app.integrations.alert_creation.office365.routes.alert import ( office365_alerts_router, ) -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/alert_creation_settings.py b/backend/app/routers/alert_creation_settings.py index a0339f96..b877e800 100644 --- a/backend/app/routers/alert_creation_settings.py +++ b/backend/app/routers/alert_creation_settings.py @@ -1,7 +1,8 @@ +from fastapi import APIRouter + from app.integrations.alert_creation_settings.routes.alert_creation_settings import ( alert_creation_settings_router, ) -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/ask_socfortress.py b/backend/app/routers/ask_socfortress.py index 8f511de9..119c415e 100644 --- a/backend/app/routers/ask_socfortress.py +++ b/backend/app/routers/ask_socfortress.py @@ -1,7 +1,8 @@ +from fastapi import APIRouter + from app.integrations.ask_socfortress.routes.ask_socfortress import ( ask_socfortress_router, ) -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/auth.py b/backend/app/routers/auth.py index f45f36dc..61a4d977 100644 --- a/backend/app/routers/auth.py +++ b/backend/app/routers/auth.py @@ -1,6 +1,7 @@ -from app.auth.routes.auth import auth_router from fastapi import APIRouter +from app.auth.routes.auth import auth_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/connectors.py b/backend/app/routers/connectors.py index 67f2c998..d40c9ef3 100644 --- a/backend/app/routers/connectors.py +++ b/backend/app/routers/connectors.py @@ -1,6 +1,7 @@ -from app.connectors.routes import connector_router from fastapi import APIRouter +from app.connectors.routes import connector_router + router = APIRouter() router.include_router(connector_router, prefix="/connectors", tags=["connectors"]) diff --git a/backend/app/routers/cortex.py b/backend/app/routers/cortex.py index 24dc3e29..03fffc0e 100644 --- a/backend/app/routers/cortex.py +++ b/backend/app/routers/cortex.py @@ -1,6 +1,7 @@ -from app.connectors.cortex.routes.analyzers import cortex_analyzer_router from fastapi import APIRouter +from app.connectors.cortex.routes.analyzers import cortex_analyzer_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/customer_provisioning.py b/backend/app/routers/customer_provisioning.py index 2bfc869b..00ac0683 100644 --- a/backend/app/routers/customer_provisioning.py +++ b/backend/app/routers/customer_provisioning.py @@ -1,8 +1,9 @@ +from fastapi import APIRouter + from app.customer_provisioning.routes.decommission import ( customer_decommissioning_router, ) from app.customer_provisioning.routes.provision import customer_provisioning_router -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/customers.py b/backend/app/routers/customers.py index fbb4d557..da741df3 100644 --- a/backend/app/routers/customers.py +++ b/backend/app/routers/customers.py @@ -1,6 +1,7 @@ -from app.customers.routes.customers import customers_router from fastapi import APIRouter +from app.customers.routes.customers import customers_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/dfir_iris.py b/backend/app/routers/dfir_iris.py index 90cc6fa8..55cba917 100644 --- a/backend/app/routers/dfir_iris.py +++ b/backend/app/routers/dfir_iris.py @@ -1,3 +1,5 @@ +from fastapi import APIRouter + from app.connectors.dfir_iris.routes.alerts import dfir_iris_alerts_router from app.connectors.dfir_iris.routes.assets import dfir_iris_assets_router from app.connectors.dfir_iris.routes.cases import dfir_iris_cases_router @@ -6,7 +8,6 @@ from app.integrations.alert_escalation.routes.general_alert import ( integration_general_alerts_router, ) -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/dnstwist.py b/backend/app/routers/dnstwist.py index ebad1bed..73549fe8 100644 --- a/backend/app/routers/dnstwist.py +++ b/backend/app/routers/dnstwist.py @@ -1,6 +1,7 @@ -from app.integrations.dnstwist.routes.analyze import dnstwist_router from fastapi import APIRouter +from app.integrations.dnstwist.routes.analyze import dnstwist_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/grafana.py b/backend/app/routers/grafana.py index 01daa565..f9044190 100644 --- a/backend/app/routers/grafana.py +++ b/backend/app/routers/grafana.py @@ -1,6 +1,7 @@ -from app.connectors.grafana.routes.dashboards import grafana_dashboards_router from fastapi import APIRouter +from app.connectors.grafana.routes.dashboards import grafana_dashboards_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/graylog.py b/backend/app/routers/graylog.py index da013662..6b0b9b1b 100644 --- a/backend/app/routers/graylog.py +++ b/backend/app/routers/graylog.py @@ -1,10 +1,11 @@ +from fastapi import APIRouter + from app.connectors.graylog.routes.collector import graylog_collector_router from app.connectors.graylog.routes.events import graylog_events_router from app.connectors.graylog.routes.management import graylog_management_router from app.connectors.graylog.routes.monitoring import graylog_monitoring_router from app.connectors.graylog.routes.pipelines import graylog_pipelines_router from app.connectors.graylog.routes.streams import graylog_streams_router -from fastapi import APIRouter router = APIRouter() diff --git a/backend/app/routers/healthcheck.py b/backend/app/routers/healthcheck.py index 05b4fba2..8a69967e 100644 --- a/backend/app/routers/healthcheck.py +++ b/backend/app/routers/healthcheck.py @@ -1,6 +1,7 @@ -from app.healthchecks.agents.routes.agents import healtcheck_agents_router from fastapi import APIRouter +from app.healthchecks.agents.routes.agents import healtcheck_agents_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/influxdb.py b/backend/app/routers/influxdb.py index 10a8886f..a9a4306c 100644 --- a/backend/app/routers/influxdb.py +++ b/backend/app/routers/influxdb.py @@ -1,6 +1,7 @@ -from app.connectors.influxdb.routes.alerts import influxdb_alerts_router from fastapi import APIRouter +from app.connectors.influxdb.routes.alerts import influxdb_alerts_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/integrations.py b/backend/app/routers/integrations.py index cfaadbee..899a4777 100644 --- a/backend/app/routers/integrations.py +++ b/backend/app/routers/integrations.py @@ -1,6 +1,7 @@ -from app.integrations.routes import integration_settings_router from fastapi import APIRouter +from app.integrations.routes import integration_settings_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/logs.py b/backend/app/routers/logs.py index 4d262204..6633fb23 100644 --- a/backend/app/routers/logs.py +++ b/backend/app/routers/logs.py @@ -1,6 +1,7 @@ -from app.utils import logs_router from fastapi import APIRouter +from app.utils import logs_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/mimecast.py b/backend/app/routers/mimecast.py index 0a960614..d0c4f2a1 100644 --- a/backend/app/routers/mimecast.py +++ b/backend/app/routers/mimecast.py @@ -1,8 +1,9 @@ +from fastapi import APIRouter + from app.integrations.mimecast.routes.mimecast import integration_mimecast_router from app.integrations.mimecast.routes.provision import ( integration_mimecast_scheduler_router, ) -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/monitoring_alert.py b/backend/app/routers/monitoring_alert.py index dcbcddaa..3c7d1c66 100644 --- a/backend/app/routers/monitoring_alert.py +++ b/backend/app/routers/monitoring_alert.py @@ -1,10 +1,11 @@ +from fastapi import APIRouter + from app.integrations.monitoring_alert.routes.monitoring_alert import ( monitoring_alerts_router, ) from app.integrations.monitoring_alert.routes.provision import ( monitoring_alerts_provision_router, ) -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/office365.py b/backend/app/routers/office365.py index 4fd14f97..fc5a659e 100644 --- a/backend/app/routers/office365.py +++ b/backend/app/routers/office365.py @@ -1,6 +1,7 @@ -from app.integrations.office365.routes.provision import integration_office365_router from fastapi import APIRouter +from app.integrations.office365.routes.provision import integration_office365_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/scheduler.py b/backend/app/routers/scheduler.py index e458c226..db97fd90 100644 --- a/backend/app/routers/scheduler.py +++ b/backend/app/routers/scheduler.py @@ -1,6 +1,7 @@ -from app.schedulers.routes.scheduler import scheduler_router from fastapi import APIRouter +from app.schedulers.routes.scheduler import scheduler_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/shuffle.py b/backend/app/routers/shuffle.py index 1b06c4af..b833dc48 100644 --- a/backend/app/routers/shuffle.py +++ b/backend/app/routers/shuffle.py @@ -1,6 +1,7 @@ -from app.connectors.shuffle.routes.workflows import shuffle_workflows_router from fastapi import APIRouter +from app.connectors.shuffle.routes.workflows import shuffle_workflows_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/smtp.py b/backend/app/routers/smtp.py index ca0ff89c..9dfe5c92 100644 --- a/backend/app/routers/smtp.py +++ b/backend/app/routers/smtp.py @@ -1,6 +1,7 @@ +from fastapi import APIRouter + from app.smtp.routes.configure import smtp_configure_router from app.smtp.routes.reports import smtp_reports_router -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/sublime.py b/backend/app/routers/sublime.py index ffa9fa42..0d9b1ae9 100644 --- a/backend/app/routers/sublime.py +++ b/backend/app/routers/sublime.py @@ -1,6 +1,7 @@ -from app.connectors.sublime.routes.alerts import sublime_alerts_router from fastapi import APIRouter +from app.connectors.sublime.routes.alerts import sublime_alerts_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/threat_intel.py b/backend/app/routers/threat_intel.py index f3501109..61633ce7 100644 --- a/backend/app/routers/threat_intel.py +++ b/backend/app/routers/threat_intel.py @@ -1,6 +1,7 @@ -from app.threat_intel.routes.socfortress import threat_intel_socfortress_router from fastapi import APIRouter +from app.threat_intel.routes.socfortress import threat_intel_socfortress_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/velociraptor.py b/backend/app/routers/velociraptor.py index a3bf3264..d1951d96 100644 --- a/backend/app/routers/velociraptor.py +++ b/backend/app/routers/velociraptor.py @@ -1,6 +1,7 @@ +from fastapi import APIRouter + from app.connectors.velociraptor.routes.artifacts import velociraptor_artifacts_router from app.connectors.velociraptor.routes.flows import velociraptor_flows_router -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/wazuh_indexer.py b/backend/app/routers/wazuh_indexer.py index e1ad9bb0..2949f828 100644 --- a/backend/app/routers/wazuh_indexer.py +++ b/backend/app/routers/wazuh_indexer.py @@ -1,6 +1,7 @@ +from fastapi import APIRouter + from app.connectors.wazuh_indexer.routes.alerts import wazuh_indexer_alerts_router from app.connectors.wazuh_indexer.routes.monitoring import wazuh_indexer_router -from fastapi import APIRouter # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/routers/wazuh_manager.py b/backend/app/routers/wazuh_manager.py index 3ee8f113..ab8c9835 100644 --- a/backend/app/routers/wazuh_manager.py +++ b/backend/app/routers/wazuh_manager.py @@ -1,6 +1,7 @@ -from app.connectors.wazuh_manager.routes.rules import wazuh_manager_rules_router from fastapi import APIRouter +from app.connectors.wazuh_manager.routes.rules import wazuh_manager_rules_router + # Instantiate the APIRouter router = APIRouter() diff --git a/backend/app/schedulers/models/scheduler.py b/backend/app/schedulers/models/scheduler.py index 14103041..fad84330 100644 --- a/backend/app/schedulers/models/scheduler.py +++ b/backend/app/schedulers/models/scheduler.py @@ -2,7 +2,8 @@ from typing import Optional from pydantic import BaseModel -from sqlmodel import Field, SQLModel +from sqlmodel import Field +from sqlmodel import SQLModel class JobMetadata(SQLModel, table=True): diff --git a/backend/app/schedulers/routes/scheduler.py b/backend/app/schedulers/routes/scheduler.py index 8096b20b..ed06ef61 100644 --- a/backend/app/schedulers/routes/scheduler.py +++ b/backend/app/schedulers/routes/scheduler.py @@ -1,11 +1,13 @@ +from fastapi import APIRouter +from fastapi import Depends +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + from app.db.db_session import get_db from app.schedulers.models.scheduler import JobMetadata from app.schedulers.scheduler import init_scheduler from app.schedulers.schema.scheduler import JobsResponse -from fastapi import APIRouter, Depends -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select scheduler_router = APIRouter() diff --git a/backend/app/schedulers/scheduler.py b/backend/app/schedulers/scheduler.py index 6d3bc9f6..a58c3f6f 100644 --- a/backend/app/schedulers/scheduler.py +++ b/backend/app/schedulers/scheduler.py @@ -1,18 +1,17 @@ -from app.db.db_session import SyncSessionLocal, sync_engine -from app.schedulers.models.scheduler import CreateSchedulerRequest, JobMetadata -from app.schedulers.services.agent_sync import agent_sync -from app.schedulers.services.invoke_mimecast import ( - invoke_mimecast_integration, - invoke_mimecast_integration_ttp, -) -from app.schedulers.services.monitoring_alert import ( - invoke_suricata_monitoring_alert, - invoke_wazuh_monitoring_alert, -) from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore from apscheduler.schedulers.asyncio import AsyncIOScheduler from loguru import logger +from app.db.db_session import SyncSessionLocal +from app.db.db_session import sync_engine +from app.schedulers.models.scheduler import CreateSchedulerRequest +from app.schedulers.models.scheduler import JobMetadata +from app.schedulers.services.agent_sync import agent_sync +from app.schedulers.services.invoke_mimecast import invoke_mimecast_integration +from app.schedulers.services.invoke_mimecast import invoke_mimecast_integration_ttp +from app.schedulers.services.monitoring_alert import invoke_suricata_monitoring_alert +from app.schedulers.services.monitoring_alert import invoke_wazuh_monitoring_alert + def init_scheduler(): """ @@ -43,9 +42,7 @@ def initialize_job_metadata(): # {"job_id": "invoke_mimecast_integration", "time_interval": 5, "function": invoke_mimecast_integration} ] for job in known_jobs: - job_metadata = ( - session.query(JobMetadata).filter_by(job_id=job["job_id"]).one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id=job["job_id"]).one_or_none() if not job_metadata: job_metadata = JobMetadata( job_id=job["job_id"], @@ -132,11 +129,7 @@ async def add_job_metadata(create_scheduler_request: CreateSchedulerRequest): create_scheduler_request (CreateSchedulerRequest): The request object containing the job details. """ with SyncSessionLocal() as session: - job_metadata = ( - session.query(JobMetadata) - .filter_by(job_id=create_scheduler_request.job_id) - .one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id=create_scheduler_request.job_id).one_or_none() if not job_metadata: job_metadata = JobMetadata( job_id=create_scheduler_request.job_id, diff --git a/backend/app/schedulers/services/agent_sync.py b/backend/app/schedulers/services/agent_sync.py index aced4751..f7cf5a21 100644 --- a/backend/app/schedulers/services/agent_sync.py +++ b/backend/app/schedulers/services/agent_sync.py @@ -2,10 +2,11 @@ from datetime import datetime import requests +from dotenv import load_dotenv + from app.db.db_session import get_sync_db_session from app.schedulers.models.scheduler import JobMetadata from app.schedulers.utils.universal import scheduler_login -from dotenv import load_dotenv load_dotenv() @@ -39,9 +40,7 @@ def agent_sync(): # Use get_sync_db_session to create and manage a synchronous session with get_sync_db_session() as session: # Synchronous ORM operations - job_metadata = ( - session.query(JobMetadata).filter_by(job_id="agent_sync").one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id="agent_sync").one_or_none() if job_metadata: job_metadata.last_success = datetime.utcnow() session.add(job_metadata) diff --git a/backend/app/schedulers/services/invoke_mimecast.py b/backend/app/schedulers/services/invoke_mimecast.py index cae26fd4..c0b58f49 100644 --- a/backend/app/schedulers/services/invoke_mimecast.py +++ b/backend/app/schedulers/services/invoke_mimecast.py @@ -1,17 +1,18 @@ from datetime import datetime -from app.db.db_session import get_db_session, get_sync_db_session -from app.integrations.mimecast.routes.mimecast import ( - invoke_mimecast_route, - mimecast_ttp_url_route, -) -from app.integrations.mimecast.schema.mimecast import MimecastRequest, MimecastResponse -from app.integrations.models.customer_integration_settings import CustomerIntegrations -from app.schedulers.models.scheduler import JobMetadata from dotenv import load_dotenv from loguru import logger from sqlalchemy import select +from app.db.db_session import get_db_session +from app.db.db_session import get_sync_db_session +from app.integrations.mimecast.routes.mimecast import invoke_mimecast_route +from app.integrations.mimecast.routes.mimecast import mimecast_ttp_url_route +from app.integrations.mimecast.schema.mimecast import MimecastRequest +from app.integrations.mimecast.schema.mimecast import MimecastResponse +from app.integrations.models.customer_integration_settings import CustomerIntegrations +from app.schedulers.models.scheduler import JobMetadata + load_dotenv() @@ -40,11 +41,7 @@ async def invoke_mimecast_integration() -> MimecastResponse: await session.close() with get_sync_db_session() as session: # Synchronous ORM operations - job_metadata = ( - session.query(JobMetadata) - .filter_by(job_id="invoke_mimecast_integration") - .one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_mimecast_integration").one_or_none() if job_metadata: job_metadata.last_success = datetime.utcnow() session.add(job_metadata) @@ -80,11 +77,7 @@ async def invoke_mimecast_integration_ttp() -> MimecastResponse: await session.close() with get_sync_db_session() as session: # Synchronous ORM operations - job_metadata = ( - session.query(JobMetadata) - .filter_by(job_id="invoke_mimecast_integration_ttp") - .one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_mimecast_integration_ttp").one_or_none() if job_metadata: job_metadata.last_success = datetime.utcnow() session.add(job_metadata) diff --git a/backend/app/schedulers/services/monitoring_alert.py b/backend/app/schedulers/services/monitoring_alert.py index 9bbd4887..73328abc 100644 --- a/backend/app/schedulers/services/monitoring_alert.py +++ b/backend/app/schedulers/services/monitoring_alert.py @@ -1,19 +1,23 @@ from datetime import datetime -from app.db.db_session import get_db_session, get_sync_db_session +from dotenv import load_dotenv +from loguru import logger +from sqlalchemy import select + +from app.db.db_session import get_db_session +from app.db.db_session import get_sync_db_session from app.db.universal_models import CustomersMeta from app.integrations.monitoring_alert.routes.monitoring_alert import ( run_suricata_analysis, - run_wazuh_analysis, ) +from app.integrations.monitoring_alert.routes.monitoring_alert import run_wazuh_analysis from app.integrations.monitoring_alert.schema.monitoring_alert import ( MonitoringWazuhAlertsRequestModel, +) +from app.integrations.monitoring_alert.schema.monitoring_alert import ( WazuhAnalysisResponse, ) from app.schedulers.models.scheduler import JobMetadata -from dotenv import load_dotenv -from loguru import logger -from sqlalchemy import select load_dotenv() @@ -41,11 +45,7 @@ async def invoke_wazuh_monitoring_alert() -> WazuhAnalysisResponse: await session.close() with get_sync_db_session() as session: # Synchronous ORM operations - job_metadata = ( - session.query(JobMetadata) - .filter_by(job_id="invoke_wazuh_monitoring_alert") - .one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_wazuh_monitoring_alert").one_or_none() if job_metadata: job_metadata.last_success = datetime.utcnow() session.add(job_metadata) @@ -83,11 +83,7 @@ async def invoke_suricata_monitoring_alert() -> WazuhAnalysisResponse: await session.close() with get_sync_db_session() as session: # Synchronous ORM operations - job_metadata = ( - session.query(JobMetadata) - .filter_by(job_id="invoke_suricata_monitoring_alert") - .one_or_none() - ) + job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_suricata_monitoring_alert").one_or_none() if job_metadata: job_metadata.last_success = datetime.utcnow() session.add(job_metadata) diff --git a/backend/app/schedulers/utils/universal.py b/backend/app/schedulers/utils/universal.py index c5a81c7d..ef9dd4b9 100644 --- a/backend/app/schedulers/utils/universal.py +++ b/backend/app/schedulers/utils/universal.py @@ -1,9 +1,10 @@ import os import requests -from app.auth.services.universal import get_scheduler_password from dotenv import load_dotenv +from app.auth.services.universal import get_scheduler_password + load_dotenv() diff --git a/backend/app/smtp/routes/configure.py b/backend/app/smtp/routes/configure.py index ad7034f2..fd540a3a 100644 --- a/backend/app/smtp/routes/configure.py +++ b/backend/app/smtp/routes/configure.py @@ -1,10 +1,13 @@ -from app.auth.models.users import SMTP, SMTPInput +from fastapi import APIRouter +from fastapi import HTTPException +from loguru import logger + +from app.auth.models.users import SMTP +from app.auth.models.users import SMTPInput from app.auth.services.universal import select_all_users from app.auth.utils import AuthHandler from app.db.db_session import session from app.smtp.schema.configure import SMTPResponse -from fastapi import APIRouter, HTTPException -from loguru import logger smtp_configure_router = APIRouter() auth_handler = AuthHandler() diff --git a/backend/app/smtp/routes/reports.py b/backend/app/smtp/routes/reports.py index f1b15e6b..89c3eddb 100644 --- a/backend/app/smtp/routes/reports.py +++ b/backend/app/smtp/routes/reports.py @@ -1,10 +1,13 @@ -from app.auth.models.users import SMTP, SMTPInput +from fastapi import APIRouter +from fastapi import HTTPException +from loguru import logger + +from app.auth.models.users import SMTP +from app.auth.models.users import SMTPInput from app.auth.services.universal import select_all_users from app.auth.utils import AuthHandler from app.db.db_session import session from app.smtp.schema.configure import SMTPResponse -from fastapi import APIRouter, HTTPException -from loguru import logger smtp_reports_router = APIRouter() auth_handler = AuthHandler() diff --git a/backend/app/smtp/services/create_report.py b/backend/app/smtp/services/create_report.py index 25efcbc2..835fca4a 100644 --- a/backend/app/smtp/services/create_report.py +++ b/backend/app/smtp/services/create_report.py @@ -8,7 +8,10 @@ from reportlab.lib.units import inch # from reportlab.pdfgen import canvas -from reportlab.platypus import Image, Paragraph, SimpleDocTemplate, Spacer +from reportlab.platypus import Image +from reportlab.platypus import Paragraph +from reportlab.platypus import SimpleDocTemplate +from reportlab.platypus import Spacer matplotlib.use( "Agg", @@ -16,6 +19,7 @@ # for scripts and web servers. This should resolve the main thread is not # in main loop issue as it bypasses the need for tkinter. import matplotlib.pyplot as plt + from app.services.wazuh_indexer.alerts import AlertsService # ! TODO: Just a template diff --git a/backend/app/smtp/services/reports.py b/backend/app/smtp/services/reports.py index 14b27517..1baceb40 100644 --- a/backend/app/smtp/services/reports.py +++ b/backend/app/smtp/services/reports.py @@ -6,7 +6,8 @@ from typing import List from app.services.smtp.create_report import create_alerts_report_pdf -from app.services.smtp.universal import EmailTemplate, UniversalEmailCredentials +from app.services.smtp.universal import EmailTemplate +from app.services.smtp.universal import UniversalEmailCredentials # ! SEND REPORT diff --git a/backend/app/threat_intel/routes/socfortress.py b/backend/app/threat_intel/routes/socfortress.py index ddac78b9..a2c6f3c5 100644 --- a/backend/app/threat_intel/routes/socfortress.py +++ b/backend/app/threat_intel/routes/socfortress.py @@ -1,14 +1,16 @@ +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Security +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + from app.auth.utils import AuthHandler from app.db.db_session import get_db -from app.threat_intel.schema.socfortress import ( - IoCResponse, - SocfortressThreatIntelRequest, -) +from app.threat_intel.schema.socfortress import IoCResponse +from app.threat_intel.schema.socfortress import SocfortressThreatIntelRequest from app.threat_intel.services.socfortress import socfortress_threat_intel_lookup from app.utils import get_connector_attribute -from fastapi import APIRouter, Depends, HTTPException, Security -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession # App specific imports diff --git a/backend/app/threat_intel/schema/socfortress.py b/backend/app/threat_intel/schema/socfortress.py index 66126ecb..24c452e8 100644 --- a/backend/app/threat_intel/schema/socfortress.py +++ b/backend/app/threat_intel/schema/socfortress.py @@ -1,6 +1,7 @@ from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel +from pydantic import Field class SocfortressThreatIntelRequest(BaseModel): diff --git a/backend/app/threat_intel/services/socfortress.py b/backend/app/threat_intel/services/socfortress.py index 8d5b1ceb..a34e060f 100644 --- a/backend/app/threat_intel/services/socfortress.py +++ b/backend/app/threat_intel/services/socfortress.py @@ -1,18 +1,18 @@ -from typing import Any, Dict +from typing import Any +from typing import Dict import httpx -from app.connectors.utils import get_connector_info_from_db -from app.db.db_session import get_db_session -from app.threat_intel.schema.socfortress import ( - IoCMapping, - IoCResponse, - SocfortressThreatIntelRequest, -) -from app.utils import get_connector_attribute from fastapi import HTTPException from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession +from app.connectors.utils import get_connector_info_from_db +from app.db.db_session import get_db_session +from app.threat_intel.schema.socfortress import IoCMapping +from app.threat_intel.schema.socfortress import IoCResponse +from app.threat_intel.schema.socfortress import SocfortressThreatIntelRequest +from app.utils import get_connector_attribute + async def get_socfortress_threat_intel_attributes( column_name: str, diff --git a/backend/app/utils.py b/backend/app/utils.py index 40fd9e2c..42a83794 100644 --- a/backend/app/utils.py +++ b/backend/app/utils.py @@ -1,26 +1,44 @@ -from datetime import datetime, timedelta +from datetime import datetime +from datetime import timedelta from enum import Enum -from typing import Any, Dict, List, Optional, Union +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union import requests +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from fastapi import Request +from fastapi import Security +from fastapi.exceptions import RequestValidationError +from loguru import logger +from pydantic import BaseModel +from pydantic import Field +from pydantic import validator +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select +from sqlalchemy.orm import joinedload + from app.auth.services.universal import find_user from app.auth.utils import AuthHandler from app.connectors.utils import get_connector_info_from_db from app.db.all_models import Connectors -from app.db.db_session import get_db, get_db_session, get_session +from app.db.db_session import get_db +from app.db.db_session import get_db_session +from app.db.db_session import get_session from app.db.universal_models import LogEntry from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationEventConfig, +) +from app.integrations.alert_creation_settings.models.alert_creation_settings import ( AlertCreationSettings, +) +from app.integrations.alert_creation_settings.models.alert_creation_settings import ( EventOrder, ) -from fastapi import APIRouter, Depends, HTTPException, Request, Security -from fastapi.exceptions import RequestValidationError -from loguru import logger -from pydantic import BaseModel, Field, validator -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.future import select -from sqlalchemy.orm import joinedload ################## ! 422 VALIDATION ERROR TYPES FOR PYDANTIC VALUE ERROR RESPONSE ! ################## @@ -229,9 +247,7 @@ async def get_user_id_from_request(self, request: Request): auth_header = request.headers.get("Authorization") if auth_header: try: - token = auth_header.split(" ")[ - 1 - ] # Better split by space and take the second part + token = auth_header.split(" ")[1] # Better split by space and take the second part except IndexError: raise HTTPException(status_code=401, detail="Invalid token") username, _ = self.auth_handler.decode_token(token) @@ -362,9 +378,7 @@ async def get_logs(session: AsyncSession = Depends(get_db)) -> LogsResponse: auth_handler_instance = AuthHandler() # Initialize your AuthHandler logger_instance = Logger(session, auth_handler_instance) - logs = ( - await logger_instance.fetch_all_logs() - ) # Assuming fetch_all_logs is an async function + logs = await logger_instance.fetch_all_logs() # Assuming fetch_all_logs is an async function if logs: return LogsResponse( logs=logs, @@ -441,12 +455,7 @@ async def get_logs_by_time_range( logs = result.scalars().all() if logs: - logs = [ - log - for log in logs - if log.timestamp - >= datetime.now() - timedelta(days=int(time_range.time_range[:-1])) - ] + logs = [log for log in logs if log.timestamp >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))] if logs != []: return LogsResponse( logs=logs, @@ -564,12 +573,7 @@ async def purge_logs_by_time_range( logs = result.scalars().all() if logs: - logs = [ - log - for log in logs - if log.timestamp - >= datetime.now() - timedelta(days=int(time_range.time_range[:-1])) - ] + logs = [log for log in logs if log.timestamp >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))] if logs != []: for log in logs: await session.delete(log) @@ -674,8 +678,7 @@ async def get_customer_alert_settings_office365( """ result = await session.execute( select(AlertCreationSettings).filter( - AlertCreationSettings.office365_organization_id - == office365_organization_id, + AlertCreationSettings.office365_organization_id == office365_organization_id, ), ) settings = result.scalars().first() diff --git a/backend/copilot.py b/backend/copilot.py index c440bc0d..f3f66a84 100644 --- a/backend/copilot.py +++ b/backend/copilot.py @@ -1,59 +1,56 @@ import os import uvicorn -from app.auth.utils import AuthHandler -from app.db.db_session import async_engine -from app.db.db_setup import ( - create_available_integrations, - create_roles, - create_tables, - ensure_admin_user, - ensure_scheduler_user, - ensure_scheduler_user_removed, -) -from app.middleware.exception_handlers import ( - custom_http_exception_handler, - validation_exception_handler, - value_error_handler, -) -from app.middleware.logger import log_requests -from app.routers import ( - agents, - alert_creation, - alert_creation_settings, - ask_socfortress, - auth, - connectors, - cortex, - customer_provisioning, - customers, - dfir_iris, - dnstwist, - grafana, - graylog, - healthcheck, - influxdb, - integrations, - logs, - mimecast, - monitoring_alert, - office365, - scheduler, - shuffle, - smtp, - sublime, - threat_intel, - velociraptor, - wazuh_indexer, - wazuh_manager, -) -from app.schedulers.scheduler import init_scheduler from dotenv import load_dotenv -from fastapi import APIRouter, FastAPI, HTTPException +from fastapi import APIRouter +from fastapi import FastAPI +from fastapi import HTTPException from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware from loguru import logger +from app.auth.utils import AuthHandler +from app.db.db_session import async_engine +from app.db.db_setup import create_available_integrations +from app.db.db_setup import create_roles +from app.db.db_setup import create_tables +from app.db.db_setup import ensure_admin_user +from app.db.db_setup import ensure_scheduler_user +from app.db.db_setup import ensure_scheduler_user_removed +from app.middleware.exception_handlers import custom_http_exception_handler +from app.middleware.exception_handlers import validation_exception_handler +from app.middleware.exception_handlers import value_error_handler +from app.middleware.logger import log_requests +from app.routers import agents +from app.routers import alert_creation +from app.routers import alert_creation_settings +from app.routers import ask_socfortress +from app.routers import auth +from app.routers import connectors +from app.routers import cortex +from app.routers import customer_provisioning +from app.routers import customers +from app.routers import dfir_iris +from app.routers import dnstwist +from app.routers import grafana +from app.routers import graylog +from app.routers import healthcheck +from app.routers import influxdb +from app.routers import integrations +from app.routers import logs +from app.routers import mimecast +from app.routers import monitoring_alert +from app.routers import office365 +from app.routers import scheduler +from app.routers import shuffle +from app.routers import smtp +from app.routers import sublime +from app.routers import threat_intel +from app.routers import velociraptor +from app.routers import wazuh_indexer +from app.routers import wazuh_manager +from app.schedulers.scheduler import init_scheduler + auth_handler = AuthHandler() # Get the `SERVER_IP` from the `.env` file load_dotenv() diff --git a/git_tasks/Vagrantfile b/git_tasks/Vagrantfile deleted file mode 100644 index e0dbb4d8..00000000 --- a/git_tasks/Vagrantfile +++ /dev/null @@ -1,3 +0,0 @@ -Vagrant.configure("2") do |config| - config.vm.box = "debian/bullseye64" -end diff --git a/git_tasks/pyproject.toml b/pyproject.toml similarity index 100% rename from git_tasks/pyproject.toml rename to pyproject.toml