Precommit fixes (#137)
* precommits * precommit adjustments
taylor_socfortress committed
Feb 9, 2024 at 10:18 UTC
c599b707aee65b5c49f5ea028cb8daa1f26eba74
214 files changed
+1875
-1786
.pre-commit-config.yaml
+1
-1
@@ -10,7 +10,7 @@ repos:
10
- id: check-json
11
- id: trailing-whitespace
12
- id: check-added-large-files
13
- - id: detect-private-key
13
+ #- id: detect-private-key # Not using since uploading nginx key
14
- id: requirements-txt-fixer
15
args: ["backend/requirements.txt", "backend/requirements.in"]
16
backend/app/agents/dfir_iris/services/cases.py
+3
-2
@@ -1,10 +1,11 @@
1
from typing import List
2
3
+from loguru import logger
4
+from sqlalchemy.ext.asyncio import AsyncSession
5
+
6
from app.agents.dfir_iris.schema.cases import AssetCaseIDResponse
7
from app.connectors.dfir_iris.services.assets import get_case_assets
8
from app.connectors.dfir_iris.services.cases import get_all_cases
6
-from loguru import logger
7
-from sqlalchemy.ext.asyncio import AsyncSession
9
10
11
async def collect_agent_soc_cases(
backend/app/agents/routes/agents.py
+17
-16
@@ -1,15 +1,21 @@
1
+from fastapi import APIRouter
2
+from fastapi import BackgroundTasks
3
+from fastapi import Depends
4
+from fastapi import HTTPException
5
+from fastapi import Security
6
+from loguru import logger
7
+from sqlalchemy import delete
8
+from sqlalchemy.ext.asyncio import AsyncSession
9
+from sqlalchemy.future import select
10
+
11
from app.agents.dfir_iris.services.cases import collect_agent_soc_cases
2
-from app.agents.schema.agents import (
3
- AgentModifyResponse,
4
- AgentsResponse,
5
- OutdatedVelociraptorAgentsResponse,
6
- OutdatedWazuhAgentsResponse,
7
- SyncedAgentsResponse,
8
-)
9
-from app.agents.services.status import (
10
- get_outdated_agents_velociraptor,
11
- get_outdated_agents_wazuh,
12
-)
12
+from app.agents.schema.agents import AgentModifyResponse
13
+from app.agents.schema.agents import AgentsResponse
14
+from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse
15
+from app.agents.schema.agents import OutdatedWazuhAgentsResponse
16
+from app.agents.schema.agents import SyncedAgentsResponse
17
+from app.agents.services.status import get_outdated_agents_velociraptor
18
+from app.agents.services.status import get_outdated_agents_wazuh
19
from app.agents.services.sync import sync_agents
20
from app.agents.velociraptor.services.agents import delete_agent_velociraptor
21
from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
@@ -23,11 +29,6 @@ from app.db.db_session import get_db
29
# App specific imports
30
# from app.db.db_session import session
31
from app.db.universal_models import Agents
26
-from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Security
27
-from loguru import logger
28
-from sqlalchemy import delete
29
-from sqlalchemy.ext.asyncio import AsyncSession
30
-from sqlalchemy.future import select
32
33
agents_router = APIRouter()
34
backend/app/agents/schema/agents.py
+3
-1
@@ -1,9 +1,11 @@
1
from typing import List
2
3
+from pydantic import BaseModel
4
+from pydantic import Field
5
+
6
from app.agents.velociraptor.schema.agents import VelociraptorAgent
7
from app.agents.wazuh.schema.agents import WazuhAgent
8
from app.db.universal_models import Agents
6
-from pydantic import BaseModel, Field
9
10
11
class AgentsResponse(BaseModel):
backend/app/agents/services/modify.py
+2
-1
@@ -1,7 +1,8 @@
1
+from fastapi import HTTPException
2
+
3
import app.agents.wazuh.services.agents as wazuh_services
4
from app.db.db_session import session
5
from app.db.universal_models import Agents
4
-from fastapi import HTTPException
6
7
8
def delete_agent_db(agent_id: str):
backend/app/agents/services/status.py
+7
-12
@@ -1,17 +1,16 @@
1
from typing import List
2
3
-from app.agents.schema.agents import (
4
- OutdatedVelociraptorAgentsResponse,
5
- OutdatedWazuhAgentsResponse,
6
-)
7
-from app.connectors.velociraptor.utils.universal import UniversalService
8
-from app.db.db_session import session
9
-from app.db.universal_models import Agents
3
from fastapi import HTTPException
4
from loguru import logger
5
from sqlalchemy.ext.asyncio import AsyncSession
6
from sqlalchemy.future import select
7
8
+from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse
9
+from app.agents.schema.agents import OutdatedWazuhAgentsResponse
10
+from app.connectors.velociraptor.utils.universal import UniversalService
11
+from app.db.db_session import session
12
+from app.db.universal_models import Agents
13
+
14
15
def get_agent(agent_id: str) -> List[Agents]:
16
"""
@@ -100,11 +99,7 @@ async def get_outdated_agents_velociraptor(
99
)
100
agents_result = await session.execute(select(Agents))
101
agents = agents_result.scalars().all()
103
- outdated_velociraptor_agents = [
104
- agent
105
- for agent in agents
106
- if agent.velociraptor_agent_version != server_version
107
- ]
102
+ outdated_velociraptor_agents = [agent for agent in agents if agent.velociraptor_agent_version != server_version]
103
104
return OutdatedVelociraptorAgentsResponse(
105
message="Outdated Velociraptor agents fetched successfully.",
backend/app/agents/services/sync.py
+8
-5
@@ -1,15 +1,18 @@
1
from typing import List
2
3
+from loguru import logger
4
+from sqlalchemy.ext.asyncio import AsyncSession
5
+from sqlalchemy.future import select
6
+
7
import app.agents.velociraptor.services.agents as velociraptor_services
8
import app.agents.wazuh.services.agents as wazuh_services
5
-from app.agents.schema.agents import SyncedAgent, SyncedAgentsResponse
9
+from app.agents.schema.agents import SyncedAgent
10
+from app.agents.schema.agents import SyncedAgentsResponse
11
from app.agents.velociraptor.schema.agents import VelociraptorAgent
7
-from app.agents.wazuh.schema.agents import WazuhAgent, WazuhAgentsList
12
+from app.agents.wazuh.schema.agents import WazuhAgent
13
+from app.agents.wazuh.schema.agents import WazuhAgentsList
14
from app.connectors.models import Connectors
15
from app.db.universal_models import Agents
10
-from loguru import logger
11
-from sqlalchemy.ext.asyncio import AsyncSession
12
-from sqlalchemy.future import select
16
17
18
async def fetch_wazuh_agents() -> WazuhAgentsList:
backend/app/agents/velociraptor/schema/agents.py
+2
-1
@@ -1,7 +1,8 @@
1
from datetime import datetime
2
from typing import Optional
3
4
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
class VelociraptorAgent(BaseModel):
backend/app/agents/velociraptor/services/agents.py
+4
-5
@@ -1,10 +1,11 @@
1
from datetime import datetime
2
3
+from fastapi import HTTPException
4
+from loguru import logger
5
+
6
from app.agents.schema.agents import AgentModifyResponse
7
from app.agents.velociraptor.schema.agents import VelociraptorAgent
8
from app.connectors.velociraptor.utils.universal import UniversalService
6
-from fastapi import HTTPException
7
-from loguru import logger
9
10
11
def create_query(query: str) -> str:
@@ -45,9 +46,7 @@ async def collect_velociraptor_agent(agent_name: str) -> VelociraptorAgent:
46
)
47
48
try:
48
- vql_last_seen_at = (
49
- f"select last_seen_at from clients(search='host:{agent_name}')"
50
- )
49
+ vql_last_seen_at = f"select last_seen_at from clients(search='host:{agent_name}')"
50
last_seen_at = await velociraptor_service._get_last_seen_timestamp(
51
vql_last_seen_at,
52
)
backend/app/agents/wazuh/schema/agents.py
+4
-2
@@ -1,7 +1,9 @@
1
from datetime import datetime
2
-from typing import List, Optional
2
+from typing import List
3
+from typing import Optional
4
4
-from pydantic import BaseModel, Field
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
8
9
class WazuhAgent(BaseModel):
backend/app/agents/wazuh/services/agents.py
+7
-11
@@ -1,12 +1,12 @@
1
-from app.agents.schema.agents import AgentModifyResponse
2
-from app.agents.wazuh.schema.agents import WazuhAgent, WazuhAgentsList
3
-from app.connectors.wazuh_manager.utils.universal import (
4
- send_delete_request,
5
- send_get_request,
6
-)
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.agents.schema.agents import AgentModifyResponse
5
+from app.agents.wazuh.schema.agents import WazuhAgent
6
+from app.agents.wazuh.schema.agents import WazuhAgentsList
7
+from app.connectors.wazuh_manager.utils.universal import send_delete_request
8
+from app.connectors.wazuh_manager.utils.universal import send_get_request
9
+
10
11
async def collect_wazuh_agents() -> WazuhAgentsList:
12
"""
@@ -29,11 +29,7 @@ async def collect_wazuh_agents() -> WazuhAgentsList:
29
try:
30
if agents_collected.get("success"):
31
wazuh_agents_list = []
32
- for agent in (
33
- agents_collected.get("data", {})
34
- .get("data", {})
35
- .get("affected_items", [])
36
- ):
32
+ for agent in agents_collected.get("data", {}).get("data", {}).get("affected_items", []):
33
os_name = agent.get("os", {}).get("name", "Unknown")
34
last_keep_alive = agent.get("lastKeepAlive", "Unknown")
35
agent_group_list = agent.get("group", [])
backend/app/agents/wazuh/services/vulnerabilities.py
+4
-5
@@ -1,13 +1,12 @@
1
from typing import List
2
3
-from app.agents.wazuh.schema.agents import (
4
- WazuhAgentVulnerabilities,
5
- WazuhAgentVulnerabilitiesResponse,
6
-)
7
-from app.connectors.wazuh_manager.utils.universal import send_get_request
3
from fastapi import HTTPException
4
from loguru import logger
5
6
+from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilities
7
+from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
8
+from app.connectors.wazuh_manager.utils.universal import send_get_request
9
+
10
11
async def collect_agent_vulnerabilities(agent_id: str):
12
"""
backend/app/auth/models/users.py
+6
-2
@@ -6,8 +6,12 @@ from enum import Enum
6
from typing import Optional
7
8
import bcrypt
9
-from pydantic import BaseModel, EmailStr, validator
10
-from sqlmodel import Field, Relationship, SQLModel
9
+from pydantic import BaseModel
10
+from pydantic import EmailStr
11
+from pydantic import validator
12
+from sqlmodel import Field
13
+from sqlmodel import Relationship
14
+from sqlmodel import SQLModel
15
16
17
class Role(SQLModel, table=True):
backend/app/auth/routes/auth.py
+19
-13
@@ -1,22 +1,28 @@
1
from datetime import timedelta
2
3
-from app.auth.models.users import (
4
- PasswordReset,
5
- PasswordResetToken,
6
- User,
7
- UserInput,
8
- UserLogin,
9
-)
10
-from app.auth.schema.auth import Token, UserLoginResponse, UserResponse
11
-from app.auth.schema.user import UserBaseResponse
12
-from app.auth.services.universal import find_user, select_all_users
13
-from app.auth.utils import AuthHandler
14
-from app.db.db_session import get_db
15
-from fastapi import APIRouter, Depends, HTTPException, Security, status
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from fastapi import status
8
from fastapi.security import OAuth2PasswordRequestForm
9
from loguru import logger
10
from sqlalchemy.ext.asyncio import AsyncSession
11
12
+from app.auth.models.users import PasswordReset
13
+from app.auth.models.users import PasswordResetToken
14
+from app.auth.models.users import User
15
+from app.auth.models.users import UserInput
16
+from app.auth.models.users import UserLogin
17
+from app.auth.schema.auth import Token
18
+from app.auth.schema.auth import UserLoginResponse
19
+from app.auth.schema.auth import UserResponse
20
+from app.auth.schema.user import UserBaseResponse
21
+from app.auth.services.universal import find_user
22
+from app.auth.services.universal import select_all_users
23
+from app.auth.utils import AuthHandler
24
+from app.db.db_session import get_db
25
+
26
ACCESS_TOKEN_EXPIRE_MINUTES = 1440
27
28
auth_router = APIRouter()
backend/app/auth/schema/user.py
+2
-1
@@ -1,6 +1,7 @@
1
from typing import List
2
3
-from pydantic import BaseModel, EmailStr
3
+from pydantic import BaseModel
4
+from pydantic import EmailStr
5
6
7
class UserBase(BaseModel):
backend/app/auth/services/universal.py
+5
-2
@@ -1,11 +1,14 @@
1
-from app.auth.models.users import Password, Role, User
2
-from app.db.db_session import async_engine
1
from loguru import logger
2
3
# ! New with Async
4
from sqlalchemy.ext.asyncio import AsyncSession
5
from sqlmodel import select
6
7
+from app.auth.models.users import Password
8
+from app.auth.models.users import Role
9
+from app.auth.models.users import User
10
+from app.db.db_session import async_engine
11
+
12
passwords_in_memory = {}
13
14
backend/app/auth/utils.py
+9
-4
@@ -1,12 +1,17 @@
1
-from datetime import datetime, timedelta
1
+from datetime import datetime
2
+from datetime import timedelta
3
4
import jwt
4
-from app.auth.services.universal import find_user, get_role
5
-from fastapi import Depends, HTTPException
6
-from fastapi.security import OAuth2PasswordBearer, SecurityScopes
5
+from fastapi import Depends
6
+from fastapi import HTTPException
7
+from fastapi.security import OAuth2PasswordBearer
8
+from fastapi.security import SecurityScopes
9
from loguru import logger
10
from passlib.context import CryptContext
11
12
+from app.auth.services.universal import find_user
13
+from app.auth.services.universal import get_role
14
+
15
16
class AuthHandler:
17
security = OAuth2PasswordBearer(
backend/app/connectors/cortex/routes/analyzers.py
+11
-8
@@ -1,15 +1,18 @@
1
from typing import List
2
3
-from app.auth.utils import AuthHandler
4
-from app.connectors.cortex.schema.analyzers import (
5
- AnalyzersResponse,
6
- RunAnalyzerBody,
7
- RunAnalyzerResponse,
8
-)
9
-from app.connectors.cortex.services.analyzers import get_analyzers, run_analyzer
10
-from fastapi import APIRouter, Depends, HTTPException, Security
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
from loguru import logger
8
9
+from app.auth.utils import AuthHandler
10
+from app.connectors.cortex.schema.analyzers import AnalyzersResponse
11
+from app.connectors.cortex.schema.analyzers import RunAnalyzerBody
12
+from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse
13
+from app.connectors.cortex.services.analyzers import get_analyzers
14
+from app.connectors.cortex.services.analyzers import run_analyzer
15
+
16
# App specific imports
17
18
backend/app/connectors/cortex/schema/analyzers.py
+8
-2
@@ -1,8 +1,14 @@
1
import ipaddress
2
import re
3
-from typing import Any, Dict, List, Optional, Tuple
3
+from typing import Any
4
+from typing import Dict
5
+from typing import List
6
+from typing import Optional
7
+from typing import Tuple
8
5
-from pydantic import BaseModel, Field, validator
9
+from pydantic import BaseModel
10
+from pydantic import Field
11
+from pydantic import validator
12
13
HASH_REGEX = re.compile(
14
r"[a-fA-F\d]{32}|[a-fA-F\d]{64}",
backend/app/connectors/cortex/services/analyzers.py
+12
-11
@@ -1,19 +1,20 @@
1
-from typing import Dict, List, Union
1
+from typing import Dict
2
+from typing import List
3
+from typing import Union
4
3
-from app.connectors.cortex.schema.analyzers import (
4
- AnalyzerJobData,
5
- AnalyzersResponse,
6
- RunAnalyzerBody,
7
- RunAnalyzerResponse,
8
-)
9
-from app.connectors.cortex.utils.universal import ( # Importing create_cortex_client; Importing from universal.py
10
- create_cortex_client,
11
- run_and_wait_for_analyzer,
12
-)
5
from cortex4py.api import Api
6
from fastapi import HTTPException
7
from loguru import logger
8
9
+from app.connectors.cortex.schema.analyzers import AnalyzerJobData
10
+from app.connectors.cortex.schema.analyzers import AnalyzersResponse
11
+from app.connectors.cortex.schema.analyzers import RunAnalyzerBody
12
+from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse
13
+from app.connectors.cortex.utils.universal import (
14
+ create_cortex_client, # Importing create_cortex_client; Importing from universal.py
15
+)
16
+from app.connectors.cortex.utils.universal import run_and_wait_for_analyzer
17
+
18
19
async def fetch_analyzers(api: Api) -> List[Dict]:
20
"""
backend/app/connectors/cortex/utils/universal.py
+6
-4
@@ -1,13 +1,15 @@
1
import time
2
-from typing import Any, Dict
2
+from typing import Any
3
+from typing import Dict
4
4
-from app.connectors.cortex.schema.analyzers import AnalyzerJobData
5
-from app.connectors.utils import get_connector_info_from_db
6
-from app.db.db_session import get_db_session
5
from cortex4py.api import Api
6
from fastapi import HTTPException
7
from loguru import logger
8
9
+from app.connectors.cortex.schema.analyzers import AnalyzerJobData
10
+from app.connectors.utils import get_connector_info_from_db
11
+from app.db.db_session import get_db_session
12
+
13
14
async def verify_cortex_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
15
"""
backend/app/connectors/dfir_iris/routes/alerts.py
+20
-20
@@ -1,26 +1,26 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
+
8
from app.auth.utils import AuthHandler
2
-from app.connectors.dfir_iris.schema.alerts import (
3
- AlertResponse,
4
- AlertsResponse,
5
- BookmarkedAlertsResponse,
6
- CaseCreationResponse,
7
- DeleteAlertResponse,
8
- DeleteMultipleAlertsRequest,
9
- FilterAlertsRequest,
10
-)
11
-from app.connectors.dfir_iris.services.alerts import (
12
- bookmark_alert,
13
- create_case,
14
- delete_alert,
15
- get_alert,
16
- get_alerts,
17
- get_bookmarked_alerts,
18
-)
9
+from app.connectors.dfir_iris.schema.alerts import AlertResponse
10
+from app.connectors.dfir_iris.schema.alerts import AlertsResponse
11
+from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse
12
+from app.connectors.dfir_iris.schema.alerts import CaseCreationResponse
13
+from app.connectors.dfir_iris.schema.alerts import DeleteAlertResponse
14
+from app.connectors.dfir_iris.schema.alerts import DeleteMultipleAlertsRequest
15
+from app.connectors.dfir_iris.schema.alerts import FilterAlertsRequest
16
+from app.connectors.dfir_iris.services.alerts import bookmark_alert
17
+from app.connectors.dfir_iris.services.alerts import create_case
18
+from app.connectors.dfir_iris.services.alerts import delete_alert
19
+from app.connectors.dfir_iris.services.alerts import get_alert
20
+from app.connectors.dfir_iris.services.alerts import get_alerts
21
+from app.connectors.dfir_iris.services.alerts import get_bookmarked_alerts
22
from app.connectors.dfir_iris.utils.universal import check_alert_exists
23
from app.db.db_session import get_db
21
-from fastapi import APIRouter, Depends, HTTPException, Security
22
-from loguru import logger
23
-from sqlalchemy.ext.asyncio import AsyncSession
24
25
# App specific imports
26
backend/app/connectors/dfir_iris/routes/assets.py
+6
-2
@@ -1,9 +1,13 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+
7
from app.auth.utils import AuthHandler
8
from app.connectors.dfir_iris.schema.assets import AssetResponse
9
from app.connectors.dfir_iris.services.assets import get_case_assets
10
from app.connectors.dfir_iris.utils.universal import check_case_exists
5
-from fastapi import APIRouter, Depends, HTTPException, Security
6
-from loguru import logger
11
12
# App specific imports
13
backend/app/connectors/dfir_iris/routes/cases.py
+23
-23
@@ -1,31 +1,31 @@
1
from datetime import timedelta
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from loguru import logger
8
+from sqlalchemy.ext.asyncio import AsyncSession
9
+
10
from app.auth.utils import AuthHandler
4
-from app.connectors.dfir_iris.schema.cases import (
5
- CaseOlderThanBody,
6
- CaseResponse,
7
- CasesBreachedResponse,
8
- ClosedCaseResponse,
9
- PurgeCaseResponse,
10
- ReopenedCaseResponse,
11
- SingleCaseBody,
12
- SingleCaseResponse,
13
- TimeUnit,
14
-)
15
-from app.connectors.dfir_iris.services.cases import (
16
- close_case,
17
- delete_single_case,
18
- get_all_cases,
19
- get_cases_older_than,
20
- get_single_case,
21
- purge_cases,
22
- reopen_case,
23
-)
11
+from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody
12
+from app.connectors.dfir_iris.schema.cases import CaseResponse
13
+from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse
14
+from app.connectors.dfir_iris.schema.cases import ClosedCaseResponse
15
+from app.connectors.dfir_iris.schema.cases import PurgeCaseResponse
16
+from app.connectors.dfir_iris.schema.cases import ReopenedCaseResponse
17
+from app.connectors.dfir_iris.schema.cases import SingleCaseBody
18
+from app.connectors.dfir_iris.schema.cases import SingleCaseResponse
19
+from app.connectors.dfir_iris.schema.cases import TimeUnit
20
+from app.connectors.dfir_iris.services.cases import close_case
21
+from app.connectors.dfir_iris.services.cases import delete_single_case
22
+from app.connectors.dfir_iris.services.cases import get_all_cases
23
+from app.connectors.dfir_iris.services.cases import get_cases_older_than
24
+from app.connectors.dfir_iris.services.cases import get_single_case
25
+from app.connectors.dfir_iris.services.cases import purge_cases
26
+from app.connectors.dfir_iris.services.cases import reopen_case
27
from app.connectors.dfir_iris.utils.universal import check_case_exists
28
from app.db.db_session import get_db
26
-from fastapi import APIRouter, Depends, HTTPException, Security
27
-from loguru import logger
28
-from sqlalchemy.ext.asyncio import AsyncSession
29
30
31
async def verify_case_exists(case_id: int) -> int:
backend/app/connectors/dfir_iris/routes/notes.py
+11
-8
@@ -1,15 +1,18 @@
1
from typing import Optional
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from loguru import logger
8
+
9
from app.auth.utils import AuthHandler
4
-from app.connectors.dfir_iris.schema.notes import (
5
- NoteCreationBody,
6
- NoteCreationResponse,
7
- NotesResponse,
8
-)
9
-from app.connectors.dfir_iris.services.notes import create_case_note, get_case_notes
10
+from app.connectors.dfir_iris.schema.notes import NoteCreationBody
11
+from app.connectors.dfir_iris.schema.notes import NoteCreationResponse
12
+from app.connectors.dfir_iris.schema.notes import NotesResponse
13
+from app.connectors.dfir_iris.services.notes import create_case_note
14
+from app.connectors.dfir_iris.services.notes import get_case_notes
15
from app.connectors.dfir_iris.utils.universal import check_case_exists
11
-from fastapi import APIRouter, Depends, HTTPException, Security
12
-from loguru import logger
16
17
18
async def verify_case_exists(case_id: int) -> int:
backend/app/connectors/dfir_iris/routes/users.py
+13
-12
@@ -1,17 +1,18 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+
7
from app.auth.utils import AuthHandler
8
from app.connectors.dfir_iris.schema.alerts import AlertResponse
3
-from app.connectors.dfir_iris.schema.users import User, UsersResponse
4
-from app.connectors.dfir_iris.services.users import (
5
- assign_user_to_alert,
6
- delete_user_from_alert,
7
- get_users,
8
-)
9
-from app.connectors.dfir_iris.utils.universal import (
10
- check_alert_exists,
11
- check_user_exists,
12
-)
13
-from fastapi import APIRouter, Depends, HTTPException, Security
14
-from loguru import logger
9
+from app.connectors.dfir_iris.schema.users import User
10
+from app.connectors.dfir_iris.schema.users import UsersResponse
11
+from app.connectors.dfir_iris.services.users import assign_user_to_alert
12
+from app.connectors.dfir_iris.services.users import delete_user_from_alert
13
+from app.connectors.dfir_iris.services.users import get_users
14
+from app.connectors.dfir_iris.utils.universal import check_alert_exists
15
+from app.connectors.dfir_iris.utils.universal import check_user_exists
16
17
18
def verify_user_exists(user_id: int) -> int:
backend/app/connectors/dfir_iris/schema/admin.py
+5
-2
@@ -1,8 +1,11 @@
1
import uuid
2
from datetime import datetime
3
-from typing import Dict, List, Optional
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
5
-from pydantic import UUID4, BaseModel
7
+from pydantic import UUID4
8
+from pydantic import BaseModel
9
10
11
class CreateCustomerData(BaseModel):
backend/app/connectors/dfir_iris/schema/alerts.py
+6
-2
@@ -1,7 +1,11 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
4
-from pydantic import BaseModel, Field
7
+from pydantic import BaseModel
8
+from pydantic import Field
9
10
11
class AlertsResponse(BaseModel):
backend/app/connectors/dfir_iris/schema/assets.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
class AssetState(BaseModel):
backend/app/connectors/dfir_iris/schema/cases.py
+8
-3
@@ -1,8 +1,13 @@
1
-from datetime import date, timedelta
1
+from datetime import date
2
+from datetime import timedelta
3
from enum import Enum
3
-from typing import Dict, List, Optional, Union
4
+from typing import Dict
5
+from typing import List
6
+from typing import Optional
7
+from typing import Union
8
5
-from pydantic import BaseModel, Field
9
+from pydantic import BaseModel
10
+from pydantic import Field
11
12
13
class CaseModel(BaseModel):
backend/app/connectors/dfir_iris/schema/notes.py
+5
-2
@@ -1,6 +1,9 @@
1
-from typing import Dict, List, Optional
1
+from typing import Dict
2
+from typing import List
3
+from typing import Optional
4
3
-from pydantic import BaseModel, Field
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
8
9
class CustomAttributes(BaseModel):
backend/app/connectors/dfir_iris/services/alerts.py
+13
-18
@@ -1,23 +1,20 @@
1
-from app.connectors.dfir_iris.schema.alerts import (
2
- AlertResponse,
3
- AlertsResponse,
4
- BookmarkedAlertsResponse,
5
- CaseCreationResponse,
6
- DeleteAlertResponse,
7
- FilterAlertsRequest,
8
-)
9
-from app.connectors.dfir_iris.utils.universal import (
10
- fetch_and_validate_data,
11
- initialize_client_and_alert,
12
-)
13
-from app.integrations.alert_creation_settings.models.alert_creation_settings import (
14
- AlertCreationSettings,
15
-)
1
from fastapi import HTTPException
2
from loguru import logger
3
from sqlalchemy.ext.asyncio import AsyncSession
4
from sqlalchemy.future import select
5
6
+from app.connectors.dfir_iris.schema.alerts import AlertResponse
7
+from app.connectors.dfir_iris.schema.alerts import AlertsResponse
8
+from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse
9
+from app.connectors.dfir_iris.schema.alerts import CaseCreationResponse
10
+from app.connectors.dfir_iris.schema.alerts import DeleteAlertResponse
11
+from app.connectors.dfir_iris.schema.alerts import FilterAlertsRequest
12
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
13
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
14
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
15
+ AlertCreationSettings,
16
+)
17
+
18
19
async def get_customer_code(session: AsyncSession, customer_id: int) -> str:
20
"""
@@ -180,9 +177,7 @@ def construct_case_creation_params(alert_details: dict) -> dict:
177
"case_tags": alert_details["alert_tags"],
178
"escalation_note": "Case created from CoPilot",
179
"iocs_import_list": [ioc["ioc_uuid"] for ioc in alert_details["iocs"]],
183
- "assets_import_list": [
184
- asset["asset_uuid"] for asset in alert_details["assets"]
185
- ],
180
+ "assets_import_list": [asset["asset_uuid"] for asset in alert_details["assets"]],
181
}
182
183
# Replace None values with the string "None"
backend/app/connectors/dfir_iris/services/assets.py
+6
-5
@@ -1,10 +1,11 @@
1
-from app.connectors.dfir_iris.schema.assets import Asset, AssetResponse, AssetState
2
-from app.connectors.dfir_iris.utils.universal import (
3
- fetch_and_validate_data,
4
- initialize_client_and_case,
5
-)
1
from fastapi import HTTPException
2
3
+from app.connectors.dfir_iris.schema.assets import Asset
4
+from app.connectors.dfir_iris.schema.assets import AssetResponse
5
+from app.connectors.dfir_iris.schema.assets import AssetState
6
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
7
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_case
8
+
9
10
async def get_case_assets(case_id: int) -> AssetResponse:
11
"""
backend/app/connectors/dfir_iris/services/cases.py
+17
-19
@@ -1,29 +1,27 @@
1
from datetime import datetime
2
-from typing import Dict, List
3
-
4
-from app.connectors.dfir_iris.schema.cases import (
5
- CaseOlderThanBody,
6
- CaseResponse,
7
- CasesBreachedResponse,
8
- ClosedCaseResponse,
9
- PurgeCaseResponse,
10
- ReopenedCaseResponse,
11
- SingleCaseBody,
12
- SingleCaseResponse,
13
-)
14
-from app.connectors.dfir_iris.utils.universal import (
15
- create_dfir_iris_client,
16
- fetch_and_parse_data,
17
-)
18
-from app.integrations.alert_creation_settings.models.alert_creation_settings import (
19
- AlertCreationSettings,
20
-)
2
+from typing import Dict
3
+from typing import List
4
+
5
from dfir_iris_client.case import Case
6
from fastapi import HTTPException
7
from loguru import logger
8
from sqlalchemy.ext.asyncio import AsyncSession
9
from sqlalchemy.future import select
10
11
+from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody
12
+from app.connectors.dfir_iris.schema.cases import CaseResponse
13
+from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse
14
+from app.connectors.dfir_iris.schema.cases import ClosedCaseResponse
15
+from app.connectors.dfir_iris.schema.cases import PurgeCaseResponse
16
+from app.connectors.dfir_iris.schema.cases import ReopenedCaseResponse
17
+from app.connectors.dfir_iris.schema.cases import SingleCaseBody
18
+from app.connectors.dfir_iris.schema.cases import SingleCaseResponse
19
+from app.connectors.dfir_iris.utils.universal import create_dfir_iris_client
20
+from app.connectors.dfir_iris.utils.universal import fetch_and_parse_data
21
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
22
+ AlertCreationSettings,
23
+)
24
+
25
26
async def get_client_and_cases() -> Dict:
27
"""
backend/app/connectors/dfir_iris/services/notes.py
+12
-13
@@ -1,19 +1,18 @@
1
-from typing import Any, Dict, List
2
-
3
-from app.connectors.dfir_iris.schema.notes import (
4
- NoteCreationBody,
5
- NoteCreationResponse,
6
- NoteDetails,
7
- NoteDetailsResponse,
8
- NotesResponse,
9
-)
10
-from app.connectors.dfir_iris.utils.universal import (
11
- fetch_and_validate_data,
12
- initialize_client_and_case,
13
-)
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+
5
from dfir_iris_client.case import Case
6
from loguru import logger
7
8
+from app.connectors.dfir_iris.schema.notes import NoteCreationBody
9
+from app.connectors.dfir_iris.schema.notes import NoteCreationResponse
10
+from app.connectors.dfir_iris.schema.notes import NoteDetails
11
+from app.connectors.dfir_iris.schema.notes import NoteDetailsResponse
12
+from app.connectors.dfir_iris.schema.notes import NotesResponse
13
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
14
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_case
15
+
16
17
async def process_notes(notes: List[Dict], case_id: int) -> List[Dict]:
18
"""
backend/app/connectors/dfir_iris/services/users.py
+3
-5
@@ -1,10 +1,8 @@
1
from app.connectors.dfir_iris.schema.alerts import AlertResponse
2
from app.connectors.dfir_iris.schema.users import UsersResponse
3
-from app.connectors.dfir_iris.utils.universal import (
4
- fetch_and_validate_data,
5
- initialize_client_and_alert,
6
- initialize_client_and_user,
7
-)
3
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
4
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
5
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_user
6
7
8
async def get_users() -> UsersResponse:
backend/app/connectors/dfir_iris/utils/universal.py
+11
-4
@@ -1,18 +1,25 @@
1
-from typing import Any, Callable, Dict, Optional, Tuple, Union
1
+from typing import Any
2
+from typing import Callable
3
+from typing import Dict
4
+from typing import Optional
5
+from typing import Tuple
6
+from typing import Union
7
8
import requests
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.db.db_session import get_db_session
9
from dfir_iris_client.admin import AdminHelper
10
from dfir_iris_client.alert import Alert
11
from dfir_iris_client.case import Case
12
from dfir_iris_client.customer import Customer
10
-from dfir_iris_client.helper.utils import assert_api_resp, get_data_from_resp
13
+from dfir_iris_client.helper.utils import assert_api_resp
14
+from dfir_iris_client.helper.utils import get_data_from_resp
15
from dfir_iris_client.session import ClientSession
16
from dfir_iris_client.users import User
17
from fastapi import HTTPException
18
from loguru import logger
19
20
+from app.connectors.utils import get_connector_info_from_db
21
+from app.db.db_session import get_db_session
22
+
23
24
async def verify_dfir_iris_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
25
"""
backend/app/connectors/event_shipper/utils/universal.py
+1
@@ -1,6 +1,7 @@
1
from typing import Optional
2
3
import asyncgelf
4
+
5
from app.connectors.utils import get_connector_info_from_db
6
from app.db.db_session import get_db_session
7
backend/app/connectors/grafana/routes/dashboards.py
+7
-6
@@ -1,11 +1,12 @@
1
+from fastapi import APIRouter
2
+from fastapi import Body
3
+from fastapi import Security
4
+from loguru import logger
5
+
6
from app.auth.utils import AuthHandler
2
-from app.connectors.grafana.schema.dashboards import (
3
- DashboardProvisionRequest,
4
- GrafanaDashboardResponse,
5
-)
7
+from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
8
+from app.connectors.grafana.schema.dashboards import GrafanaDashboardResponse
9
from app.connectors.grafana.services.dashboards import provision_dashboards
7
-from fastapi import APIRouter, Body, Security
8
-from loguru import logger
10
11
# App specific imports
12
backend/app/connectors/grafana/schema/dashboards.py
+4
-4
@@ -1,7 +1,9 @@
1
from enum import Enum
2
from typing import List
3
4
-from pydantic import BaseModel, Field, validator
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
+from pydantic import validator
7
8
9
class GrafanaDashboard(BaseModel):
@@ -81,9 +83,7 @@ class DashboardProvisionRequest(BaseModel):
83
84
@validator("dashboards", each_item=True)
85
def check_dashboard_exists(cls, e):
84
- valid_dashboards = {
85
- item.name: item for item in list(WazuhDashboard) + list(Office365Dashboard)
86
- }
86
+ valid_dashboards = {item.name: item for item in list(WazuhDashboard) + list(Office365Dashboard)}
87
if e not in valid_dashboards:
88
raise ValueError(f'Dashboard identifier "{e}" is not recognized.')
89
return e
backend/app/connectors/grafana/services/dashboards.py
+11
-23
@@ -1,18 +1,17 @@
1
import json
2
from pathlib import Path
3
4
-from app.connectors.grafana.schema.dashboards import (
5
- DashboardProvisionRequest,
6
- GrafanaDashboard,
7
- GrafanaDashboardResponse,
8
- MimecastDashboard,
9
- Office365Dashboard,
10
- WazuhDashboard,
11
-)
12
-from app.connectors.grafana.utils.universal import create_grafana_client
4
from fastapi import HTTPException
5
from loguru import logger
6
7
+from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
8
+from app.connectors.grafana.schema.dashboards import GrafanaDashboard
9
+from app.connectors.grafana.schema.dashboards import GrafanaDashboardResponse
10
+from app.connectors.grafana.schema.dashboards import MimecastDashboard
11
+from app.connectors.grafana.schema.dashboards import Office365Dashboard
12
+from app.connectors.grafana.schema.dashboards import WazuhDashboard
13
+from app.connectors.grafana.utils.universal import create_grafana_client
14
+
15
16
def get_dashboard_path(dashboard_info: tuple) -> Path:
17
"""
@@ -26,9 +25,7 @@ def get_dashboard_path(dashboard_info: tuple) -> Path:
25
"""
26
folder_name, file_name = dashboard_info
27
current_file = Path(__file__) # Path to the current file
29
- base_dir = (
30
- current_file.parent.parent
31
- ) # Move up two levels to the 'grafana' directory
28
+ base_dir = current_file.parent.parent # Move up two levels to the 'grafana' directory
29
return base_dir / "dashboards" / folder_name / file_name
30
31
@@ -147,12 +144,7 @@ async def provision_dashboards(
144
provisioned_dashboards = []
145
errors = []
146
150
- valid_dashboards = {
151
- item.name: item
152
- for item in list(WazuhDashboard)
153
- + list(Office365Dashboard)
154
- + list(MimecastDashboard)
155
- }
147
+ valid_dashboards = {item.name: item for item in list(WazuhDashboard) + list(Office365Dashboard) + list(MimecastDashboard)}
148
149
for dashboard_name in dashboard_request.dashboards:
150
dashboard_enum = valid_dashboards[dashboard_name]
@@ -175,11 +167,7 @@ async def provision_dashboards(
167
)
168
169
success = len(errors) == 0
178
- message = (
179
- "All dashboards provisioned successfully"
180
- if success
181
- else "Some dashboards failed to provision"
182
- )
170
+ message = "All dashboards provisioned successfully" if success else "Some dashboards failed to provision"
171
return GrafanaDashboardResponse(
172
provisioned_dashboards=provisioned_dashboards,
173
success=success,
backend/app/connectors/grafana/utils/universal.py
+6
-4
@@ -1,12 +1,14 @@
1
-from typing import Any, Dict
1
+from typing import Any
2
+from typing import Dict
3
3
-from app.connectors.grafana.schema.organization import GrafanaCreateOrganizationResponse
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.db.db_session import get_db_session
4
from fastapi import HTTPException
5
from grafana_client import GrafanaApi
6
from loguru import logger
7
8
+from app.connectors.grafana.schema.organization import GrafanaCreateOrganizationResponse
9
+from app.connectors.utils import get_connector_info_from_db
10
+from app.db.db_session import get_db_session
11
+
12
13
async def construct_grafana_url(
14
connector_url: str,
backend/app/connectors/graylog/routes/collector.py
+12
-14
@@ -1,19 +1,17 @@
1
-from app.auth.utils import AuthHandler
2
-from app.connectors.graylog.schema.collector import (
3
- ConfiguredInputsResponse,
4
- GraylogIndicesResponse,
5
- GraylogInputsResponse,
6
- RunningInputsResponse,
7
-)
8
-from app.connectors.graylog.services.collector import (
9
- get_indices_full,
10
- get_inputs,
11
- get_inputs_configured,
12
- get_inputs_running,
13
-)
14
-from fastapi import APIRouter, Security
1
+from fastapi import APIRouter
2
+from fastapi import Security
3
from loguru import logger
4
5
+from app.auth.utils import AuthHandler
6
+from app.connectors.graylog.schema.collector import ConfiguredInputsResponse
7
+from app.connectors.graylog.schema.collector import GraylogIndicesResponse
8
+from app.connectors.graylog.schema.collector import GraylogInputsResponse
9
+from app.connectors.graylog.schema.collector import RunningInputsResponse
10
+from app.connectors.graylog.services.collector import get_indices_full
11
+from app.connectors.graylog.services.collector import get_inputs
12
+from app.connectors.graylog.services.collector import get_inputs_configured
13
+from app.connectors.graylog.services.collector import get_inputs_running
14
+
15
# App specific imports
16
17
backend/app/connectors/graylog/routes/events.py
+9
-8
@@ -1,13 +1,14 @@
1
-from app.auth.utils import AuthHandler
2
-from app.connectors.graylog.schema.events import (
3
- AlertQuery,
4
- GraylogAlertsResponse,
5
- GraylogEventDefinitionsResponse,
6
-)
7
-from app.connectors.graylog.services.events import get_alerts, get_event_definitions
8
-from fastapi import APIRouter, Security
1
+from fastapi import APIRouter
2
+from fastapi import Security
3
from loguru import logger
4
5
+from app.auth.utils import AuthHandler
6
+from app.connectors.graylog.schema.events import AlertQuery
7
+from app.connectors.graylog.schema.events import GraylogAlertsResponse
8
+from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse
9
+from app.connectors.graylog.services.events import get_alerts
10
+from app.connectors.graylog.services.events import get_event_definitions
11
+
12
# App specific imports
13
14
backend/app/connectors/graylog/routes/management.py
+25
-27
@@ -1,34 +1,32 @@
1
from typing import List
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from loguru import logger
8
+
9
from app.auth.utils import AuthHandler
4
-from app.connectors.graylog.schema.management import (
5
- DeletedIndexBody,
6
- DeletedIndexResponse,
7
- StartInputBody,
8
- StartInputResponse,
9
- StartStreamBody,
10
- StartStreamResponse,
11
- StopInputBody,
12
- StopInputResponse,
13
- StopStreamBody,
14
- StopStreamResponse,
15
- UrlWhitelistEntryResponse,
16
-)
17
-from app.connectors.graylog.services.collector import (
18
- get_index_names,
19
- get_input_ids,
20
- get_url_whitelist_entries,
21
-)
22
-from app.connectors.graylog.services.management import (
23
- delete_index,
24
- start_input,
25
- start_stream,
26
- stop_input,
27
- stop_stream,
28
-)
10
+from app.connectors.graylog.schema.management import DeletedIndexBody
11
+from app.connectors.graylog.schema.management import DeletedIndexResponse
12
+from app.connectors.graylog.schema.management import StartInputBody
13
+from app.connectors.graylog.schema.management import StartInputResponse
14
+from app.connectors.graylog.schema.management import StartStreamBody
15
+from app.connectors.graylog.schema.management import StartStreamResponse
16
+from app.connectors.graylog.schema.management import StopInputBody
17
+from app.connectors.graylog.schema.management import StopInputResponse
18
+from app.connectors.graylog.schema.management import StopStreamBody
19
+from app.connectors.graylog.schema.management import StopStreamResponse
20
+from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse
21
+from app.connectors.graylog.services.collector import get_index_names
22
+from app.connectors.graylog.services.collector import get_input_ids
23
+from app.connectors.graylog.services.collector import get_url_whitelist_entries
24
+from app.connectors.graylog.services.management import delete_index
25
+from app.connectors.graylog.services.management import start_input
26
+from app.connectors.graylog.services.management import start_stream
27
+from app.connectors.graylog.services.management import stop_input
28
+from app.connectors.graylog.services.management import stop_stream
29
from app.connectors.graylog.services.streams import get_stream_ids
30
-from fastapi import APIRouter, Depends, HTTPException, Security
31
-from loguru import logger
30
31
graylog_management_router = APIRouter()
32
backend/app/connectors/graylog/routes/monitoring.py
+10
-12
@@ -1,17 +1,15 @@
1
-from app.auth.utils import AuthHandler
2
-from app.connectors.graylog.schema.monitoring import (
3
- GraylogEventNotificationsResponse,
4
- GraylogMessagesResponse,
5
- GraylogMetricsResponse,
6
-)
7
-from app.connectors.graylog.services.monitoring import (
8
- get_event_notifications,
9
- get_messages,
10
- get_metrics,
11
-)
12
-from fastapi import APIRouter, Security
1
+from fastapi import APIRouter
2
+from fastapi import Security
3
from loguru import logger
4
5
+from app.auth.utils import AuthHandler
6
+from app.connectors.graylog.schema.monitoring import GraylogEventNotificationsResponse
7
+from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse
8
+from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse
9
+from app.connectors.graylog.services.monitoring import get_event_notifications
10
+from app.connectors.graylog.services.monitoring import get_messages
11
+from app.connectors.graylog.services.monitoring import get_metrics
12
+
13
# App specific imports
14
15
backend/app/connectors/graylog/routes/pipelines.py
+19
-25
@@ -1,24 +1,23 @@
1
-from typing import Dict, List
1
+from typing import Dict
2
+from typing import List
3
3
-from app.auth.utils import AuthHandler
4
-from app.connectors.graylog.schema.pipelines import (
5
- GraylogPipelinesResponse,
6
- GraylogPipelinesResponseWithRuleID,
7
- Pipeline,
8
- PipelineRule,
9
- PipelineRulesResponse,
10
- PipelineWithRuleID,
11
- Stage,
12
- StageWithRuleID,
13
-)
14
-from app.connectors.graylog.services.pipelines import (
15
- get_pipeline_rule_by_id,
16
- get_pipeline_rules,
17
- get_pipelines,
18
-)
19
-from fastapi import APIRouter, Security
4
+from fastapi import APIRouter
5
+from fastapi import Security
6
from loguru import logger
7
8
+from app.auth.utils import AuthHandler
9
+from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
10
+from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponseWithRuleID
11
+from app.connectors.graylog.schema.pipelines import Pipeline
12
+from app.connectors.graylog.schema.pipelines import PipelineRule
13
+from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
14
+from app.connectors.graylog.schema.pipelines import PipelineWithRuleID
15
+from app.connectors.graylog.schema.pipelines import Stage
16
+from app.connectors.graylog.schema.pipelines import StageWithRuleID
17
+from app.connectors.graylog.services.pipelines import get_pipeline_rule_by_id
18
+from app.connectors.graylog.services.pipelines import get_pipeline_rules
19
+from app.connectors.graylog.services.pipelines import get_pipelines
20
+
21
# App specific imports
22
23
@@ -57,9 +56,7 @@ def transform_stages_with_rule_ids(
56
"""
57
new_stages = []
58
for stage in stages:
60
- rule_ids = [
61
- rule_title_to_id.get(rule_title, None) for rule_title in stage.rules
62
- ]
59
+ rule_ids = [rule_title_to_id.get(rule_title, None) for rule_title in stage.rules]
60
new_stage = StageWithRuleID(**stage.dict(), rule_ids=rule_ids)
61
new_stages.append(new_stage)
62
return new_stages
@@ -123,10 +120,7 @@ async def get_all_pipelines_with_rule_ids() -> GraylogPipelinesResponseWithRuleI
120
pipeline_rules_response.pipeline_rules,
121
)
122
126
- new_pipelines = [
127
- transform_pipeline_with_rule_ids(pipeline, rule_title_to_id)
128
- for pipeline in pipelines_response.pipelines
129
- ]
123
+ new_pipelines = [transform_pipeline_with_rule_ids(pipeline, rule_title_to_id) for pipeline in pipelines_response.pipelines]
124
125
return GraylogPipelinesResponseWithRuleID(
126
pipelines=new_pipelines,
backend/app/connectors/graylog/routes/streams.py
+4
-2
@@ -1,8 +1,10 @@
1
+from fastapi import APIRouter
2
+from fastapi import Security
3
+from loguru import logger
4
+
5
from app.auth.utils import AuthHandler
6
from app.connectors.graylog.schema.streams import GraylogStreamsResponse
7
from app.connectors.graylog.services.streams import get_streams
4
-from fastapi import APIRouter, Security
5
-from loguru import logger
8
9
# App specific imports
10
backend/app/connectors/graylog/schema/collector.py
+5
-2
@@ -1,6 +1,9 @@
1
-from typing import Dict, List, Optional
1
+from typing import Dict
2
+from typing import List
3
+from typing import Optional
4
3
-from pydantic import BaseModel, Field
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
8
9
class Document(BaseModel):
backend/app/connectors/graylog/schema/events.py
+4
-1
@@ -1,4 +1,7 @@
1
-from typing import Dict, List, Optional, Union
1
+from typing import Dict
2
+from typing import List
3
+from typing import Optional
4
+from typing import Union
5
6
from pydantic import BaseModel
7
backend/app/connectors/graylog/schema/monitoring.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
class GraylogMessages(BaseModel):
backend/app/connectors/graylog/schema/pipelines.py
+2
-1
@@ -1,4 +1,5 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
4
from pydantic import BaseModel
5
backend/app/connectors/graylog/schema/streams.py
+2
-1
@@ -1,4 +1,5 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
4
from pydantic import BaseModel
5
backend/app/connectors/graylog/services/collector.py
+16
-25
@@ -1,19 +1,19 @@
1
-from typing import List, Tuple
2
-
3
-from app.connectors.graylog.schema.collector import (
4
- ConfiguredInput,
5
- ConfiguredInputsResponse,
6
- GraylogIndexItem,
7
- GraylogIndicesResponse,
8
- GraylogInputsResponse,
9
- RunningInput,
10
- RunningInputsResponse,
11
-)
12
-from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse
13
-from app.connectors.graylog.utils.universal import send_get_request
1
+from typing import List
2
+from typing import Tuple
3
+
4
from fastapi import HTTPException
5
from loguru import logger
6
7
+from app.connectors.graylog.schema.collector import ConfiguredInput
8
+from app.connectors.graylog.schema.collector import ConfiguredInputsResponse
9
+from app.connectors.graylog.schema.collector import GraylogIndexItem
10
+from app.connectors.graylog.schema.collector import GraylogIndicesResponse
11
+from app.connectors.graylog.schema.collector import GraylogInputsResponse
12
+from app.connectors.graylog.schema.collector import RunningInput
13
+from app.connectors.graylog.schema.collector import RunningInputsResponse
14
+from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse
15
+from app.connectors.graylog.utils.universal import send_get_request
16
+
17
18
async def get_indices_full() -> GraylogIndicesResponse:
19
"""Get indices from Graylog.
@@ -33,10 +33,7 @@ async def get_indices_full() -> GraylogIndicesResponse:
33
raise HTTPException(status_code=500, detail="Failed to collect indices key")
34
35
# Convert the dictionary to a list of GraylogIndexItem
36
- indices_list = [
37
- GraylogIndexItem(index_name=name, index_info=info)
38
- for name, info in indices_data.items()
39
- ]
36
+ indices_list = [GraylogIndexItem(index_name=name, index_info=info) for name, info in indices_data.items()]
37
38
return GraylogIndicesResponse(
39
indices=indices_list,
@@ -62,10 +59,7 @@ async def fetch_configured_inputs() -> Tuple[bool, List[ConfiguredInput]]:
59
success = configured_inputs_collected.get("success", False)
60
61
if success:
65
- return True, [
66
- ConfiguredInput(**input_data)
67
- for input_data in configured_inputs_collected["data"]["inputs"]
68
- ]
62
+ return True, [ConfiguredInput(**input_data) for input_data in configured_inputs_collected["data"]["inputs"]]
63
else:
64
logger.error("Failed to fetch configured inputs")
65
return False, []
@@ -84,10 +78,7 @@ async def fetch_running_inputs() -> Tuple[bool, List[RunningInput]]:
78
success = running_inputs_collected.get("success", False)
79
80
if success:
87
- return True, [
88
- RunningInput(**input_data)
89
- for input_data in running_inputs_collected["data"]["states"]
90
- ]
81
+ return True, [RunningInput(**input_data) for input_data in running_inputs_collected["data"]["states"]]
82
else:
83
logger.error("Failed to fetch running inputs")
84
return False, []
backend/app/connectors/graylog/services/events.py
+14
-21
@@ -1,17 +1,17 @@
1
-from app.connectors.graylog.schema.events import (
2
- AlertEvent,
3
- AlertQuery,
4
- Alerts,
5
- Context,
6
- EventDefinition,
7
- GraylogAlertsResponse,
8
- GraylogEventDefinitionsResponse,
9
- Parameters,
10
-)
11
-from app.connectors.graylog.utils.universal import send_get_request, send_post_request
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.connectors.graylog.schema.events import AlertEvent
5
+from app.connectors.graylog.schema.events import AlertQuery
6
+from app.connectors.graylog.schema.events import Alerts
7
+from app.connectors.graylog.schema.events import Context
8
+from app.connectors.graylog.schema.events import EventDefinition
9
+from app.connectors.graylog.schema.events import GraylogAlertsResponse
10
+from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse
11
+from app.connectors.graylog.schema.events import Parameters
12
+from app.connectors.graylog.utils.universal import send_get_request
13
+from app.connectors.graylog.utils.universal import send_post_request
14
+
15
16
async def get_event_definitions() -> GraylogEventDefinitionsResponse:
17
"""Get event definitions from Graylog.
@@ -25,9 +25,7 @@ async def get_event_definitions() -> GraylogEventDefinitionsResponse:
25
)
26
if event_definitions_collected["success"]:
27
try:
28
- event_definitions_data = event_definitions_collected["data"][
29
- "event_definitions"
30
- ]
28
+ event_definitions_data = event_definitions_collected["data"]["event_definitions"]
29
except KeyError:
30
raise HTTPException(
31
status_code=500,
@@ -35,10 +33,7 @@ async def get_event_definitions() -> GraylogEventDefinitionsResponse:
33
)
34
35
# Convert the dictionary to a list of GraylogIndexItem
38
- event_definitions_list = [
39
- EventDefinition(**event_definition_data)
40
- for event_definition_data in event_definitions_data
41
- ]
36
+ event_definitions_list = [EventDefinition(**event_definition_data) for event_definition_data in event_definitions_data]
37
38
return GraylogEventDefinitionsResponse(
39
event_definitions=event_definitions_list,
@@ -78,9 +73,7 @@ async def get_alerts(alert_query: AlertQuery) -> GraylogAlertsResponse:
73
except KeyError:
74
raise HTTPException(status_code=500, detail="Failed to collect data key")
75
# Convert raw event data to Event objects
81
- event_objects = [
82
- AlertEvent(**event_data) for event_data in raw_alerts_data["events"]
83
- ]
76
+ event_objects = [AlertEvent(**event_data) for event_data in raw_alerts_data["events"]]
77
78
# Build the Alerts object
79
alerts = Alerts(
backend/app/connectors/graylog/services/management.py
+15
-18
@@ -1,23 +1,20 @@
1
-from app.connectors.graylog.schema.management import (
2
- DeletedIndexBody,
3
- DeletedIndexResponse,
4
- StartInputBody,
5
- StartInputResponse,
6
- StartStreamBody,
7
- StartStreamResponse,
8
- StopInputBody,
9
- StopInputResponse,
10
- StopStreamBody,
11
- StopStreamResponse,
12
-)
13
-from app.connectors.graylog.services.collector import get_index_names
14
-from app.connectors.graylog.utils.universal import (
15
- send_delete_request,
16
- send_post_request,
17
- send_put_request,
18
-)
1
from loguru import logger
2
3
+from app.connectors.graylog.schema.management import DeletedIndexBody
4
+from app.connectors.graylog.schema.management import DeletedIndexResponse
5
+from app.connectors.graylog.schema.management import StartInputBody
6
+from app.connectors.graylog.schema.management import StartInputResponse
7
+from app.connectors.graylog.schema.management import StartStreamBody
8
+from app.connectors.graylog.schema.management import StartStreamResponse
9
+from app.connectors.graylog.schema.management import StopInputBody
10
+from app.connectors.graylog.schema.management import StopInputResponse
11
+from app.connectors.graylog.schema.management import StopStreamBody
12
+from app.connectors.graylog.schema.management import StopStreamResponse
13
+from app.connectors.graylog.services.collector import get_index_names
14
+from app.connectors.graylog.utils.universal import send_delete_request
15
+from app.connectors.graylog.utils.universal import send_post_request
16
+from app.connectors.graylog.utils.universal import send_put_request
17
+
18
19
async def delete_index(index_name: DeletedIndexBody) -> DeletedIndexResponse:
20
"""Delete an index from Graylog.
backend/app/connectors/graylog/services/monitoring.py
+12
-21
@@ -1,16 +1,15 @@
1
-from app.connectors.graylog.schema.monitoring import (
2
- GraylogEventNotificationsResponse,
3
- GraylogMessages,
4
- GraylogMessagesResponse,
5
- GraylogMetricsResponse,
6
- GraylogThroughputMetrics,
7
- GraylogThroughputMetricsCollection,
8
- GraylogUncommittedJournalEntries,
9
-)
10
-from app.connectors.graylog.utils.universal import send_get_request
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.connectors.graylog.schema.monitoring import GraylogEventNotificationsResponse
5
+from app.connectors.graylog.schema.monitoring import GraylogMessages
6
+from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse
7
+from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse
8
+from app.connectors.graylog.schema.monitoring import GraylogThroughputMetrics
9
+from app.connectors.graylog.schema.monitoring import GraylogThroughputMetricsCollection
10
+from app.connectors.graylog.schema.monitoring import GraylogUncommittedJournalEntries
11
+from app.connectors.graylog.utils.universal import send_get_request
12
+
13
14
async def get_messages(page_number: int) -> GraylogMessagesResponse:
15
"""Get messages from Graylog.
@@ -109,10 +108,7 @@ def filter_and_create_throughput_metrics(merged_metrics: dict) -> list:
108
Returns:
109
list: A list of GraylogThroughputMetrics objects.
110
"""
112
- model_fields = [
113
- field_info.alias
114
- for field_info in GraylogThroughputMetricsCollection.__fields__.values()
115
- ]
111
+ model_fields = [field_info.alias for field_info in GraylogThroughputMetricsCollection.__fields__.values()]
112
throughput_metrics_list = [
113
GraylogThroughputMetrics(metric=metric_name, value=metric_data.get("value", 0))
114
for metric_name, metric_data in merged_metrics.items()
@@ -132,19 +128,14 @@ async def get_metrics() -> GraylogMetricsResponse:
128
throughput_metrics_collected = await fetch_metrics_from_graylog()
129
uncommitted_journal_entries_collected = await fetch_uncommitted_journal_entries()
130
try:
135
- if (
136
- throughput_metrics_collected["success"]
137
- and uncommitted_journal_entries_collected["success"]
138
- ):
131
+ if throughput_metrics_collected["success"] and uncommitted_journal_entries_collected["success"]:
132
merged_metrics = merge_metrics_data(throughput_metrics_collected)
133
throughput_metrics_list = filter_and_create_throughput_metrics(
134
merged_metrics,
135
)
136
137
uncommitted_journal_entries = GraylogUncommittedJournalEntries(
145
- uncommitted_journal_entries=uncommitted_journal_entries_collected[
146
- "data"
147
- ]["uncommitted_journal_entries"],
138
+ uncommitted_journal_entries=uncommitted_journal_entries_collected["data"]["uncommitted_journal_entries"],
139
)
140
141
return GraylogMetricsResponse(
backend/app/connectors/graylog/services/pipelines.py
+13
-21
@@ -1,19 +1,17 @@
1
-from app.connectors.graylog.schema.pipelines import (
2
- CreatePipeline,
3
- CreatePipelineRule,
4
- GraylogPipelinesResponse,
5
- Pipeline,
6
- PipelineRule,
7
- PipelineRulesResponse,
8
-)
9
-from app.connectors.graylog.utils.universal import send_get_request, send_post_request
10
-from app.customer_provisioning.schema.graylog import (
11
- StreamConnectionToPipelineRequest,
12
- StreamConnectionToPipelineResponse,
13
-)
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.connectors.graylog.schema.pipelines import CreatePipeline
5
+from app.connectors.graylog.schema.pipelines import CreatePipelineRule
6
+from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
7
+from app.connectors.graylog.schema.pipelines import Pipeline
8
+from app.connectors.graylog.schema.pipelines import PipelineRule
9
+from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
10
+from app.connectors.graylog.utils.universal import send_get_request
11
+from app.connectors.graylog.utils.universal import send_post_request
12
+from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest
13
+from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineResponse
14
+
15
16
async def get_pipelines() -> GraylogPipelinesResponse:
17
"""Get pipelines from Graylog.
@@ -30,10 +28,7 @@ async def get_pipelines() -> GraylogPipelinesResponse:
28
)
29
try:
30
if pipelines_collected["success"]:
33
- pipelines_list = [
34
- Pipeline(**pipeline_data)
35
- for pipeline_data in pipelines_collected["data"]
36
- ]
31
+ pipelines_list = [Pipeline(**pipeline_data) for pipeline_data in pipelines_collected["data"]]
32
return GraylogPipelinesResponse(
33
pipelines=pipelines_list,
34
success=True,
@@ -63,10 +58,7 @@ async def get_pipeline_rules() -> PipelineRulesResponse:
58
)
59
try:
60
if pipeline_rules_collected["success"]:
66
- pipeline_rules_list = [
67
- PipelineRule(**pipeline_rule_data)
68
- for pipeline_rule_data in pipeline_rules_collected["data"]
69
- ]
61
+ pipeline_rules_list = [PipelineRule(**pipeline_rule_data) for pipeline_rule_data in pipeline_rules_collected["data"]]
62
return PipelineRulesResponse(
63
pipeline_rules=pipeline_rules_list,
64
success=True,
backend/app/connectors/graylog/services/streams.py
+6
-10
@@ -1,10 +1,12 @@
1
from typing import List
2
3
-from app.connectors.graylog.schema.streams import GraylogStreamsResponse, Stream
4
-from app.connectors.graylog.utils.universal import send_get_request
3
from fastapi import HTTPException
4
from loguru import logger
5
6
+from app.connectors.graylog.schema.streams import GraylogStreamsResponse
7
+from app.connectors.graylog.schema.streams import Stream
8
+from app.connectors.graylog.utils.universal import send_get_request
9
+
10
11
async def get_streams() -> GraylogStreamsResponse:
12
"""Get streams from Graylog.
@@ -19,10 +21,7 @@ async def get_streams() -> GraylogStreamsResponse:
21
streams_collected = await send_get_request(endpoint="/api/streams")
22
try:
23
if streams_collected["success"]:
22
- streams_list = [
23
- Stream(**stream_data)
24
- for stream_data in streams_collected["data"]["streams"]
25
- ]
24
+ streams_list = [Stream(**stream_data) for stream_data in streams_collected["data"]["streams"]]
25
return GraylogStreamsResponse(
26
streams=streams_list,
27
success=True,
@@ -60,10 +59,7 @@ async def get_stream_ids() -> List[str]:
59
streams_collected = await send_get_request(endpoint="/api/streams")
60
try:
61
if streams_collected["success"]:
63
- return [
64
- stream_data["id"]
65
- for stream_data in streams_collected["data"]["streams"]
66
- ]
62
+ return [stream_data["id"] for stream_data in streams_collected["data"]["streams"]]
63
else:
64
return []
65
except KeyError as e:
backend/app/connectors/graylog/utils/universal.py
+6
-3
@@ -1,11 +1,14 @@
1
-from typing import Any, Dict, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import Optional
4
5
import requests
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.db.db_session import get_db_session
6
from fastapi import HTTPException
7
from loguru import logger
8
9
+from app.connectors.utils import get_connector_info_from_db
10
+from app.db.db_session import get_db_session
11
+
12
HEADERS = {"X-Requested-By": "CoPilot"}
13
14
backend/app/connectors/influxdb/routes/alerts.py
+4
-2
@@ -1,8 +1,10 @@
1
+from fastapi import APIRouter
2
+from fastapi import Security
3
+from loguru import logger
4
+
5
from app.auth.utils import AuthHandler
6
from app.connectors.influxdb.schema.alerts import InfluxDBAlertsResponse
7
from app.connectors.influxdb.services.alerts import get_alerts
4
-from fastapi import APIRouter, Security
5
-from loguru import logger
8
9
# App specific imports
10
backend/app/connectors/influxdb/services/alerts.py
+6
-8
@@ -1,13 +1,13 @@
1
from typing import List
2
3
-from app.connectors.influxdb.schema.alerts import InfluxDBAlert, InfluxDBAlertsResponse
4
-from app.connectors.influxdb.utils.universal import (
5
- create_influxdb_client,
6
- get_influxdb_organization,
7
-)
3
from fastapi import HTTPException
4
from loguru import logger
5
6
+from app.connectors.influxdb.schema.alerts import InfluxDBAlert
7
+from app.connectors.influxdb.schema.alerts import InfluxDBAlertsResponse
8
+from app.connectors.influxdb.utils.universal import create_influxdb_client
9
+from app.connectors.influxdb.utils.universal import get_influxdb_organization
10
+
11
# Constants
12
BUCKET_NAME = "_monitoring"
13
@@ -46,9 +46,7 @@ async def process_alert_records(result) -> List[InfluxDBAlert]:
46
for table in result:
47
for record in table.records:
48
alert = InfluxDBAlert(
49
- time=record.values.get("time").isoformat()
50
- if record.values.get("time")
51
- else None,
49
+ time=record.values.get("time").isoformat() if record.values.get("time") else None,
50
message=record.values.get("message"),
51
checkID=record.values.get("checkID"),
52
checkName=record.values.get("checkName"),
backend/app/connectors/influxdb/utils/universal.py
+5
-3
@@ -1,11 +1,13 @@
1
-from typing import Any, Dict
1
+from typing import Any
2
+from typing import Dict
3
3
-from app.connectors.utils import get_connector_info_from_db
4
-from app.db.db_session import get_db_session
4
from fastapi import HTTPException
5
from influxdb_client.client.influxdb_client_async import InfluxDBClientAsync
6
from loguru import logger
7
8
+from app.connectors.utils import get_connector_info_from_db
9
+from app.db.db_session import get_db_session
10
+
11
12
async def verify_influxdb_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
13
"""
backend/app/connectors/models.py
+5
-2
@@ -1,7 +1,10 @@
1
from datetime import datetime
2
-from typing import List, Optional
2
+from typing import List
3
+from typing import Optional
4
4
-from sqlmodel import Field, Relationship, SQLModel
5
+from sqlmodel import Field
6
+from sqlmodel import Relationship
7
+from sqlmodel import SQLModel
8
9
10
class ConnectorHistory(SQLModel, table=True):
backend/app/connectors/routes.py
+15
-12
@@ -1,21 +1,24 @@
1
from typing import Union
2
3
-from app.auth.utils import AuthHandler
4
-from app.connectors.schema import (
5
- ConnectorListResponse,
6
- ConnectorResponse,
7
- ConnectorsListResponse,
8
- UpdateConnector,
9
- VerifyConnectorResponse,
10
-)
11
-from app.connectors.services import ConnectorServices
12
-from app.db.db_session import get_db
13
-
3
## Auth Things
15
-from fastapi import APIRouter, Depends, File, HTTPException, Security, UploadFile
4
+from fastapi import APIRouter
5
+from fastapi import Depends
6
+from fastapi import File
7
+from fastapi import HTTPException
8
+from fastapi import Security
9
+from fastapi import UploadFile
10
from loguru import logger
11
from sqlalchemy.ext.asyncio import AsyncSession
12
13
+from app.auth.utils import AuthHandler
14
+from app.connectors.schema import ConnectorListResponse
15
+from app.connectors.schema import ConnectorResponse
16
+from app.connectors.schema import ConnectorsListResponse
17
+from app.connectors.schema import UpdateConnector
18
+from app.connectors.schema import VerifyConnectorResponse
19
+from app.connectors.services import ConnectorServices
20
+from app.db.db_session import get_db
21
+
22
connector_router = APIRouter()
23
24
backend/app/connectors/schema.py
+2
-1
@@ -1,5 +1,6 @@
1
from datetime import datetime
2
-from typing import List, Optional
2
+from typing import List
3
+from typing import Optional
4
5
from pydantic import BaseModel
6
backend/app/connectors/services.py
+14
-14
@@ -1,8 +1,18 @@
1
import os
2
from datetime import datetime
3
-from typing import List, Optional, Type, Union
3
+from typing import List
4
+from typing import Optional
5
+from typing import Type
6
+from typing import Union
7
8
import aiofiles
9
+from fastapi import UploadFile
10
+from loguru import logger
11
+from pydantic import BaseModel
12
+from sqlalchemy.ext.asyncio import AsyncSession
13
+from sqlalchemy.future import select
14
+from werkzeug.utils import secure_filename
15
+
16
from app.connectors.cortex.utils.universal import verify_cortex_connection
17
from app.connectors.dfir_iris.utils.universal import verify_dfir_iris_connection
18
from app.connectors.grafana.utils.universal import verify_grafana_connection
@@ -22,16 +32,8 @@ from app.integrations.utils.event_shipper import verify_event_shipper_connection
32
from app.threat_intel.services.socfortress import (
33
verifiy_socfortress_threat_intel_connector,
34
)
25
-from app.utils import (
26
- verify_alert_creation_provisioning_connection,
27
- verify_wazuh_worker_provisioning_connection,
28
-)
29
-from fastapi import UploadFile
30
-from loguru import logger
31
-from pydantic import BaseModel
32
-from sqlalchemy.ext.asyncio import AsyncSession
33
-from sqlalchemy.future import select
34
-from werkzeug.utils import secure_filename
35
+from app.utils import verify_alert_creation_provisioning_connection
36
+from app.utils import verify_wazuh_worker_provisioning_connection
37
38
UPLOAD_FOLDER = "file-store"
39
UPLOAD_FOLDER = os.path.join(
@@ -396,9 +398,7 @@ class ConnectorServices:
398
Returns:
399
bool: True if the file is allowed, False otherwise.
400
"""
399
- return (
400
- "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
401
- )
401
+ return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
402
403
@classmethod
404
async def save_file(
backend/app/connectors/shuffle/routes/workflows.py
+10
-11
@@ -1,16 +1,15 @@
1
-from app.auth.utils import AuthHandler
2
-from app.connectors.shuffle.schema.workflows import (
3
- WorkflowExecutionBodyModel,
4
- WorkflowExecutionResponseModel,
5
- WorkflowsResponse,
6
-)
7
-from app.connectors.shuffle.services.workflows import (
8
- get_workflow_executions,
9
- get_workflows,
10
-)
11
-from fastapi import APIRouter, HTTPException, Security
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from fastapi import Security
4
from loguru import logger
5
6
+from app.auth.utils import AuthHandler
7
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel
8
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionResponseModel
9
+from app.connectors.shuffle.schema.workflows import WorkflowsResponse
10
+from app.connectors.shuffle.services.workflows import get_workflow_executions
11
+from app.connectors.shuffle.services.workflows import get_workflows
12
+
13
shuffle_workflows_router = APIRouter()
14
15
backend/app/connectors/shuffle/schema/workflows.py
+6
-2
@@ -1,6 +1,10 @@
1
-from typing import Any, Dict, List, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+from typing import Optional
5
3
-from pydantic import BaseModel, Field
6
+from pydantic import BaseModel
7
+from pydantic import Field
8
9
10
class WorkflowsResponse(BaseModel):
backend/app/connectors/shuffle/services/workflows.py
+5
-6
@@ -1,14 +1,13 @@
1
from typing import List
2
3
-from app.connectors.shuffle.schema.workflows import (
4
- WorkflowExecutionBodyModel,
5
- WorkflowExecutionStatusResponseModel,
6
- WorkflowsResponse,
7
-)
8
-from app.connectors.shuffle.utils.universal import send_get_request
3
from fastapi import HTTPException
4
from loguru import logger
5
6
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel
7
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionStatusResponseModel
8
+from app.connectors.shuffle.schema.workflows import WorkflowsResponse
9
+from app.connectors.shuffle.utils.universal import send_get_request
10
+
11
12
def remove_large_images_from_actions(workflows: List) -> List:
13
"""
backend/app/connectors/shuffle/utils/universal.py
+7
-6
@@ -1,11 +1,14 @@
1
-from typing import Any, Dict, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import Optional
4
5
import requests
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.db.db_session import get_db_session
6
from fastapi import HTTPException
7
from loguru import logger
8
9
+from app.connectors.utils import get_connector_info_from_db
10
+from app.db.db_session import get_db_session
11
+
12
13
async def verify_shuffle_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
14
"""
@@ -164,9 +167,7 @@ def send_post_request(
167
return {
168
"data": response.json(),
169
"success": False if response.status_code >= 400 else True,
167
- "message": "Successfully retrieved data"
168
- if response.status_code < 400
169
- else "Failed to retrieve data",
170
+ "message": "Successfully retrieved data" if response.status_code < 400 else "Failed to retrieve data",
171
}
172
except Exception as e:
173
logger.debug(f"Response: {response}")
backend/app/connectors/sublime/models/alerts.py
+5
-2
@@ -1,7 +1,10 @@
1
import datetime
2
-from typing import List, Optional
2
+from typing import List
3
+from typing import Optional
4
4
-from sqlmodel import Field, Relationship, SQLModel
5
+from sqlmodel import Field
6
+from sqlmodel import Relationship
7
+from sqlmodel import SQLModel
8
9
10
class FlaggedRule(SQLModel, table=True):
backend/app/connectors/sublime/routes/alerts.py
+11
-9
@@ -1,15 +1,17 @@
1
-from app.auth.utils import AuthHandler
2
-from app.connectors.sublime.schema.alerts import (
3
- AlertRequestBody,
4
- AlertResponseBody,
5
- SublimeAlertsResponse,
6
-)
7
-from app.connectors.sublime.services.alerts import collect_alerts, store_sublime_alert
8
-from app.db.db_session import get_db
9
-from fastapi import APIRouter, Depends, Security
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import Security
4
from loguru import logger
5
from sqlalchemy.ext.asyncio import AsyncSession
6
7
+from app.auth.utils import AuthHandler
8
+from app.connectors.sublime.schema.alerts import AlertRequestBody
9
+from app.connectors.sublime.schema.alerts import AlertResponseBody
10
+from app.connectors.sublime.schema.alerts import SublimeAlertsResponse
11
+from app.connectors.sublime.services.alerts import collect_alerts
12
+from app.connectors.sublime.services.alerts import store_sublime_alert
13
+from app.db.db_session import get_db
14
+
15
sublime_alerts_router = APIRouter()
16
17
backend/app/connectors/sublime/schema/alerts.py
+4
-2
@@ -1,7 +1,9 @@
1
import datetime
2
-from typing import List, Optional
2
+from typing import List
3
+from typing import Optional
4
4
-from pydantic import BaseModel, Field
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
8
9
class FlaggedRule(BaseModel):
backend/app/connectors/sublime/services/alerts.py
+12
-15
@@ -1,27 +1,24 @@
1
import json
2
from typing import List
3
4
-from app.connectors.sublime.models.alerts import (
5
- FlaggedRule,
6
- Mailbox,
7
- Recipient,
8
- Sender,
9
- SublimeAlerts,
10
- TriggeredAction,
11
-)
12
-from app.connectors.sublime.schema.alerts import (
13
- AlertRequestBody,
14
- AlertResponseBody,
15
- SublimeAlertsResponse,
16
- SublimeAlertsSchema,
17
-)
18
-from app.connectors.sublime.utils.universal import send_get_request
4
from fastapi import HTTPException
5
from loguru import logger
6
from sqlalchemy.ext.asyncio import AsyncSession
7
from sqlalchemy.future import select
8
from sqlalchemy.orm import selectinload
9
10
+from app.connectors.sublime.models.alerts import FlaggedRule
11
+from app.connectors.sublime.models.alerts import Mailbox
12
+from app.connectors.sublime.models.alerts import Recipient
13
+from app.connectors.sublime.models.alerts import Sender
14
+from app.connectors.sublime.models.alerts import SublimeAlerts
15
+from app.connectors.sublime.models.alerts import TriggeredAction
16
+from app.connectors.sublime.schema.alerts import AlertRequestBody
17
+from app.connectors.sublime.schema.alerts import AlertResponseBody
18
+from app.connectors.sublime.schema.alerts import SublimeAlertsResponse
19
+from app.connectors.sublime.schema.alerts import SublimeAlertsSchema
20
+from app.connectors.sublime.utils.universal import send_get_request
21
+
22
23
def create_sublime_alert(alert_request_body: AlertRequestBody) -> SublimeAlerts:
24
"""
backend/app/connectors/sublime/utils/universal.py
+5
-2
@@ -1,9 +1,12 @@
1
-from typing import Any, Dict, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import Optional
4
5
import requests
6
+from loguru import logger
7
+
8
from app.connectors.utils import get_connector_info_from_db
9
from app.db.db_session import get_db_session
6
-from loguru import logger
10
11
12
async def verify_sublime_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
backend/app/connectors/utils.py
+6
-3
@@ -1,11 +1,14 @@
1
-from typing import Any, Dict, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import Optional
4
3
-from app.connectors.models import Connectors
4
-from app.connectors.schema import ConnectorResponse
5
from loguru import logger
6
from sqlalchemy.ext.asyncio import AsyncSession
7
from sqlalchemy.future import select
8
9
+from app.connectors.models import Connectors
10
+from app.connectors.schema import ConnectorResponse
11
+
12
13
# ! New with Async
14
async def get_connector_info_from_db(
backend/app/connectors/velociraptor/routes/artifacts.py
+23
-27
@@ -1,30 +1,30 @@
1
from typing import List
2
3
-from app.auth.utils import AuthHandler
4
-from app.connectors.velociraptor.schema.artifacts import (
5
- ArtifactsResponse,
6
- CollectArtifactBody,
7
- CollectArtifactResponse,
8
- OSPrefixEnum,
9
- OSPrefixModel,
10
- QuarantineBody,
11
- QuarantineResponse,
12
- RunCommandBody,
13
- RunCommandResponse,
14
-)
15
-from app.connectors.velociraptor.services.artifacts import (
16
- get_artifacts,
17
- quarantine_host,
18
- run_artifact_collection,
19
- run_remote_command,
20
-)
21
-from app.db.db_session import get_db
22
-from app.db.universal_models import Agents
23
-from fastapi import APIRouter, Depends, HTTPException, Security
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
from loguru import logger
8
from sqlalchemy.ext.asyncio import AsyncSession
9
from sqlalchemy.future import select
10
11
+from app.auth.utils import AuthHandler
12
+from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
13
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
14
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
15
+from app.connectors.velociraptor.schema.artifacts import OSPrefixEnum
16
+from app.connectors.velociraptor.schema.artifacts import OSPrefixModel
17
+from app.connectors.velociraptor.schema.artifacts import QuarantineBody
18
+from app.connectors.velociraptor.schema.artifacts import QuarantineResponse
19
+from app.connectors.velociraptor.schema.artifacts import RunCommandBody
20
+from app.connectors.velociraptor.schema.artifacts import RunCommandResponse
21
+from app.connectors.velociraptor.services.artifacts import get_artifacts
22
+from app.connectors.velociraptor.services.artifacts import quarantine_host
23
+from app.connectors.velociraptor.services.artifacts import run_artifact_collection
24
+from app.connectors.velociraptor.services.artifacts import run_remote_command
25
+from app.db.db_session import get_db
26
+from app.db.universal_models import Agents
27
+
28
# App specific imports
29
30
@@ -68,9 +68,7 @@ def verify_os_prefix_exists(os_prefix: str) -> str:
68
detail=f"OS prefix {os_prefix} does not exist.",
69
)
70
71
- return OSPrefixEnum[
72
- os_prefix_upper
73
- ].value # Use the uppercase version for Enum matching
71
+ return OSPrefixEnum[os_prefix_upper].value # Use the uppercase version for Enum matching
72
73
74
def get_os_prefix_from_os_name(os_name: str) -> str:
@@ -223,9 +221,7 @@ async def get_all_artifacts_for_os_prefix(
221
# Get all the artifacts names that begin with the OS prefix
222
artifacts = await get_artifacts()
223
artifacts = artifacts.artifacts
226
- artifacts_for_os_prefix = [
227
- artifact for artifact in artifacts if artifact.name.startswith(os_prefix)
228
- ]
224
+ artifacts_for_os_prefix = [artifact for artifact in artifacts if artifact.name.startswith(os_prefix)]
225
return ArtifactsResponse(
226
success=True,
227
message=f"All artifacts for OS prefix {os_prefix} retrieved",
backend/app/connectors/velociraptor/routes/flows.py
+12
-6
@@ -1,13 +1,19 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
+from sqlalchemy.future import select
8
+
9
from app.auth.utils import AuthHandler
10
from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
3
-from app.connectors.velociraptor.schema.flows import FlowResponse, RetrieveFlowRequest
4
-from app.connectors.velociraptor.services.flows import get_flow, get_flows
11
+from app.connectors.velociraptor.schema.flows import FlowResponse
12
+from app.connectors.velociraptor.schema.flows import RetrieveFlowRequest
13
+from app.connectors.velociraptor.services.flows import get_flow
14
+from app.connectors.velociraptor.services.flows import get_flows
15
from app.db.db_session import get_db
16
from app.db.universal_models import Agents
7
-from fastapi import APIRouter, Depends, HTTPException, Security
8
-from loguru import logger
9
-from sqlalchemy.ext.asyncio import AsyncSession
10
-from sqlalchemy.future import select
17
18
velociraptor_flows_router = APIRouter()
19
backend/app/connectors/velociraptor/schema/artifacts.py
+6
-2
@@ -1,7 +1,11 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
4
-from pydantic import BaseModel, Field
7
+from pydantic import BaseModel
8
+from pydantic import Field
9
10
11
class Artifacts(BaseModel):
backend/app/connectors/velociraptor/schema/flows.py
+5
-2
@@ -1,8 +1,11 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
4
from fastapi import HTTPException
5
from loguru import logger
5
-from pydantic import BaseModel, Field, root_validator
6
+from pydantic import BaseModel
7
+from pydantic import Field
8
+from pydantic import root_validator
9
10
11
class FlowSpecParameter(BaseModel):
backend/app/connectors/velociraptor/services/artifacts.py
+12
-18
@@ -1,17 +1,16 @@
1
-from app.connectors.velociraptor.schema.artifacts import (
2
- Artifacts,
3
- ArtifactsResponse,
4
- CollectArtifactBody,
5
- CollectArtifactResponse,
6
- QuarantineBody,
7
- QuarantineResponse,
8
- RunCommandBody,
9
- RunCommandResponse,
10
-)
11
-from app.connectors.velociraptor.utils.universal import UniversalService
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.connectors.velociraptor.schema.artifacts import Artifacts
5
+from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
6
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
7
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
8
+from app.connectors.velociraptor.schema.artifacts import QuarantineBody
9
+from app.connectors.velociraptor.schema.artifacts import QuarantineResponse
10
+from app.connectors.velociraptor.schema.artifacts import RunCommandBody
11
+from app.connectors.velociraptor.schema.artifacts import RunCommandResponse
12
+from app.connectors.velociraptor.utils.universal import UniversalService
13
+
14
15
def create_query(query: str) -> str:
16
"""
@@ -61,10 +60,7 @@ def get_artifact_key(analyzer_body: CollectArtifactBody) -> str:
60
f"env=dict(Command='{analyzer_body.command}'))"
61
)
62
else:
64
- return (
65
- f"collect_client(client_id='{analyzer_body.velociraptor_id}', "
66
- f"artifacts=['{analyzer_body.artifact_name}'])"
67
- )
63
+ return f"collect_client(client_id='{analyzer_body.velociraptor_id}', " f"artifacts=['{analyzer_body.artifact_name}'])"
64
65
66
async def get_artifacts() -> ArtifactsResponse:
@@ -140,9 +136,7 @@ async def run_artifact_collection(
136
message=results["message"],
137
results=results["results"],
138
)
143
- except (
144
- HTTPException
145
- ) as he: # Catch HTTPException separately to propagate the original message
139
+ except HTTPException as he: # Catch HTTPException separately to propagate the original message
140
logger.error(
141
f"HTTPException while running artifact collection on {collect_artifact_body}: {he.detail}",
142
)
backend/app/connectors/velociraptor/services/flows.py
+6
-7
@@ -1,13 +1,12 @@
1
-from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
2
-from app.connectors.velociraptor.schema.flows import (
3
- FlowClientSession,
4
- FlowResponse,
5
- RetrieveFlowRequest,
6
-)
7
-from app.connectors.velociraptor.utils.universal import UniversalService
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
5
+from app.connectors.velociraptor.schema.flows import FlowClientSession
6
+from app.connectors.velociraptor.schema.flows import FlowResponse
7
+from app.connectors.velociraptor.schema.flows import RetrieveFlowRequest
8
+from app.connectors.velociraptor.utils.universal import UniversalService
9
+
10
11
def create_query(query: str) -> str:
12
"""
backend/app/connectors/velociraptor/utils/universal.py
+11
-13
@@ -1,14 +1,18 @@
1
import json
2
from datetime import datetime
3
-from typing import Any, Dict
3
+from typing import Any
4
+from typing import Dict
5
6
import grpc
7
import pyvelociraptor
7
-from app.connectors.utils import get_connector_info_from_db
8
-from app.db.db_session import AsyncSessionLocal, get_db_session
8
from fastapi import HTTPException
9
from loguru import logger
11
-from pyvelociraptor import api_pb2, api_pb2_grpc
10
+from pyvelociraptor import api_pb2
11
+from pyvelociraptor import api_pb2_grpc
12
+
13
+from app.connectors.utils import get_connector_info_from_db
14
+from app.db.db_session import AsyncSessionLocal
15
+from app.db.db_session import get_db_session
16
17
18
async def verify_velociraptor_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
@@ -244,12 +248,8 @@ class UniversalService:
248
"""
249
# Formulate queries
250
try:
247
- vql_client_id = (
248
- f"select client_id,os_info from clients(search='host:{client_name}')"
249
- )
250
- vql_last_seen_at = (
251
- f"select last_seen_at from clients(search='host:{client_name}')"
252
- )
251
+ vql_client_id = f"select client_id,os_info from clients(search='host:{client_name}')"
252
+ vql_last_seen_at = f"select last_seen_at from clients(search='host:{client_name}')"
253
254
# Get the last seen timestamp
255
logger.info(f"Getting last seen at timestamp for {client_name}")
@@ -322,6 +322,4 @@ class UniversalService:
322
Returns:
323
bool: True if the client is offline, False otherwise.
324
"""
325
- return (
326
- datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000)
327
- ).total_seconds() > 30
325
+ return (datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000)).total_seconds() > 30
backend/app/connectors/wazuh_indexer/routes/alerts.py
+21
-21
@@ -1,28 +1,28 @@
1
from typing import List
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from loguru import logger
8
+
9
from app.auth.utils import AuthHandler
4
-from app.connectors.wazuh_indexer.schema.alerts import (
5
- AlertsByHostResponse,
6
- AlertsByRulePerHostResponse,
7
- AlertsByRuleResponse,
8
- AlertsSearchBody,
9
- AlertsSearchResponse,
10
- HostAlertsSearchBody,
11
- HostAlertsSearchResponse,
12
- IndexAlertsSearchBody,
13
- IndexAlertsSearchResponse,
14
-)
15
-from app.connectors.wazuh_indexer.services.alerts import (
16
- get_alerts,
17
- get_alerts_by_host,
18
- get_alerts_by_rule,
19
- get_alerts_by_rule_per_host,
20
- get_host_alerts,
21
- get_index_alerts,
22
-)
10
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
11
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHostResponse
12
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRuleResponse
13
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchBody
14
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchResponse
15
+from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchBody
16
+from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchResponse
17
+from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchBody
18
+from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchResponse
19
+from app.connectors.wazuh_indexer.services.alerts import get_alerts
20
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_host
21
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule
22
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule_per_host
23
+from app.connectors.wazuh_indexer.services.alerts import get_host_alerts
24
+from app.connectors.wazuh_indexer.services.alerts import get_index_alerts
25
from app.connectors.wazuh_indexer.utils.universal import collect_indices
24
-from fastapi import APIRouter, Depends, HTTPException, Security
25
-from loguru import logger
26
27
# App specific imports
28
backend/app/connectors/wazuh_indexer/routes/monitoring.py
+12
-13
@@ -1,21 +1,20 @@
1
from typing import Union
2
3
+from fastapi import APIRouter
4
+from fastapi import HTTPException
5
+from fastapi import Security
6
+
7
from app.auth.utils import AuthHandler
4
-from app.connectors.wazuh_indexer.schema.monitoring import (
5
- ClusterHealthResponse,
6
- IndicesStatsResponse,
7
- NodeAllocationResponse,
8
- ShardsResponse,
9
-)
8
+from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse
9
+from app.connectors.wazuh_indexer.schema.monitoring import IndicesStatsResponse
10
+from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocationResponse
11
+from app.connectors.wazuh_indexer.schema.monitoring import ShardsResponse
12
13
# from app.connectors.wazuh_indexer.schema import WazuhIndexerResponse, WazuhIndexerListResponse
12
-from app.connectors.wazuh_indexer.services.monitoring import (
13
- cluster_healthcheck,
14
- indices_stats,
15
- node_allocation,
16
- shards,
17
-)
18
-from fastapi import APIRouter, HTTPException, Security
14
+from app.connectors.wazuh_indexer.services.monitoring import cluster_healthcheck
15
+from app.connectors.wazuh_indexer.services.monitoring import indices_stats
16
+from app.connectors.wazuh_indexer.services.monitoring import node_allocation
17
+from app.connectors.wazuh_indexer.services.monitoring import shards
18
19
wazuh_indexer_router = APIRouter()
20
backend/app/connectors/wazuh_indexer/schema/alerts.py
+8
-5
@@ -1,7 +1,12 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
4
-from pydantic import BaseModel, Field, validator
7
+from pydantic import BaseModel
8
+from pydantic import Field
9
+from pydantic import validator
10
11
12
class Alert(BaseModel):
@@ -116,8 +121,6 @@ class AlertsByRulePerHostResponse(BaseModel):
121
122
############# ! PASSABLE MESSAGES FROM ES CLIENT ! #############
123
class SkippableWazuhIndexerClientErrors(Enum):
119
- NO_MAPPING_FOR_TIMESTAMP = (
120
- "No mapping found for [timestamp_utc] in order to sort on"
121
- )
124
+ NO_MAPPING_FOR_TIMESTAMP = "No mapping found for [timestamp_utc] in order to sort on"
125
# Add other error messages here, for example:
126
# ANOTHER_ERROR = "Another specific error message"
backend/app/connectors/wazuh_indexer/schema/indices.py
+2
-1
@@ -1,6 +1,7 @@
1
from typing import Dict
2
3
-from pydantic import BaseModel, Field
3
+from pydantic import BaseModel
4
+from pydantic import Field
5
6
7
class Indices(BaseModel):
backend/app/connectors/wazuh_indexer/schema/monitoring.py
+5
-2
@@ -1,6 +1,9 @@
1
-from typing import List, Optional, Union
1
+from typing import List
2
+from typing import Optional
3
+from typing import Union
4
3
-from pydantic import BaseModel, Field
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
8
9
class ClusterHealth(BaseModel):
backend/app/connectors/wazuh_indexer/services/alerts.py
+28
-38
@@ -1,29 +1,29 @@
1
-from typing import Dict, List, Optional, Type
2
-
3
-from app.connectors.wazuh_indexer.schema.alerts import (
4
- AlertsByHost,
5
- AlertsByHostResponse,
6
- AlertsByRule,
7
- AlertsByRulePerHost,
8
- AlertsByRulePerHostResponse,
9
- AlertsByRuleResponse,
10
- AlertsSearchBody,
11
- AlertsSearchResponse,
12
- CollectAlertsResponse,
13
- HostAlertsSearchBody,
14
- HostAlertsSearchResponse,
15
- IndexAlertsSearchBody,
16
- IndexAlertsSearchResponse,
17
- SkippableWazuhIndexerClientErrors,
18
-)
19
-from app.connectors.wazuh_indexer.utils.universal import (
20
- AlertsQueryBuilder,
21
- collect_indices,
22
- create_wazuh_indexer_client,
23
-)
1
+from typing import Dict
2
+from typing import List
3
+from typing import Optional
4
+from typing import Type
5
+
6
from fastapi import HTTPException
7
from loguru import logger
8
9
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByHost
10
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
11
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRule
12
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHost
13
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHostResponse
14
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRuleResponse
15
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchBody
16
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchResponse
17
+from app.connectors.wazuh_indexer.schema.alerts import CollectAlertsResponse
18
+from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchBody
19
+from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchResponse
20
+from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchBody
21
+from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchResponse
22
+from app.connectors.wazuh_indexer.schema.alerts import SkippableWazuhIndexerClientErrors
23
+from app.connectors.wazuh_indexer.utils.universal import AlertsQueryBuilder
24
+from app.connectors.wazuh_indexer.utils.universal import collect_indices
25
+from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
26
+
27
28
async def collect_and_aggregate_alerts(
29
field_names: List[str],
@@ -48,17 +48,11 @@ async def collect_and_aggregate_alerts(
48
alerts_response = await collect_alerts_generic(index_name, body=search_body)
49
if alerts_response.success:
50
for alert in alerts_response.alerts:
51
- composite_key = tuple(
52
- alert["_source"][field] for field in field_names
53
- )
54
- aggregated_alerts_dict[composite_key] = (
55
- aggregated_alerts_dict.get(composite_key, 0) + 1
56
- )
51
+ composite_key = tuple(alert["_source"][field] for field in field_names)
52
+ aggregated_alerts_dict[composite_key] = aggregated_alerts_dict.get(composite_key, 0) + 1
53
except HTTPException as e:
54
detail_str = str(e.detail) # Convert to string to make sure it's comparable
59
- if any(
60
- err.value in detail_str for err in SkippableWazuhIndexerClientErrors
61
- ):
55
+ if any(err.value in detail_str for err in SkippableWazuhIndexerClientErrors):
56
logger.warning(
57
f"Skipping index {index_name} due to specific error: {e.detail}",
58
)
@@ -149,9 +143,7 @@ async def get_alerts_generic(
143
)
144
alerts_summary = []
145
indices = await collect_indices()
152
- index_list = (
153
- [index_name] if index_name else indices.indices_list
154
- ) # Use the provided index_name or get all indices
146
+ index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices
147
148
for index_name in index_list:
149
try:
@@ -170,9 +162,7 @@ async def get_alerts_generic(
162
)
163
except HTTPException as e:
164
detail_str = str(e.detail) # Convert to string to make sure it's comparable
173
- if any(
174
- err.value in detail_str for err in SkippableWazuhIndexerClientErrors
175
- ):
165
+ if any(err.value in detail_str for err in SkippableWazuhIndexerClientErrors):
166
logger.warning(
167
f"Skipping index {index_name} due to specific error: {e.detail}",
168
)
backend/app/connectors/wazuh_indexer/services/monitoring.py
+18
-24
@@ -1,23 +1,21 @@
1
-from typing import Dict, Union
2
-
3
-from app.connectors.wazuh_indexer.schema.monitoring import (
4
- ClusterHealth,
5
- ClusterHealthResponse,
6
- IndicesStats,
7
- IndicesStatsResponse,
8
- NodeAllocation,
9
- NodeAllocationResponse,
10
- Shards,
11
- ShardsResponse,
12
-)
13
-from app.connectors.wazuh_indexer.utils.universal import (
14
- create_wazuh_indexer_client,
15
- format_indices_stats,
16
- format_node_allocation,
17
- format_shards,
18
-)
1
+from typing import Dict
2
+from typing import Union
3
+
4
from loguru import logger
5
6
+from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealth
7
+from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse
8
+from app.connectors.wazuh_indexer.schema.monitoring import IndicesStats
9
+from app.connectors.wazuh_indexer.schema.monitoring import IndicesStatsResponse
10
+from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocation
11
+from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocationResponse
12
+from app.connectors.wazuh_indexer.schema.monitoring import Shards
13
+from app.connectors.wazuh_indexer.schema.monitoring import ShardsResponse
14
+from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
15
+from app.connectors.wazuh_indexer.utils.universal import format_indices_stats
16
+from app.connectors.wazuh_indexer.utils.universal import format_node_allocation
17
+from app.connectors.wazuh_indexer.utils.universal import format_shards
18
+
19
20
async def cluster_healthcheck() -> Union[ClusterHealthResponse, Dict[str, str]]:
21
"""
@@ -64,9 +62,7 @@ async def node_allocation() -> Union[NodeAllocationResponse, Dict[str, bool]]:
62
raw_node_allocation_data,
63
)
64
67
- node_allocation_models = [
68
- NodeAllocation(**node) for node in formatted_node_allocation_data
69
- ]
65
+ node_allocation_models = [NodeAllocation(**node) for node in formatted_node_allocation_data]
66
67
return NodeAllocationResponse(
68
node_allocation=node_allocation_models,
@@ -97,9 +93,7 @@ async def indices_stats() -> Union[IndicesStatsResponse, Dict[str, str]]:
93
raw_indices_stats_data,
94
)
95
100
- indices_stats_models = [
101
- IndicesStats(**index) for index in formatted_indices_stats_data
102
- ]
96
+ indices_stats_models = [IndicesStats(**index) for index in formatted_indices_stats_data]
97
98
return IndicesStatsResponse(
99
indices_stats=indices_stats_models,
backend/app/connectors/wazuh_indexer/utils/universal.py
+14
-14
@@ -1,13 +1,19 @@
1
-from datetime import datetime, timedelta
2
-from typing import Any, Dict, Iterable, Tuple
1
+from datetime import datetime
2
+from datetime import timedelta
3
+from typing import Any
4
+from typing import Dict
5
+from typing import Iterable
6
+from typing import Tuple
7
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel, Indices
6
-from app.db.db_session import get_db_session
8
from elasticsearch7 import Elasticsearch
9
from fastapi import HTTPException
10
from loguru import logger
11
12
+from app.connectors.utils import get_connector_info_from_db
13
+from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
14
+from app.connectors.wazuh_indexer.schema.indices import Indices
15
+from app.db.db_session import get_db_session
16
+
17
18
async def verify_wazuh_indexer_credentials(
19
attributes: Dict[str, Any],
@@ -189,9 +195,7 @@ async def collect_indices() -> Indices:
195
indices_list = list(indices_dict.keys())
196
# Check if the index is valid
197
index_config = IndexConfigModel()
192
- indices_list = [
193
- index for index in indices_list if index_config.is_valid_index(index)
194
- ]
198
+ indices_list = [index for index in indices_list if index_config.is_valid_index(index)]
199
return Indices(
200
indices_list=indices_list,
201
success=True,
@@ -226,9 +230,7 @@ class AlertsQueryBuilder:
230
)
231
232
start = datetime.utcnow() - delta
229
- return (
230
- start.isoformat() + "Z"
231
- ) # Elasticsearch expects the time in ISO format with a Z at the end
233
+ return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end
234
235
def __init__(self):
236
self.query = {
@@ -350,9 +352,7 @@ class LogsQueryBuilder:
352
)
353
354
start = datetime.utcnow() - delta
353
- return (
354
- start.isoformat() + "Z"
355
- ) # Elasticsearch expects the time in ISO format with a Z at the end
355
+ return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end
356
357
def __init__(self):
358
self.query = {
backend/app/connectors/wazuh_manager/models/rules.py
+2
-1
@@ -1,7 +1,8 @@
1
import datetime
2
from typing import Optional
3
4
-from sqlmodel import Field, SQLModel
4
+from sqlmodel import Field
5
+from sqlmodel import SQLModel
6
7
8
class DisabledRule(SQLModel, table=True):
backend/app/connectors/wazuh_manager/routes/rules.py
+15
-12
@@ -1,24 +1,27 @@
1
# App specific imports
2
+from fastapi import APIRouter
3
+from fastapi import Depends
4
+from fastapi import HTTPException
5
+from fastapi import Security
6
+from loguru import logger
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
+from sqlalchemy.future import select
9
+
10
from app.auth.routes.auth import AuthHandler
11
from app.connectors.wazuh_manager.models.rules import DisabledRule
4
-from app.connectors.wazuh_manager.schema.rules import (
5
- AllDisabledRuleResponse,
6
- RuleDisable,
7
- RuleDisableResponse,
8
- RuleEnable,
9
- RuleEnableResponse,
10
-)
12
+from app.connectors.wazuh_manager.schema.rules import AllDisabledRuleResponse
13
+from app.connectors.wazuh_manager.schema.rules import RuleDisable
14
+from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse
15
+from app.connectors.wazuh_manager.schema.rules import RuleEnable
16
+from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse
17
18
# from app.connectors.wazuh_manager.schema.rules import RuleExclude
19
# from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse
14
-from app.connectors.wazuh_manager.services.rules import disable_rule, enable_rule
20
+from app.connectors.wazuh_manager.services.rules import disable_rule
21
+from app.connectors.wazuh_manager.services.rules import enable_rule
22
23
# from app.connectors.wazuh_manager.services.rules import exclude_rule
24
from app.db.db_session import get_db
18
-from fastapi import APIRouter, Depends, HTTPException, Security
19
-from loguru import logger
20
-from sqlalchemy.ext.asyncio import AsyncSession
21
-from sqlalchemy.future import select
25
26
# from app.connectors.wazuh_manager.schema.rules import RuleExclude
27
# from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse
backend/app/connectors/wazuh_manager/schema/rules.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
class RuleDisable(BaseModel):
backend/app/connectors/wazuh_manager/services/rules.py
+15
-14
@@ -1,25 +1,26 @@
1
import re
2
from enum import Enum
3
-from typing import Any, Dict, List, Tuple, Union
3
+from typing import Any
4
+from typing import Dict
5
+from typing import List
6
+from typing import Tuple
7
+from typing import Union
8
9
# import pcre2
10
import xmltodict
7
-from app.connectors.wazuh_manager.schema.rules import (
8
- RuleDisable,
9
- RuleDisableResponse,
10
- RuleEnable,
11
- RuleEnableResponse,
12
- RuleExclude,
13
- RuleExcludeResponse,
14
-)
15
-from app.connectors.wazuh_manager.utils.universal import (
16
- restart_service,
17
- send_get_request,
18
- send_put_request,
19
-)
11
from fastapi import HTTPException
12
from loguru import logger
13
14
+from app.connectors.wazuh_manager.schema.rules import RuleDisable
15
+from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse
16
+from app.connectors.wazuh_manager.schema.rules import RuleEnable
17
+from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse
18
+from app.connectors.wazuh_manager.schema.rules import RuleExclude
19
+from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse
20
+from app.connectors.wazuh_manager.utils.universal import restart_service
21
+from app.connectors.wazuh_manager.utils.universal import send_get_request
22
+from app.connectors.wazuh_manager.utils.universal import send_put_request
23
+
24
25
async def fetch_filename(rule_id: str) -> str:
26
"""
backend/app/connectors/wazuh_manager/utils/universal.py
+7
-3
@@ -1,10 +1,14 @@
1
-from typing import Any, Dict, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import Optional
4
5
import requests
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.db.db_session import AsyncSessionLocal, get_db_session
6
from loguru import logger
7
8
+from app.connectors.utils import get_connector_info_from_db
9
+from app.db.db_session import AsyncSessionLocal
10
+from app.db.db_session import get_db_session
11
+
12
13
async def verify_wazuh_manager_credentials(
14
attributes: Dict[str, Any],
backend/app/customer_provisioning/routes/decommission.py
+8
-4
@@ -1,12 +1,16 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
+from sqlalchemy.future import select
8
+
9
from app.auth.utils import AuthHandler
10
from app.customer_provisioning.schema.decommission import DecommissionCustomerResponse
11
from app.customer_provisioning.services.decommission import decomission_wazuh_customer
12
from app.db.db_session import get_db
13
from app.db.universal_models import CustomersMeta
6
-from fastapi import APIRouter, Depends, HTTPException, Security
7
-from loguru import logger
8
-from sqlalchemy.ext.asyncio import AsyncSession
9
-from sqlalchemy.future import select
14
15
# App specific imports
16
backend/app/customer_provisioning/routes/provision.py
+19
-14
@@ -1,21 +1,26 @@
1
-from app.auth.utils import AuthHandler
2
-from app.connectors.grafana.schema.dashboards import Office365Dashboard, WazuhDashboard
3
-from app.customer_provisioning.schema.provision import (
4
- CustomerProvisionResponse,
5
- CustomersMetaResponse,
6
- CustomerSubsctipion,
7
- GetDashboardsResponse,
8
- GetSubscriptionsResponse,
9
- ProvisionNewCustomer,
10
-)
11
-from app.customer_provisioning.services.provision import provision_wazuh_customer
12
-from app.db.db_session import get_db
13
-from app.db.universal_models import Customers, CustomersMeta
14
-from fastapi import APIRouter, Body, Depends, HTTPException, Security
1
+from fastapi import APIRouter
2
+from fastapi import Body
3
+from fastapi import Depends
4
+from fastapi import HTTPException
5
+from fastapi import Security
6
from loguru import logger
7
from sqlalchemy.ext.asyncio import AsyncSession
8
from sqlalchemy.future import select
9
10
+from app.auth.utils import AuthHandler
11
+from app.connectors.grafana.schema.dashboards import Office365Dashboard
12
+from app.connectors.grafana.schema.dashboards import WazuhDashboard
13
+from app.customer_provisioning.schema.provision import CustomerProvisionResponse
14
+from app.customer_provisioning.schema.provision import CustomersMetaResponse
15
+from app.customer_provisioning.schema.provision import CustomerSubsctipion
16
+from app.customer_provisioning.schema.provision import GetDashboardsResponse
17
+from app.customer_provisioning.schema.provision import GetSubscriptionsResponse
18
+from app.customer_provisioning.schema.provision import ProvisionNewCustomer
19
+from app.customer_provisioning.services.provision import provision_wazuh_customer
20
+from app.db.db_session import get_db
21
+from app.db.universal_models import Customers
22
+from app.db.universal_models import CustomersMeta
23
+
24
customer_provisioning_router = APIRouter()
25
26
backend/app/customer_provisioning/schema/decommission.py
+2
-1
@@ -1,6 +1,7 @@
1
from typing import List
2
3
-from pydantic import BaseModel, Field
3
+from pydantic import BaseModel
4
+from pydantic import Field
5
6
7
class DecommissionedData(BaseModel):
backend/app/customer_provisioning/schema/grafana.py
+2
-1
@@ -1,7 +1,8 @@
1
from datetime import datetime
2
from typing import Dict
3
4
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
# ! Organization ! #
backend/app/customer_provisioning/schema/graylog.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
# ! INDEX SETS ! #
backend/app/customer_provisioning/schema/provision.py
+6
-2
@@ -1,10 +1,14 @@
1
import re
2
from enum import Enum
3
-from typing import List, Optional
3
+from typing import List
4
+from typing import Optional
5
+
6
+from pydantic import BaseModel
7
+from pydantic import Field
8
+from pydantic import validator
9
10
from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
11
from app.db.universal_models import CustomersMeta
7
-from pydantic import BaseModel, Field, validator
12
13
14
class CustomerSubsctipion(Enum):
backend/app/customer_provisioning/schema/wazuh_worker.py
+2
-1
@@ -1,4 +1,5 @@
1
-from pydantic import BaseModel, Field
1
+from pydantic import BaseModel
2
+from pydantic import Field
3
4
5
class ProvisionWorkerRequest(BaseModel):
backend/app/customer_provisioning/services/decommission.py
+10
-12
@@ -1,21 +1,19 @@
1
import requests
2
+from loguru import logger
3
+from sqlalchemy.ext.asyncio import AsyncSession
4
+
5
from app.customer_provisioning.schema.decommission import DecommissionCustomerResponse
3
-from app.customer_provisioning.schema.wazuh_worker import (
4
- DecommissionWorkerRequest,
5
- DecommissionWorkerResponse,
6
-)
6
+from app.customer_provisioning.schema.wazuh_worker import DecommissionWorkerRequest
7
+from app.customer_provisioning.schema.wazuh_worker import DecommissionWorkerResponse
8
from app.customer_provisioning.services.dfir_iris import delete_customer
9
from app.customer_provisioning.services.grafana import delete_grafana_organization
9
-from app.customer_provisioning.services.graylog import delete_index_set, delete_stream
10
-from app.customer_provisioning.services.wazuh_manager import (
11
- delete_wazuh_agents,
12
- delete_wazuh_groups,
13
- gather_wazuh_agents,
14
-)
10
+from app.customer_provisioning.services.graylog import delete_index_set
11
+from app.customer_provisioning.services.graylog import delete_stream
12
+from app.customer_provisioning.services.wazuh_manager import delete_wazuh_agents
13
+from app.customer_provisioning.services.wazuh_manager import delete_wazuh_groups
14
+from app.customer_provisioning.services.wazuh_manager import gather_wazuh_agents
15
from app.db.universal_models import CustomersMeta
16
from app.utils import get_connector_attribute
17
-from loguru import logger
18
-from sqlalchemy.ext.asyncio import AsyncSession
17
18
19
async def decomission_wazuh_customer(
backend/app/customer_provisioning/services/dfir_iris.py
+6
-6
@@ -1,12 +1,12 @@
1
-from app.connectors.dfir_iris.schema.admin import CreateCustomerResponse, ListCustomers
2
-from app.connectors.dfir_iris.utils.universal import (
3
- fetch_and_validate_data,
4
- initialize_client_and_admin,
5
- initialize_client_and_customer,
6
-)
1
from fastapi import HTTPException
2
from loguru import logger
3
4
+from app.connectors.dfir_iris.schema.admin import CreateCustomerResponse
5
+from app.connectors.dfir_iris.schema.admin import ListCustomers
6
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
7
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_admin
8
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_customer
9
+
10
11
async def check_customer_exists(customer_name: str) -> bool:
12
"""
backend/app/customer_provisioning/services/grafana.py
+9
-10
@@ -1,17 +1,16 @@
1
+from fastapi import HTTPException
2
+from loguru import logger
3
+from sqlalchemy.ext.asyncio import AsyncSession
4
+
5
from app.connectors.grafana.utils.universal import create_grafana_client
6
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
3
-from app.customer_provisioning.schema.grafana import (
4
- GrafanaDatasource,
5
- GrafanaDataSourceCreationResponse,
6
- GrafanaFolderCreationResponse,
7
- GrafanaOrganizationCreation,
8
- NodesVersionResponse,
9
-)
7
+from app.customer_provisioning.schema.grafana import GrafanaDatasource
8
+from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse
9
+from app.customer_provisioning.schema.grafana import GrafanaFolderCreationResponse
10
+from app.customer_provisioning.schema.grafana import GrafanaOrganizationCreation
11
+from app.customer_provisioning.schema.grafana import NodesVersionResponse
12
from app.customer_provisioning.schema.provision import ProvisionNewCustomer
13
from app.utils import get_connector_attribute
12
-from fastapi import HTTPException
13
-from loguru import logger
14
-from sqlalchemy.ext.asyncio import AsyncSession
14
15
16
################# ! GRAFANA PROVISIONING ! #################
backend/app/customer_provisioning/services/graylog.py
+11
-14
@@ -1,23 +1,20 @@
1
import json
2
from datetime import datetime
3
4
-from app.connectors.graylog.services.pipelines import get_pipelines
5
-from app.connectors.graylog.utils.universal import (
6
- send_delete_request,
7
- send_post_request,
8
-)
9
-from app.customer_provisioning.schema.graylog import (
10
- GraylogIndexSetCreationResponse,
11
- StreamConnectionToPipelineRequest,
12
- StreamConnectionToPipelineResponse,
13
- StreamCreationResponse,
14
- TimeBasedIndexSet,
15
- WazuhEventStream,
16
-)
17
-from app.customer_provisioning.schema.provision import ProvisionNewCustomer
4
from fastapi import HTTPException
5
from loguru import logger
6
7
+from app.connectors.graylog.services.pipelines import get_pipelines
8
+from app.connectors.graylog.utils.universal import send_delete_request
9
+from app.connectors.graylog.utils.universal import send_post_request
10
+from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse
11
+from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest
12
+from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineResponse
13
+from app.customer_provisioning.schema.graylog import StreamCreationResponse
14
+from app.customer_provisioning.schema.graylog import TimeBasedIndexSet
15
+from app.customer_provisioning.schema.graylog import WazuhEventStream
16
+from app.customer_provisioning.schema.provision import ProvisionNewCustomer
17
+
18
19
######### ! GRAYLOG PROVISIONING ! ############
20
# ! INDEX SETS ! #
backend/app/customer_provisioning/services/provision.py
+21
-36
@@ -1,41 +1,32 @@
1
import requests
2
+from fastapi import HTTPException
3
+from loguru import logger
4
+from sqlalchemy.ext.asyncio import AsyncSession
5
+
6
from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
7
from app.connectors.grafana.services.dashboards import provision_dashboards
8
from app.connectors.graylog.services.management import start_stream
9
from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest
6
-from app.customer_provisioning.schema.provision import (
7
- CustomerProvisionMeta,
8
- CustomerProvisionResponse,
9
- ProvisionNewCustomer,
10
-)
11
-from app.customer_provisioning.schema.wazuh_worker import (
12
- ProvisionWorkerRequest,
13
- ProvisionWorkerResponse,
14
-)
10
+from app.customer_provisioning.schema.provision import CustomerProvisionMeta
11
+from app.customer_provisioning.schema.provision import CustomerProvisionResponse
12
+from app.customer_provisioning.schema.provision import ProvisionNewCustomer
13
+from app.customer_provisioning.schema.wazuh_worker import ProvisionWorkerRequest
14
+from app.customer_provisioning.schema.wazuh_worker import ProvisionWorkerResponse
15
from app.customer_provisioning.services.dfir_iris import create_customer
16
-from app.customer_provisioning.services.grafana import (
17
- create_grafana_datasource,
18
- create_grafana_folder,
19
- create_grafana_organization,
20
-)
21
-from app.customer_provisioning.services.graylog import (
22
- connect_stream_to_pipeline,
23
- create_event_stream,
24
- create_index_set,
25
- get_pipeline_id,
26
-)
27
-from app.customer_provisioning.services.wazuh_manager import (
28
- apply_group_configurations,
29
- create_wazuh_groups,
30
-)
16
+from app.customer_provisioning.services.grafana import create_grafana_datasource
17
+from app.customer_provisioning.services.grafana import create_grafana_folder
18
+from app.customer_provisioning.services.grafana import create_grafana_organization
19
+from app.customer_provisioning.services.graylog import connect_stream_to_pipeline
20
+from app.customer_provisioning.services.graylog import create_event_stream
21
+from app.customer_provisioning.services.graylog import create_index_set
22
+from app.customer_provisioning.services.graylog import get_pipeline_id
23
+from app.customer_provisioning.services.wazuh_manager import apply_group_configurations
24
+from app.customer_provisioning.services.wazuh_manager import create_wazuh_groups
25
from app.db.universal_models import CustomersMeta
26
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
27
AlertCreationSettings,
28
)
29
from app.utils import get_connector_attribute
36
-from fastapi import HTTPException
37
-from loguru import logger
38
-from sqlalchemy.ext.asyncio import AsyncSession
30
31
32
# ! MAIN FUNCTION ! #
@@ -61,9 +52,7 @@ async def provision_wazuh_customer(
52
# Initialize an empty dictionary to store the meta data
53
provision_meta_data = {}
54
provision_meta_data["index_set_id"] = (await create_index_set(request)).data.id
64
- provision_meta_data["stream_id"] = (
65
- await create_event_stream(request, provision_meta_data["index_set_id"])
66
- ).data.stream_id
55
+ provision_meta_data["stream_id"] = (await create_event_stream(request, provision_meta_data["index_set_id"])).data.stream_id
56
provision_meta_data["pipeline_ids"] = await get_pipeline_id(subscription="Wazuh")
57
stream_and_pipeline = StreamConnectionToPipelineRequest(
58
stream_id=provision_meta_data["stream_id"],
@@ -77,9 +66,7 @@ async def provision_wazuh_customer(
66
)
67
await create_wazuh_groups(request)
68
await apply_group_configurations(request)
80
- provision_meta_data["grafana_organization_id"] = (
81
- await create_grafana_organization(request)
82
- ).orgId
69
+ provision_meta_data["grafana_organization_id"] = (await create_grafana_organization(request)).orgId
70
provision_meta_data["wazuh_datasource_uid"] = (
71
await create_grafana_datasource(
72
request=request,
@@ -102,9 +89,7 @@ async def provision_wazuh_customer(
89
),
90
)
91
105
- provision_meta_data["iris_customer_id"] = (
106
- await create_customer(request.customer_name)
107
- ).data.customer_id
92
+ provision_meta_data["iris_customer_id"] = (await create_customer(request.customer_name)).data.customer_id
93
94
customer_provision_meta = CustomerProvisionMeta(**provision_meta_data)
95
customer_meta = await update_customer_meta_table(
backend/app/customer_provisioning/services/wazuh_manager.py
+2
-1
@@ -1,6 +1,8 @@
1
from pathlib import Path
2
from typing import List
3
4
+from loguru import logger
5
+
6
from app.connectors.wazuh_manager.utils.universal import (
7
send_delete_request as send_wazuh_delete_request,
8
)
@@ -15,7 +17,6 @@ from app.connectors.wazuh_manager.utils.universal import (
17
)
18
from app.customer_provisioning.schema.provision import ProvisionNewCustomer
19
from app.customer_provisioning.schema.wazuh_manager import WazuhAgentsTemplatePaths
18
-from loguru import logger
20
21
22
######### ! WAZUH MANAGER PROVISIONING ! ############
backend/app/customers/routes/customers.py
+25
-24
@@ -1,32 +1,33 @@
1
-from app.auth.utils import AuthHandler
2
-
3
-# App specific imports
4
-from app.customers.schema.customers import (
5
- AgentModel,
6
- AgentsResponse,
7
- CustomerFullResponse,
8
- CustomerMetaRequestBody,
9
- CustomerMetaResponse,
10
- CustomerRequestBody,
11
- CustomerResponse,
12
- CustomersResponse,
13
-)
14
-from app.db.db_session import get_db
15
-from app.db.universal_models import Agents, Customers, CustomersMeta
16
-from app.healthchecks.agents.schema.agents import (
17
- AgentHealthCheckResponse,
18
- TimeCriteriaModel,
19
-)
20
-from app.healthchecks.agents.services.agents import (
21
- velociraptor_agents_healthcheck,
22
- wazuh_agents_healthcheck,
23
-)
24
-from fastapi import APIRouter, Depends, HTTPException, Query, Security
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Query
5
+from fastapi import Security
6
from loguru import logger
7
from sqlalchemy.ext.asyncio import AsyncSession
8
from sqlalchemy.future import select
9
from starlette.status import HTTP_401_UNAUTHORIZED
10
11
+from app.auth.utils import AuthHandler
12
+
13
+# App specific imports
14
+from app.customers.schema.customers import AgentModel
15
+from app.customers.schema.customers import AgentsResponse
16
+from app.customers.schema.customers import CustomerFullResponse
17
+from app.customers.schema.customers import CustomerMetaRequestBody
18
+from app.customers.schema.customers import CustomerMetaResponse
19
+from app.customers.schema.customers import CustomerRequestBody
20
+from app.customers.schema.customers import CustomerResponse
21
+from app.customers.schema.customers import CustomersResponse
22
+from app.db.db_session import get_db
23
+from app.db.universal_models import Agents
24
+from app.db.universal_models import Customers
25
+from app.db.universal_models import CustomersMeta
26
+from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
27
+from app.healthchecks.agents.schema.agents import TimeCriteriaModel
28
+from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck
29
+from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck
30
+
31
customers_router = APIRouter()
32
33
backend/app/customers/schema/customers.py
+4
-2
@@ -1,7 +1,9 @@
1
from datetime import datetime
2
-from typing import List, Optional
2
+from typing import List
3
+from typing import Optional
4
4
-from pydantic import BaseModel, Field
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
8
9
class CustomerRequestBody(BaseModel):
backend/app/db/db_populate.py
+13
-20
@@ -1,17 +1,18 @@
1
import os
2
3
-from app.auth.models.users import Role
4
-from app.connectors.models import Connectors
5
-from app.integrations.models.customer_integration_settings import (
6
- AvailableIntegrations,
7
- AvailableIntegrationsAuthKeys,
8
-)
3
from dotenv import load_dotenv
4
from loguru import logger
5
from sqlalchemy import and_
6
from sqlalchemy.ext.asyncio import AsyncSession
7
from sqlalchemy.future import select
8
9
+from app.auth.models.users import Role
10
+from app.connectors.models import Connectors
11
+from app.integrations.models.customer_integration_settings import AvailableIntegrations
12
+from app.integrations.models.customer_integration_settings import (
13
+ AvailableIntegrationsAuthKeys,
14
+)
15
+
16
load_dotenv()
17
18
@@ -287,8 +288,7 @@ async def add_available_integrations_if_not_exist(session: AsyncSession):
288
289
for available_integration_data in available_integrations_list:
290
query = select(AvailableIntegrations).where(
290
- AvailableIntegrations.integration_name
291
- == available_integration_data["integration_name"],
291
+ AvailableIntegrations.integration_name == available_integration_data["integration_name"],
292
)
293
result = await session.execute(query)
294
existing_available_integration = result.scalars().first()
@@ -382,28 +382,21 @@ async def add_available_integrations_auth_keys_if_not_exist(session: AsyncSessio
382
Returns:
383
None
384
"""
385
- available_integrations_auth_keys_list = (
386
- await get_available_integrations_auth_keys_list(session=session)
387
- )
385
+ available_integrations_auth_keys_list = await get_available_integrations_auth_keys_list(session=session)
386
387
for available_integration_auth_keys_data in available_integrations_auth_keys_list:
388
query = select(AvailableIntegrations).where(
391
- AvailableIntegrations.integration_name
392
- == available_integration_auth_keys_data["integration_name"],
389
+ AvailableIntegrations.integration_name == available_integration_auth_keys_data["integration_name"],
390
)
391
result = await session.execute(query)
392
existing_integration = result.scalars().first()
393
394
if existing_integration:
398
- available_integration_auth_keys_data[
399
- "integration_id"
400
- ] = existing_integration.id
395
+ available_integration_auth_keys_data["integration_id"] = existing_integration.id
396
auth_key_query = select(AvailableIntegrationsAuthKeys).where(
397
and_(
403
- AvailableIntegrationsAuthKeys.integration_id
404
- == existing_integration.id,
405
- AvailableIntegrationsAuthKeys.auth_key_name
406
- == available_integration_auth_keys_data["auth_key_name"],
398
+ AvailableIntegrationsAuthKeys.integration_id == existing_integration.id,
399
+ AvailableIntegrationsAuthKeys.auth_key_name == available_integration_auth_keys_data["auth_key_name"],
400
),
401
)
402
auth_key_result = await session.execute(auth_key_query)
backend/app/db/db_session.py
+9
-4
@@ -1,6 +1,8 @@
1
# ! Old Testing without Async
2
+from sqlmodel import Session
3
+from sqlmodel import create_engine
4
+
5
from settings import SQLALCHEMY_DATABASE_URI
3
-from sqlmodel import Session, create_engine
6
7
engine = create_engine(
8
SQLALCHEMY_DATABASE_URI,
@@ -8,14 +10,17 @@ engine = create_engine(
10
)
11
session = "placeholder"
12
11
-from contextlib import asynccontextmanager, contextmanager
13
+from contextlib import asynccontextmanager
14
+from contextlib import contextmanager
15
16
from loguru import logger
14
-from settings import SQLALCHEMY_DATABASE_URI
17
from sqlalchemy import create_engine
16
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
18
+from sqlalchemy.ext.asyncio import AsyncSession
19
+from sqlalchemy.ext.asyncio import create_async_engine
20
from sqlalchemy.orm import sessionmaker
21
22
+from settings import SQLALCHEMY_DATABASE_URI
23
+
24
# create async engine for SQLite using aiosqlite
25
async_engine = create_async_engine(SQLALCHEMY_DATABASE_URI, echo=False)
26
sync_engine = create_engine(
backend/app/db/db_setup.py
+8
-11
@@ -1,20 +1,17 @@
1
-from app.auth.services.universal import (
2
- create_admin_user,
3
- create_scheduler_user,
4
- remove_scheduler_user,
5
-)
6
-from app.db.db_populate import (
7
- add_available_integrations_auth_keys_if_not_exist,
8
- add_available_integrations_if_not_exist,
9
- add_connectors_if_not_exist,
10
- add_roles_if_not_exist,
11
-)
1
from loguru import logger
2
from sqlalchemy.ext.asyncio import AsyncSession
3
4
# ! New with Async
5
from sqlmodel import SQLModel
6
7
+from app.auth.services.universal import create_admin_user
8
+from app.auth.services.universal import create_scheduler_user
9
+from app.auth.services.universal import remove_scheduler_user
10
+from app.db.db_populate import add_available_integrations_auth_keys_if_not_exist
11
+from app.db.db_populate import add_available_integrations_if_not_exist
12
+from app.db.db_populate import add_connectors_if_not_exist
13
+from app.db.db_populate import add_roles_if_not_exist
14
+
15
16
async def create_tables(async_engine):
17
"""
backend/app/db/universal_models.py
+12
-29
@@ -1,7 +1,9 @@
1
from datetime import datetime
2
from typing import Optional
3
4
-from sqlmodel import Field, Relationship, SQLModel
4
+from sqlmodel import Field
5
+from sqlmodel import Relationship
6
+from sqlmodel import SQLModel
7
8
9
class Customers(SQLModel, table=True):
@@ -71,22 +73,12 @@ class CustomersMeta(SQLModel, table=True):
73
self.customer_meta_grafana_org_id = customer_meta.customer_meta_grafana_org_id
74
self.customer_meta_wazuh_group = customer_meta.customer_meta_wazuh_group
75
self.customer_meta_index_retention = customer_meta.customer_meta_index_retention
74
- self.customer_meta_wazuh_registration_port = (
75
- customer_meta.customer_meta_wazuh_registration_port
76
- )
77
- self.customer_meta_wazuh_log_ingestion_port = (
78
- customer_meta.customer_meta_wazuh_log_ingestion_port
79
- )
76
+ self.customer_meta_wazuh_registration_port = customer_meta.customer_meta_wazuh_registration_port
77
+ self.customer_meta_wazuh_log_ingestion_port = customer_meta.customer_meta_wazuh_log_ingestion_port
78
self.customer_meta_wazuh_api_port = customer_meta.customer_meta_wazuh_api_port
81
- self.customer_meta_wazuh_auth_password = (
82
- customer_meta.customer_meta_wazuh_auth_password
83
- )
84
- self.customer_meta_iris_customer_id = (
85
- customer_meta.customer_meta_iris_customer_id
86
- )
87
- self.customer_meta_office365_organization_id = (
88
- customer_meta.customer_meta_office365_organization_id
89
- )
79
+ self.customer_meta_wazuh_auth_password = customer_meta.customer_meta_wazuh_auth_password
80
+ self.customer_meta_iris_customer_id = customer_meta.customer_meta_iris_customer_id
81
+ self.customer_meta_office365_organization_id = customer_meta.customer_meta_office365_organization_id
82
83
84
class Agents(SQLModel, table=True):
@@ -111,9 +103,7 @@ class Agents(SQLModel, table=True):
103
def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code):
104
# Check if agent_last_seen is 'Unknown' and set wazuh_last_seen accordingly
105
if wazuh_agent.agent_last_seen == "Unknown":
114
- wazuh_last_seen_value = (
115
- "1970-01-01T00:00:00+00:00" # default datetime value
116
- )
106
+ wazuh_last_seen_value = "1970-01-01T00:00:00+00:00" # default datetime value
107
else:
108
wazuh_last_seen_value = wazuh_agent.agent_last_seen_as_datetime
109
@@ -125,19 +115,14 @@ class Agents(SQLModel, table=True):
115
label=wazuh_agent.agent_label,
116
wazuh_last_seen=wazuh_last_seen_value,
117
wazuh_agent_version=wazuh_agent.wazuh_agent_version,
128
- velociraptor_id=velociraptor_agent.client_id
129
- if velociraptor_agent.client_id
130
- else "n/a",
118
+ velociraptor_id=velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a",
119
velociraptor_last_seen=velociraptor_agent.client_last_seen_as_datetime,
120
velociraptor_agent_version=velociraptor_agent.client_version,
121
customer_code=customer_code,
122
)
123
124
def update_from_model(self, wazuh_agent, velociraptor_agent, customer_code):
137
- if (
138
- wazuh_agent.agent_last_seen == "Unknown"
139
- or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00"
140
- ):
125
+ if wazuh_agent.agent_last_seen == "Unknown" or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00":
126
wazuh_last_seen_value = datetime.strptime(
127
"1970-01-01T00:00:00+00:00",
128
"%Y-%m-%dT%H:%M:%S%z",
@@ -152,9 +137,7 @@ class Agents(SQLModel, table=True):
137
self.label = wazuh_agent.agent_label
138
self.wazuh_last_seen = wazuh_last_seen_value
139
self.wazuh_agent_version = wazuh_agent.wazuh_agent_version
155
- self.velociraptor_id = (
156
- velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a"
157
- )
140
+ self.velociraptor_id = velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a"
141
self.velociraptor_last_seen = velociraptor_agent.client_last_seen_as_datetime
142
self.velociraptor_agent_version = velociraptor_agent.client_version
143
self.customer_code = customer_code
backend/app/healthchecks/agents/routes/agents.py
+18
-17
@@ -1,24 +1,25 @@
1
-from app.auth.utils import AuthHandler
2
-from app.db.db_session import get_db
3
-from app.db.universal_models import Agents
4
-from app.healthchecks.agents.schema.agents import (
5
- AgentHealthCheckResponse,
6
- HostLogsSearchBody,
7
- HostLogsSearchResponse,
8
- TimeCriteriaModel,
9
-)
10
-from app.healthchecks.agents.services.agents import (
11
- host_logs,
12
- velociraptor_agent_healthcheck,
13
- velociraptor_agents_healthcheck,
14
- wazuh_agent_healthcheck,
15
- wazuh_agents_healthcheck,
16
-)
17
-from fastapi import APIRouter, Depends, HTTPException, Query, Security
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Query
5
+from fastapi import Security
6
from loguru import logger
7
from sqlalchemy.ext.asyncio import AsyncSession
8
from sqlalchemy.future import select
9
10
+from app.auth.utils import AuthHandler
11
+from app.db.db_session import get_db
12
+from app.db.universal_models import Agents
13
+from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
14
+from app.healthchecks.agents.schema.agents import HostLogsSearchBody
15
+from app.healthchecks.agents.schema.agents import HostLogsSearchResponse
16
+from app.healthchecks.agents.schema.agents import TimeCriteriaModel
17
+from app.healthchecks.agents.services.agents import host_logs
18
+from app.healthchecks.agents.services.agents import velociraptor_agent_healthcheck
19
+from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck
20
+from app.healthchecks.agents.services.agents import wazuh_agent_healthcheck
21
+from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck
22
+
23
healtcheck_agents_router = APIRouter()
24
25
backend/app/healthchecks/agents/schema/agents.py
+8
-3
@@ -1,7 +1,12 @@
1
from datetime import datetime
2
-from typing import Any, Dict, List, Optional
3
-
4
-from pydantic import BaseModel, Field, validator
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
+
7
+from pydantic import BaseModel
8
+from pydantic import Field
9
+from pydantic import validator
10
11
12
class AgentModel(BaseModel):
backend/app/healthchecks/agents/services/agents.py
+20
-27
@@ -1,24 +1,23 @@
1
-from datetime import datetime, timedelta
2
-from typing import Optional, Type
3
-
4
-from app.connectors.wazuh_indexer.utils.universal import (
5
- LogsQueryBuilder,
6
- collect_indices,
7
- create_wazuh_indexer_client,
8
-)
9
-from app.healthchecks.agents.schema.agents import (
10
- AgentHealthCheckResponse,
11
- AgentModel,
12
- CollectLogsResponse,
13
- ExtendedAgentModel,
14
- HostLogsSearchBody,
15
- HostLogsSearchResponse,
16
- LogsSearchBody,
17
- TimeCriteriaModel,
18
-)
1
+from datetime import datetime
2
+from datetime import timedelta
3
+from typing import Optional
4
+from typing import Type
5
+
6
from fastapi import HTTPException
7
from loguru import logger
8
9
+from app.connectors.wazuh_indexer.utils.universal import LogsQueryBuilder
10
+from app.connectors.wazuh_indexer.utils.universal import collect_indices
11
+from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
12
+from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
13
+from app.healthchecks.agents.schema.agents import AgentModel
14
+from app.healthchecks.agents.schema.agents import CollectLogsResponse
15
+from app.healthchecks.agents.schema.agents import ExtendedAgentModel
16
+from app.healthchecks.agents.schema.agents import HostLogsSearchBody
17
+from app.healthchecks.agents.schema.agents import HostLogsSearchResponse
18
+from app.healthchecks.agents.schema.agents import LogsSearchBody
19
+from app.healthchecks.agents.schema.agents import TimeCriteriaModel
20
+
21
22
def is_wazuh_agent_unhealthy(
23
agent: AgentModel,
@@ -44,9 +43,7 @@ def is_wazuh_agent_unhealthy(
43
return ExtendedAgentModel(**agent.dict(), unhealthy_wazuh_agent=True)
44
45
# Calculate the total time delta based on the criteria
47
- total_minutes = (
48
- time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60
49
- )
46
+ total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60
47
time_delta = timedelta(minutes=total_minutes)
48
49
is_unhealthy = (current_time - wazuh_last_seen) > time_delta
@@ -77,9 +74,7 @@ def is_velociraptor_agent_unhealthy(
74
return ExtendedAgentModel(**agent.dict(), unhealthy_velociraptor_agent=True)
75
76
# Calculate the total time delta based on the criteria
80
- total_minutes = (
81
- time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60
82
- )
77
+ total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60
78
time_delta = timedelta(minutes=total_minutes)
79
80
is_unhealthy = (current_time - velociraptor_last_seen) > time_delta
@@ -277,9 +272,7 @@ async def get_logs_generic(
272
)
273
logs_summary = []
274
indices = await collect_indices()
280
- index_list = (
281
- [index_name] if index_name else indices.indices_list
282
- ) # Use the provided index_name or get all indices
275
+ index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices
276
277
for index_name in index_list:
278
try:
backend/app/integrations/alert_creation/general/routes/alert.py
+12
-15
@@ -1,16 +1,17 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from loguru import logger
5
+from sqlalchemy.ext.asyncio import AsyncSession
6
+from sqlalchemy.future import select
7
+
8
from app.db.db_session import get_db
2
-from app.integrations.alert_creation.general.schema.alert import (
3
- CreateAlertRequest,
4
- CreateAlertResponse,
5
-)
9
+from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest
10
+from app.integrations.alert_creation.general.schema.alert import CreateAlertResponse
11
from app.integrations.alert_creation.general.services.alert import create_alert
12
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
13
AlertCreationSettings,
14
)
10
-from fastapi import APIRouter, Depends, HTTPException
11
-from loguru import logger
12
-from sqlalchemy.ext.asyncio import AsyncSession
13
-from sqlalchemy.future import select
15
16
general_alerts_router = APIRouter()
17
@@ -35,15 +36,12 @@ async def is_rule_id_valid(
36
37
result = await session.execute(
38
select(AlertCreationSettings).where(
38
- AlertCreationSettings.customer_code
39
- == create_alert_request.agent_labels_customer,
39
+ AlertCreationSettings.customer_code == create_alert_request.agent_labels_customer,
40
),
41
)
42
settings = result.scalars().first()
43
44
- if settings and str(create_alert_request.rule_id) in (
45
- settings.excluded_wazuh_rules or ""
46
- ).split(","):
44
+ if settings and str(create_alert_request.rule_id) in (settings.excluded_wazuh_rules or "").split(","):
45
return False
46
47
return True
@@ -69,8 +67,7 @@ async def is_customer_code_valid(
67
68
result = await session.execute(
69
select(AlertCreationSettings).where(
72
- AlertCreationSettings.customer_code
73
- == create_alert_request.agent_labels_customer,
70
+ AlertCreationSettings.customer_code == create_alert_request.agent_labels_customer,
71
),
72
)
73
settings = result.scalars().first()
backend/app/integrations/alert_creation/general/schema/alert.py
+8
-3
@@ -1,7 +1,12 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
3
-
4
-from pydantic import BaseModel, Extra, Field
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
+
7
+from pydantic import BaseModel
8
+from pydantic import Extra
9
+from pydantic import Field
10
11
12
class ValidIocFields(Enum):
backend/app/integrations/alert_creation/general/services/alert.py
+20
-29
@@ -1,33 +1,29 @@
1
-from typing import Optional, Set
1
+from typing import Optional
2
+from typing import Set
3
+
4
+from fastapi import HTTPException
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
8
from app.agents.routes.agents import get_agent
9
from app.agents.schema.agents import AgentsResponse
5
-from app.connectors.dfir_iris.utils.universal import (
6
- fetch_and_validate_data,
7
- initialize_client_and_alert,
8
-)
9
-from app.integrations.alert_creation.general.schema.alert import (
10
- CreateAlertRequest,
11
- CreateAlertResponse,
12
- IrisAlertContext,
13
- IrisAlertPayload,
14
- IrisAsset,
15
- IrisIoc,
16
- ValidIocFields,
17
-)
10
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
11
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
12
+from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest
13
+from app.integrations.alert_creation.general.schema.alert import CreateAlertResponse
14
+from app.integrations.alert_creation.general.schema.alert import IrisAlertContext
15
+from app.integrations.alert_creation.general.schema.alert import IrisAlertPayload
16
+from app.integrations.alert_creation.general.schema.alert import IrisAsset
17
+from app.integrations.alert_creation.general.schema.alert import IrisIoc
18
+from app.integrations.alert_creation.general.schema.alert import ValidIocFields
19
from app.integrations.alert_creation.general.services.alert_multi_exclude import (
20
AlertDetailsService,
21
)
21
-from app.integrations.utils.alerts import (
22
- get_asset_type_id,
23
- send_to_shuffle,
24
- validate_ioc_type,
25
-)
22
+from app.integrations.utils.alerts import get_asset_type_id
23
+from app.integrations.utils.alerts import send_to_shuffle
24
+from app.integrations.utils.alerts import validate_ioc_type
25
from app.integrations.utils.schema import ShufflePayload
26
from app.utils import get_customer_alert_settings
28
-from fastapi import HTTPException
29
-from loguru import logger
30
-from sqlalchemy.ext.asyncio import AsyncSession
27
28
29
def valid_ioc_fields() -> Set[str]:
@@ -57,13 +53,8 @@ async def construct_alert_source_link(
53
The alert source link.
54
"""
55
# Check if the alert has a process id and that it is not "No process ID found"
60
- if (
61
- hasattr(alert_details, "process_id")
62
- and alert_details.process_id != "No process ID found"
63
- ):
64
- query_string = (
65
- f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20"
66
- )
56
+ if hasattr(alert_details, "process_id") and alert_details.process_id != "No process ID found":
57
+ query_string = f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20"
58
else:
59
query_string = f"%22query%22:%22_id:%5C%22{alert_details.id}%5C%22%20AND%20"
60
backend/app/integrations/alert_creation/general/services/alert_multi_exclude.py
+10
-10
@@ -1,13 +1,17 @@
1
-from typing import Any, Dict, List, Tuple
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+from typing import Tuple
5
+
6
+from elasticsearch7 import NotFoundError
7
+from loguru import logger
8
+from sqlalchemy.ext.asyncio import AsyncSession
9
10
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
11
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
12
AlertCreationEventConfig,
13
)
14
from app.utils import get_customer_alert_event_configs
8
-from elasticsearch7 import NotFoundError
9
-from loguru import logger
10
-from sqlalchemy.ext.asyncio import AsyncSession
15
16
17
class AlertDetailsService:
@@ -65,9 +69,7 @@ class AlertDetailsService:
69
result = self.es.search(index="_all", body=query)
70
71
# Extract (index, id) pairs from the result
68
- index_id_pairs = [
69
- (hit["_index"], hit["_id"]) for hit in result["hits"]["hits"]
70
- ]
72
+ index_id_pairs = [(hit["_index"], hit["_id"]) for hit in result["hits"]["hits"]]
73
logger.info(
74
f"Found {len(index_id_pairs)} alerts with syslog_level of 'ALERT' within the last 1 hour.",
75
)
@@ -237,9 +239,7 @@ class AlertDetailsService:
239
logger.info(f"Total alert timeline hits: {total_hits}")
240
241
# Build and sort the list of events
240
- events = [
241
- event["_source"] for event in alert_timeline_events["hits"]["hits"]
242
- ]
242
+ events = [event["_source"] for event in alert_timeline_events["hits"]["hits"]]
243
events.sort(key=lambda x: x["timestamp_utc"])
244
245
# return self.process_events(events)
backend/app/integrations/alert_creation/office365/routes/alert.py
+18
-12
@@ -1,13 +1,28 @@
1
# from app.alerts.office365.services.threat_intel import create_threat_intel_alert
2
+from fastapi import APIRouter
3
+from fastapi import Depends
4
+from fastapi import HTTPException
5
+from loguru import logger
6
+from sqlalchemy import select
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
+
9
from app.db.db_session import get_db
10
from app.integrations.alert_creation.office365.schema.exchange import (
11
Office365ExchangeAlertBase,
12
+)
13
+from app.integrations.alert_creation.office365.schema.exchange import (
14
Office365ExchangeAlertRequest,
15
+)
16
+from app.integrations.alert_creation.office365.schema.exchange import (
17
Office365ExchangeAlertResponse,
18
+)
19
+from app.integrations.alert_creation.office365.schema.exchange import (
20
ValidOffice365Workloads,
21
)
22
from app.integrations.alert_creation.office365.schema.threat_intel import (
23
Office365ThreatIntelAlertRequest,
24
+)
25
+from app.integrations.alert_creation.office365.schema.threat_intel import (
26
Office365ThreatIntelAlertResponse,
27
)
28
from app.integrations.alert_creation.office365.services.exchange import (
@@ -19,10 +34,6 @@ from app.integrations.alert_creation.office365.services.threat_intel import (
34
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
35
AlertCreationSettings,
36
)
22
-from fastapi import APIRouter, Depends, HTTPException
23
-from loguru import logger
24
-from sqlalchemy import select
25
-from sqlalchemy.ext.asyncio import AsyncSession
37
38
office365_alerts_router = APIRouter()
39
@@ -47,8 +58,7 @@ async def is_office365_organization_id_valid(
58
59
result = await session.execute(
60
select(AlertCreationSettings).where(
50
- AlertCreationSettings.office365_organization_id
51
- == create_alert_request.data_office365_OrganizationId,
61
+ AlertCreationSettings.office365_organization_id == create_alert_request.data_office365_OrganizationId,
62
),
63
)
64
settings = result.scalars().first()
@@ -71,9 +81,7 @@ async def create_office365_exchange_alert(
81
session: AsyncSession = Depends(get_db),
82
):
83
logger.info(f"create_alert_request: {create_alert_request}")
74
- if create_alert_request.data_office365_Workload not in [
75
- workload.value for workload in ValidOffice365Workloads
76
- ]:
84
+ if create_alert_request.data_office365_Workload not in [workload.value for workload in ValidOffice365Workloads]:
85
logger.info(f"Invalid workload: {create_alert_request.data_office365_Workload}")
86
raise HTTPException(status_code=400, detail="Invalid workload")
87
logger.info(f"Workload is valid: {create_alert_request.data_office365_Workload}")
@@ -91,9 +99,7 @@ async def create_office365_threat_intel_alert(
99
session: AsyncSession = Depends(get_db),
100
):
101
logger.info(f"create_alert_request: {create_alert_request}")
94
- if create_alert_request.data_office365_Workload not in [
95
- workload.value for workload in ValidOffice365Workloads
96
- ]:
102
+ if create_alert_request.data_office365_Workload not in [workload.value for workload in ValidOffice365Workloads]:
103
logger.info(f"Invalid workload: {create_alert_request.data_office365_Workload}")
104
raise HTTPException(status_code=400, detail="Invalid workload")
105
logger.info(f"Workload is valid: {create_alert_request.data_office365_Workload}")
backend/app/integrations/alert_creation/office365/schema/exchange.py
+8
-3
@@ -1,7 +1,12 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
3
-
4
-from pydantic import BaseModel, Extra, Field
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
+
7
+from pydantic import BaseModel
8
+from pydantic import Extra
9
+from pydantic import Field
10
11
12
class ValidOffice365Workloads(Enum):
backend/app/integrations/alert_creation/office365/schema/threat_intel.py
+7
-2
@@ -1,7 +1,12 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
4
-from pydantic import BaseModel, Extra, Field
7
+from pydantic import BaseModel
8
+from pydantic import Extra
9
+from pydantic import Field
10
11
12
class ValidOffice365Workloads(Enum):
backend/app/integrations/alert_creation/office365/services/exchange.py
+15
-12
@@ -1,23 +1,26 @@
1
-from typing import Optional, Set
1
+from typing import Optional
2
+from typing import Set
3
3
-from app.connectors.dfir_iris.utils.universal import (
4
- fetch_and_validate_data,
5
- initialize_client_and_alert,
6
-)
4
+from loguru import logger
5
+from sqlalchemy.ext.asyncio import AsyncSession
6
+
7
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
8
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
9
from app.integrations.alert_creation.general.schema.alert import ValidIocFields
10
+from app.integrations.alert_creation.office365.schema.exchange import IrisAlertContext
11
+from app.integrations.alert_creation.office365.schema.exchange import IrisAlertPayload
12
+from app.integrations.alert_creation.office365.schema.exchange import IrisAsset
13
+from app.integrations.alert_creation.office365.schema.exchange import IrisIoc
14
from app.integrations.alert_creation.office365.schema.exchange import (
9
- IrisAlertContext,
10
- IrisAlertPayload,
11
- IrisAsset,
12
- IrisIoc,
15
Office365ExchangeAlertRequest,
16
+)
17
+from app.integrations.alert_creation.office365.schema.exchange import (
18
Office365ExchangeAlertResponse,
19
)
16
-from app.integrations.utils.alerts import send_to_shuffle, validate_ioc_type
20
+from app.integrations.utils.alerts import send_to_shuffle
21
+from app.integrations.utils.alerts import validate_ioc_type
22
from app.integrations.utils.schema import ShufflePayload
23
from app.utils import get_customer_alert_settings_office365
19
-from loguru import logger
20
-from sqlalchemy.ext.asyncio import AsyncSession
24
25
26
def valid_ioc_fields() -> Set[str]:
backend/app/integrations/alert_creation/office365/services/threat_intel.py
+17
-10
@@ -1,23 +1,30 @@
1
-from typing import Optional, Set
1
+from typing import Optional
2
+from typing import Set
3
3
-from app.connectors.dfir_iris.utils.universal import (
4
- fetch_and_validate_data,
5
- initialize_client_and_alert,
6
-)
4
+from loguru import logger
5
+from sqlalchemy.ext.asyncio import AsyncSession
6
+
7
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
8
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
9
from app.integrations.alert_creation.general.schema.alert import ValidIocFields
10
from app.integrations.alert_creation.office365.schema.threat_intel import (
11
IrisAlertContext,
12
+)
13
+from app.integrations.alert_creation.office365.schema.threat_intel import (
14
IrisAlertPayload,
11
- IrisAsset,
12
- IrisIoc,
15
+)
16
+from app.integrations.alert_creation.office365.schema.threat_intel import IrisAsset
17
+from app.integrations.alert_creation.office365.schema.threat_intel import IrisIoc
18
+from app.integrations.alert_creation.office365.schema.threat_intel import (
19
Office365ThreatIntelAlertRequest,
20
+)
21
+from app.integrations.alert_creation.office365.schema.threat_intel import (
22
Office365ThreatIntelAlertResponse,
23
)
16
-from app.integrations.utils.alerts import send_to_shuffle, validate_ioc_type
24
+from app.integrations.utils.alerts import send_to_shuffle
25
+from app.integrations.utils.alerts import validate_ioc_type
26
from app.integrations.utils.schema import ShufflePayload
27
from app.utils import get_customer_alert_settings_office365
19
-from loguru import logger
20
-from sqlalchemy.ext.asyncio import AsyncSession
28
29
30
def valid_ioc_fields() -> Set[str]:
backend/app/integrations/alert_creation_settings/models/alert_creation_settings.py
+5
-2
@@ -1,6 +1,9 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from sqlmodel import Field, Relationship, SQLModel
4
+from sqlmodel import Field
5
+from sqlmodel import Relationship
6
+from sqlmodel import SQLModel
7
8
9
class Condition(SQLModel, table=True):
backend/app/integrations/alert_creation_settings/routes/alert_creation_settings.py
+23
-15
@@ -1,24 +1,39 @@
1
from typing import List
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from loguru import logger
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
+from sqlalchemy.future import select
9
+from sqlalchemy.orm import joinedload
10
+
11
from app.db.db_session import get_db
12
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
13
AlertCreationEventConfig,
14
+)
15
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
16
AlertCreationSettings,
17
+)
18
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
19
EventOrder,
20
)
21
from app.integrations.alert_creation_settings.schema.alert_creation_settings import (
22
AlertCreationEventConfigResponse,
23
+)
24
+from app.integrations.alert_creation_settings.schema.alert_creation_settings import (
25
AlertCreationSettingsCreate,
26
+)
27
+from app.integrations.alert_creation_settings.schema.alert_creation_settings import (
28
AlertCreationSettingsResponse,
29
+)
30
+from app.integrations.alert_creation_settings.schema.alert_creation_settings import (
31
EventOrderCreate,
32
+)
33
+from app.integrations.alert_creation_settings.schema.alert_creation_settings import (
34
EventOrderResponse,
35
)
36
from app.utils import get_customer_alert_event_configs
17
-from fastapi import APIRouter, Depends, HTTPException
18
-from loguru import logger
19
-from sqlalchemy.ext.asyncio import AsyncSession
20
-from sqlalchemy.future import select
21
-from sqlalchemy.orm import joinedload
37
38
alert_creation_settings_router = APIRouter()
39
@@ -79,8 +94,7 @@ async def create_alert_creation_settings(
94
95
result = await session.execute(
96
select(AlertCreationSettings).where(
82
- AlertCreationSettings.customer_code
83
- == alert_creation_settings.customer_code,
97
+ AlertCreationSettings.customer_code == alert_creation_settings.customer_code,
98
),
99
)
100
settings = result.scalars().first()
@@ -215,9 +229,7 @@ async def add_event_order(
229
230
# Query the EventOrder instance again to ensure event_configs are loaded
231
result = await session.execute(
218
- select(EventOrder)
219
- .options(joinedload(EventOrder.event_configs))
220
- .where(EventOrder.id == event_order_db.id),
232
+ select(EventOrder).options(joinedload(EventOrder.event_configs)).where(EventOrder.id == event_order_db.id),
233
)
234
event_order_db = result.scalars().first()
235
@@ -266,11 +278,7 @@ async def update_event_orders(
278
for event_order in event_orders:
279
# Check if an EventOrder with the given order_label already exists
280
existing_order = next(
269
- (
270
- order
271
- for order in settings.event_orders
272
- if order.order_label == event_order.order_label
273
- ),
281
+ (order for order in settings.event_orders if order.order_label == event_order.order_label),
282
None,
283
)
284
backend/app/integrations/alert_creation_settings/schema/alert_creation_settings.py
+2
-1
@@ -1,4 +1,5 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
4
from pydantic import BaseModel
5
backend/app/integrations/alert_escalation/routes/general_alert.py
+8
-7
@@ -1,13 +1,14 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import Security
4
+from loguru import logger
5
+from sqlalchemy.ext.asyncio import AsyncSession
6
+
7
from app.auth.utils import AuthHandler
8
from app.db.db_session import get_db
3
-from app.integrations.alert_escalation.schema.general_alert import (
4
- CreateAlertRequest,
5
- CreateAlertResponse,
6
-)
9
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
10
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
11
from app.integrations.alert_escalation.services.general_alert import create_alert
8
-from fastapi import APIRouter, Depends, Security
9
-from loguru import logger
10
-from sqlalchemy.ext.asyncio import AsyncSession
12
13
integration_general_alerts_router = APIRouter()
14
backend/app/integrations/alert_escalation/schema/general_alert.py
+8
-3
@@ -1,7 +1,12 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
3
-
4
-from pydantic import BaseModel, Extra, Field
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
+
7
+from pydantic import BaseModel
8
+from pydantic import Extra
9
+from pydantic import Field
10
11
12
class ValidIocFields(Enum):
backend/app/integrations/alert_escalation/services/general_alert.py
+21
-25
@@ -1,34 +1,33 @@
1
-from typing import Optional, Set
1
+from typing import Optional
2
+from typing import Set
3
+
4
+from fastapi import HTTPException
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
+from sqlalchemy.future import select
8
9
# from app.integrations.alert_escalation.utils.universal import get_agent_data
10
from app.agents.routes.agents import get_agent
11
from app.agents.schema.agents import AgentsResponse
6
-from app.connectors.dfir_iris.utils.universal import (
7
- fetch_and_validate_data,
8
- initialize_client_and_alert,
9
-)
12
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
13
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
14
from app.connectors.utils import get_connector_info_from_db
15
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
16
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
17
AlertCreationSettings,
18
)
15
-from app.integrations.alert_escalation.schema.general_alert import (
16
- CreateAlertRequest,
17
- CreateAlertResponse,
18
- GenericAlertModel,
19
- GenericSourceModel,
20
- IrisAlertContext,
21
- IrisAlertPayload,
22
- IrisAsset,
23
- IrisIoc,
24
- ValidIocFields,
25
-)
26
-from app.integrations.utils.alerts import get_asset_type_id, validate_ioc_type
19
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
20
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
21
+from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel
22
+from app.integrations.alert_escalation.schema.general_alert import GenericSourceModel
23
+from app.integrations.alert_escalation.schema.general_alert import IrisAlertContext
24
+from app.integrations.alert_escalation.schema.general_alert import IrisAlertPayload
25
+from app.integrations.alert_escalation.schema.general_alert import IrisAsset
26
+from app.integrations.alert_escalation.schema.general_alert import IrisIoc
27
+from app.integrations.alert_escalation.schema.general_alert import ValidIocFields
28
+from app.integrations.utils.alerts import get_asset_type_id
29
+from app.integrations.utils.alerts import validate_ioc_type
30
from app.utils import get_customer_alert_settings
28
-from fastapi import HTTPException
29
-from loguru import logger
30
-from sqlalchemy.ext.asyncio import AsyncSession
31
-from sqlalchemy.future import select
31
32
33
async def is_customer_code_valid(customer_code: str, session: AsyncSession) -> bool:
@@ -85,10 +84,7 @@ async def construct_alert_source_link(
84
The alert source link.
85
"""
86
# Check if the alert has a process id and that it is not "No process ID found"
88
- if (
89
- hasattr(alert_details, "process_id")
90
- and alert_details._source.process_id != "No process ID found"
91
- ):
87
+ if hasattr(alert_details, "process_id") and alert_details._source.process_id != "No process ID found":
88
query_string = f"%22query%22:%22process_id:%5C%22{alert_details._source.process_id}%5C%22%20AND%20"
89
else:
90
query_string = f"%22query%22:%22_id:%5C%22{alert_details._id}%5C%22%20AND%20"
backend/app/integrations/ask_socfortress/routes/ask_socfortress.py
+9
-3
@@ -1,16 +1,22 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
+
8
from app.auth.utils import AuthHandler
9
from app.db.db_session import get_db
10
from app.integrations.ask_socfortress.schema.ask_socfortress import (
11
AskSocfortressRequest,
12
+)
13
+from app.integrations.ask_socfortress.schema.ask_socfortress import (
14
AskSocfortressSigmaResponse,
15
)
16
from app.integrations.ask_socfortress.services.ask_socfortress import (
17
ask_socfortress_lookup,
18
)
19
from app.utils import get_connector_attribute
11
-from fastapi import APIRouter, Depends, HTTPException, Security
12
-from loguru import logger
13
-from sqlalchemy.ext.asyncio import AsyncSession
20
21
# App specific imports
22
backend/app/integrations/ask_socfortress/schema/ask_socfortress.py
+2
-1
@@ -1,4 +1,5 @@
1
-from pydantic import BaseModel, Field
1
+from pydantic import BaseModel
2
+from pydantic import Field
3
4
5
class AskSocfortressRequest(BaseModel):
backend/app/integrations/ask_socfortress/services/ask_socfortress.py
+14
-9
@@ -1,23 +1,28 @@
1
-from typing import Any, Dict, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import Optional
4
5
import httpx
6
+from fastapi import HTTPException
7
+from loguru import logger
8
+from sqlalchemy.ext.asyncio import AsyncSession
9
+
10
from app.connectors.utils import get_connector_info_from_db
11
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
12
from app.db.db_session import get_db_session
7
-from app.integrations.alert_escalation.schema.general_alert import (
8
- CreateAlertRequest,
9
- GenericAlertModel,
10
- GenericSourceModel,
11
-)
13
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
14
+from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel
15
+from app.integrations.alert_escalation.schema.general_alert import GenericSourceModel
16
from app.integrations.ask_socfortress.schema.ask_socfortress import (
17
AskSocfortressRequest,
18
+)
19
+from app.integrations.ask_socfortress.schema.ask_socfortress import (
20
AskSocfortressSigmaRequest,
21
+)
22
+from app.integrations.ask_socfortress.schema.ask_socfortress import (
23
AskSocfortressSigmaResponse,
24
)
25
from app.utils import get_connector_attribute
18
-from fastapi import HTTPException
19
-from loguru import logger
20
-from sqlalchemy.ext.asyncio import AsyncSession
26
27
28
async def get_single_alert_details(
backend/app/integrations/dnstwist/routes/analyze.py
+7
-6
@@ -1,12 +1,13 @@
1
import regex
2
-from app.integrations.dnstwist.schema.analyze import (
3
- DomainAnalysisResponse,
4
- DomainRequestBody,
5
-)
6
-from app.integrations.dnstwist.services.analyze import analyze_domain
7
-from fastapi import APIRouter, Depends, HTTPException
2
+from fastapi import APIRouter
3
+from fastapi import Depends
4
+from fastapi import HTTPException
5
from loguru import logger
6
7
+from app.integrations.dnstwist.schema.analyze import DomainAnalysisResponse
8
+from app.integrations.dnstwist.schema.analyze import DomainRequestBody
9
+from app.integrations.dnstwist.services.analyze import analyze_domain
10
+
11
dnstwist_router = APIRouter()
12
13
backend/app/integrations/dnstwist/schema/analyze.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
class DomainData(BaseModel):
backend/app/integrations/dnstwist/services/analyze.py
+3
-4
@@ -1,10 +1,9 @@
1
import dnstwist
2
-from app.integrations.dnstwist.schema.analyze import (
3
- DomainAnalysisResponse,
4
- DomainRequestBody,
5
-)
2
from loguru import logger
3
4
+from app.integrations.dnstwist.schema.analyze import DomainAnalysisResponse
5
+from app.integrations.dnstwist.schema.analyze import DomainRequestBody
6
+
7
8
def analyze_domain(domain: DomainRequestBody) -> DomainAnalysisResponse:
9
"""
backend/app/integrations/mimecast/routes/mimecast.py
+14
-14
@@ -1,20 +1,20 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import Security
4
+from loguru import logger
5
+from sqlalchemy.ext.asyncio import AsyncSession
6
+
7
from app.auth.utils import AuthHandler
8
from app.db.db_session import get_db
3
-from app.integrations.mimecast.schema.mimecast import (
4
- MimecastAuthKeys,
5
- MimecastRequest,
6
- MimecastResponse,
7
- MimecastTTPURLSRequest,
8
-)
9
-from app.integrations.mimecast.services.mimecast import get_ttp_urls, invoke_mimecast
9
+from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys
10
+from app.integrations.mimecast.schema.mimecast import MimecastRequest
11
+from app.integrations.mimecast.schema.mimecast import MimecastResponse
12
+from app.integrations.mimecast.schema.mimecast import MimecastTTPURLSRequest
13
+from app.integrations.mimecast.services.mimecast import get_ttp_urls
14
+from app.integrations.mimecast.services.mimecast import invoke_mimecast
15
from app.integrations.routes import find_customer_integration
11
-from app.integrations.utils.utils import (
12
- extract_mimecast_auth_keys,
13
- get_customer_integration_response,
14
-)
15
-from fastapi import APIRouter, Depends, Security
16
-from loguru import logger
17
-from sqlalchemy.ext.asyncio import AsyncSession
16
+from app.integrations.utils.utils import extract_mimecast_auth_keys
17
+from app.integrations.utils.utils import get_customer_integration_response
18
19
integration_mimecast_router = APIRouter()
20
backend/app/integrations/mimecast/routes/provision.py
+6
-6
@@ -1,15 +1,15 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from sqlalchemy.ext.asyncio import AsyncSession
4
+
5
from app.db.db_session import get_db
6
from app.integrations.mimecast.schema.mimecast import MimecastScheduledResponse
3
-from app.integrations.mimecast.schema.provision import (
4
- ProvisionMimecastRequest,
5
- ProvisionMimecastResponse,
6
-)
7
+from app.integrations.mimecast.schema.provision import ProvisionMimecastRequest
8
+from app.integrations.mimecast.schema.provision import ProvisionMimecastResponse
9
from app.integrations.mimecast.services.provision import provision_mimecast
10
from app.integrations.utils.utils import get_customer_integration_response
11
from app.schedulers.models.scheduler import CreateSchedulerRequest
12
from app.schedulers.scheduler import add_scheduler_jobs
11
-from fastapi import APIRouter, Depends
12
-from sqlalchemy.ext.asyncio import AsyncSession
13
14
integration_mimecast_scheduler_router = APIRouter()
15
backend/app/integrations/mimecast/schema/mimecast.py
+12
-12
@@ -2,11 +2,17 @@ import base64
2
import hashlib
3
import hmac
4
import uuid
5
-from datetime import datetime, timedelta
5
+from datetime import datetime
6
+from datetime import timedelta
7
from enum import Enum
7
-from typing import Dict, List, Optional
8
+from typing import Dict
9
+from typing import List
10
+from typing import Optional
11
9
-from pydantic import BaseModel, Field, HttpUrl, root_validator
12
+from pydantic import BaseModel
13
+from pydantic import Field
14
+from pydantic import HttpUrl
15
+from pydantic import root_validator
16
17
18
class PipelineRuleTitles(Enum):
@@ -141,9 +147,7 @@ class MimecastHeaders(BaseModel):
147
)
148
149
class Config:
144
- allow_population_by_field_name = (
145
- True # This allows field population by both alias and field name
146
- )
150
+ allow_population_by_field_name = True # This allows field population by both alias and field name
151
152
153
class MimecastTTPURLSRequest(BaseModel):
@@ -199,9 +203,7 @@ class MimecastTTPURLSRequest(BaseModel):
203
elif unit == "w":
204
lower_bound = now - timedelta(weeks=amount)
205
202
- values["lower_bound"] = (
203
- lower_bound.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
204
- )
206
+ values["lower_bound"] = lower_bound.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
207
values["upper_bound"] = now.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
208
return values
209
@@ -248,9 +250,7 @@ class DataItem(BaseModel):
250
scanResult: str = Field(..., description="Scan result.")
251
252
class Config:
251
- allow_population_by_field_name = (
252
- True # This allows field population by both alias and field name
253
- )
253
+ allow_population_by_field_name = True # This allows field population by both alias and field name
254
255
256
class RequestBody(BaseModel):
backend/app/integrations/mimecast/schema/provision.py
+7
-2
@@ -1,6 +1,11 @@
1
-from typing import Any, Dict, List, Optional
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+from typing import Optional
5
3
-from pydantic import BaseModel, Field, root_validator
6
+from pydantic import BaseModel
7
+from pydantic import Field
8
+from pydantic import root_validator
9
10
11
class ProvisionMimecastRequest(BaseModel):
backend/app/integrations/mimecast/services/mimecast.py
+11
-12
@@ -12,21 +12,20 @@ from zipfile import ZipFile
12
13
import aiofiles
14
import requests
15
-from app.integrations.mimecast.schema.mimecast import (
16
- DataItem,
17
- MimecastAPIEndpointResponse,
18
- MimecastAuthKeys,
19
- MimecastRequest,
20
- MimecastResponse,
21
- MimecastTTPURLSRequest,
22
- RequestBody,
23
- TtpURLResponseBody,
24
-)
15
+from fastapi import HTTPException
16
+from loguru import logger
17
+
18
+from app.integrations.mimecast.schema.mimecast import DataItem
19
+from app.integrations.mimecast.schema.mimecast import MimecastAPIEndpointResponse
20
+from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys
21
+from app.integrations.mimecast.schema.mimecast import MimecastRequest
22
+from app.integrations.mimecast.schema.mimecast import MimecastResponse
23
+from app.integrations.mimecast.schema.mimecast import MimecastTTPURLSRequest
24
+from app.integrations.mimecast.schema.mimecast import RequestBody
25
+from app.integrations.mimecast.schema.mimecast import TtpURLResponseBody
26
from app.integrations.utils.collection import send_post_request
27
from app.integrations.utils.event_shipper import event_shipper
28
from app.integrations.utils.schema import EventShipperPayload
28
-from fastapi import HTTPException
29
-from loguru import logger
29
30
31
async def get_checkpoint_filename(customer_code: str):
backend/app/integrations/mimecast/services/provision.py
+20
-29
@@ -1,40 +1,33 @@
1
import json
2
from datetime import datetime
3
4
-from app.connectors.grafana.schema.dashboards import (
5
- DashboardProvisionRequest,
6
- MimecastDashboard,
7
-)
4
+from loguru import logger
5
+from sqlalchemy import and_
6
+from sqlalchemy import update
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
+
9
+from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
10
+from app.connectors.grafana.schema.dashboards import MimecastDashboard
11
from app.connectors.grafana.services.dashboards import provision_dashboards
12
from app.connectors.grafana.utils.universal import create_grafana_client
13
from app.connectors.graylog.services.management import start_stream
14
from app.connectors.graylog.utils.universal import send_post_request
12
-from app.customer_provisioning.schema.grafana import (
13
- GrafanaDatasource,
14
- GrafanaDataSourceCreationResponse,
15
-)
16
-from app.customer_provisioning.schema.graylog import (
17
- GraylogIndexSetCreationResponse,
18
- StreamCreationResponse,
19
- TimeBasedIndexSet,
20
-)
21
-from app.customer_provisioning.services.grafana import (
22
- create_grafana_folder,
23
- get_opensearch_version,
24
-)
25
-from app.customers.routes.customers import get_customer, get_customer_meta
26
-from app.integrations.mimecast.schema.provision import (
27
- MimecastEventStream,
28
- ProvisionMimecastRequest,
29
- ProvisionMimecastResponse,
30
-)
15
+from app.customer_provisioning.schema.grafana import GrafanaDatasource
16
+from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse
17
+from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse
18
+from app.customer_provisioning.schema.graylog import StreamCreationResponse
19
+from app.customer_provisioning.schema.graylog import TimeBasedIndexSet
20
+from app.customer_provisioning.services.grafana import create_grafana_folder
21
+from app.customer_provisioning.services.grafana import get_opensearch_version
22
+from app.customers.routes.customers import get_customer
23
+from app.customers.routes.customers import get_customer_meta
24
+from app.integrations.mimecast.schema.provision import MimecastEventStream
25
+from app.integrations.mimecast.schema.provision import ProvisionMimecastRequest
26
+from app.integrations.mimecast.schema.provision import ProvisionMimecastResponse
27
from app.integrations.models.customer_integration_settings import CustomerIntegrations
28
from app.integrations.routes import create_integration_meta
29
from app.integrations.schema import CustomerIntegrationsMetaSchema
30
from app.utils import get_connector_attribute
35
-from loguru import logger
36
-from sqlalchemy import and_, update
37
-from sqlalchemy.ext.asyncio import AsyncSession
31
32
33
################## ! GRAYLOG ! ##################
@@ -223,9 +216,7 @@ async def create_grafana_datasource(
216
grafana_client = await create_grafana_client("Grafana")
217
# Switch to the newly created organization
218
grafana_client.user.switch_actual_user_organisation(
226
- (
227
- await get_customer_meta(customer_code, session)
228
- ).customer_meta.customer_meta_grafana_org_id,
219
+ (await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id,
220
)
221
datasource_payload = GrafanaDatasource(
222
name="MIMECAST",
backend/app/integrations/models/customer_integration_settings.py
+5
-2
@@ -1,7 +1,10 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
4
from sqlalchemy import Text
4
-from sqlmodel import Field, Relationship, SQLModel
5
+from sqlmodel import Field
6
+from sqlmodel import Relationship
7
+from sqlmodel import SQLModel
8
9
10
class AvailableIntegrations(SQLModel, table=True):
backend/app/integrations/monitoring_alert/models/monitoring_alert.py
+2
-1
@@ -1,6 +1,7 @@
1
from typing import Optional
2
3
-from sqlmodel import Field, SQLModel
3
+from sqlmodel import Field
4
+from sqlmodel import SQLModel
5
6
7
class MonitoringAlerts(SQLModel, table=True):
backend/app/integrations/monitoring_alert/routes/monitoring_alert.py
+17
-9
@@ -1,22 +1,32 @@
1
from typing import List
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from loguru import logger
8
+from sqlalchemy.ext.asyncio import AsyncSession
9
+from sqlalchemy.future import select
10
+
11
from app.auth.utils import AuthHandler
12
from app.db.db_session import get_db
13
from app.db.universal_models import CustomersMeta
14
from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
15
+from app.integrations.monitoring_alert.schema.monitoring_alert import GraylogPostRequest
16
from app.integrations.monitoring_alert.schema.monitoring_alert import (
8
- GraylogPostRequest,
17
GraylogPostResponse,
18
+)
19
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
20
MonitoringAlertsRequestModel,
21
+)
22
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
23
MonitoringWazuhAlertsRequestModel,
24
+)
25
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
26
WazuhAnalysisResponse,
27
)
28
from app.integrations.monitoring_alert.services.suricata import analyze_suricata_alerts
29
from app.integrations.monitoring_alert.services.wazuh import analyze_wazuh_alerts
16
-from fastapi import APIRouter, Depends, HTTPException, Security
17
-from loguru import logger
18
-from sqlalchemy.ext.asyncio import AsyncSession
19
-from sqlalchemy.future import select
30
31
monitoring_alerts_router = APIRouter()
32
@@ -147,8 +157,7 @@ async def run_wazuh_analysis(
157
158
monitoring_alerts = await session.execute(
159
select(MonitoringAlerts).where(
150
- (MonitoringAlerts.customer_code == request.customer_code)
151
- & (MonitoringAlerts.alert_source == "WAZUH"),
160
+ (MonitoringAlerts.customer_code == request.customer_code) & (MonitoringAlerts.alert_source == "WAZUH"),
161
),
162
)
163
monitoring_alerts = monitoring_alerts.scalars().all()
@@ -196,8 +205,7 @@ async def run_suricata_analysis(
205
206
monitoring_alerts = await session.execute(
207
select(MonitoringAlerts).where(
199
- (MonitoringAlerts.customer_code == request.customer_code)
200
- & (MonitoringAlerts.alert_source == "SURICATA"),
208
+ (MonitoringAlerts.customer_code == request.customer_code) & (MonitoringAlerts.alert_source == "SURICATA"),
209
),
210
)
211
monitoring_alerts = monitoring_alerts.scalars().all()
backend/app/integrations/monitoring_alert/routes/provision.py
+13
-11
@@ -1,21 +1,29 @@
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from loguru import logger
4
+
5
from app.connectors.graylog.routes.events import get_all_event_definitions
6
from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse
7
+from app.integrations.monitoring_alert.schema.provision import AvailableMonitoringAlerts
8
from app.integrations.monitoring_alert.schema.provision import (
4
- AvailableMonitoringAlerts,
9
AvailableMonitoringAlertsResponse,
10
+)
11
+from app.integrations.monitoring_alert.schema.provision import (
12
ProvisionMonitoringAlertRequest,
13
+)
14
+from app.integrations.monitoring_alert.schema.provision import (
15
ProvisionWazuhMonitoringAlertResponse,
16
)
17
from app.integrations.monitoring_alert.services.provision import (
18
provision_suricata_monitoring_alert,
19
+)
20
+from app.integrations.monitoring_alert.services.provision import (
21
provision_wazuh_monitoring_alert,
22
)
23
from app.integrations.utils.event_shipper import event_shipper
24
from app.integrations.utils.schema import EventShipperPayload
25
from app.schedulers.models.scheduler import CreateSchedulerRequest
26
from app.schedulers.scheduler import add_scheduler_jobs
17
-from fastapi import APIRouter, HTTPException
18
-from loguru import logger
27
28
monitoring_alerts_provision_router = APIRouter()
29
@@ -79,10 +87,7 @@ async def check_if_event_definition_exists(event_definition: str) -> bool:
87
logger.info(
88
f"Event definitions collected: {event_definitions_response.event_definitions}",
89
)
82
- if event_definition in [
83
- event_definition.title
84
- for event_definition in event_definitions_response.event_definitions
85
- ]:
90
+ if event_definition in [event_definition.title for event_definition in event_definitions_response.event_definitions]:
91
raise HTTPException(
92
status_code=400,
93
detail=f"Event definition {event_definition} already exists",
@@ -99,10 +104,7 @@ async def get_available_monitoring_alerts_route() -> AvailableMonitoringAlertsRe
104
"""
105
Get the available monitoring alerts.
106
"""
102
- alerts = [
103
- {"name": alert.name.replace("_", " "), "value": alert.value}
104
- for alert in AvailableMonitoringAlerts
105
- ]
107
+ alerts = [{"name": alert.name.replace("_", " "), "value": alert.value} for alert in AvailableMonitoringAlerts]
108
return AvailableMonitoringAlertsResponse(
109
success=True,
110
message="Alerts retrieved successfully",
backend/app/integrations/monitoring_alert/schema/monitoring_alert.py
+10
-3
@@ -1,8 +1,15 @@
1
from enum import Enum
2
-from typing import Any, Dict, List, Optional
2
+from typing import Any
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
4
-from app.integrations.alert_creation.general.schema.alert import IrisAsset, IrisIoc
5
-from pydantic import BaseModel, Extra, Field
7
+from pydantic import BaseModel
8
+from pydantic import Extra
9
+from pydantic import Field
10
+
11
+from app.integrations.alert_creation.general.schema.alert import IrisAsset
12
+from app.integrations.alert_creation.general.schema.alert import IrisIoc
13
14
15
class MonitoringAlertsRequestModel(BaseModel):
backend/app/integrations/monitoring_alert/schema/provision.py
+6
-2
@@ -1,8 +1,12 @@
1
from enum import Enum
2
-from typing import Dict, List, Optional
2
+from typing import Dict
3
+from typing import List
4
+from typing import Optional
5
6
from fastapi import HTTPException
5
-from pydantic import BaseModel, Field, validator
7
+from pydantic import BaseModel
8
+from pydantic import Field
9
+from pydantic import validator
10
11
12
class AvailableMonitoringAlerts(str, Enum):
backend/app/integrations/monitoring_alert/services/provision.py
+29
-13
@@ -1,27 +1,49 @@
1
import os
2
from typing import Optional
3
4
+from dotenv import load_dotenv
5
+from fastapi import HTTPException
6
+from loguru import logger
7
+
8
from app.connectors.graylog.routes.monitoring import get_all_event_notifications
9
from app.connectors.graylog.schema.management import UrlWhitelistEntryResponse
10
from app.connectors.graylog.schema.monitoring import GraylogEventNotificationsResponse
11
from app.connectors.graylog.services.collector import get_url_whitelist_entries
8
-from app.connectors.graylog.utils.universal import send_post_request, send_put_request
12
+from app.connectors.graylog.utils.universal import send_post_request
13
+from app.connectors.graylog.utils.universal import send_put_request
14
from app.integrations.monitoring_alert.schema.provision import (
15
GraylogAlertProvisionConfig,
16
+)
17
+from app.integrations.monitoring_alert.schema.provision import (
18
GraylogAlertProvisionFieldSpecItem,
19
+)
20
+from app.integrations.monitoring_alert.schema.provision import (
21
GraylogAlertProvisionModel,
22
+)
23
+from app.integrations.monitoring_alert.schema.provision import (
24
GraylogAlertProvisionNotification,
25
+)
26
+from app.integrations.monitoring_alert.schema.provision import (
27
GraylogAlertProvisionNotificationSettings,
28
+)
29
+from app.integrations.monitoring_alert.schema.provision import (
30
GraylogAlertProvisionProvider,
31
+)
32
+from app.integrations.monitoring_alert.schema.provision import (
33
GraylogAlertWebhookNotificationModel,
34
+)
35
+from app.integrations.monitoring_alert.schema.provision import (
36
GraylogUrlWhitelistEntries,
37
+)
38
+from app.integrations.monitoring_alert.schema.provision import (
39
GraylogUrlWhitelistEntryConfig,
40
+)
41
+from app.integrations.monitoring_alert.schema.provision import (
42
ProvisionMonitoringAlertRequest,
43
+)
44
+from app.integrations.monitoring_alert.schema.provision import (
45
ProvisionWazuhMonitoringAlertResponse,
46
)
22
-from dotenv import load_dotenv
23
-from fastapi import HTTPException
24
-from loguru import logger
47
48
load_dotenv()
49
import uuid
@@ -72,10 +94,7 @@ async def check_if_url_whitelist_entry_exists(url: str) -> bool:
94
logger.info(
95
f"Url whitelist entries collected: {url_whitelist_entries_response.url_whitelist_entries}",
96
)
75
- if url in [
76
- url_whitelist_entry.value
77
- for url_whitelist_entry in url_whitelist_entries_response.url_whitelist_entries.entries
78
- ]:
97
+ if url in [url_whitelist_entry.value for url_whitelist_entry in url_whitelist_entries_response.url_whitelist_entries.entries]:
98
logger.info(f"Url whitelist entry {url} already exists")
99
return True
100
return False
@@ -103,9 +122,7 @@ async def get_notification_id(notification_title: str) -> Optional[str]:
122
logger.info(
123
f"Event notifications collected: {event_notifications_response.event_notifications}",
124
)
106
- for (
107
- event_notification
108
- ) in event_notifications_response.event_notifications.notifications:
125
+ for event_notification in event_notifications_response.event_notifications.notifications:
126
if event_notification.title == notification_title:
127
return event_notification.id
128
return None
@@ -184,8 +201,7 @@ async def check_if_event_notification_exists(event_notification: str) -> bool:
201
f"Event notifications collected: {event_notifications_response.event_notifications}",
202
)
203
if event_notification in [
187
- event_notification.title
188
- for event_notification in event_notifications_response.event_notifications.notifications
204
+ event_notification.title for event_notification in event_notifications_response.event_notifications.notifications
205
]:
206
return True
207
return False
backend/app/integrations/monitoring_alert/services/suricata.py
+27
-29
@@ -1,20 +1,21 @@
1
import json
2
-from typing import Optional, Set
2
+from typing import Optional
3
+from typing import Set
4
+
5
+from fastapi import HTTPException
6
+from loguru import logger
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
9
from app.agents.routes.agents import get_agent
10
from app.agents.schema.agents import AgentsResponse
6
-from app.connectors.dfir_iris.utils.universal import (
7
- fetch_and_validate_data,
8
- initialize_client_and_alert,
9
-)
11
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
12
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
13
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
14
from app.db.universal_models import CustomersMeta
12
-from app.integrations.alert_creation.general.schema.alert import (
13
- CreateAlertRequest,
14
- IrisAsset,
15
- IrisIoc,
16
- ValidIocFields,
17
-)
15
+from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest
16
+from app.integrations.alert_creation.general.schema.alert import IrisAsset
17
+from app.integrations.alert_creation.general.schema.alert import IrisIoc
18
+from app.integrations.alert_creation.general.schema.alert import ValidIocFields
19
from app.integrations.alert_creation.general.services.alert_multi_exclude import (
20
AlertDetailsService,
21
)
@@ -27,19 +28,25 @@ from app.integrations.alert_escalation.services.general_alert import (
28
from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
29
from app.integrations.monitoring_alert.schema.monitoring_alert import (
30
FilterAlertsRequest,
30
- SuricataAlertModel,
31
+)
32
+from app.integrations.monitoring_alert.schema.monitoring_alert import SuricataAlertModel
33
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
34
SuricataIrisAlertContext,
32
- SuricataIrisAsset,
35
+)
36
+from app.integrations.monitoring_alert.schema.monitoring_alert import SuricataIrisAsset
37
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
38
WazuhAnalysisResponse,
39
+)
40
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
41
WazuhIrisAlertContext,
42
+)
43
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
44
WazuhIrisAlertPayload,
45
)
46
from app.integrations.monitoring_alert.utils.db_operations import remove_alert_id
38
-from app.integrations.utils.alerts import get_asset_type_id, validate_ioc_type
47
+from app.integrations.utils.alerts import get_asset_type_id
48
+from app.integrations.utils.alerts import validate_ioc_type
49
from app.utils import get_customer_alert_settings
40
-from fastapi import HTTPException
41
-from loguru import logger
42
-from sqlalchemy.ext.asyncio import AsyncSession
50
51
52
def valid_ioc_fields() -> Set[str]:
@@ -69,13 +76,8 @@ async def construct_alert_source_link(
76
The alert source link.
77
"""
78
# Check if the alert has a process id and that it is not "No process ID found"
72
- if (
73
- hasattr(alert_details, "process_id")
74
- and alert_details.process_id != "No process ID found"
75
- ):
76
- query_string = (
77
- f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20"
78
- )
79
+ if hasattr(alert_details, "process_id") and alert_details.process_id != "No process ID found":
80
+ query_string = f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20"
81
else:
82
query_string = f"%22query%22:%22_id:%5C%22{alert_details.id}%5C%22%20AND%20"
83
@@ -218,11 +220,7 @@ async def check_if_open_alert_exists_in_iris(alert_details: SuricataAlertModel)
220
lambda: alert_client.filter_alerts(**params),
221
)
222
logger.info(f"Alert exists: {alert_exists['data']['alerts']}")
221
- return (
222
- alert_exists["data"]["alerts"][0]["alert_id"]
223
- if alert_exists["data"]["alerts"]
224
- else []
225
- )
223
+ return alert_exists["data"]["alerts"][0]["alert_id"] if alert_exists["data"]["alerts"] else []
224
225
226
def construct_params(request: FilterAlertsRequest) -> dict:
backend/app/integrations/monitoring_alert/services/wazuh.py
+24
-28
@@ -1,20 +1,21 @@
1
import json
2
-from typing import Optional, Set
2
+from typing import Optional
3
+from typing import Set
4
+
5
+from fastapi import HTTPException
6
+from loguru import logger
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
9
from app.agents.routes.agents import get_agent
10
from app.agents.schema.agents import AgentsResponse
6
-from app.connectors.dfir_iris.utils.universal import (
7
- fetch_and_validate_data,
8
- initialize_client_and_alert,
9
-)
11
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
12
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
13
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
14
from app.db.universal_models import CustomersMeta
12
-from app.integrations.alert_creation.general.schema.alert import (
13
- CreateAlertRequest,
14
- IrisAsset,
15
- IrisIoc,
16
- ValidIocFields,
17
-)
15
+from app.integrations.alert_creation.general.schema.alert import CreateAlertRequest
16
+from app.integrations.alert_creation.general.schema.alert import IrisAsset
17
+from app.integrations.alert_creation.general.schema.alert import IrisIoc
18
+from app.integrations.alert_creation.general.schema.alert import ValidIocFields
19
from app.integrations.alert_creation.general.services.alert_multi_exclude import (
20
AlertDetailsService,
21
)
@@ -27,17 +28,21 @@ from app.integrations.alert_escalation.services.general_alert import (
28
from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
29
from app.integrations.monitoring_alert.schema.monitoring_alert import (
30
FilterAlertsRequest,
30
- WazuhAlertModel,
31
+)
32
+from app.integrations.monitoring_alert.schema.monitoring_alert import WazuhAlertModel
33
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
34
WazuhAnalysisResponse,
35
+)
36
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
37
WazuhIrisAlertContext,
38
+)
39
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
40
WazuhIrisAlertPayload,
41
)
42
from app.integrations.monitoring_alert.utils.db_operations import remove_alert_id
36
-from app.integrations.utils.alerts import get_asset_type_id, validate_ioc_type
43
+from app.integrations.utils.alerts import get_asset_type_id
44
+from app.integrations.utils.alerts import validate_ioc_type
45
from app.utils import get_customer_alert_settings
38
-from fastapi import HTTPException
39
-from loguru import logger
40
-from sqlalchemy.ext.asyncio import AsyncSession
46
47
48
def valid_ioc_fields() -> Set[str]:
@@ -67,13 +72,8 @@ async def construct_alert_source_link(
72
The alert source link.
73
"""
74
# Check if the alert has a process id and that it is not "No process ID found"
70
- if (
71
- hasattr(alert_details, "process_id")
72
- and alert_details.process_id != "No process ID found"
73
- ):
74
- query_string = (
75
- f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20"
76
- )
75
+ if hasattr(alert_details, "process_id") and alert_details.process_id != "No process ID found":
76
+ query_string = f"%22query%22:%22process_id:%5C%22{alert_details.process_id}%5C%22%20AND%20"
77
else:
78
query_string = f"%22query%22:%22_id:%5C%22{alert_details.id}%5C%22%20AND%20"
79
@@ -214,11 +214,7 @@ async def check_if_open_alert_exists_in_iris(alert_details: WazuhAlertModel) ->
214
lambda: alert_client.filter_alerts(**params),
215
)
216
logger.info(f"Alert exists: {alert_exists['data']['alerts']}")
217
- return (
218
- alert_exists["data"]["alerts"][0]["alert_id"]
219
- if alert_exists["data"]["alerts"]
220
- else []
221
- )
217
+ return alert_exists["data"]["alerts"][0]["alert_id"] if alert_exists["data"]["alerts"] else []
218
219
220
def construct_params(request: FilterAlertsRequest) -> dict:
backend/app/integrations/monitoring_alert/utils/db_operations.py
+2
-1
@@ -1,8 +1,9 @@
1
-from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
1
from loguru import logger
2
from sqlalchemy.ext.asyncio import AsyncSession
3
from sqlalchemy.future import select
4
5
+from app.integrations.monitoring_alert.models.monitoring_alert import MonitoringAlerts
6
+
7
8
async def remove_alert_id(alert_id: str, session: AsyncSession) -> None:
9
"""
backend/app/integrations/office365/routes/provision.py
+13
-12
@@ -1,20 +1,21 @@
1
from typing import Dict
2
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from fastapi import Security
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
+
9
from app.auth.utils import AuthHandler
10
from app.db.db_session import get_db
5
-from app.integrations.office365.schema.provision import (
6
- ProvisionOffice365AuthKeys,
7
- ProvisionOffice365Request,
8
- ProvisionOffice365Response,
9
-)
11
+from app.integrations.office365.schema.provision import ProvisionOffice365AuthKeys
12
+from app.integrations.office365.schema.provision import ProvisionOffice365Request
13
+from app.integrations.office365.schema.provision import ProvisionOffice365Response
14
from app.integrations.office365.services.provision import provision_office365
11
-from app.integrations.routes import (
12
- find_customer_integration,
13
- get_customer_integrations_by_customer_code,
14
-)
15
-from app.integrations.schema import CustomerIntegrations, CustomerIntegrationsResponse
16
-from fastapi import APIRouter, Depends, HTTPException, Security
17
-from sqlalchemy.ext.asyncio import AsyncSession
15
+from app.integrations.routes import find_customer_integration
16
+from app.integrations.routes import get_customer_integrations_by_customer_code
17
+from app.integrations.schema import CustomerIntegrations
18
+from app.integrations.schema import CustomerIntegrationsResponse
19
20
integration_office365_router = APIRouter()
21
backend/app/integrations/office365/schema/provision.py
+5
-2
@@ -1,7 +1,10 @@
1
from enum import Enum
2
-from typing import Any, Dict
2
+from typing import Any
3
+from typing import Dict
4
4
-from pydantic import BaseModel, Field, root_validator
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+from pydantic import root_validator
8
9
10
class PipelineRuleTitles(Enum):
backend/app/integrations/office365/services/provision.py
+46
-80
@@ -4,65 +4,50 @@ from datetime import datetime
4
from typing import List
5
6
import requests
7
-from app.connectors.grafana.schema.dashboards import (
8
- DashboardProvisionRequest,
9
- Office365Dashboard,
10
-)
7
+from dotenv import load_dotenv
8
+from fastapi import HTTPException
9
+from loguru import logger
10
+from sqlalchemy import and_
11
+from sqlalchemy import update
12
+from sqlalchemy.ext.asyncio import AsyncSession
13
+
14
+from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
15
+from app.connectors.grafana.schema.dashboards import Office365Dashboard
16
from app.connectors.grafana.services.dashboards import provision_dashboards
17
from app.connectors.grafana.utils.universal import create_grafana_client
13
-from app.connectors.graylog.schema.pipelines import (
14
- CreatePipeline,
15
- CreatePipelineRule,
16
- GraylogPipelinesResponse,
17
- PipelineRulesResponse,
18
-)
18
+from app.connectors.graylog.schema.pipelines import CreatePipeline
19
+from app.connectors.graylog.schema.pipelines import CreatePipelineRule
20
+from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
21
+from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
22
from app.connectors.graylog.services.management import start_stream
20
-from app.connectors.graylog.services.pipelines import (
21
- connect_stream_to_pipeline,
22
- create_pipeline_graylog,
23
- create_pipeline_rule,
24
- get_pipeline_id,
25
- get_pipeline_rules,
26
- get_pipelines,
27
-)
23
+from app.connectors.graylog.services.pipelines import connect_stream_to_pipeline
24
+from app.connectors.graylog.services.pipelines import create_pipeline_graylog
25
+from app.connectors.graylog.services.pipelines import create_pipeline_rule
26
+from app.connectors.graylog.services.pipelines import get_pipeline_id
27
+from app.connectors.graylog.services.pipelines import get_pipeline_rules
28
+from app.connectors.graylog.services.pipelines import get_pipelines
29
from app.connectors.graylog.utils.universal import send_post_request
29
-from app.connectors.wazuh_manager.utils.universal import (
30
- send_get_request,
31
- send_put_request,
32
-)
33
-from app.customer_provisioning.schema.grafana import (
34
- GrafanaDatasource,
35
- GrafanaDataSourceCreationResponse,
36
-)
37
-from app.customer_provisioning.schema.graylog import (
38
- GraylogIndexSetCreationResponse,
39
- Office365EventStream,
40
- StreamConnectionToPipelineRequest,
41
- StreamCreationResponse,
42
- TimeBasedIndexSet,
43
-)
44
-from app.customer_provisioning.services.grafana import (
45
- create_grafana_folder,
46
- get_opensearch_version,
47
-)
48
-from app.customers.routes.customers import get_customer, get_customer_meta
30
+from app.connectors.wazuh_manager.utils.universal import send_get_request
31
+from app.connectors.wazuh_manager.utils.universal import send_put_request
32
+from app.customer_provisioning.schema.grafana import GrafanaDatasource
33
+from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse
34
+from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse
35
+from app.customer_provisioning.schema.graylog import Office365EventStream
36
+from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest
37
+from app.customer_provisioning.schema.graylog import StreamCreationResponse
38
+from app.customer_provisioning.schema.graylog import TimeBasedIndexSet
39
+from app.customer_provisioning.services.grafana import create_grafana_folder
40
+from app.customer_provisioning.services.grafana import get_opensearch_version
41
+from app.customers.routes.customers import get_customer
42
+from app.customers.routes.customers import get_customer_meta
43
from app.integrations.models.customer_integration_settings import CustomerIntegrations
50
-from app.integrations.office365.schema.provision import (
51
- PipelineRuleTitles,
52
- PipelineTitles,
53
- ProvisionOffice365AuthKeys,
54
- ProvisionOffice365Response,
55
-)
56
-from app.integrations.utils.schema import (
57
- PraecoAlertConfig,
58
- PraecoProvisionAlertResponse,
59
-)
44
+from app.integrations.office365.schema.provision import PipelineRuleTitles
45
+from app.integrations.office365.schema.provision import PipelineTitles
46
+from app.integrations.office365.schema.provision import ProvisionOffice365AuthKeys
47
+from app.integrations.office365.schema.provision import ProvisionOffice365Response
48
+from app.integrations.utils.schema import PraecoAlertConfig
49
+from app.integrations.utils.schema import PraecoProvisionAlertResponse
50
from app.utils import get_connector_attribute
61
-from dotenv import load_dotenv
62
-from fastapi import HTTPException
63
-from loguru import logger
64
-from sqlalchemy import and_, update
65
-from sqlalchemy.ext.asyncio import AsyncSession
51
52
load_dotenv()
53
@@ -488,9 +473,7 @@ async def pipeline_rules_exists(pipeline_rules: PipelineRulesResponse) -> List[s
473
return [
474
rule_title.value
475
for rule_title in PipelineRuleTitles
491
- if not any(
492
- rule.title == rule_title.value for rule in pipeline_rules.pipeline_rules
493
- )
476
+ if not any(rule.title == rule_title.value for rule in pipeline_rules.pipeline_rules)
477
]
478
479
@@ -601,12 +584,7 @@ async def create_wazuh_alert_rule(rule_title: str) -> None:
584
Creates the 'WAZUH CREATE FIELD SYSLOG LEVEL - ALERT' pipeline rule.
585
"""
586
rule_source = (
604
- f'rule "{rule_title}"\n'
605
- "when\n"
606
- " to_long($message.rule_level) > 11\n"
607
- "then\n"
608
- ' set_field("syslog_level", "ALERT");\n'
609
- "end"
587
+ f'rule "{rule_title}"\n' "when\n" " to_long($message.rule_level) > 11\n" "then\n" ' set_field("syslog_level", "ALERT");\n' "end"
588
)
589
await create_pipeline_rule(
590
CreatePipelineRule(
@@ -636,9 +614,7 @@ async def pipeline_exists(pipelines: GraylogPipelinesResponse) -> List[str]:
614
return [
615
pipeline_title.value
616
for pipeline_title in PipelineTitles
639
- if not any(
640
- pipeline.title == pipeline_title.value for pipeline in pipelines.pipelines
641
- )
617
+ if not any(pipeline.title == pipeline_title.value for pipeline in pipelines.pipelines)
618
]
619
620
@@ -699,9 +675,7 @@ async def create_grafana_datasource(
675
grafana_client = await create_grafana_client("Grafana")
676
# Switch to the newly created organization
677
grafana_client.user.switch_actual_user_organisation(
702
- (
703
- await get_customer_meta(customer_code, session)
704
- ).customer_meta.customer_meta_grafana_org_id,
678
+ (await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id,
679
)
680
datasource_payload = GrafanaDatasource(
681
name="O365",
@@ -784,9 +758,7 @@ async def provision_office365(
758
await check_pipeline()
759
760
# Create Index Set
787
- index_set_id = (
788
- await create_index_set(customer_code=customer_code, session=session)
789
- ).data.id
761
+ index_set_id = (await create_index_set(customer_code=customer_code, session=session)).data.id
762
logger.info(f"Index set: {index_set_id}")
763
# Create event stream
764
stream_id = (
@@ -810,23 +782,17 @@ async def provision_office365(
782
await start_stream(stream_id=stream_id)
783
784
# Grafana Deployment
813
- office365_datasource_uid = (
814
- await create_grafana_datasource(customer_code=customer_code, session=session)
815
- ).datasource.uid
785
+ office365_datasource_uid = (await create_grafana_datasource(customer_code=customer_code, session=session)).datasource.uid
786
grafana_o365_folder_id = (
787
await create_grafana_folder(
818
- organization_id=(
819
- await get_customer_meta(customer_code, session)
820
- ).customer_meta.customer_meta_grafana_org_id,
788
+ organization_id=(await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id,
789
folder_title="OFFICE 365",
790
)
791
).id
792
await provision_dashboards(
793
DashboardProvisionRequest(
794
dashboards=[dashboard.name for dashboard in Office365Dashboard],
827
- organizationId=(
828
- await get_customer_meta(customer_code, session)
829
- ).customer_meta.customer_meta_grafana_org_id,
795
+ organizationId=(await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id,
796
folderId=grafana_o365_folder_id,
797
datasourceUid=office365_datasource_uid,
798
),
backend/app/integrations/routes.py
+43
-50
@@ -1,42 +1,49 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
+
4
+from fastapi import APIRouter
5
+from fastapi import Depends
6
+from fastapi import HTTPException
7
+from fastapi import Security
8
+from loguru import logger
9
+from sqlalchemy import delete
10
+from sqlalchemy import update
11
+from sqlalchemy.exc import NoResultFound
12
+from sqlalchemy.ext.asyncio import AsyncSession
13
+from sqlalchemy.future import select
14
+from sqlalchemy.orm import joinedload
15
16
from app.auth.utils import AuthHandler
17
from app.db.db_session import get_db
5
-from app.db.universal_models import Customers, CustomersMeta
18
+from app.db.universal_models import Customers
19
+from app.db.universal_models import CustomersMeta
20
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
21
AlertCreationSettings,
22
)
23
+from app.integrations.models.customer_integration_settings import AvailableIntegrations
24
+from app.integrations.models.customer_integration_settings import CustomerIntegrations
25
from app.integrations.models.customer_integration_settings import (
10
- AvailableIntegrations,
11
- CustomerIntegrations,
26
CustomerIntegrationsMeta,
13
- IntegrationAuthKeys,
14
- IntegrationConfig,
15
- IntegrationService,
16
- IntegrationSubscription,
27
)
18
-from app.integrations.schema import (
19
- AuthKey,
20
- AvailableIntegrationsResponse,
21
- CreateIntegrationAuthKeys,
22
- CreateIntegrationService,
23
- CustomerIntegrationCreate,
24
- CustomerIntegrationCreateResponse,
25
- CustomerIntegrationDeleteResponse,
26
- CustomerIntegrationsMetaResponse,
27
- CustomerIntegrationsMetaSchema,
28
- CustomerIntegrationsResponse,
29
- DeleteCustomerIntegration,
30
- IntegrationWithAuthKeys,
31
- UpdateCustomerIntegration,
28
+from app.integrations.models.customer_integration_settings import IntegrationAuthKeys
29
+from app.integrations.models.customer_integration_settings import IntegrationConfig
30
+from app.integrations.models.customer_integration_settings import IntegrationService
31
+from app.integrations.models.customer_integration_settings import (
32
+ IntegrationSubscription,
33
)
33
-from fastapi import APIRouter, Depends, HTTPException, Security
34
-from loguru import logger
35
-from sqlalchemy import delete, update
36
-from sqlalchemy.exc import NoResultFound
37
-from sqlalchemy.ext.asyncio import AsyncSession
38
-from sqlalchemy.future import select
39
-from sqlalchemy.orm import joinedload
34
+from app.integrations.schema import AuthKey
35
+from app.integrations.schema import AvailableIntegrationsResponse
36
+from app.integrations.schema import CreateIntegrationAuthKeys
37
+from app.integrations.schema import CreateIntegrationService
38
+from app.integrations.schema import CustomerIntegrationCreate
39
+from app.integrations.schema import CustomerIntegrationCreateResponse
40
+from app.integrations.schema import CustomerIntegrationDeleteResponse
41
+from app.integrations.schema import CustomerIntegrationsMetaResponse
42
+from app.integrations.schema import CustomerIntegrationsMetaSchema
43
+from app.integrations.schema import CustomerIntegrationsResponse
44
+from app.integrations.schema import DeleteCustomerIntegration
45
+from app.integrations.schema import IntegrationWithAuthKeys
46
+from app.integrations.schema import UpdateCustomerIntegration
47
48
integration_settings_router = APIRouter()
49
@@ -61,9 +68,7 @@ async def fetch_available_integrations(session: AsyncSession):
68
69
integrations_with_auth_keys = []
70
for integration in unique_integrations:
64
- auth_keys = [
65
- AuthKey(auth_key_name=key.auth_key_name) for key in integration.auth_keys
66
- ]
71
+ auth_keys = [AuthKey(auth_key_name=key.auth_key_name) for key in integration.auth_keys]
72
integration_data = IntegrationWithAuthKeys(
73
id=integration.id,
74
integration_name=integration.integration_name,
@@ -97,9 +102,7 @@ async def validate_integration_auth_keys(
102
Validate if the integration auth keys are valid.
103
"""
104
available_integrations = await fetch_available_integrations(session)
100
- integration = [
101
- ai for ai in available_integrations if ai.integration_name == integration_name
102
- ][0]
105
+ integration = [ai for ai in available_integrations if ai.integration_name == integration_name][0]
106
available_auth_keys = [ak.auth_key_name for ak in integration.auth_keys]
107
# loop through the `available_auth_keys` and check if the `integration_auth_keys` contains the `auth_key_name`
108
for auth_key in available_auth_keys:
@@ -120,9 +123,7 @@ async def validate_integration_auth_key_update(
123
"""
124
logger.info(f"integration_auth_key: {integration_auth_key}")
125
available_integrations = await fetch_available_integrations(session)
123
- integration = [
124
- ai for ai in available_integrations if ai.integration_name == integration_name
125
- ][0]
126
+ integration = [ai for ai in available_integrations if ai.integration_name == integration_name][0]
127
available_auth_keys = [ak.auth_key_name for ak in integration.auth_keys]
128
for auth_key in integration_auth_key:
129
if auth_key.auth_key_name not in available_auth_keys:
@@ -500,20 +501,14 @@ def process_customer_integrations(customer_integrations_data):
501
"""
502
processed_customer_integrations = []
503
for ci in customer_integrations_data:
503
- first_service_id = (
504
- ci.integration_subscriptions[0].integration_service_id
505
- if ci.integration_subscriptions
506
- else None
507
- )
504
+ first_service_id = ci.integration_subscriptions[0].integration_service_id if ci.integration_subscriptions else None
505
customer_integration_obj = CustomerIntegrations(
506
id=ci.id,
507
customer_code=ci.customer_code,
508
customer_name=ci.customer_name,
509
integration_subscriptions=ci.integration_subscriptions,
510
integration_service_id=first_service_id,
514
- integration_service_name=ci.integration_subscriptions[
515
- 0
516
- ].integration_service.service_name
511
+ integration_service_name=ci.integration_subscriptions[0].integration_service.service_name
512
if ci.integration_subscriptions
513
else None,
514
deployed=ci.deployed,
@@ -936,10 +931,8 @@ async def delete_integration_meta(
931
"""
932
try:
933
stmt = delete(CustomerIntegrationsMeta).where(
939
- CustomerIntegrationsMeta.customer_code
940
- == customer_integration_meta.customer_code,
941
- CustomerIntegrationsMeta.integration_name
942
- == customer_integration_meta.integration_name,
934
+ CustomerIntegrationsMeta.customer_code == customer_integration_meta.customer_code,
935
+ CustomerIntegrationsMeta.integration_name == customer_integration_meta.integration_name,
936
)
937
await session.execute(stmt)
938
await session.commit()
backend/app/integrations/schema.py
+4
-2
@@ -1,6 +1,8 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Field
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
7
8
class AuthKey(BaseModel):
backend/app/integrations/utils/alerts.py
+6
-3
@@ -1,16 +1,19 @@
1
import ipaddress
2
import re
3
from abc import ABC
4
-from typing import Dict, Optional, Union
4
+from typing import Dict
5
+from typing import Optional
6
+from typing import Union
7
8
import httpx
9
import regex
8
-from app.integrations.utils.schema import ShufflePayload
9
-from app.utils import get_customer_alert_settings
10
from fastapi import HTTPException
11
from loguru import logger
12
from sqlalchemy.ext.asyncio import AsyncSession
13
14
+from app.integrations.utils.schema import ShufflePayload
15
+from app.utils import get_customer_alert_settings
16
+
17
18
#################### ! DFIR IRIS ASSET VALIDATOR ! ####################
19
class AssetValidator(ABC):
backend/app/integrations/utils/collection.py
+3
-1
@@ -1,5 +1,7 @@
1
import asyncio
2
-from typing import Any, Dict, Optional
2
+from typing import Any
3
+from typing import Dict
4
+from typing import Optional
5
6
import httpx
7
from loguru import logger
backend/app/integrations/utils/event_shipper.py
+7
-7
@@ -1,15 +1,15 @@
1
import asyncio
2
-from typing import Any, Dict
2
+from typing import Any
3
+from typing import Dict
4
+
5
+from fastapi import HTTPException
6
+from loguru import logger
7
8
from app.connectors.event_shipper.utils.universal import create_gelf_logger
9
from app.connectors.utils import get_connector_info_from_db
10
from app.db.db_session import get_db_session
7
-from app.integrations.utils.schema import (
8
- EventShipperPayload,
9
- EventShipperPayloadResponse,
10
-)
11
-from fastapi import HTTPException
12
-from loguru import logger
11
+from app.integrations.utils.schema import EventShipperPayload
12
+from app.integrations.utils.schema import EventShipperPayloadResponse
13
14
15
async def get_gelf_logger():
backend/app/integrations/utils/schema.py
+5
-2
@@ -1,6 +1,9 @@
1
-from typing import List, Optional
1
+from typing import List
2
+from typing import Optional
3
3
-from pydantic import BaseModel, Extra, Field
4
+from pydantic import BaseModel
5
+from pydantic import Extra
6
+from pydantic import Field
7
8
9
class WazuhOSInfo(BaseModel):
backend/app/integrations/utils/utils.py
+4
-2
@@ -1,11 +1,13 @@
1
from typing import Dict
2
3
-from app.integrations.routes import get_customer_integrations_by_customer_code
4
-from app.integrations.schema import CustomerIntegrations, CustomerIntegrationsResponse
3
from fastapi import HTTPException
4
from loguru import logger
5
from sqlalchemy.ext.asyncio import AsyncSession
6
7
+from app.integrations.routes import get_customer_integrations_by_customer_code
8
+from app.integrations.schema import CustomerIntegrations
9
+from app.integrations.schema import CustomerIntegrationsResponse
10
+
11
12
async def get_customer_integration_response(
13
customer_code: str,
backend/app/middleware/exception_handlers.py
+9
-4
@@ -1,11 +1,16 @@
1
-from app.auth.utils import AuthHandler
2
-from app.db.db_session import async_engine # Make sure to import the async engine
3
-from app.utils import ErrorType, Logger, ValidationErrorItem, ValidationErrorResponse
4
-from fastapi import HTTPException, Request
1
+from fastapi import HTTPException
2
+from fastapi import Request
3
from fastapi.exceptions import RequestValidationError
4
from fastapi.responses import JSONResponse
5
from sqlalchemy.ext.asyncio import AsyncSession
6
7
+from app.auth.utils import AuthHandler
8
+from app.db.db_session import async_engine # Make sure to import the async engine
9
+from app.utils import ErrorType
10
+from app.utils import Logger
11
+from app.utils import ValidationErrorItem
12
+from app.utils import ValidationErrorResponse
13
+
14
15
# Utility function to get user_id from request
16
async def get_user_id_from_request(request: Request, logger_instance):
backend/app/middleware/logger.py
+7
-11
@@ -1,9 +1,11 @@
1
+from fastapi import HTTPException
2
+from fastapi import Request
3
+from fastapi.responses import JSONResponse
4
+from sqlalchemy.ext.asyncio import AsyncSession
5
+
6
from app.auth.utils import AuthHandler
7
from app.db.db_session import async_engine
8
from app.utils import Logger
4
-from fastapi import HTTPException, Request
5
-from fastapi.responses import JSONResponse
6
-from sqlalchemy.ext.asyncio import AsyncSession
9
10
EXCLUDED_PATHS = ["/auth/token", "/auth/register"]
11
INTERNAL_SERVER_ERROR = 500
@@ -54,20 +56,14 @@ async def handle_exception(e, user_id, request, logger_instance):
56
JSONResponse: The response containing the error message.
57
"""
58
try:
57
- user_id = (
58
- await logger_instance.get_user_id_from_request(request)
59
- if user_id is None
60
- else user_id
61
- )
59
+ user_id = await logger_instance.get_user_id_from_request(request) if user_id is None else user_id
60
except HTTPException as http_exc:
61
return JSONResponse(
62
status_code=http_exc.status_code,
63
content={"message": str(http_exc), "success": False},
64
)
65
await logger_instance.log_error(user_id, request, e)
68
- status_code = (
69
- e.status_code if isinstance(e, HTTPException) else INTERNAL_SERVER_ERROR
70
- )
66
+ status_code = e.status_code if isinstance(e, HTTPException) else INTERNAL_SERVER_ERROR
67
return JSONResponse(
68
status_code=status_code,
69
content={"message": str(e), "success": False},
backend/app/routers/agents.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.agents.routes.agents import agents_router
1
from fastapi import APIRouter
2
3
+from app.agents.routes.agents import agents_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/alert_creation.py
+2
-1
@@ -1,8 +1,9 @@
1
+from fastapi import APIRouter
2
+
3
from app.integrations.alert_creation.general.routes.alert import general_alerts_router
4
from app.integrations.alert_creation.office365.routes.alert import (
5
office365_alerts_router,
6
)
5
-from fastapi import APIRouter
7
8
# Instantiate the APIRouter
9
router = APIRouter()
backend/app/routers/alert_creation_settings.py
+2
-1
@@ -1,7 +1,8 @@
1
+from fastapi import APIRouter
2
+
3
from app.integrations.alert_creation_settings.routes.alert_creation_settings import (
4
alert_creation_settings_router,
5
)
4
-from fastapi import APIRouter
6
7
# Instantiate the APIRouter
8
router = APIRouter()
backend/app/routers/ask_socfortress.py
+2
-1
@@ -1,7 +1,8 @@
1
+from fastapi import APIRouter
2
+
3
from app.integrations.ask_socfortress.routes.ask_socfortress import (
4
ask_socfortress_router,
5
)
4
-from fastapi import APIRouter
6
7
# Instantiate the APIRouter
8
router = APIRouter()
backend/app/routers/auth.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.auth.routes.auth import auth_router
1
from fastapi import APIRouter
2
3
+from app.auth.routes.auth import auth_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/connectors.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.routes import connector_router
1
from fastapi import APIRouter
2
3
+from app.connectors.routes import connector_router
4
+
5
router = APIRouter()
6
7
router.include_router(connector_router, prefix="/connectors", tags=["connectors"])
backend/app/routers/cortex.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.cortex.routes.analyzers import cortex_analyzer_router
1
from fastapi import APIRouter
2
3
+from app.connectors.cortex.routes.analyzers import cortex_analyzer_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/customer_provisioning.py
+2
-1
@@ -1,8 +1,9 @@
1
+from fastapi import APIRouter
2
+
3
from app.customer_provisioning.routes.decommission import (
4
customer_decommissioning_router,
5
)
6
from app.customer_provisioning.routes.provision import customer_provisioning_router
5
-from fastapi import APIRouter
7
8
# Instantiate the APIRouter
9
router = APIRouter()
backend/app/routers/customers.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.customers.routes.customers import customers_router
1
from fastapi import APIRouter
2
3
+from app.customers.routes.customers import customers_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/dfir_iris.py
+2
-1
@@ -1,3 +1,5 @@
1
+from fastapi import APIRouter
2
+
3
from app.connectors.dfir_iris.routes.alerts import dfir_iris_alerts_router
4
from app.connectors.dfir_iris.routes.assets import dfir_iris_assets_router
5
from app.connectors.dfir_iris.routes.cases import dfir_iris_cases_router
@@ -6,7 +8,6 @@ from app.connectors.dfir_iris.routes.users import dfir_iris_users_router
8
from app.integrations.alert_escalation.routes.general_alert import (
9
integration_general_alerts_router,
10
)
9
-from fastapi import APIRouter
11
12
# Instantiate the APIRouter
13
router = APIRouter()
backend/app/routers/dnstwist.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.integrations.dnstwist.routes.analyze import dnstwist_router
1
from fastapi import APIRouter
2
3
+from app.integrations.dnstwist.routes.analyze import dnstwist_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/grafana.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.grafana.routes.dashboards import grafana_dashboards_router
1
from fastapi import APIRouter
2
3
+from app.connectors.grafana.routes.dashboards import grafana_dashboards_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/graylog.py
+2
-1
@@ -1,10 +1,11 @@
1
+from fastapi import APIRouter
2
+
3
from app.connectors.graylog.routes.collector import graylog_collector_router
4
from app.connectors.graylog.routes.events import graylog_events_router
5
from app.connectors.graylog.routes.management import graylog_management_router
6
from app.connectors.graylog.routes.monitoring import graylog_monitoring_router
7
from app.connectors.graylog.routes.pipelines import graylog_pipelines_router
8
from app.connectors.graylog.routes.streams import graylog_streams_router
7
-from fastapi import APIRouter
9
10
router = APIRouter()
11
backend/app/routers/healthcheck.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.healthchecks.agents.routes.agents import healtcheck_agents_router
1
from fastapi import APIRouter
2
3
+from app.healthchecks.agents.routes.agents import healtcheck_agents_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/influxdb.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.influxdb.routes.alerts import influxdb_alerts_router
1
from fastapi import APIRouter
2
3
+from app.connectors.influxdb.routes.alerts import influxdb_alerts_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/integrations.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.integrations.routes import integration_settings_router
1
from fastapi import APIRouter
2
3
+from app.integrations.routes import integration_settings_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/logs.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.utils import logs_router
1
from fastapi import APIRouter
2
3
+from app.utils import logs_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/mimecast.py
+2
-1
@@ -1,8 +1,9 @@
1
+from fastapi import APIRouter
2
+
3
from app.integrations.mimecast.routes.mimecast import integration_mimecast_router
4
from app.integrations.mimecast.routes.provision import (
5
integration_mimecast_scheduler_router,
6
)
5
-from fastapi import APIRouter
7
8
# Instantiate the APIRouter
9
router = APIRouter()
backend/app/routers/monitoring_alert.py
+2
-1
@@ -1,10 +1,11 @@
1
+from fastapi import APIRouter
2
+
3
from app.integrations.monitoring_alert.routes.monitoring_alert import (
4
monitoring_alerts_router,
5
)
6
from app.integrations.monitoring_alert.routes.provision import (
7
monitoring_alerts_provision_router,
8
)
7
-from fastapi import APIRouter
9
10
# Instantiate the APIRouter
11
router = APIRouter()
backend/app/routers/office365.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.integrations.office365.routes.provision import integration_office365_router
1
from fastapi import APIRouter
2
3
+from app.integrations.office365.routes.provision import integration_office365_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/scheduler.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.schedulers.routes.scheduler import scheduler_router
1
from fastapi import APIRouter
2
3
+from app.schedulers.routes.scheduler import scheduler_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/shuffle.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.shuffle.routes.workflows import shuffle_workflows_router
1
from fastapi import APIRouter
2
3
+from app.connectors.shuffle.routes.workflows import shuffle_workflows_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/smtp.py
+2
-1
@@ -1,6 +1,7 @@
1
+from fastapi import APIRouter
2
+
3
from app.smtp.routes.configure import smtp_configure_router
4
from app.smtp.routes.reports import smtp_reports_router
3
-from fastapi import APIRouter
5
6
# Instantiate the APIRouter
7
router = APIRouter()
backend/app/routers/sublime.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.sublime.routes.alerts import sublime_alerts_router
1
from fastapi import APIRouter
2
3
+from app.connectors.sublime.routes.alerts import sublime_alerts_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/threat_intel.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.threat_intel.routes.socfortress import threat_intel_socfortress_router
1
from fastapi import APIRouter
2
3
+from app.threat_intel.routes.socfortress import threat_intel_socfortress_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/routers/velociraptor.py
+2
-1
@@ -1,6 +1,7 @@
1
+from fastapi import APIRouter
2
+
3
from app.connectors.velociraptor.routes.artifacts import velociraptor_artifacts_router
4
from app.connectors.velociraptor.routes.flows import velociraptor_flows_router
3
-from fastapi import APIRouter
5
6
# Instantiate the APIRouter
7
router = APIRouter()
backend/app/routers/wazuh_indexer.py
+2
-1
@@ -1,6 +1,7 @@
1
+from fastapi import APIRouter
2
+
3
from app.connectors.wazuh_indexer.routes.alerts import wazuh_indexer_alerts_router
4
from app.connectors.wazuh_indexer.routes.monitoring import wazuh_indexer_router
3
-from fastapi import APIRouter
5
6
# Instantiate the APIRouter
7
router = APIRouter()
backend/app/routers/wazuh_manager.py
+2
-1
@@ -1,6 +1,7 @@
1
-from app.connectors.wazuh_manager.routes.rules import wazuh_manager_rules_router
1
from fastapi import APIRouter
2
3
+from app.connectors.wazuh_manager.routes.rules import wazuh_manager_rules_router
4
+
5
# Instantiate the APIRouter
6
router = APIRouter()
7
backend/app/schedulers/models/scheduler.py
+2
-1
@@ -2,7 +2,8 @@ from datetime import datetime
2
from typing import Optional
3
4
from pydantic import BaseModel
5
-from sqlmodel import Field, SQLModel
5
+from sqlmodel import Field
6
+from sqlmodel import SQLModel
7
8
9
class JobMetadata(SQLModel, table=True):
backend/app/schedulers/routes/scheduler.py
+6
-4
@@ -1,11 +1,13 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from loguru import logger
4
+from sqlalchemy.ext.asyncio import AsyncSession
5
+from sqlalchemy.future import select
6
+
7
from app.db.db_session import get_db
8
from app.schedulers.models.scheduler import JobMetadata
9
from app.schedulers.scheduler import init_scheduler
10
from app.schedulers.schema.scheduler import JobsResponse
5
-from fastapi import APIRouter, Depends
6
-from loguru import logger
7
-from sqlalchemy.ext.asyncio import AsyncSession
8
-from sqlalchemy.future import select
11
12
scheduler_router = APIRouter()
13
backend/app/schedulers/scheduler.py
+12
-19
@@ -1,18 +1,17 @@
1
-from app.db.db_session import SyncSessionLocal, sync_engine
2
-from app.schedulers.models.scheduler import CreateSchedulerRequest, JobMetadata
3
-from app.schedulers.services.agent_sync import agent_sync
4
-from app.schedulers.services.invoke_mimecast import (
5
- invoke_mimecast_integration,
6
- invoke_mimecast_integration_ttp,
7
-)
8
-from app.schedulers.services.monitoring_alert import (
9
- invoke_suricata_monitoring_alert,
10
- invoke_wazuh_monitoring_alert,
11
-)
1
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
2
from apscheduler.schedulers.asyncio import AsyncIOScheduler
3
from loguru import logger
4
5
+from app.db.db_session import SyncSessionLocal
6
+from app.db.db_session import sync_engine
7
+from app.schedulers.models.scheduler import CreateSchedulerRequest
8
+from app.schedulers.models.scheduler import JobMetadata
9
+from app.schedulers.services.agent_sync import agent_sync
10
+from app.schedulers.services.invoke_mimecast import invoke_mimecast_integration
11
+from app.schedulers.services.invoke_mimecast import invoke_mimecast_integration_ttp
12
+from app.schedulers.services.monitoring_alert import invoke_suricata_monitoring_alert
13
+from app.schedulers.services.monitoring_alert import invoke_wazuh_monitoring_alert
14
+
15
16
def init_scheduler():
17
"""
@@ -43,9 +42,7 @@ def initialize_job_metadata():
42
# {"job_id": "invoke_mimecast_integration", "time_interval": 5, "function": invoke_mimecast_integration}
43
]
44
for job in known_jobs:
46
- job_metadata = (
47
- session.query(JobMetadata).filter_by(job_id=job["job_id"]).one_or_none()
48
- )
45
+ job_metadata = session.query(JobMetadata).filter_by(job_id=job["job_id"]).one_or_none()
46
if not job_metadata:
47
job_metadata = JobMetadata(
48
job_id=job["job_id"],
@@ -132,11 +129,7 @@ async def add_job_metadata(create_scheduler_request: CreateSchedulerRequest):
129
create_scheduler_request (CreateSchedulerRequest): The request object containing the job details.
130
"""
131
with SyncSessionLocal() as session:
135
- job_metadata = (
136
- session.query(JobMetadata)
137
- .filter_by(job_id=create_scheduler_request.job_id)
138
- .one_or_none()
139
- )
132
+ job_metadata = session.query(JobMetadata).filter_by(job_id=create_scheduler_request.job_id).one_or_none()
133
if not job_metadata:
134
job_metadata = JobMetadata(
135
job_id=create_scheduler_request.job_id,
backend/app/schedulers/services/agent_sync.py
+3
-4
@@ -2,10 +2,11 @@ import os
2
from datetime import datetime
3
4
import requests
5
+from dotenv import load_dotenv
6
+
7
from app.db.db_session import get_sync_db_session
8
from app.schedulers.models.scheduler import JobMetadata
9
from app.schedulers.utils.universal import scheduler_login
8
-from dotenv import load_dotenv
10
11
load_dotenv()
12
@@ -39,9 +40,7 @@ def agent_sync():
40
# Use get_sync_db_session to create and manage a synchronous session
41
with get_sync_db_session() as session:
42
# Synchronous ORM operations
42
- job_metadata = (
43
- session.query(JobMetadata).filter_by(job_id="agent_sync").one_or_none()
44
- )
43
+ job_metadata = session.query(JobMetadata).filter_by(job_id="agent_sync").one_or_none()
44
if job_metadata:
45
job_metadata.last_success = datetime.utcnow()
46
session.add(job_metadata)
backend/app/schedulers/services/invoke_mimecast.py
+11
-18
@@ -1,17 +1,18 @@
1
from datetime import datetime
2
3
-from app.db.db_session import get_db_session, get_sync_db_session
4
-from app.integrations.mimecast.routes.mimecast import (
5
- invoke_mimecast_route,
6
- mimecast_ttp_url_route,
7
-)
8
-from app.integrations.mimecast.schema.mimecast import MimecastRequest, MimecastResponse
9
-from app.integrations.models.customer_integration_settings import CustomerIntegrations
10
-from app.schedulers.models.scheduler import JobMetadata
3
from dotenv import load_dotenv
4
from loguru import logger
5
from sqlalchemy import select
6
7
+from app.db.db_session import get_db_session
8
+from app.db.db_session import get_sync_db_session
9
+from app.integrations.mimecast.routes.mimecast import invoke_mimecast_route
10
+from app.integrations.mimecast.routes.mimecast import mimecast_ttp_url_route
11
+from app.integrations.mimecast.schema.mimecast import MimecastRequest
12
+from app.integrations.mimecast.schema.mimecast import MimecastResponse
13
+from app.integrations.models.customer_integration_settings import CustomerIntegrations
14
+from app.schedulers.models.scheduler import JobMetadata
15
+
16
load_dotenv()
17
18
@@ -40,11 +41,7 @@ async def invoke_mimecast_integration() -> MimecastResponse:
41
await session.close()
42
with get_sync_db_session() as session:
43
# Synchronous ORM operations
43
- job_metadata = (
44
- session.query(JobMetadata)
45
- .filter_by(job_id="invoke_mimecast_integration")
46
- .one_or_none()
47
- )
44
+ job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_mimecast_integration").one_or_none()
45
if job_metadata:
46
job_metadata.last_success = datetime.utcnow()
47
session.add(job_metadata)
@@ -80,11 +77,7 @@ async def invoke_mimecast_integration_ttp() -> MimecastResponse:
77
await session.close()
78
with get_sync_db_session() as session:
79
# Synchronous ORM operations
83
- job_metadata = (
84
- session.query(JobMetadata)
85
- .filter_by(job_id="invoke_mimecast_integration_ttp")
86
- .one_or_none()
87
- )
80
+ job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_mimecast_integration_ttp").one_or_none()
81
if job_metadata:
82
job_metadata.last_success = datetime.utcnow()
83
session.add(job_metadata)
backend/app/schedulers/services/monitoring_alert.py
+11
-15
@@ -1,19 +1,23 @@
1
from datetime import datetime
2
3
-from app.db.db_session import get_db_session, get_sync_db_session
3
+from dotenv import load_dotenv
4
+from loguru import logger
5
+from sqlalchemy import select
6
+
7
+from app.db.db_session import get_db_session
8
+from app.db.db_session import get_sync_db_session
9
from app.db.universal_models import CustomersMeta
10
from app.integrations.monitoring_alert.routes.monitoring_alert import (
11
run_suricata_analysis,
7
- run_wazuh_analysis,
12
)
13
+from app.integrations.monitoring_alert.routes.monitoring_alert import run_wazuh_analysis
14
from app.integrations.monitoring_alert.schema.monitoring_alert import (
15
MonitoringWazuhAlertsRequestModel,
16
+)
17
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
18
WazuhAnalysisResponse,
19
)
20
from app.schedulers.models.scheduler import JobMetadata
14
-from dotenv import load_dotenv
15
-from loguru import logger
16
-from sqlalchemy import select
21
22
load_dotenv()
23
@@ -41,11 +45,7 @@ async def invoke_wazuh_monitoring_alert() -> WazuhAnalysisResponse:
45
await session.close()
46
with get_sync_db_session() as session:
47
# Synchronous ORM operations
44
- job_metadata = (
45
- session.query(JobMetadata)
46
- .filter_by(job_id="invoke_wazuh_monitoring_alert")
47
- .one_or_none()
48
- )
48
+ job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_wazuh_monitoring_alert").one_or_none()
49
if job_metadata:
50
job_metadata.last_success = datetime.utcnow()
51
session.add(job_metadata)
@@ -83,11 +83,7 @@ async def invoke_suricata_monitoring_alert() -> WazuhAnalysisResponse:
83
await session.close()
84
with get_sync_db_session() as session:
85
# Synchronous ORM operations
86
- job_metadata = (
87
- session.query(JobMetadata)
88
- .filter_by(job_id="invoke_suricata_monitoring_alert")
89
- .one_or_none()
90
- )
86
+ job_metadata = session.query(JobMetadata).filter_by(job_id="invoke_suricata_monitoring_alert").one_or_none()
87
if job_metadata:
88
job_metadata.last_success = datetime.utcnow()
89
session.add(job_metadata)
backend/app/schedulers/utils/universal.py
+2
-1
@@ -1,9 +1,10 @@
1
import os
2
3
import requests
4
-from app.auth.services.universal import get_scheduler_password
4
from dotenv import load_dotenv
5
6
+from app.auth.services.universal import get_scheduler_password
7
+
8
load_dotenv()
9
10
backend/app/smtp/routes/configure.py
+6
-3
@@ -1,10 +1,13 @@
1
-from app.auth.models.users import SMTP, SMTPInput
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from loguru import logger
4
+
5
+from app.auth.models.users import SMTP
6
+from app.auth.models.users import SMTPInput
7
from app.auth.services.universal import select_all_users
8
from app.auth.utils import AuthHandler
9
from app.db.db_session import session
10
from app.smtp.schema.configure import SMTPResponse
6
-from fastapi import APIRouter, HTTPException
7
-from loguru import logger
11
12
smtp_configure_router = APIRouter()
13
auth_handler = AuthHandler()
backend/app/smtp/routes/reports.py
+6
-3
@@ -1,10 +1,13 @@
1
-from app.auth.models.users import SMTP, SMTPInput
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from loguru import logger
4
+
5
+from app.auth.models.users import SMTP
6
+from app.auth.models.users import SMTPInput
7
from app.auth.services.universal import select_all_users
8
from app.auth.utils import AuthHandler
9
from app.db.db_session import session
10
from app.smtp.schema.configure import SMTPResponse
6
-from fastapi import APIRouter, HTTPException
7
-from loguru import logger
11
12
smtp_reports_router = APIRouter()
13
auth_handler = AuthHandler()
backend/app/smtp/services/create_report.py
+5
-1
@@ -8,7 +8,10 @@ from reportlab.lib.styles import getSampleStyleSheet
8
from reportlab.lib.units import inch
9
10
# from reportlab.pdfgen import canvas
11
-from reportlab.platypus import Image, Paragraph, SimpleDocTemplate, Spacer
11
+from reportlab.platypus import Image
12
+from reportlab.platypus import Paragraph
13
+from reportlab.platypus import SimpleDocTemplate
14
+from reportlab.platypus import Spacer
15
16
matplotlib.use(
17
"Agg",
@@ -16,6 +19,7 @@ matplotlib.use(
19
# for scripts and web servers. This should resolve the main thread is not
20
# in main loop issue as it bypasses the need for tkinter.
21
import matplotlib.pyplot as plt
22
+
23
from app.services.wazuh_indexer.alerts import AlertsService
24
25
# ! TODO: Just a template
backend/app/smtp/services/reports.py
+2
-1
@@ -6,7 +6,8 @@ from email.mime.text import MIMEText
6
from typing import List
7
8
from app.services.smtp.create_report import create_alerts_report_pdf
9
-from app.services.smtp.universal import EmailTemplate, UniversalEmailCredentials
9
+from app.services.smtp.universal import EmailTemplate
10
+from app.services.smtp.universal import UniversalEmailCredentials
11
12
# ! SEND REPORT
13
backend/app/threat_intel/routes/socfortress.py
+9
-7
@@ -1,14 +1,16 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
+
8
from app.auth.utils import AuthHandler
9
from app.db.db_session import get_db
3
-from app.threat_intel.schema.socfortress import (
4
- IoCResponse,
5
- SocfortressThreatIntelRequest,
6
-)
10
+from app.threat_intel.schema.socfortress import IoCResponse
11
+from app.threat_intel.schema.socfortress import SocfortressThreatIntelRequest
12
from app.threat_intel.services.socfortress import socfortress_threat_intel_lookup
13
from app.utils import get_connector_attribute
9
-from fastapi import APIRouter, Depends, HTTPException, Security
10
-from loguru import logger
11
-from sqlalchemy.ext.asyncio import AsyncSession
14
15
# App specific imports
16
backend/app/threat_intel/schema/socfortress.py
+2
-1
@@ -1,6 +1,7 @@
1
from typing import Optional
2
3
-from pydantic import BaseModel, Field
3
+from pydantic import BaseModel
4
+from pydantic import Field
5
6
7
class SocfortressThreatIntelRequest(BaseModel):
backend/app/threat_intel/services/socfortress.py
+9
-9
@@ -1,18 +1,18 @@
1
-from typing import Any, Dict
1
+from typing import Any
2
+from typing import Dict
3
4
import httpx
4
-from app.connectors.utils import get_connector_info_from_db
5
-from app.db.db_session import get_db_session
6
-from app.threat_intel.schema.socfortress import (
7
- IoCMapping,
8
- IoCResponse,
9
- SocfortressThreatIntelRequest,
10
-)
11
-from app.utils import get_connector_attribute
5
from fastapi import HTTPException
6
from loguru import logger
7
from sqlalchemy.ext.asyncio import AsyncSession
8
9
+from app.connectors.utils import get_connector_info_from_db
10
+from app.db.db_session import get_db_session
11
+from app.threat_intel.schema.socfortress import IoCMapping
12
+from app.threat_intel.schema.socfortress import IoCResponse
13
+from app.threat_intel.schema.socfortress import SocfortressThreatIntelRequest
14
+from app.utils import get_connector_attribute
15
+
16
17
async def get_socfortress_threat_intel_attributes(
18
column_name: str,
backend/app/utils.py
+33
-30
@@ -1,26 +1,44 @@
1
-from datetime import datetime, timedelta
1
+from datetime import datetime
2
+from datetime import timedelta
3
from enum import Enum
3
-from typing import Any, Dict, List, Optional, Union
4
+from typing import Any
5
+from typing import Dict
6
+from typing import List
7
+from typing import Optional
8
+from typing import Union
9
10
import requests
11
+from fastapi import APIRouter
12
+from fastapi import Depends
13
+from fastapi import HTTPException
14
+from fastapi import Request
15
+from fastapi import Security
16
+from fastapi.exceptions import RequestValidationError
17
+from loguru import logger
18
+from pydantic import BaseModel
19
+from pydantic import Field
20
+from pydantic import validator
21
+from sqlalchemy.ext.asyncio import AsyncSession
22
+from sqlalchemy.future import select
23
+from sqlalchemy.orm import joinedload
24
+
25
from app.auth.services.universal import find_user
26
from app.auth.utils import AuthHandler
27
from app.connectors.utils import get_connector_info_from_db
28
from app.db.all_models import Connectors
10
-from app.db.db_session import get_db, get_db_session, get_session
29
+from app.db.db_session import get_db
30
+from app.db.db_session import get_db_session
31
+from app.db.db_session import get_session
32
from app.db.universal_models import LogEntry
33
from app.integrations.alert_creation_settings.models.alert_creation_settings import (
34
AlertCreationEventConfig,
35
+)
36
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
37
AlertCreationSettings,
38
+)
39
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
40
EventOrder,
41
)
17
-from fastapi import APIRouter, Depends, HTTPException, Request, Security
18
-from fastapi.exceptions import RequestValidationError
19
-from loguru import logger
20
-from pydantic import BaseModel, Field, validator
21
-from sqlalchemy.ext.asyncio import AsyncSession
22
-from sqlalchemy.future import select
23
-from sqlalchemy.orm import joinedload
42
43
44
################## ! 422 VALIDATION ERROR TYPES FOR PYDANTIC VALUE ERROR RESPONSE ! ##################
@@ -229,9 +247,7 @@ class Logger:
247
auth_header = request.headers.get("Authorization")
248
if auth_header:
249
try:
232
- token = auth_header.split(" ")[
233
- 1
234
- ] # Better split by space and take the second part
250
+ token = auth_header.split(" ")[1] # Better split by space and take the second part
251
except IndexError:
252
raise HTTPException(status_code=401, detail="Invalid token")
253
username, _ = self.auth_handler.decode_token(token)
@@ -362,9 +378,7 @@ async def get_logs(session: AsyncSession = Depends(get_db)) -> LogsResponse:
378
auth_handler_instance = AuthHandler() # Initialize your AuthHandler
379
logger_instance = Logger(session, auth_handler_instance)
380
365
- logs = (
366
- await logger_instance.fetch_all_logs()
367
- ) # Assuming fetch_all_logs is an async function
381
+ logs = await logger_instance.fetch_all_logs() # Assuming fetch_all_logs is an async function
382
if logs:
383
return LogsResponse(
384
logs=logs,
@@ -441,12 +455,7 @@ async def get_logs_by_time_range(
455
logs = result.scalars().all()
456
457
if logs:
444
- logs = [
445
- log
446
- for log in logs
447
- if log.timestamp
448
- >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))
449
- ]
458
+ logs = [log for log in logs if log.timestamp >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))]
459
if logs != []:
460
return LogsResponse(
461
logs=logs,
@@ -564,12 +573,7 @@ async def purge_logs_by_time_range(
573
logs = result.scalars().all()
574
575
if logs:
567
- logs = [
568
- log
569
- for log in logs
570
- if log.timestamp
571
- >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))
572
- ]
576
+ logs = [log for log in logs if log.timestamp >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))]
577
if logs != []:
578
for log in logs:
579
await session.delete(log)
@@ -674,8 +678,7 @@ async def get_customer_alert_settings_office365(
678
"""
679
result = await session.execute(
680
select(AlertCreationSettings).filter(
677
- AlertCreationSettings.office365_organization_id
678
- == office365_organization_id,
681
+ AlertCreationSettings.office365_organization_id == office365_organization_id,
682
),
683
)
684
settings = result.scalars().first()
backend/copilot.py
+45
-48
@@ -1,59 +1,56 @@
1
import os
2
3
import uvicorn
4
-from app.auth.utils import AuthHandler
5
-from app.db.db_session import async_engine
6
-from app.db.db_setup import (
7
- create_available_integrations,
8
- create_roles,
9
- create_tables,
10
- ensure_admin_user,
11
- ensure_scheduler_user,
12
- ensure_scheduler_user_removed,
13
-)
14
-from app.middleware.exception_handlers import (
15
- custom_http_exception_handler,
16
- validation_exception_handler,
17
- value_error_handler,
18
-)
19
-from app.middleware.logger import log_requests
20
-from app.routers import (
21
- agents,
22
- alert_creation,
23
- alert_creation_settings,
24
- ask_socfortress,
25
- auth,
26
- connectors,
27
- cortex,
28
- customer_provisioning,
29
- customers,
30
- dfir_iris,
31
- dnstwist,
32
- grafana,
33
- graylog,
34
- healthcheck,
35
- influxdb,
36
- integrations,
37
- logs,
38
- mimecast,
39
- monitoring_alert,
40
- office365,
41
- scheduler,
42
- shuffle,
43
- smtp,
44
- sublime,
45
- threat_intel,
46
- velociraptor,
47
- wazuh_indexer,
48
- wazuh_manager,
49
-)
50
-from app.schedulers.scheduler import init_scheduler
4
from dotenv import load_dotenv
52
-from fastapi import APIRouter, FastAPI, HTTPException
5
+from fastapi import APIRouter
6
+from fastapi import FastAPI
7
+from fastapi import HTTPException
8
from fastapi.exceptions import RequestValidationError
9
from fastapi.middleware.cors import CORSMiddleware
10
from loguru import logger
11
12
+from app.auth.utils import AuthHandler
13
+from app.db.db_session import async_engine
14
+from app.db.db_setup import create_available_integrations
15
+from app.db.db_setup import create_roles
16
+from app.db.db_setup import create_tables
17
+from app.db.db_setup import ensure_admin_user
18
+from app.db.db_setup import ensure_scheduler_user
19
+from app.db.db_setup import ensure_scheduler_user_removed
20
+from app.middleware.exception_handlers import custom_http_exception_handler
21
+from app.middleware.exception_handlers import validation_exception_handler
22
+from app.middleware.exception_handlers import value_error_handler
23
+from app.middleware.logger import log_requests
24
+from app.routers import agents
25
+from app.routers import alert_creation
26
+from app.routers import alert_creation_settings
27
+from app.routers import ask_socfortress
28
+from app.routers import auth
29
+from app.routers import connectors
30
+from app.routers import cortex
31
+from app.routers import customer_provisioning
32
+from app.routers import customers
33
+from app.routers import dfir_iris
34
+from app.routers import dnstwist
35
+from app.routers import grafana
36
+from app.routers import graylog
37
+from app.routers import healthcheck
38
+from app.routers import influxdb
39
+from app.routers import integrations
40
+from app.routers import logs
41
+from app.routers import mimecast
42
+from app.routers import monitoring_alert
43
+from app.routers import office365
44
+from app.routers import scheduler
45
+from app.routers import shuffle
46
+from app.routers import smtp
47
+from app.routers import sublime
48
+from app.routers import threat_intel
49
+from app.routers import velociraptor
50
+from app.routers import wazuh_indexer
51
+from app.routers import wazuh_manager
52
+from app.schedulers.scheduler import init_scheduler
53
+
54
auth_handler = AuthHandler()
55
# Get the `SERVER_IP` from the `.env` file
56
load_dotenv()
git_tasks/Vagrantfile
deleted
-3
@@ -1,3 +0,0 @@
1
-Vagrant.configure("2") do |config|
2
- config.vm.box = "debian/bullseye64"
3
-end