fixed my wrong commit
Davide Di Modica committed
Oct 16, 2023 at 18:05 UTC
550fe69d2844e8156f5a19033ae0403cf6bf8a86
115 files changed
+9652
backend/app/agents/routes/agents.py
new
+120
@@ -0,0 +1,120 @@
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from loguru import logger
4
+from starlette.status import HTTP_401_UNAUTHORIZED
5
+
6
+from app.agents.schema.agents import AgentModifyResponse
7
+from app.agents.schema.agents import AgentsResponse
8
+from app.agents.schema.agents import AgentUpdateCustomerCodeBody
9
+from app.agents.schema.agents import AgentUpdateCustomerCodeResponse
10
+from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse
11
+from app.agents.schema.agents import OutdatedWazuhAgentsResponse
12
+from app.agents.schema.agents import SyncedAgentsResponse
13
+from app.agents.services.modify import mark_agent_criticality
14
+from app.agents.services.status import get_outdated_agents_velociraptor
15
+from app.agents.services.status import get_outdated_agents_wazuh
16
+from app.agents.services.sync import sync_agents
17
+from app.agents.velociraptor.services.agents import delete_agent_velociraptor
18
+from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
19
+from app.agents.wazuh.services.vulnerabilities import collect_agent_vulnerabilities
20
+
21
+# App specific imports
22
+from app.db.db_session import session
23
+from app.db.universal_models import Agents
24
+
25
+agents_router = APIRouter()
26
+
27
+
28
+def verify_admin(user):
29
+ if not user.is_admin:
30
+ raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Unauthorized")
31
+
32
+
33
+@agents_router.get("", response_model=AgentsResponse, description="Get all disabled rules")
34
+async def get_agents() -> AgentsResponse:
35
+ logger.info("Fetching all agents")
36
+ agents = session.query(Agents).all()
37
+ return AgentsResponse(agents=agents, success=True, message="Agents fetched successfully")
38
+
39
+
40
+@agents_router.get("/{agent_id}", response_model=AgentsResponse, description="Get agent by agent_id")
41
+async def get_agent(agent_id: str) -> AgentsResponse:
42
+ logger.info(f"Fetching agent with agent_id: {agent_id}")
43
+ agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()
44
+ if not agent:
45
+ raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
46
+ return AgentsResponse(agents=[agent], success=True, message="Agent fetched successfully")
47
+
48
+
49
+@agents_router.get("/hostname/{hostname}", response_model=AgentsResponse, description="Get agent by hostname")
50
+async def get_agent_by_hostname(hostname: str) -> AgentsResponse:
51
+ logger.info(f"Fetching agent with hostname: {hostname}")
52
+ agent = session.query(Agents).filter(Agents.hostname == hostname).first()
53
+ if not agent:
54
+ raise HTTPException(status_code=404, detail=f"Agent with hostname {hostname} not found")
55
+ return AgentsResponse(agents=[agent], success=True, message="Agent fetched successfully")
56
+
57
+
58
+@agents_router.post("/sync", response_model=SyncedAgentsResponse, description="Sync agents from Wazuh Manager")
59
+async def sync_all_agents() -> SyncedAgentsResponse:
60
+ logger.info("Syncing agents from Wazuh Manager")
61
+ return sync_agents()
62
+
63
+
64
+@agents_router.post("/{agent_id}/critical", response_model=AgentModifyResponse, description="Mark agent as critical")
65
+async def mark_agent_as_critical(agent_id: str) -> AgentModifyResponse:
66
+ logger.info(f"Marking agent {agent_id} as critical")
67
+ return mark_agent_criticality(agent_id, True)
68
+
69
+
70
+@agents_router.post("/{agent_id}/noncritical", response_model=AgentModifyResponse, description="Mark agent as not critical")
71
+async def mark_agent_as_not_critical(agent_id: str) -> AgentModifyResponse:
72
+ logger.info(f"Marking agent {agent_id} as not critical")
73
+ return mark_agent_criticality(agent_id, False)
74
+
75
+
76
+@agents_router.get("/{agent_id}/vulnerabilities", response_model=WazuhAgentVulnerabilitiesResponse, description="Get agent vulnerabilities")
77
+async def get_agent_vulnerabilities(agent_id: str) -> WazuhAgentVulnerabilitiesResponse:
78
+ logger.info(f"Fetching agent {agent_id} vulnerabilities")
79
+ return collect_agent_vulnerabilities(agent_id)
80
+
81
+
82
+@agents_router.get("/wazuh/outdated", response_model=OutdatedWazuhAgentsResponse, description="Get all outdated Wazuh agents")
83
+async def get_outdated_wazuh_agents() -> OutdatedWazuhAgentsResponse:
84
+ logger.info("Fetching all outdated Wazuh agents")
85
+ return get_outdated_agents_wazuh()
86
+
87
+
88
+@agents_router.get(
89
+ "/velociraptor/outdated",
90
+ response_model=OutdatedVelociraptorAgentsResponse,
91
+ description="Get all outdated Velociraptor agents",
92
+)
93
+async def get_outdated_velociraptor_agents() -> OutdatedVelociraptorAgentsResponse:
94
+ logger.info("Fetching all outdated Velociraptor agents")
95
+ return get_outdated_agents_velociraptor()
96
+
97
+
98
+@agents_router.delete("/{agent_id}/delete", response_model=AgentModifyResponse, description="Delete agent")
99
+async def delete_agent(agent_id: str) -> AgentModifyResponse:
100
+ logger.info(f"Deleting agent {agent_id}")
101
+ # delete_agent_db(agent_id)
102
+ # delete_agent_wazuh(agent_id)
103
+ client_id = session.query(Agents).filter(Agents.agent_id == agent_id).first().velociraptor_id
104
+ delete_agent_velociraptor(client_id)
105
+ return {"success": True, "message": f"Agent {agent_id} deleted from database and Wazuh"}
106
+
107
+
108
+@agents_router.put(
109
+ "/{agent_id}/update-customer-code",
110
+ response_model=AgentUpdateCustomerCodeResponse,
111
+ description="Update agent customer code",
112
+)
113
+async def update_agent_customer_code(agent_id: str, body: AgentUpdateCustomerCodeBody) -> AgentUpdateCustomerCodeResponse:
114
+ logger.info(f"Updating agent {agent_id} customer code to {body.customer_code}")
115
+ agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()
116
+ if not agent:
117
+ raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
118
+ agent.customer_code = body.customer_code
119
+ session.commit()
120
+ return {"success": True, "message": f"Agent {agent_id} customer code updated to {body.customer_code}"}
backend/app/agents/schema/agents.py
new
+50
@@ -0,0 +1,50 @@
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
9
+
10
+
11
+class AgentsResponse(BaseModel):
12
+ agents: List[Agents]
13
+ success: bool
14
+ message: str
15
+
16
+
17
+class SyncedAgent(WazuhAgent, VelociraptorAgent):
18
+ pass
19
+
20
+
21
+class SyncedAgentsResponse(BaseModel):
22
+ agents_added: List[SyncedAgent]
23
+ success: bool
24
+ message: str
25
+
26
+
27
+class AgentModifyResponse(BaseModel):
28
+ success: bool
29
+ message: str
30
+
31
+
32
+class OutdatedWazuhAgentsResponse(BaseModel):
33
+ outdated_wazuh_agents: List[Agents]
34
+ success: bool
35
+ message: str
36
+
37
+
38
+class OutdatedVelociraptorAgentsResponse(BaseModel):
39
+ outdated_velociraptor_agents: List[Agents]
40
+ success: bool
41
+ message: str
42
+
43
+
44
+class AgentUpdateCustomerCodeBody(BaseModel):
45
+ customer_code: str = Field(None, description="Customer code to be updated")
46
+
47
+
48
+class AgentUpdateCustomerCodeResponse(BaseModel):
49
+ success: bool
50
+ message: str
backend/app/agents/services/modify.py
new
+34
@@ -0,0 +1,34 @@
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
6
+
7
+
8
+def mark_agent_criticality(agent_id: str, critical: bool):
9
+ """Mark agent as critical or not critical."""
10
+ agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()
11
+ if not agent:
12
+ raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
13
+ agent.critical_asset = critical
14
+ session.commit()
15
+ return {"success": True, "message": f"Agent {agent_id} marked as critical: {critical}"}
16
+
17
+
18
+def delete_agent_db(agent_id: str):
19
+ """Delete agent from database."""
20
+ agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()
21
+ if not agent:
22
+ raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
23
+ session.delete(agent)
24
+ session.commit()
25
+ return {"success": True, "message": f"Agent {agent_id} deleted from database"}
26
+
27
+
28
+def delete_agent_wazuh(agent_id: str):
29
+ """Delete agent from Wazuh service."""
30
+ try:
31
+ wazuh_services.delete_agent(agent_id)
32
+ return {"success": True, "message": f"Agent {agent_id} deleted from Wazuh"}
33
+ except Exception as e:
34
+ raise HTTPException(status_code=500, detail=f"Failed to delete agent {agent_id} from Wazuh: {e}")
backend/app/agents/services/status.py
new
+61
@@ -0,0 +1,61 @@
1
+from typing import List
2
+
3
+from loguru import logger
4
+
5
+from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse
6
+from app.agents.schema.agents import OutdatedWazuhAgentsResponse
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
10
+
11
+
12
+def get_agent(agent_id: str) -> List[Agents]:
13
+ """
14
+ Retrieves a specific agent from the database using its ID.
15
+
16
+ Args:
17
+ agent_id (str): The ID of the agent to retrieve.
18
+
19
+ Returns:
20
+ AgentMetadata: The agent object if found, otherwise None.
21
+ """
22
+ return session.query(Agents).filter(Agents.agent_id == agent_id).first()
23
+
24
+
25
+def get_outdated_agents_wazuh() -> OutdatedWazuhAgentsResponse:
26
+ """
27
+ Retrieves all agents with outdated Wazuh agent versions from the database.
28
+
29
+ Returns:
30
+ List[dict]: A list of dictionaries where each dictionary represents the serialized data of an outdated agent.
31
+ """
32
+ wazuh_manager = get_agent("000")
33
+ if wazuh_manager is None:
34
+ logger.error("Wazuh Manager with agent_id '000' not found.")
35
+ return {"message": "Wazuh Manager with agent_id '000' not found.", "success": False}
36
+
37
+ outdated_wazuh_agents = (
38
+ session.query(Agents).filter(Agents.agent_id != "000", Agents.wazuh_agent_version != wazuh_manager.wazuh_agent_version).all()
39
+ )
40
+ return {"message": "Outdated Wazuh agents fetched successfully.", "success": True, "outdated_wazuh_agents": outdated_wazuh_agents}
41
+
42
+
43
+def get_outdated_agents_velociraptor() -> OutdatedVelociraptorAgentsResponse:
44
+ """
45
+ Retrieves all agents with outdated Velociraptor client versions from the database.
46
+
47
+ Returns:
48
+ List[dict]: A list of dictionaries where each dictionary represents the serialized data of an outdated agent.
49
+ """
50
+ outdated_velociraptor_agents = []
51
+ vql_server_version = "select * from config"
52
+ server_version = UniversalService()._get_server_version(vql_server_version)
53
+ agents = session.query(Agents).all()
54
+ for agent in agents:
55
+ if agent.velociraptor_agent_version != server_version:
56
+ outdated_velociraptor_agents.append(agent)
57
+ return {
58
+ "message": "Outdated Velociraptor agents fetched successfully.",
59
+ "success": True,
60
+ "outdated_velociraptor_agents": outdated_velociraptor_agents,
61
+ }
backend/app/agents/services/sync.py
new
+78
@@ -0,0 +1,78 @@
1
+from typing import List
2
+
3
+from loguru import logger
4
+
5
+import app.agents.velociraptor.services.agents as velociraptor_services
6
+import app.agents.wazuh.services.agents as wazuh_services
7
+from app.agents.schema.agents import SyncedAgent
8
+from app.agents.schema.agents import SyncedAgentsResponse
9
+from app.agents.velociraptor.schema.agents import VelociraptorAgent
10
+from app.agents.wazuh.schema.agents import WazuhAgent
11
+from app.agents.wazuh.schema.agents import WazuhAgentsList
12
+from app.db.db_session import session
13
+from app.db.universal_models import Agents
14
+
15
+
16
+def fetch_wazuh_agents() -> WazuhAgentsList:
17
+ """Fetch agents from Wazuh service."""
18
+ collected_wazuh_agents = wazuh_services.collect_wazuh_agents()
19
+ return WazuhAgentsList(
20
+ agents=collected_wazuh_agents.agents,
21
+ success=collected_wazuh_agents.success,
22
+ message=collected_wazuh_agents.message,
23
+ )
24
+
25
+
26
+def fetch_velociraptor_agent(agent_name: str) -> VelociraptorAgent:
27
+ """Fetch agent details from Velociraptor service."""
28
+ return velociraptor_services.collect_velociraptor_agent(agent_name)
29
+
30
+
31
+def add_agent_to_db(agent: WazuhAgent, client: VelociraptorAgent, customer_code: str):
32
+ """Add new agent to database."""
33
+ new_agent = Agents.create_from_model(agent, client, customer_code)
34
+ session.add(new_agent)
35
+ session.commit()
36
+ logger.info(f"Agent {agent.agent_name} added to the database")
37
+
38
+
39
+def update_agent_in_db(existing_agent: Agents, agent: WazuhAgent, client: VelociraptorAgent, customer_code: str):
40
+ """Update existing agent in database."""
41
+ existing_agent.update_from_model(agent, client, customer_code)
42
+ session.commit()
43
+ logger.info(f"Agent {agent.agent_name} updated in the database")
44
+
45
+
46
+def extract_customer_code(customer_code: str):
47
+ """Extract customer code from agent label."""
48
+ parts = customer_code.split("_")
49
+ return parts[1] if len(parts) > 1 else None
50
+
51
+
52
+def sync_agents() -> SyncedAgentsResponse:
53
+ """Synchronize agents from Wazuh and Velociraptor services."""
54
+ wazuh_agents_list = fetch_wazuh_agents()
55
+ logger.info(f"Collected Wazuh Agents: {wazuh_agents_list}")
56
+
57
+ agents_added_list: List[WazuhAgent] = []
58
+
59
+ for wazuh_agent in wazuh_agents_list.agents:
60
+ logger.info(f"Collecting Velociraptor Agent for {wazuh_agent.agent_name}")
61
+
62
+ velociraptor_agent = fetch_velociraptor_agent(wazuh_agent.agent_name)
63
+
64
+ customer_code = extract_customer_code(wazuh_agent.agent_label)
65
+
66
+ existing_agent = session.query(Agents).filter(Agents.hostname == wazuh_agent.agent_name).first()
67
+
68
+ if existing_agent:
69
+ update_agent_in_db(existing_agent, wazuh_agent, velociraptor_agent, customer_code)
70
+ else:
71
+ add_agent_to_db(wazuh_agent, velociraptor_agent, customer_code)
72
+
73
+ # Combine the wazuh agent and velociraptor agent into one object
74
+ synced_agent = SyncedAgent(**wazuh_agent.dict(), **velociraptor_agent.dict())
75
+ agents_added_list.append(synced_agent)
76
+
77
+ logger.info(f"Agents Added List: {agents_added_list}")
78
+ return SyncedAgentsResponse(success=True, message="Agents synced successfully", agents_added=agents_added_list)
backend/app/agents/velociraptor/schema/agents.py
new
+19
@@ -0,0 +1,19 @@
1
+from datetime import datetime
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
+
7
+
8
+class VelociraptorAgent(BaseModel):
9
+ client_id: Optional[str] = Field("n/a", alias="velociraptor_id")
10
+ client_last_seen: str = Field(..., alias="velociraptor_last_seen")
11
+ client_version: str = Field(..., alias="velociraptor_agent_version")
12
+
13
+ @property
14
+ def client_last_seen_as_datetime(self):
15
+ dt = datetime.strptime(self.client_last_seen, "%Y-%m-%dT%H:%M:%S%z")
16
+ return dt.replace(tzinfo=None)
17
+
18
+ class Config:
19
+ allow_population_by_field_name = True
backend/app/agents/velociraptor/services/agents.py
new
+67
@@ -0,0 +1,67 @@
1
+from datetime import datetime
2
+
3
+from loguru import logger
4
+
5
+from app.agents.schema.agents import AgentsResponse
6
+from app.agents.velociraptor.schema.agents import VelociraptorAgent
7
+from app.connectors.velociraptor.services.artifacts import ArtifactsService
8
+from app.connectors.velociraptor.utils.universal import UniversalService
9
+
10
+
11
+def collect_velociraptor_agent(agent_name: str) -> VelociraptorAgent:
12
+ """
13
+ Retrieves the client ID, last_seen_at and client version based on the agent name from Velociraptor.
14
+
15
+ Args:
16
+ agent_name (str): The name of the agent.
17
+
18
+ Returns:
19
+ str: The client ID if found, None otherwise.
20
+ str: The last seen at timestamp if found, Default timsetamp otherwise.
21
+ """
22
+ logger.info(f"Collecting agent {agent_name} from Velociraptor")
23
+ try:
24
+ client_id = UniversalService().get_client_id(agent_name)["results"][0]["client_id"]
25
+ except (KeyError, IndexError, TypeError) as e:
26
+ logger.error(f"Failed to get client ID for {agent_name}. Error: {e}")
27
+ return VelociraptorAgent(client_id="Unknown", client_last_seen="Unknown", client_version="Unknown")
28
+
29
+ try:
30
+ vql_last_seen_at = f"select last_seen_at from clients(search='host:{agent_name}')"
31
+ last_seen_at = UniversalService()._get_last_seen_timestamp(vql_last_seen_at)
32
+ client_last_seen = datetime.fromtimestamp(
33
+ int(last_seen_at) / 1000000,
34
+ ).strftime(
35
+ "%Y-%m-%dT%H:%M:%S+00:00",
36
+ ) # Converting to string format
37
+ except Exception as e:
38
+ logger.error(f"Failed to get or convert last seen at for {agent_name}. Error: {e}")
39
+ client_last_seen = "1970-01-01T00:00:00+00:00"
40
+
41
+ try:
42
+ vql_client_version = f"select * from clients(search='host:{agent_name}')"
43
+ client_version = UniversalService()._get_client_version(vql_client_version)
44
+ except Exception as e:
45
+ logger.error(f"Failed to get client version for {agent_name}. Error: {e}")
46
+ client_version = "Unknown"
47
+
48
+ return VelociraptorAgent(client_id=client_id, client_last_seen=client_last_seen, client_version=client_version)
49
+
50
+
51
+def delete_agent_velociraptor(client_id: str) -> AgentsResponse:
52
+ """
53
+ Deletes an agent from Velociraptor.
54
+
55
+ Args:
56
+ client_id (str): The client ID of the agent to delete.
57
+
58
+ Returns:
59
+ AgentsResponse: The response object.
60
+ """
61
+ logger.info(f"Deleting agent {client_id} from Velociraptor")
62
+ try:
63
+ ArtifactsService().delete_client(client_id=client_id)
64
+ return AgentsResponse(success=True, message="Agent deleted successfully")
65
+ except Exception as e:
66
+ logger.error(f"Failed to delete agent {client_id}. Error: {e}")
67
+ return AgentsResponse(success=False, message="Failed to delete agent")
backend/app/agents/velociraptor/utils/universal.py
new
+20
@@ -0,0 +1,20 @@
1
+from datetime import datetime
2
+
3
+from loguru import logger
4
+
5
+
6
+def parse_date(date_string: str) -> datetime:
7
+ """
8
+ Parses a date string into a datetime object.
9
+
10
+ Args:
11
+ date_string (str): The date string to parse.
12
+
13
+ Returns:
14
+ datetime: The parsed datetime object.
15
+ """
16
+ try:
17
+ return datetime.strptime(date_string, "%Y-%m-%dT%H:%M:%S+00:00")
18
+ except ValueError:
19
+ logger.info(f"Invalid format for date: {date_string}. Using the epoch time as default.")
20
+ return datetime.strptime("1970-01-01T00:00:00+00:00", "%Y-%m-%dT%H:%M:%S+00:00")
backend/app/agents/wazuh/schema/agents.py
new
+57
@@ -0,0 +1,57 @@
1
+from datetime import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+
8
+
9
+class WazuhAgent(BaseModel):
10
+ agent_id: str = Field(..., alias="agent_id")
11
+ agent_name: str = Field(..., alias="hostname")
12
+ agent_ip: str = Field(..., alias="ip_address")
13
+ agent_os: str = Field(..., alias="os")
14
+ agent_label: str = Field(..., alias="label")
15
+ agent_last_seen: str = Field(..., alias="wazuh_last_seen")
16
+ wazuh_agent_version: str = Field(..., alias="wazuh_agent_version")
17
+
18
+ @property
19
+ def agent_last_seen_as_datetime(self):
20
+ dt = datetime.strptime(self.agent_last_seen, "%Y-%m-%dT%H:%M:%S%z")
21
+ return dt.replace(tzinfo=None)
22
+
23
+ class Config:
24
+ allow_population_by_field_name = True
25
+
26
+
27
+class WazuhAgentsList(BaseModel):
28
+ agents: List[WazuhAgent]
29
+ success: bool
30
+ message: str
31
+
32
+ class Config:
33
+ allow_population_by_field_name = True
34
+
35
+
36
+class WazuhAgentVulnerabilities(BaseModel):
37
+ severity: Optional[str]
38
+ updated: Optional[str]
39
+ version: Optional[str]
40
+ type: Optional[str]
41
+ name: Optional[str]
42
+ external_references: Optional[List[str]]
43
+ condition: Optional[str]
44
+ detection_time: Optional[str]
45
+ cvss3_score: Optional[float]
46
+ published: Optional[str]
47
+ architecture: Optional[str]
48
+ cve: Optional[str]
49
+ status: Optional[str]
50
+ title: Optional[str]
51
+ cvss2_score: Optional[float]
52
+
53
+
54
+class WazuhAgentVulnerabilitiesResponse(BaseModel):
55
+ vulnerabilities: Optional[List[WazuhAgentVulnerabilities]]
56
+ success: bool
57
+ message: str
backend/app/agents/wazuh/services/agents.py
new
+51
@@ -0,0 +1,51 @@
1
+from loguru import logger
2
+
3
+from app.agents.schema.agents import AgentModifyResponse
4
+from app.agents.wazuh.schema.agents import WazuhAgent
5
+from app.agents.wazuh.schema.agents import WazuhAgentsList
6
+from app.connectors.wazuh_manager.utils.universal import send_delete_request
7
+from app.connectors.wazuh_manager.utils.universal import send_get_request
8
+
9
+
10
+def collect_wazuh_agents() -> WazuhAgentsList:
11
+ logger.info("Collecting all agents from Wazuh Manager")
12
+ agents_collected = send_get_request(endpoint="/agents", params={"limit": 1000})
13
+ logger.info(f"Agents collected: {agents_collected}")
14
+ if agents_collected["success"]:
15
+ wazuh_agents_list = []
16
+ for agent in agents_collected["data"]["data"]["affected_items"]:
17
+ os_name = agent.get("os", {}).get("name", "Unknown")
18
+ last_keep_alive = agent.get("lastKeepAlive", "Unknown")
19
+ agent_group_list = agent.get("group", [])
20
+ agent_group = agent_group_list[0] if agent_group_list else "Unknown"
21
+
22
+ wazuh_agent = WazuhAgent(
23
+ agent_id=agent["id"],
24
+ agent_name=agent["name"],
25
+ agent_ip=agent["ip"],
26
+ agent_os=os_name,
27
+ agent_label=agent_group,
28
+ agent_last_seen=last_keep_alive,
29
+ wazuh_agent_version=agent["version"] if "version" in agent else "n/a",
30
+ )
31
+ wazuh_agents_list.append(wazuh_agent)
32
+
33
+ return WazuhAgentsList(agents=wazuh_agents_list, success=True, message="Agents collected successfully")
34
+ else:
35
+ return WazuhAgentsList(agents=[], success=False, message="Failed to collect agents")
36
+
37
+
38
+def delete_agent(agent_id: str) -> AgentModifyResponse:
39
+ """Delete agent from Wazuh Manager."""
40
+ logger.info(f"Deleting agent {agent_id} from Wazuh Manager")
41
+ params = {
42
+ "purge": True,
43
+ "agents_list": [agent_id],
44
+ "status": "all",
45
+ "older_than": "0s",
46
+ }
47
+ agent_deleted = send_delete_request(endpoint="/agents", params=params)
48
+ if agent_deleted["success"]:
49
+ return AgentModifyResponse(success=True, message="Agent deleted successfully")
50
+ else:
51
+ return AgentModifyResponse(success=False, message="Failed to delete agent")
backend/app/agents/wazuh/services/vulnerabilities.py
new
+25
@@ -0,0 +1,25 @@
1
+from typing import List
2
+
3
+from loguru import logger
4
+
5
+from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilities
6
+from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
7
+from app.connectors.wazuh_manager.utils.universal import send_get_request
8
+
9
+
10
+def collect_agent_vulnerabilities(agent_id: str):
11
+ """Collect agent vulnerabilities from Wazuh Manager."""
12
+ logger.info(f"Collecting agent {agent_id} vulnerabilities from Wazuh Manager")
13
+ agent_vulnerabilities = send_get_request(endpoint=f"/vulnerability/{agent_id}")
14
+ if agent_vulnerabilities["success"]:
15
+ processed_vulnerabilities = process_agent_vulnerabilities(agent_vulnerabilities["data"])
16
+ return WazuhAgentVulnerabilitiesResponse(
17
+ vulnerabilities=processed_vulnerabilities,
18
+ success=True,
19
+ message="Vulnerabilities collected successfully",
20
+ )
21
+
22
+
23
+def process_agent_vulnerabilities(agent_vulnerabilities: dict) -> List[WazuhAgentVulnerabilities]:
24
+ vulnerabilities = agent_vulnerabilities.get("data", {}).get("affected_items", [])
25
+ return [WazuhAgentVulnerabilities(**vuln) for vuln in vulnerabilities]
backend/app/auth/models/users.py
new
+77
@@ -0,0 +1,77 @@
1
+import datetime
2
+from enum import Enum
3
+from typing import Optional
4
+
5
+from pydantic import EmailStr
6
+from pydantic import validator
7
+from sqlmodel import Field
8
+from sqlmodel import Relationship
9
+from sqlmodel import SQLModel
10
+
11
+
12
+class Role(SQLModel, table=True):
13
+ id: Optional[int] = Field(primary_key=True)
14
+ name: str = Field(max_length=256)
15
+ description: str = Field(max_length=256)
16
+
17
+ user: Optional["User"] = Relationship(back_populates="role")
18
+
19
+
20
+class User(SQLModel, table=True):
21
+ id: Optional[int] = Field(primary_key=True)
22
+ username: str = Field(index=True)
23
+ password: str = Field(max_length=256, min_length=6)
24
+ email: EmailStr
25
+ created_at: datetime.datetime = datetime.datetime.now()
26
+ role_id: Optional[int] = Field(foreign_key="role.id")
27
+
28
+ smtp: "SMTP" = Relationship(back_populates="user")
29
+ role: Optional["Role"] = Relationship(back_populates="user")
30
+
31
+
32
+# Enum class for role_id 1,2
33
+class RoleEnum(int, Enum):
34
+ admin = 1
35
+ analyst = 2
36
+
37
+
38
+class UserInput(SQLModel):
39
+ username: str
40
+ password: str = Field(
41
+ max_length=256,
42
+ min_length=8,
43
+ regex="^(?=.*[a-z])(?=.*[A-Z])(?=.*\d)[a-zA-Z\d]{8,}$",
44
+ description="Password must be at least 8 characters long and contain at least one uppercase letter, one lowercase letter, and one number",
45
+ )
46
+ email: EmailStr
47
+ role_id: RoleEnum = Field(RoleEnum.analyst, description="Role ID 1: admin, 2: analyst", foreign_key="role.id")
48
+
49
+
50
+class UserLogin(SQLModel):
51
+ username: str
52
+ password: str
53
+
54
+
55
+class SMTP(SQLModel, table=True):
56
+ id: Optional[int] = Field(primary_key=True)
57
+ email: EmailStr
58
+ smtp_password: str = Field(max_length=256)
59
+ smtp_server: str = Field(max_length=256)
60
+ smtp_port: int
61
+ user_id: int = Field(foreign_key="user.id")
62
+
63
+ user: "User" = Relationship(back_populates="smtp")
64
+
65
+
66
+class SMTPInput(SQLModel):
67
+ email: EmailStr
68
+ smtp_password: str = Field(max_length=256)
69
+ smtp_password2: str = Field(max_length=256)
70
+ smtp_server: str = Field(max_length=256)
71
+ smtp_port: int
72
+
73
+ @validator("smtp_password2")
74
+ def password_match(cls, v, values, **kwargs):
75
+ if "smtp_password" in values and v != values["smtp_password"]:
76
+ raise ValueError("passwords don't match")
77
+ return v
backend/app/auth/routes/auth.py
new
+73
@@ -0,0 +1,73 @@
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 status
7
+from fastapi.security import OAuth2PasswordRequestForm
8
+
9
+from app.auth.models.users import User
10
+from app.auth.models.users import UserInput
11
+from app.auth.models.users import UserLogin
12
+from app.auth.schema.auth import Token
13
+from app.auth.schema.auth import UserLoginResponse
14
+from app.auth.schema.auth import UserResponse
15
+from app.auth.services.universal import find_user
16
+from app.auth.services.universal import select_all_users
17
+from app.auth.utils import AuthHandler
18
+from app.db.db_session import session
19
+
20
+ACCESS_TOKEN_EXPIRE_MINUTES = 1440
21
+
22
+user_router = APIRouter()
23
+auth_handler = AuthHandler()
24
+
25
+
26
+@user_router.post("/token", response_model=Token)
27
+async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
28
+ user = auth_handler.authenticate_user(form_data.username, form_data.password)
29
+ if not user:
30
+ raise HTTPException(
31
+ status_code=status.HTTP_401_UNAUTHORIZED,
32
+ detail="Incorrect username or password",
33
+ headers={"WWW-Authenticate": "Bearer"},
34
+ )
35
+ access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
36
+ access_token = auth_handler.encode_token(user.username, access_token_expires)
37
+ return {"access_token": access_token, "token_type": "bearer"}
38
+
39
+
40
+@user_router.get("/refresh", response_model=Token)
41
+async def refresh_token(current_user: User = Depends(auth_handler.get_current_user)):
42
+ access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
43
+ access_token = auth_handler.encode_token(current_user.username, access_token_expires)
44
+ return {"access_token": access_token, "token_type": "bearer"}
45
+
46
+
47
+@user_router.post("/register", response_model=UserResponse, status_code=201, description="Register new user")
48
+def register(user: UserInput):
49
+ users = select_all_users()
50
+ if any(x.username == user.username for x in users):
51
+ raise HTTPException(status_code=400, detail="Username is taken")
52
+ hashed_pwd = auth_handler.get_password_hash(user.password)
53
+ u = User(username=user.username, password=hashed_pwd, email=user.email, role_id=user.role_id)
54
+ session.add(u)
55
+ session.commit()
56
+ return {"message": "User created successfully", "success": True}
57
+
58
+
59
+@user_router.post("/login", response_model=UserLoginResponse, description="Login user", deprecated=True)
60
+def login(user: UserLogin):
61
+ user_found = find_user(user.username)
62
+ if not user_found:
63
+ raise HTTPException(status_code=401, detail="Invalid username and/or password")
64
+ verified = auth_handler.verify_password(user.password, user_found.password)
65
+ if not verified:
66
+ raise HTTPException(status_code=401, detail="Invalid username and/or password")
67
+ token = auth_handler.encode_token(user_found.username)
68
+ return {"token": token, "success": True, "message": "Login successful"}
69
+
70
+
71
+# @user_router.get("/users/me", description="Get current user")
72
+# def get_current_user(user: User = Depends(auth_handler.get_current_user)):
73
+# return user
backend/app/auth/schema/auth.py
new
+21
@@ -0,0 +1,21 @@
1
+from pydantic import BaseModel
2
+
3
+
4
+class UserResponse(BaseModel):
5
+ message: str
6
+ success: bool
7
+
8
+
9
+class UserLoginResponse(BaseModel):
10
+ token: str
11
+ message: str
12
+ success: bool
13
+
14
+
15
+class Token(BaseModel):
16
+ access_token: str
17
+ token_type: str
18
+
19
+
20
+class TokenData(BaseModel):
21
+ username: str | None = None
backend/app/auth/services/universal.py
new
+29
@@ -0,0 +1,29 @@
1
+from sqlmodel import Session
2
+from sqlmodel import select
3
+
4
+from app.auth.models.users import Role
5
+from app.auth.models.users import User
6
+from app.db.db_session import engine
7
+
8
+
9
+def select_all_users():
10
+ with Session(engine) as session:
11
+ statement = select(User)
12
+ res = session.exec(statement).all()
13
+ return res
14
+
15
+
16
+def find_user(name):
17
+ with Session(engine) as session:
18
+ statement = select(User).where(User.username == name)
19
+ return session.exec(statement).first()
20
+
21
+
22
+def get_role(name):
23
+ with Session(engine) as session:
24
+ statement = select(User).where(User.username == name)
25
+ res = session.exec(statement).first()
26
+ # Get the role name
27
+ statement = select(Role).where(Role.id == res.role_id)
28
+ role = session.exec(statement).first()
29
+ return role.name
backend/app/auth/utils.py
new
+122
@@ -0,0 +1,122 @@
1
+from datetime import datetime
2
+from datetime import timedelta
3
+
4
+import jwt
5
+from fastapi import Depends
6
+from fastapi import HTTPException
7
+from fastapi.security import OAuth2PasswordBearer
8
+from fastapi.security import SecurityScopes
9
+from passlib.context import CryptContext
10
+
11
+from app.auth.services.universal import find_user
12
+from app.auth.services.universal import get_role
13
+
14
+
15
+class AuthHandler:
16
+ security = OAuth2PasswordBearer(
17
+ tokenUrl="auth/token",
18
+ scopes={"admin": "Admin users", "analyst": "SOC Analysts"},
19
+ )
20
+ pwd_context = CryptContext(schemes=["bcrypt"])
21
+ secret = "bL4unrkoxtFs1MT6A7Ns2yMLkduyuqrkTxDV9CjlbNc="
22
+
23
+ def get_password_hash(self, password):
24
+ return self.pwd_context.hash(password)
25
+
26
+ def verify_password(self, plain_password, hashed_password):
27
+ return self.pwd_context.verify(plain_password, hashed_password)
28
+
29
+ def authenticate_user(self, username: str, password: str):
30
+ user = find_user(username)
31
+ if not user or not self.verify_password(password, user.password):
32
+ return False
33
+ return user
34
+
35
+ def encode_token(self, username: str, access_token_expires: timedelta = timedelta(minutes=60)):
36
+ payload = {
37
+ "exp": datetime.utcnow() + access_token_expires,
38
+ "iat": datetime.utcnow(),
39
+ "sub": username,
40
+ "scopes": [get_role(username)],
41
+ }
42
+ return jwt.encode(payload, self.secret, algorithm="HS256")
43
+
44
+ def decode_token(self, token):
45
+ try:
46
+ payload = jwt.decode(token, self.secret, algorithms=["HS256"])
47
+ return payload["sub"], payload.get("scopes", [])
48
+ except jwt.ExpiredSignatureError:
49
+ raise HTTPException(status_code=401, detail="Expired signature")
50
+ except jwt.InvalidTokenError:
51
+ raise HTTPException(status_code=401, detail="Invalid token")
52
+
53
+ def get_current_user(self, security_scopes: SecurityScopes, token: str = Depends(security)):
54
+ if security_scopes.scopes:
55
+ authenticate_value = f'Bearer scope="{security_scopes.scope_str}"'
56
+ else:
57
+ authenticate_value = "Bearer"
58
+
59
+ credentials_exception = HTTPException(
60
+ status_code=401,
61
+ detail="Could not validate credentials",
62
+ headers={"WWW-Authenticate": authenticate_value},
63
+ )
64
+
65
+ try:
66
+ username, token_scopes = self.decode_token(token)
67
+ except Exception as e:
68
+ raise HTTPException(
69
+ status_code=401,
70
+ detail=f"Could not decode token: {e}",
71
+ headers={"WWW-Authenticate": authenticate_value},
72
+ )
73
+
74
+ if username is None:
75
+ raise HTTPException(
76
+ status_code=401,
77
+ detail="Username not found in token",
78
+ headers={"WWW-Authenticate": authenticate_value},
79
+ )
80
+ user = find_user(username)
81
+ if user is None:
82
+ raise HTTPException(
83
+ status_code=401,
84
+ detail="User not found",
85
+ headers={"WWW-Authenticate": authenticate_value},
86
+ )
87
+
88
+ for scope in security_scopes.scopes:
89
+ if scope not in token_scopes:
90
+ raise HTTPException(
91
+ status_code=401,
92
+ detail="Not enough permissions",
93
+ headers={"WWW-Authenticate": authenticate_value},
94
+ )
95
+
96
+ return user
97
+
98
+ def return_username_for_logging(self, token: str = Depends(security)):
99
+ username, token_scopes = self.decode_token(token)
100
+ return username
101
+
102
+ def require_any_scope(self, *required_scopes: str):
103
+ async def _require_any_scope(token: str = Depends(self.security)):
104
+ if not token:
105
+ raise HTTPException(
106
+ status_code=401,
107
+ detail="Not authenticated",
108
+ headers={"WWW-Authenticate": "Bearer"},
109
+ )
110
+
111
+ username, token_scopes = self.decode_token(token)
112
+
113
+ if not any(scope in token_scopes for scope in required_scopes):
114
+ raise HTTPException(
115
+ status_code=401,
116
+ detail="Not enough permissions, you don't have any of the required scopes.",
117
+ headers={"WWW-Authenticate": "Bearer"},
118
+ )
119
+
120
+ return username
121
+
122
+ return _require_any_scope
backend/app/auth/utils_backup.py
new
+53
@@ -0,0 +1,53 @@
1
+#### ! COMMENTING OUT FOR PRECOMMIT TO PASS ####
2
+# import datetime
3
+# from typing import Optional
4
+#### ! COMMENTING OUT FOR PRECOMMIT TO PASS ####
5
+# import jwt
6
+# from fastapi import HTTPException
7
+# from fastapi import Security
8
+# from fastapi.security import HTTPAuthorizationCredentials
9
+# from fastapi.security import HTTPBearer
10
+# from fastapi.security import OAuth2PasswordBearer
11
+# from passlib.context import CryptContext
12
+# from starlette import status
13
+
14
+# from app.auth.models.users import User
15
+# from app.auth.services.universal import find_user
16
+
17
+
18
+# class AuthHandler:
19
+# security = HTTPBearer()
20
+# pwd_context = CryptContext(schemes=["bcrypt"])
21
+# secret = "supersecret"
22
+
23
+# def get_password_hash(self, password):
24
+# return self.pwd_context.hash(password)
25
+
26
+# def verify_password(self, pwd, hashed_pwd):
27
+# return self.pwd_context.verify(pwd, hashed_pwd)
28
+
29
+# def encode_token(self, user_id):
30
+# payload = {"exp": datetime.datetime.utcnow() + datetime.timedelta(hours=8), "iat": datetime.datetime.utcnow(), "sub": user_id}
31
+# return jwt.encode(payload, self.secret, algorithm="HS256")
32
+
33
+# def decode_token(self, token):
34
+# try:
35
+# payload = jwt.decode(token, self.secret, algorithms=["HS256"])
36
+# return payload["sub"]
37
+# except jwt.ExpiredSignatureError:
38
+# raise HTTPException(status_code=401, detail="Expired signature")
39
+# except jwt.InvalidTokenError:
40
+# raise HTTPException(status_code=401, detail="Invalid token")
41
+
42
+# def auth_wrapper(self, auth: HTTPAuthorizationCredentials = Security(security)):
43
+# return self.decode_token(auth.credentials)
44
+
45
+# def get_current_user(self, auth: HTTPAuthorizationCredentials = Security(security)):
46
+# credentials_exception = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials")
47
+# username = self.decode_token(auth.credentials)
48
+# if username is None:
49
+# raise credentials_exception
50
+# user = find_user(username)
51
+# if user is None:
52
+# raise credentials_exception
53
+# return user
backend/app/connectors/cortex/routes/analyzers.py
new
+44
@@ -0,0 +1,44 @@
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
+
8
+from app.connectors.cortex.schema.analyzers import AnalyzersResponse
9
+from app.connectors.cortex.schema.analyzers import RunAnalyzerBody
10
+from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse
11
+from app.connectors.cortex.services.analyzers import get_analyzers
12
+from app.connectors.cortex.services.analyzers import run_analyzer
13
+
14
+# App specific imports
15
+
16
+
17
+cortex_analyzer_router = APIRouter()
18
+
19
+
20
+def get_available_analyzers() -> List[str]:
21
+ return get_analyzers().analyzers
22
+
23
+
24
+def verify_analyzer_exists(run_analyzer_body: RunAnalyzerBody) -> RunAnalyzerBody:
25
+ available_analyzers = get_available_analyzers()
26
+ if run_analyzer_body.analyzer_name not in available_analyzers:
27
+ raise HTTPException(status_code=400, detail=f"Analyzer {run_analyzer_body.analyzer_name} does not exist.")
28
+ return run_analyzer_body
29
+
30
+
31
+@cortex_analyzer_router.get("", response_model=AnalyzersResponse, description="Get all analyzers")
32
+async def get_all_analyzers() -> AnalyzersResponse:
33
+ logger.info("Fetching all analyzers")
34
+ return get_analyzers()
35
+
36
+
37
+@cortex_analyzer_router.post("/run", response_model=RunAnalyzerResponse, description="Run an analyzer")
38
+async def run_analyzer_route(run_analyzer_body: RunAnalyzerBody = Depends(verify_analyzer_exists)) -> RunAnalyzerResponse:
39
+ is_valid, data_type = RunAnalyzerBody.is_valid_datatype(run_analyzer_body.analyzer_data)
40
+ if not is_valid:
41
+ raise HTTPException(status_code=400, detail=f"Invalid data type: {data_type}")
42
+
43
+ logger.info(f"Running analyzer {run_analyzer_body.analyzer_name} with data {run_analyzer_body.analyzer_data} of type {data_type}")
44
+ return run_analyzer(run_analyzer_body, data_type)
backend/app/connectors/cortex/schema/analyzers.py
new
+76
@@ -0,0 +1,76 @@
1
+import ipaddress
2
+import re
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
+
9
+from pydantic import BaseModel
10
+from pydantic import Field
11
+from pydantic import validator
12
+
13
+HASH_REGEX = re.compile(r"[a-fA-F\d]{32}|[a-fA-F\d]{64}") # Update this regex to match your specific hash format
14
+DOMAIN_REGEX = re.compile(
15
+ r"^(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.)+[a-z]{2,6}$",
16
+) # Update this regex to match your specific domain format
17
+
18
+
19
+class AnalyzersResponse(BaseModel):
20
+ analyzers: List[str]
21
+ message: str
22
+ success: bool
23
+
24
+
25
+class RunAnalyzerBody(BaseModel):
26
+ analyzer_name: str = Field(..., description="Name of the analyzer to be run.")
27
+ analyzer_data: str = Field(..., description="The Indicator of Compromise (IoC) to be analyzed.")
28
+ data_type: Optional[str] = Field(default=None, description="Data type determined after validation")
29
+
30
+ @validator("analyzer_data", pre=True, always=True)
31
+ def validate_and_set_data_type(cls, value: str, values: dict) -> str:
32
+ is_valid, data_type = cls.is_valid_datatype(value)
33
+ if not is_valid:
34
+ raise ValueError(f"Invalid data type: {data_type}")
35
+ values["data_type"] = data_type
36
+ return value
37
+
38
+ @classmethod
39
+ def is_valid_datatype(cls, value: str) -> Tuple[bool, str]:
40
+ if cls._is_valid_ipv4(value):
41
+ return True, "ip"
42
+ elif cls._is_valid_hash(value):
43
+ return True, "hash"
44
+ elif cls._is_valid_domain(value):
45
+ return True, "domain"
46
+ else:
47
+ return False, "Unknown"
48
+
49
+ @staticmethod
50
+ def _is_valid_ipv4(value: str) -> bool:
51
+ try:
52
+ ipaddress.IPv4Address(value)
53
+ return True
54
+ except ValueError:
55
+ return False
56
+
57
+ @staticmethod
58
+ def _is_valid_hash(value: str) -> bool:
59
+ return bool(HASH_REGEX.match(value))
60
+
61
+ @staticmethod
62
+ def _is_valid_domain(value: str) -> bool:
63
+ return bool(DOMAIN_REGEX.match(value))
64
+
65
+
66
+class RunAnalyzerResponse(BaseModel):
67
+ report: Dict[str, Any]
68
+ message: str
69
+ success: bool
70
+
71
+
72
+class AnalyzerJobData(BaseModel):
73
+ data: str = Field(..., description="The Indicator of Compromise (IoC) to be analyzed.")
74
+ dataType: str = Field(..., description="The type of the IoC (e.g., 'IP', 'hash', 'domain').")
75
+ tlp: int = Field(1, description="Traffic Light Protocol (TLP) level.")
76
+ message: str = Field("custom message sent to analyzer", description="Custom message.")
backend/app/connectors/cortex/services/analyzers.py
new
+71
@@ -0,0 +1,71 @@
1
+# analyzers.py
2
+
3
+from typing import Dict
4
+from typing import List
5
+from typing import Union
6
+
7
+from cortex4py.api import Api
8
+from fastapi import HTTPException
9
+from loguru import logger
10
+
11
+from app.connectors.cortex.schema.analyzers import AnalyzerJobData
12
+from app.connectors.cortex.schema.analyzers import AnalyzersResponse
13
+from app.connectors.cortex.schema.analyzers import RunAnalyzerBody
14
+from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse
15
+from app.connectors.cortex.utils.universal import (
16
+ create_cortex_client, # Importing create_cortex_client
17
+)
18
+from app.connectors.cortex.utils.universal import (
19
+ run_and_wait_for_analyzer, # Importing from universal.py
20
+)
21
+
22
+
23
+def fetch_analyzers(api: Api) -> List[Dict]:
24
+ return api.analyzers.find_all({}, range="all")
25
+
26
+
27
+def extract_analyzer_names(analyzers: List[Dict]) -> List[str]:
28
+ try:
29
+ return [analyzer.name for analyzer in analyzers]
30
+ except Exception as e:
31
+ logger.error(f"Error processing analyzers: {e}")
32
+ raise HTTPException(status_code=500, detail=f"Error processing analyzers: {e}")
33
+
34
+
35
+def init_cortex_client() -> Union[Api, None]:
36
+ return create_cortex_client("Cortex")
37
+
38
+
39
+def handle_api_initialization(api: Union[Api, None]) -> Api:
40
+ if api is None:
41
+ logger.error("API initialization failed")
42
+ raise HTTPException(status_code=500, detail="API initialization failed")
43
+ return api
44
+
45
+
46
+def get_analyzers() -> AnalyzersResponse:
47
+ api = init_cortex_client()
48
+ handle_api_initialization(api)
49
+
50
+ analyzers = fetch_analyzers(api)
51
+ analyzer_names = extract_analyzer_names(analyzers)
52
+
53
+ return AnalyzersResponse(success=True, message="Successfully fetched analyzers", analyzers=analyzer_names)
54
+
55
+
56
+def run_analyzer(run_analyzer_body: RunAnalyzerBody, data_type: str) -> RunAnalyzerResponse:
57
+ api = init_cortex_client()
58
+ handle_api_initialization(api)
59
+
60
+ analyzer_name = run_analyzer_body.analyzer_name
61
+ analyzer_data = run_analyzer_body.analyzer_data
62
+ logger.info(f"Running analyzer {analyzer_name} with data {analyzer_data} of type {data_type}")
63
+ job_data = AnalyzerJobData(data=analyzer_data, dataType=data_type)
64
+
65
+ result = run_and_wait_for_analyzer(analyzer_name=analyzer_name, job_data=job_data)
66
+
67
+ if result is None:
68
+ logger.error(f"Failed to run analyzer {analyzer_name}")
69
+ raise HTTPException(status_code=500, detail=f"Failed to run analyzer {analyzer_name}")
70
+
71
+ return RunAnalyzerResponse(success=True, message="Successfully ran analyzer", report=result)
backend/app/connectors/cortex/utils/universal.py
new
+112
@@ -0,0 +1,112 @@
1
+import time
2
+import traceback
3
+from typing import Any
4
+from typing import Dict
5
+
6
+from cortex4py.api import Api
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
+
12
+
13
+def verify_cortex_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
14
+ """
15
+ Verifies the connection to Cortex service.
16
+
17
+ Returns:
18
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
19
+ """
20
+ logger.info(f"Verifying the Cortex connection to {attributes['connector_url']}")
21
+
22
+ try:
23
+ api = Api(attributes["connector_url"], attributes["connector_api_key"], verify_cert=False)
24
+ # Get Cortex Status
25
+ status = api.status
26
+ if status:
27
+ logger.debug("Cortex connection successful")
28
+ return {"connectionSuccessful": True, "message": "Cortex connection successful"}
29
+ else:
30
+ logger.error(f"Connection to {attributes['connector_url']} failed with error.")
31
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error."}
32
+ except Exception as e:
33
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")
34
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
35
+
36
+
37
+def verify_cortex_connection(connector_name: str) -> str:
38
+ """
39
+ Returns the authentication token for the Cortex service.
40
+
41
+ Returns:
42
+ str: Authentication token for the Cortex service.
43
+ """
44
+ attributes = get_connector_info_from_db(connector_name)
45
+ if attributes is None:
46
+ logger.error("No Cortex connector found in the database")
47
+ return None
48
+ return verify_cortex_credentials(attributes)
49
+
50
+
51
+def create_cortex_client(connector_name: str) -> Api:
52
+ """
53
+ Returns an Cortex client for the Wazuh Indexer service.
54
+
55
+ Returns:
56
+ Cortex: Cortex client for the Cortex service.
57
+ """
58
+ attributes = get_connector_info_from_db(connector_name)
59
+ if attributes is None:
60
+ logger.error("No Wazuh Indexer connector found in the database")
61
+ return None
62
+ return Api(attributes["connector_url"], attributes["connector_api_key"], verify_cert=False)
63
+
64
+
65
+def run_and_wait_for_analyzer(analyzer_name: str, job_data: AnalyzerJobData) -> Dict[str, Any]:
66
+ api = create_cortex_client("Cortex") # Create Api object
67
+ if api is None:
68
+ return {"success": False, "message": "API initialization failed"}
69
+ try:
70
+ job = api.analyzers.run_by_name(analyzer_name, job_data.dict(), force=1)
71
+ return monitor_analyzer_job(api, job)
72
+ except Exception as e:
73
+ logger.error(f"Error running analyzer {analyzer_name}: {e}")
74
+ logger.debug(f"job_data dict: {job_data.dict()}")
75
+ logger.debug(f"Exception details: {traceback.format_exc()}")
76
+ logger.debug(f"Error running analyzer {analyzer_name}: {e}", exc_info=True)
77
+ return {"success": False, "message": f"Error running analyzer {analyzer_name}: {e}"}
78
+
79
+
80
+def monitor_analyzer_job(api: Api, job: Any) -> Dict[str, Any]:
81
+ r_json = job.json()
82
+ job_id = r_json["id"]
83
+ logger.info(f"Job ID is: {job_id}")
84
+
85
+ job_state = r_json["status"]
86
+ timer = 0
87
+
88
+ while job_state != "Success":
89
+ if timer == 60:
90
+ logger.error("Job failed to complete after 5 minutes.")
91
+ return {"success": False, "message": "Job timed out"}
92
+
93
+ timer += 1
94
+ logger.info(f"Timer is: {timer}")
95
+
96
+ if job_state == "Failure":
97
+ error_message = r_json["errorMessage"]
98
+ logger.error(f"Cortex Failure: {error_message}")
99
+ return {"success": False, "message": f"Analyzer failed: {error_message}"}
100
+
101
+ time.sleep(5)
102
+ followup_request = api.jobs.get_by_id(job_id)
103
+ r_json = followup_request.json()
104
+ job_state = r_json["status"]
105
+
106
+ return retrieve_final_report(api, job_id)
107
+
108
+
109
+def retrieve_final_report(api: Api, job_id: str) -> Dict[str, Any]:
110
+ report = api.jobs.get_report(job_id).report
111
+ final_report = report["full"]
112
+ return {"success": True, "message": "Analyzer ran successfully", "report": final_report}
backend/app/connectors/dfir_iris/routes/alerts.py
new
+47
@@ -0,0 +1,47 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from loguru import logger
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.services.alerts import bookmark_alert
10
+from app.connectors.dfir_iris.services.alerts import get_alerts
11
+from app.connectors.dfir_iris.services.alerts import get_bookmarked_alerts
12
+from app.connectors.dfir_iris.utils.universal import check_alert_exists
13
+
14
+# App specific imports
15
+
16
+
17
+def verify_alert_exists(alert_id: str) -> str:
18
+ if not check_alert_exists(alert_id):
19
+ raise HTTPException(status_code=400, detail=f"Alert {alert_id} does not exist.")
20
+ return alert_id
21
+
22
+
23
+dfir_iris_alerts_router = APIRouter()
24
+
25
+
26
+@dfir_iris_alerts_router.get("", response_model=AlertsResponse, description="Get all alerts")
27
+async def get_all_alerts() -> AlertsResponse:
28
+ logger.info("Fetching all alerts")
29
+ return get_alerts()
30
+
31
+
32
+@dfir_iris_alerts_router.get("/bookmark", response_model=BookmarkedAlertsResponse, description="Get all bookmarked alerts")
33
+async def get_all_bookmarked_alerts() -> BookmarkedAlertsResponse:
34
+ logger.info("Fetching all bookmarked alerts")
35
+ return get_bookmarked_alerts()
36
+
37
+
38
+@dfir_iris_alerts_router.post("/bookmark/{alert_id}", response_model=AlertResponse, description="Bookmark an alert")
39
+async def bookmark_alert_route(alert_id: str = Depends(verify_alert_exists)) -> AlertResponse:
40
+ logger.info(f"Bookmarking alert {alert_id}")
41
+ return bookmark_alert(alert_id, bookmarked=True)
42
+
43
+
44
+@dfir_iris_alerts_router.delete("/bookmark/{alert_id}", response_model=AlertResponse, description="Unbookmark an alert")
45
+async def unbookmark_alert_route(alert_id: str = Depends(verify_alert_exists)) -> AlertResponse:
46
+ logger.info(f"Unbookmarking alert {alert_id}")
47
+ return bookmark_alert(alert_id, bookmarked=False)
backend/app/connectors/dfir_iris/routes/assets.py
new
+25
@@ -0,0 +1,25 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from loguru import logger
5
+
6
+from app.connectors.dfir_iris.schema.assets import AssetResponse
7
+from app.connectors.dfir_iris.services.assets import get_case_assets
8
+from app.connectors.dfir_iris.utils.universal import check_case_exists
9
+
10
+# App specific imports
11
+
12
+
13
+def verify_case_exists(case_id: int) -> int:
14
+ if not check_case_exists(case_id):
15
+ raise HTTPException(status_code=400, detail=f"Case {case_id} does not exist.")
16
+ return case_id
17
+
18
+
19
+assets_router = APIRouter()
20
+
21
+
22
+@assets_router.get("/{case_id}", response_model=AssetResponse, description="Get all assets for a case")
23
+async def get_case_assets_route(case_id: int = Depends(verify_case_exists)) -> AssetResponse:
24
+ logger.info(f"Fetching assets for case {case_id}")
25
+ return get_case_assets(case_id)
backend/app/connectors/dfir_iris/routes/cases.py
new
+56
@@ -0,0 +1,56 @@
1
+from datetime import timedelta
2
+
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
6
+from loguru import logger
7
+
8
+from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody
9
+from app.connectors.dfir_iris.schema.cases import CaseResponse
10
+from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse
11
+from app.connectors.dfir_iris.schema.cases import SingleCaseBody
12
+from app.connectors.dfir_iris.schema.cases import SingleCaseResponse
13
+from app.connectors.dfir_iris.schema.cases import TimeUnit
14
+from app.connectors.dfir_iris.services.cases import get_all_cases
15
+from app.connectors.dfir_iris.services.cases import get_cases_older_than
16
+from app.connectors.dfir_iris.services.cases import get_single_case
17
+from app.connectors.dfir_iris.utils.universal import check_case_exists
18
+
19
+
20
+def verify_case_exists(case_id: int) -> int:
21
+ if not check_case_exists(case_id):
22
+ raise HTTPException(status_code=400, detail=f"Case {case_id} does not exist.")
23
+ return case_id
24
+
25
+
26
+cases_router = APIRouter()
27
+
28
+
29
+def get_timedelta(older_than: int, time_unit: TimeUnit) -> CaseOlderThanBody:
30
+ delta = None
31
+ if time_unit == TimeUnit.HOURS:
32
+ delta = timedelta(hours=older_than)
33
+ elif time_unit == TimeUnit.DAYS:
34
+ delta = timedelta(days=older_than)
35
+ elif time_unit == TimeUnit.WEEKS:
36
+ delta = timedelta(weeks=older_than)
37
+ return CaseOlderThanBody(older_than=delta, time_unit=time_unit)
38
+
39
+
40
+@cases_router.get("", response_model=CaseResponse, description="Get all cases")
41
+async def get_cases_route() -> CaseResponse:
42
+ logger.info("Fetching all cases")
43
+ return get_all_cases()
44
+
45
+
46
+@cases_router.post("/older_than", response_model=CasesBreachedResponse, description="Get all cases older than a specified date")
47
+async def get_cases_older_than_route(case_older_than_body: CaseOlderThanBody = Depends(get_timedelta)) -> CaseResponse:
48
+ logger.info(f"Fetching all cases older than {case_older_than_body.older_than} ({case_older_than_body.time_unit.value})")
49
+ return get_cases_older_than(case_older_than_body)
50
+
51
+
52
+@cases_router.get("/{case_id}", response_model=SingleCaseResponse, description="Get a single case")
53
+async def get_single_case_route(case_id: int = Depends(verify_case_exists)) -> SingleCaseResponse:
54
+ logger.info(f"Fetching case {case_id}")
55
+ single_case_body = SingleCaseBody(case_id=case_id)
56
+ return get_single_case(single_case_body.case_id)
backend/app/connectors/dfir_iris/routes/notes.py
new
+35
@@ -0,0 +1,35 @@
1
+from typing import Optional
2
+
3
+from fastapi import APIRouter
4
+from fastapi import Depends
5
+from fastapi import HTTPException
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 NotesResponse
11
+from app.connectors.dfir_iris.services.notes import create_case_note
12
+from app.connectors.dfir_iris.services.notes import get_case_notes
13
+from app.connectors.dfir_iris.utils.universal import check_case_exists
14
+
15
+
16
+def verify_case_exists(case_id: int) -> int:
17
+ if not check_case_exists(case_id):
18
+ raise HTTPException(status_code=400, detail=f"Case {case_id} does not exist.")
19
+ return case_id
20
+
21
+
22
+notes_router = APIRouter()
23
+
24
+
25
+@notes_router.get("/{case_id}", response_model=NotesResponse, description="Get all notes for a case")
26
+async def get_case_notes_route(case_id: int = Depends(verify_case_exists), search_term: Optional[str] = "%") -> NotesResponse:
27
+ logger.info(f"Fetching notes for case {case_id}")
28
+ return get_case_notes(case_id, search_term)
29
+
30
+
31
+@notes_router.post("/{case_id}", response_model=NoteCreationResponse, description="Create a note for a case")
32
+async def create_case_note_route(case_id: int, note_creation_body: NoteCreationBody) -> NoteCreationResponse:
33
+ verify_case_exists(case_id)
34
+ logger.info(f"Creating a note for case {case_id}")
35
+ return create_case_note(case_id, note_creation_body)
backend/app/connectors/dfir_iris/routes/users.py
new
+39
@@ -0,0 +1,39 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from loguru import logger
5
+
6
+from app.connectors.dfir_iris.schema.alerts import AlertResponse
7
+from app.connectors.dfir_iris.schema.users import User
8
+from app.connectors.dfir_iris.schema.users import UsersResponse
9
+from app.connectors.dfir_iris.services.users import assign_user_to_alert
10
+from app.connectors.dfir_iris.services.users import get_users
11
+from app.connectors.dfir_iris.utils.universal import check_alert_exists
12
+from app.connectors.dfir_iris.utils.universal import check_user_exists
13
+
14
+
15
+def verify_user_exists(user_id: int) -> int:
16
+ if not check_user_exists(user_id):
17
+ raise HTTPException(status_code=400, detail=f"User {user_id} does not exist.")
18
+ return user_id
19
+
20
+
21
+def verify_alert_exists(alert_id: str) -> str:
22
+ if not check_alert_exists(alert_id):
23
+ raise HTTPException(status_code=400, detail=f"Alert {alert_id} does not exist.")
24
+ return alert_id
25
+
26
+
27
+dfir_iris_users_router = APIRouter()
28
+
29
+
30
+@dfir_iris_users_router.get("", response_model=UsersResponse, description="Get all users")
31
+async def get_all_users() -> UsersResponse:
32
+ logger.info("Fetching all users")
33
+ return get_users()
34
+
35
+
36
+@dfir_iris_users_router.post("/assign/{alert_id}/{user_id}", response_model=AlertResponse, description="Assign a user to an alert")
37
+async def assign_user_to_alert_route(alert_id: str = Depends(verify_alert_exists), user_id: int = Depends(verify_user_exists)) -> User:
38
+ logger.info(f"Assigning user {user_id} to alert {alert_id}")
39
+ return assign_user_to_alert(alert_id, user_id)
backend/app/connectors/dfir_iris/schema/alerts.py
new
+25
@@ -0,0 +1,25 @@
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+from typing import Optional
5
+
6
+from pydantic import BaseModel
7
+from pydantic import Field
8
+
9
+
10
+class AlertsResponse(BaseModel):
11
+ alerts: Optional[List[Dict[str, Any]]] = Field([], description="The alerts returned from the search.")
12
+ message: str
13
+ success: bool
14
+
15
+
16
+class AlertResponse(BaseModel):
17
+ alert: Optional[Dict[str, Any]] = Field({}, description="The alert returned from the search.")
18
+ message: str
19
+ success: bool
20
+
21
+
22
+class BookmarkedAlertsResponse(BaseModel):
23
+ bookmarked_alerts: Optional[List[Dict[str, Any]]] = Field([], description="The alerts returned from the search.")
24
+ message: str
25
+ success: bool
backend/app/connectors/dfir_iris/schema/assets.py
new
+40
@@ -0,0 +1,40 @@
1
+from typing import List
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+
6
+
7
+class AssetState(BaseModel):
8
+ object_last_update: str
9
+ object_state: int
10
+
11
+
12
+class Asset(BaseModel):
13
+ analysis_status: str
14
+ analysis_status_id: int
15
+ asset_compromise_status_id: int
16
+ asset_description: str
17
+ asset_domain: str
18
+ asset_icon_compromised: str
19
+ asset_icon_not_compromised: str
20
+ asset_id: int
21
+ asset_ip: str
22
+ asset_name: str
23
+ asset_tags: str
24
+ asset_type: str
25
+ asset_type_id: int
26
+ asset_uuid: str
27
+ ioc_links: Optional[None]
28
+ link: List
29
+
30
+
31
+class AssetData(BaseModel):
32
+ assets: List[Asset]
33
+ state: AssetState
34
+
35
+
36
+class AssetResponse(BaseModel):
37
+ assets: List[Asset]
38
+ state: AssetState
39
+ message: str
40
+ success: bool
backend/app/connectors/dfir_iris/schema/cases.py
new
+97
@@ -0,0 +1,97 @@
1
+from datetime import timedelta
2
+from enum import Enum
3
+from typing import Dict
4
+from typing import List
5
+from typing import Optional
6
+from typing import Union
7
+
8
+from pydantic import BaseModel
9
+from pydantic import Field
10
+
11
+
12
+class CaseModel(BaseModel):
13
+ access_level: int
14
+ case_close_date: str
15
+ case_description: str
16
+ case_id: int
17
+ case_name: str
18
+ case_open_date: str
19
+ case_soc_id: str
20
+ case_uuid: str
21
+ classification: Optional[str]
22
+ classification_id: Optional[int]
23
+ client_name: str
24
+ opened_by: str
25
+ opened_by_user_id: int
26
+ owner: str
27
+ owner_id: int
28
+ state_id: int
29
+ state_name: str
30
+
31
+
32
+class CaseResponse(BaseModel):
33
+ cases: List[CaseModel]
34
+ message: str
35
+ success: bool
36
+
37
+
38
+class ModificationHistoryItem(BaseModel):
39
+ action: str
40
+ user: str
41
+ user_id: int
42
+
43
+
44
+class SingleCaseModel(BaseModel):
45
+ case_description: str
46
+ case_id: int
47
+ case_name: str
48
+ case_soc_id: str
49
+ case_tags: Optional[str]
50
+ case_uuid: str
51
+ classification: Optional[Union[str, None]]
52
+ classification_id: Optional[Union[int, None]]
53
+ close_date: Optional[Union[str, None]]
54
+ custom_attributes: Optional[Union[str, None]]
55
+ customer_id: int
56
+ customer_name: str
57
+ initial_date: str
58
+ modification_history: Dict[str, ModificationHistoryItem]
59
+ open_by_user: str
60
+ open_by_user_id: int
61
+ open_date: str
62
+ owner: str
63
+ owner_id: int
64
+ protagonists: List[str]
65
+ reviewer: Optional[Union[str, None]]
66
+ reviewer_id: Optional[Union[int, None]]
67
+ state_id: int
68
+ state_name: str
69
+ status_id: int
70
+ status_name: str
71
+
72
+
73
+class SingleCaseBody(BaseModel):
74
+ case_id: int
75
+
76
+
77
+class SingleCaseResponse(BaseModel):
78
+ case: SingleCaseModel
79
+ message: str
80
+ success: bool
81
+
82
+
83
+class TimeUnit(str, Enum):
84
+ HOURS = "hours"
85
+ DAYS = "days"
86
+ WEEKS = "weeks"
87
+
88
+
89
+class CaseOlderThanBody(BaseModel):
90
+ older_than: timedelta = Field(..., description="Amount of time to filter cases by")
91
+ time_unit: TimeUnit
92
+
93
+
94
+class CasesBreachedResponse(BaseModel):
95
+ cases_breached: List[CaseModel]
96
+ message: str
97
+ success: bool
backend/app/connectors/dfir_iris/schema/notes.py
new
+68
@@ -0,0 +1,68 @@
1
+from typing import Dict
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+
8
+
9
+class CustomAttributes(BaseModel):
10
+ # Define additional fields if custom_attributes contains specific keys
11
+ pass
12
+
13
+
14
+class NoteDetails(BaseModel):
15
+ custom_attributes: CustomAttributes
16
+ group_id: int
17
+ group_title: str
18
+ group_uuid: str
19
+ note_content: str
20
+ note_creationdate: str
21
+ note_id: int
22
+ note_lastupdate: str
23
+ note_title: str
24
+ note_uuid: str
25
+
26
+
27
+class NoteDetailsResponse(BaseModel):
28
+ note_details: NoteDetails
29
+ message: str
30
+ success: bool
31
+
32
+
33
+class NoteItem(BaseModel):
34
+ note_details: NoteDetails
35
+ note_id: int
36
+ note_title: str
37
+
38
+
39
+class NotesResponse(BaseModel):
40
+ notes: List[NoteItem]
41
+ message: str
42
+ success: bool
43
+
44
+
45
+class NotesQueryParams(BaseModel):
46
+ case_id: int
47
+ search_term: Optional[str] = Field("%", description="Search term to filter notes by. Defaults to wildcard search (%).")
48
+
49
+
50
+class NoteCreationBody(BaseModel):
51
+ note_title: str = Field(..., description="Title of the note to be created.")
52
+ note_content: str = Field(..., description="Content of the note to be created.")
53
+
54
+
55
+class NoteAttributes(BaseModel):
56
+ custom_attributes: Dict[str, str] = Field(...)
57
+ note_content: str = Field(...)
58
+ note_creationdate: str = Field(...)
59
+ note_id: int = Field(...)
60
+ note_lastupdate: str = Field(...)
61
+ note_title: str = Field(...)
62
+ note_uuid: str = Field(...)
63
+
64
+
65
+class NoteCreationResponse(BaseModel):
66
+ message: str = Field(...)
67
+ note: NoteAttributes = Field(...)
68
+ success: bool = Field(...)
backend/app/connectors/dfir_iris/schema/users.py
new
+17
@@ -0,0 +1,17 @@
1
+from typing import List
2
+
3
+from pydantic import BaseModel
4
+
5
+
6
+class User(BaseModel):
7
+ user_active: bool
8
+ user_id: int
9
+ user_login: str
10
+ user_name: str
11
+ user_uuid: str
12
+
13
+
14
+class UsersResponse(BaseModel):
15
+ message: str
16
+ success: bool
17
+ users: List[User]
backend/app/connectors/dfir_iris/services/alerts.py
new
+30
@@ -0,0 +1,30 @@
1
+from app.connectors.dfir_iris.schema.alerts import AlertResponse
2
+from app.connectors.dfir_iris.schema.alerts import AlertsResponse
3
+from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse
4
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
5
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
6
+
7
+
8
+def get_alerts() -> AlertsResponse:
9
+ client, alert = initialize_client_and_alert("DFIR-IRIS")
10
+ result = fetch_and_validate_data(client, alert.filter_alerts)
11
+ return AlertsResponse(success=True, message="Successfully fetched alerts", alerts=result["data"]["alerts"])
12
+
13
+
14
+def bookmark_alert(alert_id: str, bookmarked: bool) -> AlertResponse:
15
+ client, alert = initialize_client_and_alert("DFIR-IRIS")
16
+ if bookmarked:
17
+ result = fetch_and_validate_data(client, alert.update_alert, alert_id, {"alert_tags": "bookmarked"})
18
+ return AlertResponse(success=True, message="Successfully bookmarked alert", alert=result["data"])
19
+ result = fetch_and_validate_data(client, alert.update_alert, alert_id, {"alert_tags": ""})
20
+ return AlertResponse(success=True, message="Successfully removed bookmark from alert", alert=result["data"])
21
+
22
+
23
+def get_bookmarked_alerts() -> BookmarkedAlertsResponse:
24
+ alerts = get_alerts().alerts
25
+ bookmarked_alerts = []
26
+ for alert in alerts:
27
+ if alert["alert_tags"] is not None and "bookmarked" in alert["alert_tags"]:
28
+ bookmarked_alerts.append(alert)
29
+
30
+ return BookmarkedAlertsResponse(success=True, message="Successfully fetched bookmarked alerts", bookmarked_alerts=bookmarked_alerts)
backend/app/connectors/dfir_iris/services/assets.py
new
+20
@@ -0,0 +1,20 @@
1
+from app.connectors.dfir_iris.schema.assets import Asset
2
+from app.connectors.dfir_iris.schema.assets import AssetResponse
3
+from app.connectors.dfir_iris.schema.assets import AssetState
4
+from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
5
+from app.connectors.dfir_iris.utils.universal import initialize_client_and_case
6
+
7
+
8
+def get_case_assets(case_id: int) -> AssetResponse:
9
+ client, case = initialize_client_and_case("DFIR-IRIS")
10
+ result = fetch_and_validate_data(client, case.list_assets, case_id)
11
+
12
+ asset_list = result["data"]["assets"]
13
+ state_data = result["data"]["state"]
14
+
15
+ return AssetResponse(
16
+ success=True,
17
+ message="Successfully fetched assets for case",
18
+ assets=[Asset(**asset) for asset in asset_list], # List[Asset]
19
+ state=AssetState(**state_data), # AssetState
20
+ )
backend/app/connectors/dfir_iris/services/cases.py
new
+97
@@ -0,0 +1,97 @@
1
+from datetime import datetime
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
+
9
+from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody
10
+from app.connectors.dfir_iris.schema.cases import CaseResponse
11
+from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse
12
+from app.connectors.dfir_iris.schema.cases import SingleCaseBody
13
+from app.connectors.dfir_iris.schema.cases import SingleCaseResponse
14
+from app.connectors.dfir_iris.utils.universal import create_dfir_iris_client
15
+from app.connectors.dfir_iris.utils.universal import fetch_and_parse_data
16
+
17
+
18
+def get_client_and_cases() -> Dict:
19
+ """
20
+ Initialize the client session and fetch all cases.
21
+
22
+ Returns:
23
+ Dictionary containing the success status and either the case data or an error message.
24
+ """
25
+ dfir_iris_client = create_dfir_iris_client("DFIR-IRIS")
26
+ case = Case(session=dfir_iris_client)
27
+ logger.info("Fetching all cases after getting session")
28
+ result = fetch_and_parse_data(dfir_iris_client, case.list_cases)
29
+ return result
30
+
31
+
32
+def filter_open_cases(cases: List[Dict]) -> List[Dict]:
33
+ """
34
+ Filters out cases that are still open.
35
+
36
+ Args:
37
+ cases (List): List of all cases.
38
+
39
+ Returns:
40
+ List of cases that are still open.
41
+ """
42
+ return [case for case in cases if case["case_close_date"] == ""]
43
+
44
+
45
+def filter_cases_older_than(cases: List[Dict], older_than: datetime) -> List[Dict]:
46
+ """
47
+ Filters out cases that are older than the specified time.
48
+
49
+ Args:
50
+ cases (List): List of all cases.
51
+ older_than (datetime): The datetime to filter by.
52
+
53
+ Returns:
54
+ List of cases that are older than the specified time.
55
+ """
56
+ current_time = datetime.now()
57
+ filtered_cases = []
58
+ for case in cases:
59
+ case_open_date = (
60
+ datetime.strptime(case["case_open_date"], "%m/%d/%Y")
61
+ if not isinstance(case["case_open_date"], datetime)
62
+ else case["case_open_date"]
63
+ )
64
+ if case_open_date < current_time - older_than:
65
+ case["case_open_date"] = case_open_date.strftime("%m/%d/%Y") # Convert back to string to match the model
66
+ filtered_cases.append(case)
67
+ return filtered_cases
68
+
69
+
70
+def get_all_cases() -> CaseResponse:
71
+ result = get_client_and_cases()
72
+ if not result["success"]:
73
+ logger.error(f"Failed to get all cases: {result['message']}")
74
+ return HTTPException(status_code=500, detail=f"Failed to get all cases: {result['message']}")
75
+ return CaseResponse(success=True, message="Successfully fetched all cases", cases=result["data"])
76
+
77
+
78
+def get_cases_older_than(case_older_than_body: CaseOlderThanBody) -> CasesBreachedResponse:
79
+ result = get_client_and_cases()
80
+ if not result["success"]:
81
+ logger.error(f"Failed to get all cases: {result['message']}")
82
+ return HTTPException(status_code=500, detail=f"Failed to get all cases: {result['message']}")
83
+
84
+ open_cases = filter_open_cases(result["data"])
85
+ breached_cases = filter_cases_older_than(open_cases, case_older_than_body.older_than)
86
+ return CasesBreachedResponse(
87
+ success=True,
88
+ message=f"Successfully fetched all cases older than {case_older_than_body.older_than}",
89
+ cases_breached=breached_cases,
90
+ )
91
+
92
+
93
+def get_single_case(case_id: SingleCaseBody) -> SingleCaseResponse:
94
+ dfir_iris_client = create_dfir_iris_client("DFIR-IRIS")
95
+ case = Case(session=dfir_iris_client)
96
+ result = fetch_and_parse_data(dfir_iris_client, case.get_case, case_id)
97
+ return SingleCaseResponse(success=True, message="Successfully fetched single case", case=result["data"])
backend/app/connectors/dfir_iris/services/notes.py
new
+59
@@ -0,0 +1,59 @@
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
+def process_notes(notes: List[Dict], case_id: int) -> List[Dict]:
18
+ processed_notes = []
19
+ for note in notes:
20
+ note_details = get_case_note_details(note["note_id"], case_id)
21
+ logger.info(f"Note details: {note_details}")
22
+ note["note_details"] = note_details.note_details
23
+ processed_notes.append(note)
24
+ return processed_notes
25
+
26
+
27
+def get_case_notes(case_id: int, search_term: str) -> NotesResponse:
28
+ client, case = initialize_client_and_case("DFIR-IRIS")
29
+ result = fetch_and_validate_data(client, case.search_notes, search_term, case_id)
30
+ processed_notes = process_notes(result["data"], case_id)
31
+ return NotesResponse(success=True, message="Successfully fetched notes for case", notes=processed_notes)
32
+
33
+
34
+def get_case_note_details(note_id: int, case_id: int) -> NoteDetailsResponse:
35
+ client, case = initialize_client_and_case("DFIR-IRIS")
36
+ result = fetch_and_validate_data(client, case.get_note, note_id, case_id)
37
+ note_details = NoteDetails(**result["data"])
38
+ return NoteDetailsResponse(success=True, message="Successfully fetched note details", note_details=note_details)
39
+
40
+
41
+def perform_note_creation(client: Any, case: Case, note_creation_body: NoteCreationBody, case_id: int) -> Dict:
42
+ result = fetch_and_validate_data(client, case.add_notes_group, note_creation_body.note_title, case_id)
43
+ note_id = result["data"]["group_id"]
44
+ custom_attributes = {}
45
+ return fetch_and_validate_data(
46
+ client,
47
+ case.add_note,
48
+ note_creation_body.note_title,
49
+ note_creation_body.note_content,
50
+ note_id,
51
+ custom_attributes,
52
+ case_id,
53
+ )
54
+
55
+
56
+def create_case_note(case_id: int, note_creation_body: NoteCreationBody) -> NoteCreationResponse:
57
+ client, case = initialize_client_and_case("DFIR-IRIS")
58
+ result = perform_note_creation(client, case, note_creation_body, case_id)
59
+ return NoteCreationResponse(success=True, message="Successfully created note", note=result["data"])
backend/app/connectors/dfir_iris/services/users.py
new
+17
@@ -0,0 +1,17 @@
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 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
+def get_users() -> UsersResponse:
9
+ client, user = initialize_client_and_user("DFIR-IRIS")
10
+ result = fetch_and_validate_data(client, user.list_users)
11
+ return UsersResponse(success=True, message="Successfully fetched users", users=result["data"])
12
+
13
+
14
+def assign_user_to_alert(alert_id: str, user_id: int) -> AlertResponse:
15
+ client, alert = initialize_client_and_alert("DFIR-IRIS")
16
+ result = fetch_and_validate_data(client, alert.update_alert, alert_id, {"alert_owner_id": user_id})
17
+ return AlertResponse(success=True, message="Successfully assigned user to alert", alert=result["data"])
backend/app/connectors/dfir_iris/utils/universal.py
new
+197
@@ -0,0 +1,197 @@
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
9
+from dfir_iris_client.alert import Alert
10
+from dfir_iris_client.case import Case
11
+from dfir_iris_client.helper.utils import assert_api_resp
12
+from dfir_iris_client.helper.utils import get_data_from_resp
13
+from dfir_iris_client.session import ClientSession
14
+from dfir_iris_client.users import User
15
+from fastapi import HTTPException
16
+from loguru import logger
17
+
18
+from app.connectors.utils import get_connector_info_from_db
19
+
20
+
21
+def verify_dfir_iris_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
22
+ """
23
+ Verifies the connection to DFIR-IRIS service.
24
+
25
+ Returns:
26
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
27
+ """
28
+ logger.info(f"Verifying the DFIR-IRIS connection to {attributes['connector_url']}")
29
+
30
+ try:
31
+ headers = {
32
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
33
+ }
34
+ dfir_iris = requests.get(
35
+ f"{attributes['connector_url']}/api/ping",
36
+ headers=headers,
37
+ verify=False,
38
+ )
39
+ # See if 200 is returned
40
+ if dfir_iris.status_code == 200:
41
+ logger.info(
42
+ f"Connection to {attributes['connector_url']} successful",
43
+ )
44
+ logger.debug("DFIR-IRIS connection successful")
45
+ return {"connectionSuccessful": True, "message": "DFIR-IRIS connection successful"}
46
+ except Exception as e:
47
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")
48
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
49
+
50
+
51
+def verify_dfir_iris_connection(connector_name: str) -> str:
52
+ """
53
+ Returns the authentication token for the DFIR-IRIS service.
54
+
55
+ Returns:
56
+ str: Authentication token for the DFIR-IRIS service.
57
+ """
58
+ attributes = get_connector_info_from_db(connector_name)
59
+ if attributes is None:
60
+ logger.error("No DFIR-IRIS connector found in the database")
61
+ return None
62
+ return verify_dfir_iris_credentials(attributes)
63
+
64
+
65
+def create_dfir_iris_client(connector_name: str) -> ClientSession:
66
+ """
67
+ Creates a session with DFIR-IRIS.
68
+
69
+ This method creates a session with DFIR-IRIS and returns a dictionary with a success status and the session object.
70
+ If a session cannot be established, an error is logged and a dictionary with "success" set to False and an error message is
71
+ returned.
72
+
73
+ Returns:
74
+ dict: A dictionary containing the success status and either the session object or an error message.
75
+ """
76
+ try:
77
+ attributes = get_connector_info_from_db(connector_name)
78
+ logger.info("Creating session with DFIR-IRIS.")
79
+ return ClientSession(
80
+ host=attributes["connector_url"],
81
+ apikey=attributes["connector_api_key"],
82
+ agent="iris-client",
83
+ ssl_verify=False,
84
+ timeout=120,
85
+ proxy=None,
86
+ )
87
+ except Exception as e:
88
+ logger.error(f"Error creating session with DFIR-IRIS: {e}")
89
+ return HTTPException(status_code=500, detail=f"Error creating session with DFIR-IRIS: {e}")
90
+
91
+
92
+def fetch_and_parse_data(session: ClientSession, action: Callable, *args) -> Dict[str, Union[bool, Optional[Dict]]]:
93
+ """
94
+ Fetches and parses data from DFIR-IRIS using a specified action.
95
+
96
+ Args:
97
+ session (ClientSession): The DFIR-IRIS session object.
98
+ action (Callable): The function to execute to fetch data from DFIR-IRIS. This function should accept *args.
99
+ args: The arguments to pass to the action function.
100
+
101
+ Returns:
102
+ dict: A dictionary containing the success status and either the fetched data or None if the operation was unsuccessful.
103
+ """
104
+ try:
105
+ logger.info(f"Executing {action.__name__}... on args: {args}")
106
+ status = action(*args)
107
+ assert_api_resp(status, soft_fail=False)
108
+ data = get_data_from_resp(status)
109
+ logger.info(f"Successfully executed {action.__name__}")
110
+ return {"success": True, "data": data}
111
+ except Exception as err:
112
+ logger.error(f"Failed to execute {action.__name__}: {err}")
113
+ return HTTPException(status_code=500, detail=f"Failed to execute {action.__name__}: {err}")
114
+
115
+
116
+def initialize_client_and_case(service_name: str) -> Tuple[Any, Case]:
117
+ dfir_iris_client = create_dfir_iris_client(service_name)
118
+ case = Case(session=dfir_iris_client)
119
+ return dfir_iris_client, case
120
+
121
+
122
+def initialize_client_and_alert(service_name: str) -> Tuple[Any, Alert]:
123
+ dfir_iris_client = create_dfir_iris_client(service_name)
124
+ alert = Alert(session=dfir_iris_client)
125
+ return dfir_iris_client, alert
126
+
127
+
128
+def initialize_client_and_user(service_name: str) -> Tuple[Any, Alert]:
129
+ dfir_iris_client = create_dfir_iris_client(service_name)
130
+ user = User(session=dfir_iris_client)
131
+ return dfir_iris_client, user
132
+
133
+
134
+def handle_error(error_message: str, status_code: int = 500):
135
+ logger.error(error_message)
136
+ raise HTTPException(status_code=status_code, detail=error_message)
137
+
138
+
139
+def fetch_and_validate_data(client: Any, func: Callable, *args: Any) -> Dict:
140
+ result = fetch_and_parse_data(client, func, *args)
141
+ if not result["success"]:
142
+ handle_error(f"Failed to fetch data: {result['message']}")
143
+ return result
144
+
145
+
146
+def check_case_exists(case_id: int) -> bool:
147
+ try:
148
+ logger.info(f"Checking if case {case_id} exists")
149
+ dfir_iris_client = create_dfir_iris_client("DFIR-IRIS")
150
+ case = Case(session=dfir_iris_client)
151
+ data = case.get_case(case_id)
152
+ assert_api_resp(data, soft_fail=False)
153
+ result = get_data_from_resp(data)
154
+ if result is None:
155
+ logger.info(f"Case {case_id} does not exist")
156
+ return False
157
+ logger.info(f"Case {case_id} exists")
158
+ return True
159
+ except Exception as e:
160
+ logger.error(f"Failed to check if case {case_id} exists: {e}")
161
+ return False
162
+
163
+
164
+def check_alert_exists(alert_id: str) -> bool:
165
+ try:
166
+ logger.info(f"Checking if alert {alert_id} exists")
167
+ dfir_iris_client = create_dfir_iris_client("DFIR-IRIS")
168
+ alert = Alert(session=dfir_iris_client)
169
+ data = alert.get_alert(alert_id)
170
+ assert_api_resp(data, soft_fail=False)
171
+ result = get_data_from_resp(data)
172
+ if result is None:
173
+ logger.info(f"Alert {alert_id} does not exist")
174
+ return False
175
+ logger.info(f"Alert {alert_id} exists")
176
+ return True
177
+ except Exception as e:
178
+ logger.error(f"Failed to check if alert {alert_id} exists: {e}")
179
+ return False
180
+
181
+
182
+def check_user_exists(user_id: int) -> bool:
183
+ try:
184
+ logger.info(f"Checking if user {user_id} exists")
185
+ dfir_iris_client = create_dfir_iris_client("DFIR-IRIS")
186
+ user = User(session=dfir_iris_client)
187
+ data = user.get_user(user_id)
188
+ assert_api_resp(data, soft_fail=False)
189
+ result = get_data_from_resp(data)
190
+ if result is None:
191
+ logger.info(f"User {user_id} does not exist")
192
+ return False
193
+ logger.info(f"User {user_id} exists")
194
+ return True
195
+ except Exception as e:
196
+ logger.error(f"Failed to check if user {user_id} exists: {e}")
197
+ return False
backend/app/connectors/graylog/routes/collector.py
new
+40
@@ -0,0 +1,40 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.connectors.graylog.schema.collector import ConfiguredInputsResponse
5
+from app.connectors.graylog.schema.collector import GraylogIndicesResponse
6
+from app.connectors.graylog.schema.collector import GraylogInputsResponse
7
+from app.connectors.graylog.schema.collector import RunningInputsResponse
8
+from app.connectors.graylog.services.collector import get_indices_full
9
+from app.connectors.graylog.services.collector import get_inputs
10
+from app.connectors.graylog.services.collector import get_inputs_configured
11
+from app.connectors.graylog.services.collector import get_inputs_running
12
+
13
+# App specific imports
14
+
15
+
16
+graylog_collector_router = APIRouter()
17
+
18
+
19
+@graylog_collector_router.get("/indices", response_model=GraylogIndicesResponse, description="Get all indices")
20
+async def get_all_indices() -> GraylogIndicesResponse:
21
+ logger.info("Fetching all graylog indices")
22
+ return get_indices_full()
23
+
24
+
25
+@graylog_collector_router.get("/inputs", response_model=GraylogInputsResponse, description="Get all inputs")
26
+async def get_all_inputs() -> GraylogInputsResponse:
27
+ logger.info("Fetching all graylog inputs")
28
+ return get_inputs()
29
+
30
+
31
+@graylog_collector_router.get("/inputs/running", response_model=RunningInputsResponse, description="Get all running inputs")
32
+async def get_all_running_inputs() -> RunningInputsResponse:
33
+ logger.info("Fetching all graylog running inputs")
34
+ return get_inputs_running()
35
+
36
+
37
+@graylog_collector_router.get("/inputs/configured", response_model=ConfiguredInputsResponse, description="Get all configured inputs")
38
+async def get_all_configured_inputs() -> ConfiguredInputsResponse:
39
+ logger.info("Fetching all graylog configured inputs")
40
+ return get_inputs_configured()
backend/app/connectors/graylog/routes/events.py
new
+25
@@ -0,0 +1,25 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.connectors.graylog.schema.events import AlertQuery
5
+from app.connectors.graylog.schema.events import GraylogAlertsResponse
6
+from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse
7
+from app.connectors.graylog.services.events import get_alerts
8
+from app.connectors.graylog.services.events import get_event_definitions
9
+
10
+# App specific imports
11
+
12
+
13
+graylog_events_router = APIRouter()
14
+
15
+
16
+@graylog_events_router.get("/event/definitions", response_model=GraylogEventDefinitionsResponse, description="Get all event definitions")
17
+async def get_all_event_definitions() -> GraylogEventDefinitionsResponse:
18
+ logger.info("Fetching all graylog event definitions")
19
+ return get_event_definitions()
20
+
21
+
22
+@graylog_events_router.post("/event/alerts", response_model=GraylogAlertsResponse, description="Get all alerts")
23
+async def get_all_alerts(alert_query: AlertQuery) -> GraylogAlertsResponse:
24
+ logger.info("Fetching all graylog alerts")
25
+ return get_alerts(alert_query)
backend/app/connectors/graylog/routes/management.py
new
+110
@@ -0,0 +1,110 @@
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
+
8
+from app.connectors.graylog.schema.management import DeletedIndexBody
9
+from app.connectors.graylog.schema.management import DeletedIndexResponse
10
+from app.connectors.graylog.schema.management import StartInputBody
11
+from app.connectors.graylog.schema.management import StartInputResponse
12
+from app.connectors.graylog.schema.management import StartStreamBody
13
+from app.connectors.graylog.schema.management import StartStreamResponse
14
+from app.connectors.graylog.schema.management import StopInputBody
15
+from app.connectors.graylog.schema.management import StopInputResponse
16
+from app.connectors.graylog.schema.management import StopStreamBody
17
+from app.connectors.graylog.schema.management import StopStreamResponse
18
+from app.connectors.graylog.services.collector import get_index_names
19
+from app.connectors.graylog.services.collector import get_input_ids
20
+from app.connectors.graylog.services.management import delete_index
21
+from app.connectors.graylog.services.management import start_input
22
+from app.connectors.graylog.services.management import start_stream
23
+from app.connectors.graylog.services.management import stop_input
24
+from app.connectors.graylog.services.management import stop_stream
25
+from app.connectors.graylog.services.streams import get_stream_ids
26
+
27
+graylog_management_router = APIRouter()
28
+
29
+
30
+def get_managed_index_names() -> List[str]:
31
+ return get_index_names()
32
+
33
+
34
+def get_managed_input_ids() -> List[str]:
35
+ return get_input_ids()
36
+
37
+
38
+def get_managed_stream_ids() -> List[str]:
39
+ return get_stream_ids()
40
+
41
+
42
+def verify_index_name(deleted_index_body: DeletedIndexBody) -> DeletedIndexBody:
43
+ # Remove any extra spaces from index_name
44
+ deleted_index_body.index_name = deleted_index_body.index_name.strip()
45
+
46
+ managed_index_names = get_managed_index_names()
47
+ if deleted_index_body.index_name not in managed_index_names:
48
+ raise HTTPException(
49
+ status_code=400,
50
+ detail=f"Index name '{deleted_index_body.index_name}' is not managed by Graylog or no longer exists.",
51
+ )
52
+ return deleted_index_body
53
+
54
+
55
+def verify_input_id(stop_input_body: StopInputBody) -> StopInputBody:
56
+ # Remove any extra spaces from input_id
57
+ stop_input_body.input_id = stop_input_body.input_id.strip()
58
+
59
+ managed_input_ids = get_managed_input_ids()
60
+ if stop_input_body.input_id not in managed_input_ids:
61
+ raise HTTPException(status_code=400, detail=f"Input ID '{stop_input_body.input_id}' is not managed by Graylog or no longer exists.")
62
+ return stop_input_body
63
+
64
+
65
+def verify_stream_id(stop_stream_body: StopStreamBody) -> StopStreamBody:
66
+ # Remove any extra spaces from stream_id
67
+ stop_stream_body.stream_id = stop_stream_body.stream_id.strip()
68
+
69
+ managed_stream_ids = get_managed_stream_ids()
70
+ if stop_stream_body.stream_id not in managed_stream_ids:
71
+ raise HTTPException(
72
+ status_code=400,
73
+ detail=f"Stream ID '{stop_stream_body.stream_id}' is not managed by Graylog or no longer exists.",
74
+ )
75
+ return stop_stream_body
76
+
77
+
78
+@graylog_management_router.delete("/index", response_model=DeletedIndexResponse, description="Delete index")
79
+async def delete_index_route(deleted_index_body: DeletedIndexBody = Depends(verify_index_name)) -> DeletedIndexResponse:
80
+ logger.info(f"Deleting index {deleted_index_body.index_name}")
81
+
82
+ return delete_index(deleted_index_body.index_name)
83
+
84
+
85
+@graylog_management_router.post("/input/stop", response_model=StopInputResponse, description="Stop input")
86
+async def stop_input_route(stop_input_body: StopInputBody = Depends(verify_input_id)) -> StopInputResponse:
87
+ logger.info(f"Stopping input {stop_input_body.input_id}")
88
+
89
+ return stop_input(stop_input_body.input_id)
90
+
91
+
92
+@graylog_management_router.post("/input/start", response_model=StartInputResponse, description="Start input")
93
+async def start_input_route(start_input_body: StartInputBody = Depends(verify_input_id)) -> StartInputResponse:
94
+ logger.info(f"Starting input {start_input_body.input_id}")
95
+
96
+ return start_input(start_input_body.input_id)
97
+
98
+
99
+@graylog_management_router.post("/stream/stop", response_model=StopStreamResponse, description="Stop stream")
100
+async def stop_stream_route(stop_stream_body: StopStreamBody = Depends(verify_stream_id)) -> StopStreamResponse:
101
+ logger.info(f"Stopping stream {stop_stream_body.stream_id}")
102
+
103
+ return stop_stream(stop_stream_body.stream_id)
104
+
105
+
106
+@graylog_management_router.post("/stream/start", response_model=StartStreamResponse, description="Start stream")
107
+async def start_stream_route(start_stream_body: StartStreamBody = Depends(verify_stream_id)) -> StartStreamResponse:
108
+ logger.info(f"Starting stream {start_stream_body.stream_id}")
109
+
110
+ return start_stream(start_stream_body.stream_id)
backend/app/connectors/graylog/routes/monitoring.py
new
+25
@@ -0,0 +1,25 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse
5
+from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse
6
+from app.connectors.graylog.services.monitoring import get_messages
7
+from app.connectors.graylog.services.monitoring import get_metrics
8
+
9
+# App specific imports
10
+
11
+
12
+graylog_monitoring_router = APIRouter()
13
+
14
+
15
+@graylog_monitoring_router.get("/messages", response_model=GraylogMessagesResponse, description="Get all messages")
16
+async def get_all_messages(page_number: int = 1) -> GraylogMessagesResponse:
17
+ logger.info("Fetching all graylog messages")
18
+ logger.info(f"Page number: {page_number}")
19
+ return get_messages(page_number)
20
+
21
+
22
+@graylog_monitoring_router.get("/metrics", response_model=GraylogMetricsResponse, description="Get all metrics")
23
+async def get_all_metrics() -> GraylogMetricsResponse:
24
+ logger.info("Fetching all graylog metrics")
25
+ return get_metrics()
backend/app/connectors/graylog/routes/pipelines.py
new
+24
@@ -0,0 +1,24 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
5
+from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
6
+from app.connectors.graylog.services.pipelines import get_pipeline_rules
7
+from app.connectors.graylog.services.pipelines import get_pipelines
8
+
9
+# App specific imports
10
+
11
+
12
+graylog_pipelines_router = APIRouter()
13
+
14
+
15
+@graylog_pipelines_router.get("/pipelines", response_model=GraylogPipelinesResponse, description="Get all pipelines")
16
+async def get_all_pipelines() -> GraylogPipelinesResponse:
17
+ logger.info("Fetching all graylog pipelines")
18
+ return get_pipelines()
19
+
20
+
21
+@graylog_pipelines_router.get("/pipeline/rules", response_model=PipelineRulesResponse, description="Get all pipeline rules")
22
+async def get_all_pipeline_rules() -> PipelineRulesResponse:
23
+ logger.info("Fetching all graylog pipeline rules")
24
+ return get_pipeline_rules()
backend/app/connectors/graylog/routes/streams.py
new
+16
@@ -0,0 +1,16 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.connectors.graylog.schema.streams import GraylogStreamsResponse
5
+from app.connectors.graylog.services.streams import get_streams
6
+
7
+# App specific imports
8
+
9
+
10
+graylog_streams_router = APIRouter()
11
+
12
+
13
+@graylog_streams_router.get("/streams", response_model=GraylogStreamsResponse, description="Get all streams")
14
+async def get_all_streams() -> GraylogStreamsResponse:
15
+ logger.info("Fetching all graylog streams")
16
+ return get_streams()
backend/app/connectors/graylog/schema/collector.py
new
+152
@@ -0,0 +1,152 @@
1
+from typing import Dict
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+
8
+
9
+class Document(BaseModel):
10
+ count: int
11
+ deleted: int
12
+
13
+
14
+class Operation(BaseModel):
15
+ time_seconds: int
16
+ total: int
17
+
18
+
19
+class ShardInfo(BaseModel):
20
+ documents: Document
21
+ flush: Operation
22
+ get: Operation
23
+ index: Operation
24
+ merge: Operation
25
+ open_search_contexts: int
26
+ refresh: Operation
27
+ search_fetch: Operation
28
+ search_query: Operation
29
+ segments: int
30
+ store_size_bytes: int
31
+
32
+
33
+class Routing(BaseModel):
34
+ active: bool
35
+ id: int
36
+ node_hostname: str
37
+ node_id: str
38
+ node_name: str
39
+ primary: bool
40
+ relocating_to: Optional[None] # Assuming this is always None based on your example
41
+ state: str
42
+
43
+
44
+class IndexInfo(BaseModel):
45
+ all_shards: ShardInfo
46
+ primary_shards: ShardInfo
47
+ reopened: bool
48
+ routing: List[Routing]
49
+
50
+
51
+class GraylogIndexItem(BaseModel):
52
+ index_name: str
53
+ index_info: IndexInfo
54
+
55
+
56
+class GraylogIndicesResponse(BaseModel):
57
+ indices: List[GraylogIndexItem]
58
+ message: str
59
+ success: bool
60
+
61
+
62
+class ConfiguredInputAttributes(BaseModel):
63
+ recv_buffer_size: int
64
+ tcp_keepalive: bool
65
+ use_null_delimiter: bool
66
+ number_worker_threads: int
67
+ tls_client_auth_cert_file: Optional[str]
68
+ force_rdns: Optional[bool]
69
+ bind_address: str
70
+ tls_cert_file: Optional[str]
71
+ store_full_message: Optional[bool]
72
+ expand_structured_data: Optional[bool]
73
+ port: int
74
+ tls_key_file: Optional[str]
75
+ tls_enable: bool
76
+ tls_key_password: Optional[str]
77
+ max_message_size: int
78
+ tls_client_auth: str
79
+ override_source: Optional[str]
80
+ charset_name: Optional[str]
81
+ allow_override_date: Optional[bool]
82
+
83
+
84
+class ConfiguredInput(BaseModel):
85
+ title: str
86
+ global_field: bool = Field(alias="global")
87
+ name: str
88
+ content_pack: Optional[str]
89
+ created_at: str
90
+ type: str
91
+ creator_user_id: str
92
+ attributes: ConfiguredInputAttributes
93
+ static_fields: Dict[str, str]
94
+ node: str
95
+ id: str
96
+
97
+
98
+class MessageInputAttributes(BaseModel):
99
+ recv_buffer_size: int
100
+ tcp_keepalive: bool
101
+ use_null_delimiter: bool
102
+ number_worker_threads: int
103
+ tls_client_auth_cert_file: Optional[str]
104
+ bind_address: str
105
+ tls_cert_file: Optional[str]
106
+ port: int
107
+ tls_key_file: Optional[str]
108
+ tls_enable: bool
109
+ tls_key_password: Optional[str]
110
+ max_message_size: int
111
+ tls_client_auth: str
112
+
113
+
114
+class MessageInput(BaseModel):
115
+ title: str
116
+ global_field: bool = Field(alias="global")
117
+ name: str
118
+ content_pack: Optional[str]
119
+ created_at: str
120
+ type: str
121
+ creator_user_id: str
122
+ attributes: MessageInputAttributes
123
+ static_fields: Dict[str, str]
124
+ node: str
125
+ id: str
126
+
127
+
128
+class RunningInput(BaseModel):
129
+ id: str
130
+ state: str
131
+ started_at: str
132
+ detailed_message: Optional[str]
133
+ message_input: MessageInput
134
+
135
+
136
+class ConfiguredInputsResponse(BaseModel):
137
+ configured_inputs: List[ConfiguredInput]
138
+ message: str
139
+ success: bool
140
+
141
+
142
+class RunningInputsResponse(BaseModel):
143
+ running_inputs: List[RunningInput]
144
+ message: str
145
+ success: bool
146
+
147
+
148
+class GraylogInputsResponse(BaseModel):
149
+ configured_inputs: List[ConfiguredInput]
150
+ running_inputs: List[RunningInput]
151
+ message: str
152
+ success: bool
backend/app/connectors/graylog/schema/events.py
new
+155
@@ -0,0 +1,155 @@
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
+
8
+
9
+class Provider(BaseModel):
10
+ require_values: bool
11
+ template: str
12
+ type: str
13
+
14
+
15
+class FieldSpecItem(BaseModel):
16
+ data_type: str
17
+ providers: List[Provider]
18
+
19
+
20
+class Conditions(BaseModel):
21
+ expression: Optional[str]
22
+
23
+
24
+class Config(BaseModel):
25
+ conditions: Conditions
26
+ execute_every_ms: int
27
+ group_by: List[str]
28
+ query: str
29
+ query_parameters: List[str]
30
+ search_within_ms: int
31
+ series: List[str]
32
+ streams: List[str]
33
+ type: str
34
+
35
+
36
+class NotificationSettings(BaseModel):
37
+ backlog_size: int
38
+ grace_period_ms: int
39
+
40
+
41
+class Storage(BaseModel):
42
+ streams: List[str]
43
+ type: str
44
+
45
+
46
+class EventDefinition(BaseModel):
47
+ _scope: str
48
+ alert: bool
49
+ config: Config
50
+ description: str
51
+ field_spec: Dict[str, FieldSpecItem]
52
+ id: str
53
+ key_spec: List[str]
54
+ notification_settings: NotificationSettings
55
+ notifications: Optional[List[Dict[str, Union[str, None]]]]
56
+ priority: int
57
+ storage: List[Storage]
58
+ title: str
59
+
60
+
61
+class GraylogEventDefinitionsResponse(BaseModel):
62
+ event_definitions: List[EventDefinition]
63
+ message: str
64
+ success: bool
65
+
66
+
67
+class AlertQuery(BaseModel):
68
+ query: Optional[str] = ""
69
+ page: int = 1
70
+ per_page: int = 100
71
+ filter: Optional[Dict[str, Union[str, List[str]]]] = {"alerts": "only", "event_definitions": []}
72
+ timerange: Optional[Dict[str, Union[int, str]]] = {"range": 86400, "type": "relative"}
73
+
74
+
75
+class SimplifiedEventDefinition(BaseModel):
76
+ description: str
77
+ id: str
78
+ title: str
79
+
80
+
81
+class Stream(BaseModel):
82
+ description: str
83
+ id: str
84
+ title: str
85
+
86
+
87
+class Context(BaseModel):
88
+ event_definitions: Dict[str, SimplifiedEventDefinition]
89
+ streams: Dict[str, Stream]
90
+
91
+
92
+class Fields(BaseModel):
93
+ test: str
94
+
95
+
96
+class Event(BaseModel):
97
+ alert: bool
98
+ event_definition_id: str
99
+ event_definition_type: str
100
+ fields: Fields
101
+ group_by_fields: Dict[str, str]
102
+ id: str
103
+ key: Optional[str]
104
+ key_tuple: List[str]
105
+ message: str
106
+ origin_context: str
107
+ priority: int
108
+ source: str
109
+ source_streams: List[str]
110
+ streams: List[str]
111
+ timerange_end: Optional[str]
112
+ timerange_start: Optional[str]
113
+ timestamp: str
114
+ timestamp_processing: str
115
+
116
+
117
+class AlertEvent(BaseModel):
118
+ event: Event
119
+ index_name: str
120
+ index_type: str
121
+
122
+
123
+class Filter(BaseModel):
124
+ alerts: str
125
+ event_definitions: List[str]
126
+
127
+
128
+class Timerange(BaseModel):
129
+ range: int
130
+ type: str
131
+
132
+
133
+class Parameters(BaseModel):
134
+ page: int
135
+ per_page: int
136
+ query: str
137
+ sort_by: str
138
+ sort_direction: str
139
+ timerange: Timerange
140
+ filter: Filter
141
+
142
+
143
+class Alerts(BaseModel):
144
+ context: Context
145
+ duration: int
146
+ events: List[AlertEvent]
147
+ parameters: Parameters
148
+ total_events: int
149
+ used_indices: List[str]
150
+
151
+
152
+class GraylogAlertsResponse(BaseModel):
153
+ alerts: Alerts
154
+ message: str
155
+ success: bool
backend/app/connectors/graylog/schema/management.py
new
+46
@@ -0,0 +1,46 @@
1
+from pydantic import BaseModel
2
+
3
+
4
+class DeletedIndexBody(BaseModel):
5
+ index_name: str
6
+
7
+
8
+class DeletedIndexResponse(BaseModel):
9
+ success: bool
10
+ message: str
11
+
12
+
13
+class StopInputBody(BaseModel):
14
+ input_id: str
15
+
16
+
17
+class StopInputResponse(BaseModel):
18
+ success: bool
19
+ message: str
20
+
21
+
22
+class StartInputBody(BaseModel):
23
+ input_id: str
24
+
25
+
26
+class StartInputResponse(BaseModel):
27
+ success: bool
28
+ message: str
29
+
30
+
31
+class StopStreamBody(BaseModel):
32
+ stream_id: str
33
+
34
+
35
+class StopStreamResponse(BaseModel):
36
+ success: bool
37
+ message: str
38
+
39
+
40
+class StartStreamBody(BaseModel):
41
+ stream_id: str
42
+
43
+
44
+class StartStreamResponse(BaseModel):
45
+ success: bool
46
+ message: str
backend/app/connectors/graylog/schema/monitoring.py
new
+53
@@ -0,0 +1,53 @@
1
+from typing import List
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
+
7
+
8
+class GraylogMessages(BaseModel):
9
+ caller: str
10
+ content: str
11
+ node_id: str
12
+ timestamp: str
13
+
14
+
15
+class GraylogTotalMessages(BaseModel):
16
+ total: int
17
+
18
+
19
+class GraylogMessagesResponse(BaseModel):
20
+ graylog_messages: List[GraylogMessages]
21
+ success: bool
22
+ message: str
23
+ total_messages: int
24
+
25
+
26
+class GraylogThroughputMetrics(BaseModel):
27
+ metric: str
28
+ value: float
29
+
30
+
31
+class GraylogThroughputMetricsCollection(BaseModel):
32
+ graylog2_buffers_input_usage: Optional[str] = Field(alias="org.graylog2.buffers.input.usage")
33
+ graylog2_buffers_output_usage: Optional[str] = Field(alias="org.graylog2.buffers.output.usage")
34
+ graylog2_buffers_process_usage: Optional[str] = Field(alias="org.graylog2.buffers.process.usage")
35
+ graylog2_throughput_input_1_sec_rate: Optional[str] = Field(alias="org.graylog2.throughput.input.1-sec-rate")
36
+ graylog2_throughput_output_1_sec_rate: Optional[str] = Field(alias="org.graylog2.throughput.output.1-sec-rate")
37
+ graylog2_throughput_output: Optional[str] = Field(alias="org.graylog2.throughput.output")
38
+ graylog2_throughput_input: Optional[str] = Field(alias="org.graylog2.throughput.input")
39
+
40
+
41
+class GraylogThroughputMetricsList(BaseModel):
42
+ throughput_metrics: List[GraylogThroughputMetrics]
43
+
44
+
45
+class GraylogUncommittedJournalEntries(BaseModel):
46
+ uncommitted_journal_entries: int
47
+
48
+
49
+class GraylogMetricsResponse(BaseModel):
50
+ throughput_metrics: List[GraylogThroughputMetrics]
51
+ uncommitted_journal_entries: int
52
+ message: str
53
+ success: bool
backend/app/connectors/graylog/schema/pipelines.py
new
+44
@@ -0,0 +1,44 @@
1
+from typing import List
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+
6
+
7
+class Stage(BaseModel):
8
+ match: str
9
+ rules: List[str]
10
+ stage: int
11
+
12
+
13
+class Pipeline(BaseModel):
14
+ created_at: str
15
+ description: str
16
+ errors: Optional[None]
17
+ id: str
18
+ modified_at: str
19
+ source: str
20
+ stages: List[Stage]
21
+ title: str
22
+
23
+
24
+class GraylogPipelinesResponse(BaseModel):
25
+ message: str
26
+ pipelines: List[Pipeline]
27
+ success: bool
28
+
29
+
30
+class PipelineRule(BaseModel):
31
+ created_at: str
32
+ description: str
33
+ errors: Optional[None]
34
+ id: str
35
+ modified_at: str
36
+ source: str
37
+ title: str
38
+
39
+
40
+# Define the main response model
41
+class PipelineRulesResponse(BaseModel):
42
+ message: str
43
+ pipeline_rules: List[PipelineRule]
44
+ success: bool
backend/app/connectors/graylog/schema/streams.py
new
+38
@@ -0,0 +1,38 @@
1
+from typing import List
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+
6
+
7
+class Rule(BaseModel):
8
+ description: Optional[str]
9
+ field: str
10
+ id: str
11
+ inverted: bool
12
+ stream_id: str
13
+ type: int
14
+ value: str
15
+
16
+
17
+class Stream(BaseModel):
18
+ content_pack: Optional[str]
19
+ created_at: str
20
+ creator_user_id: str
21
+ description: str
22
+ disabled: bool
23
+ id: str
24
+ index_set_id: str
25
+ is_default: bool
26
+ is_editable: bool
27
+ matching_type: str
28
+ outputs: list
29
+ remove_matches_from_default_stream: bool
30
+ rules: List[Rule]
31
+ title: str
32
+
33
+
34
+class GraylogStreamsResponse(BaseModel):
35
+ message: str
36
+ streams: List[Stream]
37
+ total: int
38
+ success: bool
backend/app/connectors/graylog/services/collector.py
new
+126
@@ -0,0 +1,126 @@
1
+from typing import List
2
+from typing import Tuple
3
+
4
+from loguru import logger
5
+
6
+from app.connectors.graylog.schema.collector import ConfiguredInput
7
+from app.connectors.graylog.schema.collector import ConfiguredInputsResponse
8
+from app.connectors.graylog.schema.collector import GraylogIndexItem
9
+from app.connectors.graylog.schema.collector import GraylogIndicesResponse
10
+from app.connectors.graylog.schema.collector import GraylogInputsResponse
11
+from app.connectors.graylog.schema.collector import RunningInput
12
+from app.connectors.graylog.schema.collector import RunningInputsResponse
13
+from app.connectors.graylog.utils.universal import send_get_request
14
+
15
+
16
+def get_indices_full() -> GraylogIndicesResponse:
17
+ """Get indices from Graylog."""
18
+ logger.info("Getting indices from Graylog")
19
+ indices_collected = send_get_request(endpoint="/api/system/indexer/indices")
20
+ if indices_collected["success"]:
21
+ indices_data = indices_collected["data"]["all"]["indices"]
22
+
23
+ # Convert the dictionary to a list of GraylogIndexItem
24
+ indices_list = [GraylogIndexItem(index_name=name, index_info=info) for name, info in indices_data.items()]
25
+
26
+ return GraylogIndicesResponse(indices=indices_list, success=True, message="Indices collected successfully")
27
+ else:
28
+ return GraylogIndicesResponse(indices=[], success=False, message="Failed to collect indices")
29
+
30
+
31
+def fetch_configured_inputs() -> Tuple[bool, List[ConfiguredInput]]:
32
+ configured_inputs_collected = send_get_request(endpoint="/api/system/inputs")
33
+ success = configured_inputs_collected.get("success", False)
34
+
35
+ if success:
36
+ return True, [ConfiguredInput(**input_data) for input_data in configured_inputs_collected["data"]["inputs"]]
37
+ else:
38
+ logger.error("Failed to fetch configured inputs")
39
+ return False, []
40
+
41
+
42
+def fetch_running_inputs() -> Tuple[bool, List[RunningInput]]:
43
+ running_inputs_collected = send_get_request(endpoint="/api/system/inputstates")
44
+ success = running_inputs_collected.get("success", False)
45
+
46
+ if success:
47
+ return True, [RunningInput(**input_data) for input_data in running_inputs_collected["data"]["states"]]
48
+ else:
49
+ logger.error("Failed to fetch running inputs")
50
+ return False, []
51
+
52
+
53
+def get_inputs() -> GraylogInputsResponse:
54
+ """Get inputs from Graylog."""
55
+ logger.info("Getting inputs from Graylog")
56
+
57
+ config_success, configured_inputs_list = fetch_configured_inputs()
58
+ run_success, running_inputs_list = fetch_running_inputs()
59
+
60
+ if config_success and run_success:
61
+ logger.info("Successfully fetched both configured and running inputs")
62
+ return GraylogInputsResponse(
63
+ configured_inputs=configured_inputs_list,
64
+ running_inputs=running_inputs_list,
65
+ success=True,
66
+ message="Successfully retrieved inputs",
67
+ )
68
+ else:
69
+ logger.error("Failed to fetch one or both types of inputs")
70
+ return GraylogInputsResponse(configured_inputs=[], running_inputs=[], success=False, message="Failed to collect inputs")
71
+
72
+
73
+def get_inputs_running() -> RunningInputsResponse:
74
+ """Get running inputs from Graylog."""
75
+ logger.info("Getting running inputs from Graylog")
76
+ run_success, running_inputs_list = fetch_running_inputs()
77
+ if run_success:
78
+ return RunningInputsResponse(running_inputs=running_inputs_list, success=True, message="Successfully retrieved running inputs")
79
+
80
+
81
+def get_inputs_configured() -> ConfiguredInputsResponse:
82
+ """Get configured inputs from Graylog."""
83
+ logger.info("Getting configured inputs from Graylog")
84
+ config_success, configured_inputs_list = fetch_configured_inputs()
85
+ if config_success:
86
+ return ConfiguredInputsResponse(
87
+ configured_inputs=configured_inputs_list,
88
+ success=True,
89
+ message="Successfully retrieved configured inputs",
90
+ )
91
+
92
+
93
+def get_index_names() -> List[str]:
94
+ """
95
+ Gets the names of all the indices in Graylog.
96
+
97
+ Returns:
98
+ List[str]: A list of all the index names.
99
+ """
100
+ logger.info("Getting index names from Graylog")
101
+
102
+ indices_collected = get_indices_full()
103
+
104
+ if indices_collected.success:
105
+ # Access the index_name attribute directly
106
+ return [index.index_name for index in indices_collected.indices]
107
+ else:
108
+ return []
109
+
110
+
111
+def get_input_ids() -> List[str]:
112
+ """
113
+ Gets the IDs of all the inputs in Graylog.
114
+
115
+ Returns:
116
+ List[str]: A list of all the input IDs.
117
+ """
118
+ logger.info("Getting input IDs from Graylog")
119
+
120
+ success, inputs_collected = fetch_configured_inputs()
121
+
122
+ if success:
123
+ # Access the input_id attribute directly
124
+ return [input.id for input in inputs_collected]
125
+ else:
126
+ return []
backend/app/connectors/graylog/services/events.py
new
+59
@@ -0,0 +1,59 @@
1
+from loguru import logger
2
+
3
+from app.connectors.graylog.schema.events import AlertEvent
4
+from app.connectors.graylog.schema.events import AlertQuery
5
+from app.connectors.graylog.schema.events import Alerts
6
+from app.connectors.graylog.schema.events import Context
7
+from app.connectors.graylog.schema.events import EventDefinition
8
+from app.connectors.graylog.schema.events import GraylogAlertsResponse
9
+from app.connectors.graylog.schema.events import GraylogEventDefinitionsResponse
10
+from app.connectors.graylog.schema.events import Parameters
11
+from app.connectors.graylog.utils.universal import send_get_request
12
+from app.connectors.graylog.utils.universal import send_post_request
13
+
14
+
15
+def get_event_definitions() -> GraylogEventDefinitionsResponse:
16
+ """Get event definitions from Graylog."""
17
+ logger.info("Getting event definitions from Graylog")
18
+ event_definitions_collected = send_get_request(endpoint="/api/events/definitions")
19
+ if event_definitions_collected["success"]:
20
+ event_definitions_data = event_definitions_collected["data"]["event_definitions"]
21
+
22
+ # Convert the dictionary to a list of GraylogIndexItem
23
+ event_definitions_list = [EventDefinition(**event_definition_data) for event_definition_data in event_definitions_data]
24
+
25
+ return GraylogEventDefinitionsResponse(
26
+ event_definitions=event_definitions_list,
27
+ success=True,
28
+ message="Event definitions collected successfully",
29
+ )
30
+ else:
31
+ return GraylogEventDefinitionsResponse(event_definitions=[], success=False, message="Failed to collect event definitions")
32
+
33
+
34
+def get_alerts(alert_query: AlertQuery) -> GraylogAlertsResponse:
35
+ logger.info("Getting alerts from Graylog")
36
+ response = send_post_request(endpoint="/api/events/search", data=alert_query.dict())
37
+
38
+ if response["success"]:
39
+ raw_alerts_data = response["data"]
40
+ # Convert raw event data to Event objects
41
+ event_objects = [AlertEvent(**event_data) for event_data in raw_alerts_data["events"]]
42
+
43
+ # Build the Alerts object
44
+ alerts = Alerts(
45
+ context=Context(**raw_alerts_data["context"]),
46
+ duration=raw_alerts_data["duration"],
47
+ events=event_objects,
48
+ parameters=Parameters(**raw_alerts_data["parameters"]),
49
+ total_events=raw_alerts_data["total_events"],
50
+ used_indices=raw_alerts_data["used_indices"],
51
+ )
52
+
53
+ # Build the final GraylogAlertsResponse
54
+ final_response = GraylogAlertsResponse(alerts=alerts, message="Successfully collected alerts", success=True)
55
+
56
+ logger.info(f"Events collected: {event_objects}")
57
+ return final_response
58
+ else:
59
+ return GraylogAlertsResponse(alerts=Alerts(events=[]), success=False, message="Failed to collect alerts")
backend/app/connectors/graylog/services/management.py
new
+73
@@ -0,0 +1,73 @@
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
+def delete_index(index_name: DeletedIndexBody) -> DeletedIndexResponse:
20
+ """Delete an index from Graylog."""
21
+ logger.info(f"Deleting index {index_name} from Graylog")
22
+ send_delete_request(endpoint=f"/api/system/indexer/indices/{index_name}")
23
+ # Check if the index still exists
24
+ index_names = get_index_names()
25
+ logger.info(f"Index names: {index_names}")
26
+ if index_name in index_names:
27
+ return DeletedIndexResponse(
28
+ success=False,
29
+ message=f"Failed to delete index {index_name}. If the index is still in use, it cannot be deleted.",
30
+ )
31
+ else:
32
+ return DeletedIndexResponse(success=True, message=f"Successfully deleted index {index_name}")
33
+
34
+
35
+def stop_input(input_id: StopInputBody) -> StopInputResponse:
36
+ """Stop an input in Graylog."""
37
+ logger.info(f"Stopping input {input_id} in Graylog")
38
+ response = send_delete_request(endpoint=f"/api/system/inputstates/{input_id}")
39
+ if response["success"]:
40
+ return StopInputResponse(success=True, message=f"Successfully stopped input {input_id}")
41
+ else:
42
+ return StopInputResponse(success=False, message=f"Failed to stop input {input_id}")
43
+
44
+
45
+def start_input(input_id: StartInputBody) -> StartInputResponse:
46
+ """Start an input in Graylog."""
47
+ logger.info(f"Starting input {input_id} in Graylog")
48
+ response = send_put_request(endpoint=f"/api/system/inputstates/{input_id}")
49
+ if response["success"]:
50
+ return StartInputResponse(success=True, message=f"Successfully started input {input_id}")
51
+ else:
52
+ return StartInputResponse(success=False, message=f"Failed to start input {input_id}")
53
+
54
+
55
+def stop_stream(stream_id: StopStreamBody) -> StopStreamResponse:
56
+ """Stop a stream in Graylog."""
57
+ logger.info(f"Stopping stream {stream_id} in Graylog")
58
+ response = send_post_request(endpoint=f"/api/streams/{stream_id}/pause")
59
+ logger.info(f"Response: {response}")
60
+ if response["success"]:
61
+ return StopStreamResponse(success=True, message=f"Successfully stopped stream {stream_id}")
62
+ else:
63
+ return StopStreamResponse(success=False, message=f"Failed to stop stream {stream_id}")
64
+
65
+
66
+def start_stream(stream_id: StartStreamBody) -> StartStreamResponse:
67
+ """Start a stream in Graylog."""
68
+ logger.info(f"Starting stream {stream_id} in Graylog")
69
+ response = send_post_request(endpoint=f"/api/streams/{stream_id}/resume")
70
+ if response["success"]:
71
+ return StartStreamResponse(success=True, message=f"Successfully started stream {stream_id}")
72
+ else:
73
+ return StartStreamResponse(success=False, message=f"Failed to start stream {stream_id}")
backend/app/connectors/graylog/services/monitoring.py
new
+86
@@ -0,0 +1,86 @@
1
+from loguru import logger
2
+
3
+from app.connectors.graylog.schema.monitoring import GraylogMessages
4
+from app.connectors.graylog.schema.monitoring import GraylogMessagesResponse
5
+from app.connectors.graylog.schema.monitoring import GraylogMetricsResponse
6
+from app.connectors.graylog.schema.monitoring import GraylogThroughputMetrics
7
+from app.connectors.graylog.schema.monitoring import GraylogThroughputMetricsCollection
8
+from app.connectors.graylog.schema.monitoring import GraylogUncommittedJournalEntries
9
+from app.connectors.graylog.utils.universal import send_get_request
10
+
11
+
12
+def get_messages(page_number: int) -> GraylogMessagesResponse:
13
+ """Get messages from Graylog."""
14
+ logger.info("Getting messages from Graylog")
15
+ params = {"page": page_number}
16
+ messages_collected = send_get_request(endpoint="/api/system/messages", params=params)
17
+ if messages_collected["success"]:
18
+ graylog_messages_list = []
19
+ for message in messages_collected["data"]["messages"]:
20
+ graylog_message = GraylogMessages(
21
+ caller=message["caller"],
22
+ content=message["content"],
23
+ node_id=message["node_id"],
24
+ timestamp=message["timestamp"],
25
+ )
26
+ graylog_messages_list.append(graylog_message)
27
+ return GraylogMessagesResponse(
28
+ graylog_messages=graylog_messages_list,
29
+ success=True,
30
+ message="Messages collected successfully",
31
+ total_messages=messages_collected["data"]["total"],
32
+ )
33
+ else:
34
+ return GraylogMessagesResponse(graylog_messages=[], success=False, message="Failed to collect messages")
35
+
36
+
37
+def fetch_metrics_from_graylog() -> dict:
38
+ return send_get_request(endpoint="/api/system/metrics")
39
+
40
+
41
+def fetch_uncommitted_journal_entries() -> dict:
42
+ return send_get_request(endpoint="/api/system/journal")
43
+
44
+
45
+def merge_metrics_data(throughput_metrics_collected: dict) -> dict:
46
+ throughput_metrics = throughput_metrics_collected["data"]["gauges"]
47
+ input_output_metrics = throughput_metrics_collected["data"]["counters"]
48
+ return {**throughput_metrics, **input_output_metrics}
49
+
50
+
51
+def filter_and_create_throughput_metrics(merged_metrics: dict) -> list:
52
+ model_fields = [field_info.alias for field_info in GraylogThroughputMetricsCollection.__fields__.values()]
53
+ throughput_metrics_list = [
54
+ GraylogThroughputMetrics(metric=metric_name, value=metric_data.get("value", 0))
55
+ for metric_name, metric_data in merged_metrics.items()
56
+ if metric_name in model_fields
57
+ ]
58
+ return throughput_metrics_list
59
+
60
+
61
+def get_metrics() -> GraylogMetricsResponse:
62
+ logger.info("Getting metrics from Graylog")
63
+ throughput_metrics_collected = fetch_metrics_from_graylog()
64
+ uncommitted_journal_entries_collected = fetch_uncommitted_journal_entries()
65
+
66
+ if throughput_metrics_collected["success"] and uncommitted_journal_entries_collected["success"]:
67
+ merged_metrics = merge_metrics_data(throughput_metrics_collected)
68
+ throughput_metrics_list = filter_and_create_throughput_metrics(merged_metrics)
69
+
70
+ uncommitted_journal_entries = GraylogUncommittedJournalEntries(
71
+ uncommitted_journal_entries=uncommitted_journal_entries_collected["data"]["uncommitted_journal_entries"],
72
+ )
73
+
74
+ return GraylogMetricsResponse(
75
+ throughput_metrics=throughput_metrics_list,
76
+ uncommitted_journal_entries=uncommitted_journal_entries.uncommitted_journal_entries,
77
+ success=True,
78
+ message="Metrics collected successfully",
79
+ )
80
+ else:
81
+ return GraylogMetricsResponse(
82
+ throughput_metrics=[],
83
+ uncommitted_journal_entries=0,
84
+ success=False,
85
+ message="Failed to collect metrics",
86
+ )
backend/app/connectors/graylog/services/pipelines.py
new
+29
@@ -0,0 +1,29 @@
1
+from loguru import logger
2
+
3
+from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse
4
+from app.connectors.graylog.schema.pipelines import Pipeline
5
+from app.connectors.graylog.schema.pipelines import PipelineRule
6
+from app.connectors.graylog.schema.pipelines import PipelineRulesResponse
7
+from app.connectors.graylog.utils.universal import send_get_request
8
+
9
+
10
+def get_pipelines() -> GraylogPipelinesResponse:
11
+ """Get pipelines from Graylog."""
12
+ logger.info("Getting pipelines from Graylog")
13
+ pipelines_collected = send_get_request(endpoint="/api/system/pipelines/pipeline")
14
+ if pipelines_collected["success"]:
15
+ pipelines_list = [Pipeline(**pipeline_data) for pipeline_data in pipelines_collected["data"]]
16
+ return GraylogPipelinesResponse(pipelines=pipelines_list, success=True, message="Pipelines collected successfully")
17
+ else:
18
+ return GraylogPipelinesResponse(pipelines=[], success=False, message="Failed to collect pipelines")
19
+
20
+
21
+def get_pipeline_rules() -> PipelineRulesResponse:
22
+ """Get pipeline rules from Graylog."""
23
+ logger.info("Getting pipeline rules from Graylog")
24
+ pipeline_rules_collected = send_get_request(endpoint="/api/system/pipelines/rule")
25
+ if pipeline_rules_collected["success"]:
26
+ pipeline_rules_list = [PipelineRule(**pipeline_rule_data) for pipeline_rule_data in pipeline_rules_collected["data"]]
27
+ return PipelineRulesResponse(pipeline_rules=pipeline_rules_list, success=True, message="Pipeline rules collected successfully")
28
+ else:
29
+ return PipelineRulesResponse(pipeline_rules=[], success=False, message="Failed to collect pipeline rules")
backend/app/connectors/graylog/services/streams.py
new
+33
@@ -0,0 +1,33 @@
1
+from typing import List
2
+
3
+from loguru import logger
4
+
5
+from app.connectors.graylog.schema.streams import GraylogStreamsResponse
6
+from app.connectors.graylog.schema.streams import Stream
7
+from app.connectors.graylog.utils.universal import send_get_request
8
+
9
+
10
+def get_streams() -> GraylogStreamsResponse:
11
+ """Get streams from Graylog."""
12
+ logger.info("Getting streams from Graylog")
13
+ streams_collected = send_get_request(endpoint="/api/streams")
14
+ if streams_collected["success"]:
15
+ streams_list = [Stream(**stream_data) for stream_data in streams_collected["data"]["streams"]]
16
+ return GraylogStreamsResponse(
17
+ streams=streams_list,
18
+ success=True,
19
+ message="Streams collected successfully",
20
+ total=streams_collected["data"]["total"],
21
+ )
22
+ else:
23
+ return GraylogStreamsResponse(streams=[], success=False, message="Failed to collect streams", total=0)
24
+
25
+
26
+def get_stream_ids() -> List[str]:
27
+ """Get stream IDs from Graylog."""
28
+ logger.info("Getting stream IDs from Graylog")
29
+ streams_collected = send_get_request(endpoint="/api/streams")
30
+ if streams_collected["success"]:
31
+ return [stream_data["id"] for stream_data in streams_collected["data"]["streams"]]
32
+ else:
33
+ return []
backend/app/connectors/graylog/utils/universal.py
new
+207
@@ -0,0 +1,207 @@
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
+
10
+HEADERS = {"X-Requested-By": "CoPilot"}
11
+
12
+
13
+def verify_graylog_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
14
+ """
15
+ Verifies the connection to Graylog service.
16
+
17
+ Returns:
18
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
19
+ """
20
+ logger.info(
21
+ f"Verifying the graylog connection to {attributes['connector_url']}",
22
+ )
23
+ try:
24
+ graylog_roles = requests.get(
25
+ f"{attributes['connector_url']}/api/authz/roles/user/{attributes['connector_username']}",
26
+ auth=(
27
+ attributes["connector_username"],
28
+ attributes["connector_password"],
29
+ ),
30
+ verify=False,
31
+ )
32
+ if graylog_roles.status_code == 200:
33
+ logger.info(
34
+ f"Connection to {attributes['connector_url']} successful",
35
+ )
36
+ return {"connectionSuccessful": True, "message": "Graylog connection successful"}
37
+ else:
38
+ logger.error(
39
+ f"Connection to {attributes['connector_url']} failed with error: {graylog_roles.text}",
40
+ )
41
+ return {
42
+ "connectionSuccessful": False,
43
+ "message": f"Connection to {attributes['connector_url']} failed with error: {graylog_roles.text}",
44
+ }
45
+ except Exception as e:
46
+ logger.error(
47
+ f"Connection to {attributes['connector_url']} failed with error: {e}",
48
+ )
49
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
50
+
51
+
52
+def verify_graylog_connection(connector_name: str) -> str:
53
+ """
54
+ Returns if connection to Graylog service is successful.
55
+ """
56
+ logger.info("Getting Graylog authentication token")
57
+ attributes = get_connector_info_from_db(connector_name)
58
+ if attributes is None:
59
+ logger.error("No Graylog connector found in the database")
60
+ return None
61
+ return verify_graylog_credentials(attributes)
62
+
63
+
64
+def send_get_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
65
+ """
66
+ Sends a GET request to the Graylog service.
67
+
68
+ Args:
69
+ endpoint (str): The endpoint to send the GET request to.
70
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the GET request. Defaults to None.
71
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylogr".
72
+
73
+ Returns:
74
+ Dict[str, Any]: The response from the GET request.
75
+ """
76
+ logger.info(f"Sending GET request to {endpoint}")
77
+ attributes = get_connector_info_from_db(connector_name)
78
+ if attributes is None:
79
+ logger.error("No Graylog connector found in the database")
80
+ return None
81
+ try:
82
+ response = requests.get(
83
+ f"{attributes['connector_url']}{endpoint}",
84
+ headers=HEADERS,
85
+ auth=(
86
+ attributes["connector_username"],
87
+ attributes["connector_password"],
88
+ ),
89
+ params=params,
90
+ verify=False,
91
+ )
92
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
93
+ except Exception as e:
94
+ logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
95
+ return {"success": False, "message": f"Failed to send GET request to {endpoint} with error: {e}"}
96
+
97
+
98
+def send_post_request(endpoint: str, data: Dict[str, Any] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
99
+ """
100
+ Sends a POST request to the Graylog service.
101
+
102
+ Args:
103
+ endpoint (str): The endpoint to send the POST request to.
104
+ data (Dict[str, Any]): The data to send with the POST request.
105
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylog".
106
+
107
+ Returns:
108
+ Dict[str, Any]: The response from the POST request.
109
+ """
110
+ logger.info(f"Sending POST request to {endpoint}")
111
+ attributes = get_connector_info_from_db(connector_name)
112
+ if attributes is None:
113
+ logger.error("No Graylog connector found in the database")
114
+ return {"success": False, "message": "No Graylog connector found in the database"}
115
+
116
+ try:
117
+ response = requests.post(
118
+ f"{attributes['connector_url']}{endpoint}",
119
+ headers=HEADERS,
120
+ auth=(
121
+ attributes["connector_username"],
122
+ attributes["connector_password"],
123
+ ),
124
+ json=data,
125
+ verify=False,
126
+ )
127
+
128
+ if response.status_code == 204:
129
+ return {"data": None, "success": True, "message": "Successfully completed request with no content"}
130
+ else:
131
+ return {
132
+ "data": response.json(),
133
+ "success": False if response.status_code >= 400 else True,
134
+ "message": "Successfully retrieved data" if response.status_code < 400 else "Failed to retrieve data",
135
+ }
136
+ except Exception as e:
137
+ logger.debug(f"Response: {response}")
138
+ logger.error(f"Failed to send POST request to {endpoint} with error: {e}")
139
+ return {"success": False, "message": f"Failed to send POST request to {endpoint} with error: {e}"}
140
+
141
+
142
+def send_delete_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
143
+ """
144
+ Sends a DELETE request to the Graylog service.
145
+
146
+ Args:
147
+ endpoint (str): The endpoint to send the DELETE request to.
148
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the DELETE request. Defaults to None.
149
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylog".
150
+
151
+ Returns:
152
+ Dict[str, Any]: The response from the DELETE request.
153
+ """
154
+ logger.info(f"Sending DELETE request to {endpoint}")
155
+ attributes = get_connector_info_from_db(connector_name)
156
+ if attributes is None:
157
+ logger.error("No Graylog connector found in the database")
158
+ return None
159
+ try:
160
+ response = requests.delete(
161
+ f"{attributes['connector_url']}{endpoint}",
162
+ headers=HEADERS,
163
+ auth=(
164
+ attributes["connector_username"],
165
+ attributes["connector_password"],
166
+ ),
167
+ params=params,
168
+ verify=False,
169
+ )
170
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
171
+ except Exception as e:
172
+ logger.error(f"Failed to send DELETE request to {endpoint} with error: {e}")
173
+ return {"success": False, "message": f"Failed to send DELETE request to {endpoint} with error: {e}"}
174
+
175
+
176
+def send_put_request(endpoint: str, data: Optional[Dict[str, Any]] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
177
+ """
178
+ Sends a PUT request to the Graylog service.
179
+
180
+ Args:
181
+ endpoint (str): The endpoint to send the PUT request to.
182
+ data (Optional[Dict[str, Any]]): The data to send with the PUT request.
183
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylog".
184
+
185
+ Returns:
186
+ Dict[str, Any]: The response from the PUT request.
187
+ """
188
+ logger.info(f"Sending PUT request to {endpoint}")
189
+ attributes = get_connector_info_from_db(connector_name)
190
+ if attributes is None:
191
+ logger.error("No Graylog connector found in the database")
192
+ return None
193
+ try:
194
+ response = requests.put(
195
+ f"{attributes['connector_url']}{endpoint}",
196
+ headers=HEADERS,
197
+ auth=(
198
+ attributes["connector_username"],
199
+ attributes["connector_password"],
200
+ ),
201
+ json=data,
202
+ verify=False,
203
+ )
204
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
205
+ except Exception as e:
206
+ logger.error(f"Failed to send PUT request to {endpoint} with error: {e}")
207
+ return {"success": False, "message": f"Failed to send PUT request to {endpoint} with error: {e}"}
backend/app/connectors/models.py
new
+95
@@ -0,0 +1,95 @@
1
+from datetime import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from sqlmodel import Field
6
+from sqlmodel import Relationship
7
+from sqlmodel import SQLModel
8
+
9
+
10
+class ConnectorHistory(SQLModel, table=True):
11
+ """
12
+ Model representing the history logs of each connector.
13
+
14
+ :ivar id: Unique integer ID of the history log.
15
+ :ivar connector_id: Foreign key linking to the Connectors table.
16
+ :ivar change_timestamp: Timestamp when the change was made.
17
+ :ivar change_description: Description of the change.
18
+ :ivar connector: Relationship to the Connectors model.
19
+ """
20
+
21
+ id: Optional[int] = Field(default=None, primary_key=True)
22
+ connector_id: int = Field(foreign_key="connectors.id")
23
+ change_timestamp: datetime = Field(default=datetime.utcnow())
24
+ change_description: str
25
+
26
+ # Relationship
27
+ connector: Optional["Connectors"] = Relationship(back_populates="history_logs")
28
+
29
+
30
+class Connectors(SQLModel, table=True):
31
+ """
32
+ Model representing each connector and its attributes.
33
+
34
+ :ivar id: Unique integer ID of the connector.
35
+ :ivar connector_name: Name of the connector.
36
+ :ivar connector_type: Type or version of the connector.
37
+ :ivar connector_url: URL endpoint of the connector.
38
+ :ivar connector_last_updated: Timestamp when the connector was last updated.
39
+ :ivar connector_username: Optional username for the connector.
40
+ :ivar connector_password: Optional password for the connector.
41
+ :ivar connector_api_key: Optional API key for the connector.
42
+ :ivar connector_description: Description of what the connector does.
43
+ :ivar connector_supports: Information on what the connector supports.
44
+ :ivar connector_configured: Boolean indicating if the connector is configured.
45
+ :ivar connector_verified: Boolean indicating if the connector is verified.
46
+ :ivar connector_accepts_api_key: Boolean indicating if the connector accepts API keys.
47
+ :ivar connector_accepts_username_password: Boolean indicating if the connector accepts username and password.
48
+ :ivar connector_accepts_file: Boolean indicating if the connector accepts files.
49
+ :ivar history_logs: Relationship to the ConnectorHistory model.
50
+ """
51
+
52
+ id: Optional[int] = Field(default=None, primary_key=True)
53
+ connector_name: str = Field()
54
+ connector_type: str = Field()
55
+ connector_url: str = Field()
56
+ connector_last_updated: datetime = Field(default=datetime.utcnow())
57
+ connector_username: Optional[str] = Field(default=None)
58
+ connector_password: Optional[str] = Field(default=None)
59
+ connector_api_key: Optional[str] = Field(default=None)
60
+
61
+ # Fields moved from ConnectorsAvailable
62
+ connector_description: Optional[str] = Field(default=None)
63
+ connector_supports: Optional[str] = Field(default=None)
64
+ connector_configured: bool = Field(default=False)
65
+ connector_verified: bool = Field(default=False)
66
+ connector_accepts_api_key: bool = Field(default=False)
67
+ connector_accepts_username_password: bool = Field(default=False)
68
+ connector_accepts_file: bool = Field(default=False)
69
+
70
+ # Relationship
71
+ history_logs: List[ConnectorHistory] = Relationship(back_populates="connector")
72
+
73
+
74
+# Example usage
75
+# new_connector = Connectors(
76
+# connector_name="Wazuh-Indexer",
77
+# connector_type="4.4.1",
78
+# connector_url="https://ashwix01.socfortress.local:9200",
79
+# connector_username="admin",
80
+# connector_password="password_here",
81
+# connector_api_key="api_key_here",
82
+# # Fields from ConnectorsAvailable
83
+# connector_description="Description here",
84
+# connector_supports="Supports list here",
85
+# connector_configured=True,
86
+# connector_verified=True,
87
+# connector_accepts_api_key=True,
88
+# connector_accepts_username_password=True,
89
+# connector_accepts_file=False
90
+# )
91
+
92
+# new_log = ConnectorHistory(
93
+# connector_id=1, # This should be the ID of the corresponding connector
94
+# change_description="Changed the API key."
95
+# )
backend/app/connectors/routes.py
new
+159
@@ -0,0 +1,159 @@
1
+from typing import Union
2
+
3
+## Auth Things
4
+from fastapi import APIRouter
5
+from fastapi import File
6
+from fastapi import HTTPException
7
+from fastapi import Security
8
+from fastapi import UploadFile
9
+from loguru import logger
10
+
11
+from app.auth.utils import AuthHandler
12
+from app.connectors.schema import ConnectorListResponse
13
+from app.connectors.schema import ConnectorResponse
14
+from app.connectors.schema import ConnectorsListResponse
15
+from app.connectors.schema import UpdateConnector
16
+from app.connectors.schema import VerifyConnectorResponse
17
+from app.connectors.services import ConnectorServices
18
+
19
+connector_router = APIRouter()
20
+
21
+
22
+@connector_router.get(
23
+ "",
24
+ response_model=ConnectorsListResponse,
25
+ description="Fetch all available connectors",
26
+ dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
27
+)
28
+async def get_connectors() -> ConnectorListResponse:
29
+ """
30
+ Fetch all available connectors from the database.
31
+
32
+ This endpoint retrieves all the connectors stored in the database and returns them
33
+ along with a success status and message.
34
+
35
+ Returns:
36
+ ConnectorListResponse: A Pydantic model containing a list of connectors and additional metadata.
37
+
38
+ Raises:
39
+ HTTPException: An exception with a 404 status code is raised if no connectors are found.
40
+ """
41
+
42
+ connectors = ConnectorServices.fetch_all_connectors()
43
+ if connectors:
44
+ return {"connectors": connectors, "success": True, "message": "Connectors fetched successfully"}
45
+ else:
46
+ raise HTTPException(status_code=404, detail="No connectors found")
47
+
48
+
49
+@connector_router.get(
50
+ "/{connector_id}",
51
+ response_model=ConnectorListResponse,
52
+ description="Fetch a specific connector",
53
+ dependencies=[Security(AuthHandler().require_any_scope("admin", "test"))],
54
+)
55
+async def get_connector(connector_id: int) -> Union[ConnectorResponse, HTTPException]:
56
+ """
57
+ Fetch a specific connector by its ID.
58
+
59
+ This endpoint retrieves a connector identified by `connector_id` from the database.
60
+
61
+ Args:
62
+ connector_id (int): The unique identifier for the connector to fetch.
63
+
64
+ Returns:
65
+ ConnectorResponse: A Pydantic model representing the fetched connector.
66
+
67
+ Raises:
68
+ HTTPException: An exception with a 404 status code is raised if the connector is not found.
69
+ """
70
+ connector = ConnectorServices.fetch_connector_by_id(connector_id)
71
+ if connector is not None:
72
+ return {"connector": connector, "success": True, "message": "Connector fetched successfully"}
73
+ else:
74
+ raise HTTPException(status_code=404, detail=f"No connector found for ID: {connector_id}".format(connector_id=connector_id))
75
+
76
+
77
+@connector_router.post(
78
+ "/verify/{connector_id}",
79
+ response_model=VerifyConnectorResponse,
80
+ description="Verify a connector. Makes an API call to the connector to verify it is working.",
81
+)
82
+async def verify_connector(connector_id: int) -> Union[VerifyConnectorResponse, HTTPException]:
83
+ """
84
+ Verify a connector by its ID.
85
+
86
+ This endpoint verifies a connector identified by `connector_id` by making an API call to the connector.
87
+
88
+ Args:
89
+ connector_id (int): The unique identifier for the connector to verify.
90
+
91
+ Returns:
92
+ ConnectorResponse: A Pydantic model representing the verified connector.
93
+
94
+ Raises:
95
+ HTTPException: An exception with a 404 status code is raised if the connector is not found.
96
+ """
97
+ connector = ConnectorServices.verify_connector_by_id(connector_id)
98
+ if connector is not None:
99
+ logger.info(f"Connector verified successfully: {connector}")
100
+ return connector
101
+ else:
102
+ raise HTTPException(status_code=404, detail=f"No connector found for ID: {connector_id}".format(connector_id=connector_id))
103
+
104
+
105
+@connector_router.put("/{connector_id}", response_model=ConnectorListResponse, description="Update a connector")
106
+async def update_connector(connector_id: int, connector: UpdateConnector) -> ConnectorListResponse:
107
+ """
108
+ Update a connector by its ID.
109
+
110
+ This endpoint updates a connector identified by `connector_id` in the database.
111
+
112
+ Args:
113
+ connector_id (int): The unique identifier for the connector to update.
114
+ connector (ConnectorListResponse): The updated connector data.
115
+
116
+ Returns:
117
+ ConnectorListResponse: A Pydantic model representing the updated connector.
118
+
119
+ Raises:
120
+ HTTPException: An exception with a 404 status code is raised if the connector is not found.
121
+ """
122
+ updated_connector = ConnectorServices.update_connector_by_id(connector_id, connector)
123
+ if updated_connector is not None:
124
+ return {"connector": updated_connector, "success": True, "message": "Connector updated successfully"}
125
+ else:
126
+ raise HTTPException(status_code=404, detail=f"No connector found for ID: {connector_id}".format(connector_id=connector_id))
127
+
128
+
129
+@connector_router.post("/upload/{connector_id}", description="Upload a YAML file for a specific connector")
130
+async def upload_yaml_file(connector_id: int, file: UploadFile = File(...)) -> dict:
131
+ """
132
+ Upload a YAML file for a specific connector ID.
133
+
134
+ This endpoint allows you to upload a `.yaml` file for a specific connector
135
+ identified by `connector_id`.
136
+
137
+ Args:
138
+ connector_id (int): The unique identifier for the connector.
139
+ file (UploadFile): The `.yaml` file to be uploaded.
140
+
141
+ Returns:
142
+ dict: A dictionary with a success message and other information.
143
+
144
+ Raises:
145
+ HTTPException: An exception with a 400 status code is raised if the file format is incorrect or connector ID is not 6.
146
+ """
147
+ if connector_id != 6:
148
+ raise HTTPException(status_code=400, detail="Only the Velociraptor connector is allowed for YAML file uploads.")
149
+ if not file.filename.endswith(".yaml"):
150
+ raise HTTPException(status_code=400, detail="Only .yaml files are allowed.")
151
+ try:
152
+ save_file_result = ConnectorServices.save_file(file)
153
+ if save_file_result:
154
+ return {"success": True, "message": "File uploaded successfully"}
155
+ else:
156
+ raise HTTPException(status_code=500, detail="Failed to upload file")
157
+ except Exception as e:
158
+ logger.error(f"Failed to upload file: {e}")
159
+ raise HTTPException(status_code=500, detail="Failed to upload file")
backend/app/connectors/schema.py
new
+61
@@ -0,0 +1,61 @@
1
+from datetime import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+
7
+
8
+class ConnectorHistoryResponse(BaseModel):
9
+ id: Optional[int]
10
+ connector_id: int
11
+ change_timestamp: datetime
12
+ change_description: str
13
+
14
+ class Config:
15
+ orm_mode = True
16
+
17
+
18
+class ConnectorResponse(BaseModel):
19
+ id: Optional[int]
20
+ connector_name: str
21
+ connector_type: str
22
+ connector_url: str
23
+ connector_last_updated: datetime
24
+ connector_username: Optional[str]
25
+ connector_password: Optional[str]
26
+ connector_api_key: Optional[str]
27
+ connector_description: Optional[str]
28
+ connector_supports: Optional[str]
29
+ connector_configured: bool
30
+ connector_verified: bool
31
+ connector_accepts_api_key: bool
32
+ connector_accepts_username_password: bool
33
+ connector_accepts_file: bool
34
+ history_logs: Optional[List[ConnectorHistoryResponse]]
35
+
36
+ class Config:
37
+ orm_mode = True
38
+
39
+
40
+class ConnectorsListResponse(BaseModel):
41
+ connectors: List[ConnectorResponse]
42
+ success: bool
43
+ message: str
44
+
45
+
46
+class ConnectorListResponse(BaseModel):
47
+ connector: ConnectorResponse
48
+ success: bool
49
+ message: str
50
+
51
+
52
+class VerifyConnectorResponse(BaseModel):
53
+ connectionSuccessful: bool
54
+ message: str
55
+
56
+
57
+class UpdateConnector(BaseModel):
58
+ connector_url: str
59
+ connector_username: Optional[str]
60
+ connector_password: Optional[str]
61
+ connector_api_key: Optional[str]
backend/app/connectors/services.py
new
+284
@@ -0,0 +1,284 @@
1
+import os
2
+from contextlib import contextmanager
3
+from datetime import datetime
4
+from typing import Generator
5
+from typing import List
6
+from typing import Optional
7
+from typing import Type
8
+
9
+from fastapi import UploadFile
10
+from loguru import logger
11
+from pydantic import BaseModel
12
+from sqlmodel import Session
13
+from sqlmodel 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.graylog.utils.universal import verify_graylog_connection
19
+from app.connectors.models import Connectors
20
+from app.connectors.schema import ConnectorResponse
21
+from app.connectors.shuffle.utils.universal import verify_shuffle_connection
22
+from app.connectors.sublime.utils.universal import verify_sublime_connection
23
+from app.connectors.velociraptor.utils.universal import verify_velociraptor_connection
24
+from app.connectors.wazuh_indexer.utils.universal import verify_wazuh_indexer_connection
25
+from app.connectors.wazuh_manager.utils.universal import verify_wazuh_manager_connection
26
+from app.db.db_session import engine # Import the shared engine
27
+
28
+UPLOAD_FOLDER = "file-store"
29
+UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), UPLOAD_FOLDER)
30
+ALLOWED_EXTENSIONS = set(["yaml"]) # replace with your allowed file extensions
31
+
32
+
33
+# Create an interface for connector services
34
+class ConnectorServiceInterface(BaseModel):
35
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
36
+ raise NotImplementedError
37
+
38
+
39
+# Wazuh Manager Service
40
+class WazuhManagerService(ConnectorServiceInterface):
41
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
42
+ return verify_wazuh_manager_connection(connector.connector_name)
43
+
44
+
45
+# Wazuh Indexer Service
46
+class WazuhIndexerService(ConnectorServiceInterface):
47
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
48
+ return verify_wazuh_indexer_connection(connector.connector_name)
49
+
50
+
51
+# Velociraptor Service
52
+class VelociraptorService(ConnectorServiceInterface):
53
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
54
+ return verify_velociraptor_connection(connector.connector_name)
55
+
56
+
57
+# Graylog Service
58
+class GraylogService(ConnectorServiceInterface):
59
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
60
+ return verify_graylog_connection(connector.connector_name)
61
+
62
+
63
+# DFIR-IRIS Service
64
+class DfirIrisService(ConnectorServiceInterface):
65
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
66
+ return verify_dfir_iris_connection(connector.connector_name)
67
+
68
+
69
+# Cortex Service
70
+class CortexService(ConnectorServiceInterface):
71
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
72
+ return verify_cortex_connection(connector.connector_name)
73
+
74
+
75
+# Shuffle Service
76
+class ShuffleService(ConnectorServiceInterface):
77
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
78
+ return verify_shuffle_connection(connector.connector_name)
79
+
80
+
81
+# Sublime Service
82
+class SublimeService(ConnectorServiceInterface):
83
+ def verify_authentication(self, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
84
+ return verify_sublime_connection(connector.connector_name)
85
+
86
+
87
+# Factory function to create a service instance based on connector name
88
+def get_connector_service(connector_name: str) -> Type[ConnectorServiceInterface]:
89
+ service_map = {
90
+ "Wazuh-Manager": WazuhManagerService,
91
+ "Wazuh-Indexer": WazuhIndexerService,
92
+ "Velociraptor": VelociraptorService,
93
+ "Graylog": GraylogService,
94
+ "DFIR-IRIS": DfirIrisService,
95
+ "Cortex": CortexService,
96
+ "Shuffle": ShuffleService,
97
+ "Sublime": SublimeService,
98
+ }
99
+ return service_map.get(connector_name, None)
100
+
101
+
102
+class ConnectorServices:
103
+ """
104
+ Service class for handling operations related to connectors.
105
+ """
106
+
107
+ @staticmethod
108
+ @contextmanager
109
+ def get_session() -> Generator[Session, None, None]:
110
+ """
111
+ Get a new session for database interaction.
112
+
113
+ This method is a context manager, which ensures that the session is closed
114
+ once the operations within the context are completed.
115
+
116
+ Yields:
117
+ Session: The database session object.
118
+ """
119
+ session = Session(engine)
120
+ try:
121
+ yield session
122
+ finally:
123
+ session.close()
124
+
125
+ @classmethod
126
+ def fetch_all_connectors(cls) -> List[ConnectorResponse]:
127
+ """
128
+ Fetch all connectors from the database.
129
+
130
+ This method retrieves all connector records from the database, converts them
131
+ to Pydantic models, and returns them as a list.
132
+
133
+ Returns:
134
+ List[ConnectorResponse]: A list of connectors in their Pydantic representation.
135
+ """
136
+ # Get a new session
137
+ with cls.get_session() as session:
138
+ query = select(Connectors)
139
+ connectors = session.exec(query).all()
140
+
141
+ # Convert the SQLModel object to a Pydantic model
142
+ connector_responses = [ConnectorResponse.from_orm(connector) for connector in connectors]
143
+ return connector_responses
144
+
145
+ @classmethod
146
+ def fetch_connector_by_id(cls, connector_id: int) -> Optional[ConnectorResponse]:
147
+ """
148
+ Fetch a connector by its ID from the database.
149
+
150
+ Given a connector ID, this method retrieves the corresponding connector
151
+ record from the database, if it exists.
152
+
153
+ Args:
154
+ connector_id (int): The ID of the connector to fetch.
155
+
156
+ Returns:
157
+ Optional[ConnectorResponse]: The connector in its Pydantic representation, or None if not found.
158
+ """
159
+ # Get a new session
160
+ with cls.get_session() as session:
161
+ query = select(Connectors).where(Connectors.id == connector_id)
162
+ connector = session.exec(query).first()
163
+
164
+ if not connector:
165
+ logger.info(f"No connector found for ID: {connector_id}")
166
+ return None
167
+
168
+ try:
169
+ # Convert the SQLModel object to a Pydantic model
170
+ connector_response = ConnectorResponse.from_orm(connector)
171
+ return connector_response
172
+ except Exception as e:
173
+ logger.exception(f"Failed to create ConnectorResponse object: {e}")
174
+ return None
175
+
176
+ @classmethod
177
+ def verify_connector_by_id(cls, connector_id: int) -> Optional[ConnectorResponse]:
178
+ """
179
+ Verify a connector by making an API call to it.
180
+
181
+ Given a connector ID, this method retrieves the corresponding connector
182
+ record from the database, if it exists, and makes an API call to the connector.
183
+
184
+ Args:
185
+ connector_id (int): The ID of the connector to verify.
186
+
187
+ Returns:
188
+ Optional[ConnectorResponse]: The connector in its Pydantic representation, or None if not found.
189
+ """
190
+ # Get a new session
191
+ with cls.get_session() as session:
192
+ query = select(Connectors).where(Connectors.id == connector_id)
193
+ connector = session.exec(query).first()
194
+
195
+ if not connector:
196
+ logger.info(f"No connector found for ID: {connector_id}")
197
+ return None
198
+
199
+ try:
200
+ # Convert the SQLModel object to a Pydantic model
201
+ connector_response = ConnectorResponse.from_orm(connector)
202
+
203
+ # Get the appropriate service for this connector
204
+ ServiceClass = get_connector_service(connector_response.connector_name)
205
+
206
+ if ServiceClass is not None:
207
+ service_instance = ServiceClass()
208
+ connector_response = service_instance.verify_authentication(connector_response)
209
+ else:
210
+ logger.error(f"Connector type {connector_response.connector_name} is not supported")
211
+ return None
212
+
213
+ return connector_response
214
+ except Exception as e:
215
+ logger.exception(f"Failed to create ConnectorResponse object: {e}")
216
+ return None
217
+
218
+ @classmethod
219
+ def update_connector_by_id(cls, connector_id: int, connector: ConnectorResponse) -> Optional[ConnectorResponse]:
220
+ """
221
+ Update a connector by its ID in the database.
222
+
223
+ Given a connector ID and a Pydantic representation of a connector, this method
224
+ updates the corresponding connector record in the database, if it exists.
225
+
226
+ Args:
227
+ connector_id (int): The ID of the connector to update.
228
+ connector (ConnectorResponse): The updated connector in its Pydantic representation.
229
+
230
+ Returns:
231
+ Optional[ConnectorResponse]: The updated connector in its Pydantic representation, or None if not found.
232
+ """
233
+ # Get a new session
234
+ with cls.get_session() as session:
235
+ query = select(Connectors).where(Connectors.id == connector_id)
236
+ connector_record = session.exec(query).first()
237
+
238
+ if not connector_record:
239
+ logger.info(f"No connector found for ID: {connector_id}")
240
+ return None
241
+
242
+ try:
243
+ # Update the connector record
244
+ connector_record.connector_url = connector.connector_url
245
+ connector_record.connector_username = connector.connector_username
246
+ connector_record.connector_password = connector.connector_password
247
+ connector_record.connector_api_key = connector.connector_api_key
248
+ connector_record.connector_last_updated = datetime.now()
249
+
250
+ # Commit the changes to the database
251
+ session.add(connector_record)
252
+ session.commit()
253
+
254
+ # Convert the SQLModel object to a Pydantic model
255
+ connector_response = ConnectorResponse.from_orm(connector_record)
256
+ return connector_response
257
+ except Exception as e:
258
+ logger.exception(f"Failed to update connector: {e}")
259
+ return None
260
+
261
+ @staticmethod
262
+ def allowed_file(filename):
263
+ return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
264
+
265
+ @classmethod
266
+ def save_file(cls, file: UploadFile):
267
+ if file and cls.allowed_file(file.filename):
268
+ filename = secure_filename(file.filename)
269
+ file_path = os.path.join(UPLOAD_FOLDER, filename)
270
+
271
+ # Save the file
272
+ with open(file_path, "wb") as buffer:
273
+ buffer.write(file.file.read())
274
+
275
+ # Update connector
276
+ connector = cls.fetch_connector_by_id(6)
277
+ connector.connector_configured = True
278
+ connector.connector_api_key = file_path
279
+ cls.update_connector_by_id(6, connector)
280
+
281
+ connector_response = ConnectorResponse.from_orm(connector)
282
+ return connector_response
283
+ else:
284
+ return False
backend/app/connectors/shuffle/routes/workflows.py
new
+45
@@ -0,0 +1,45 @@
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from loguru import logger
4
+
5
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel
6
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionResponseModel
7
+from app.connectors.shuffle.schema.workflows import WorkflowsResponse
8
+from app.connectors.shuffle.services.workflows import get_workflow_executions
9
+from app.connectors.shuffle.services.workflows import get_workflows
10
+
11
+shuffle_workflows_router = APIRouter()
12
+
13
+
14
+@shuffle_workflows_router.get("", response_model=WorkflowsResponse, description="Get all workflows")
15
+async def get_all_workflows() -> WorkflowsResponse:
16
+ logger.info("Fetching all workflows")
17
+ return get_workflows()
18
+
19
+
20
+@shuffle_workflows_router.get("/executions", response_model=WorkflowExecutionResponseModel, description="Get all workflow executions")
21
+async def get_all_workflow_executions() -> WorkflowExecutionResponseModel:
22
+ logger.info("Fetching all workflow executions")
23
+
24
+ # Initialize an empty list for storing workflow details
25
+ workflow_details = []
26
+
27
+ # Get the workflow response by awaiting the asynchronous function get_workflows()
28
+ workflow_response = await get_all_workflows()
29
+
30
+ # Access the workflows attribute from the response
31
+ workflows = workflow_response.workflows
32
+
33
+ # Check if workflows is not None before proceeding
34
+ if workflows:
35
+ for workflow in workflows:
36
+ workflow_details.append(
37
+ {
38
+ "workflow_id": workflow["id"],
39
+ "workflow_name": workflow["name"],
40
+ "status": get_workflow_executions(WorkflowExecutionBodyModel(workflow_id=workflow["id"])),
41
+ },
42
+ )
43
+ return WorkflowExecutionResponseModel(success=True, message="Successfully fetched workflow executions", workflows=workflow_details)
44
+ else:
45
+ raise HTTPException(status_code=404, detail="No workflows found")
backend/app/connectors/shuffle/schema/workflows.py
new
+39
@@ -0,0 +1,39 @@
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+from typing import Optional
5
+
6
+from pydantic import BaseModel
7
+from pydantic import Field
8
+
9
+
10
+class WorkflowsResponse(BaseModel):
11
+ message: str
12
+ success: bool
13
+ workflows: Optional[List[Dict[str, Any]]] = Field([], description="The alerts returned from the search.")
14
+
15
+
16
+class WorkflowStatusExecutionModel(BaseModel):
17
+ executions: Optional[str] = Field(None, description="Status of workflow executions")
18
+ message: str = Field(..., description="Status message")
19
+ success: bool = Field(..., description="Success status")
20
+
21
+
22
+class WorkflowExecutionBodyModel(BaseModel):
23
+ workflow_id: str = Field(..., description="Unique identifier for the workflow")
24
+
25
+
26
+class WorkflowExecutionStatusResponseModel(BaseModel):
27
+ last_run: Optional[str] = Field(..., description="Status of workflow executions")
28
+
29
+
30
+class WorkflowExecutionModel(BaseModel):
31
+ status: WorkflowExecutionStatusResponseModel = Field(..., description="Status object")
32
+ workflow_id: str = Field(..., description="Unique identifier for the workflow")
33
+ workflow_name: str = Field(..., description="Name of the workflow")
34
+
35
+
36
+class WorkflowExecutionResponseModel(BaseModel):
37
+ message: str = Field(..., description="Response message")
38
+ success: bool = Field(..., description="Success status")
39
+ workflows: List[WorkflowExecutionModel] = Field(..., description="List of workflow objects")
backend/app/connectors/shuffle/services/workflows.py
new
+33
@@ -0,0 +1,33 @@
1
+from loguru import logger
2
+
3
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionBodyModel
4
+from app.connectors.shuffle.schema.workflows import WorkflowExecutionStatusResponseModel
5
+from app.connectors.shuffle.schema.workflows import WorkflowsResponse
6
+from app.connectors.shuffle.utils.universal import send_get_request
7
+
8
+
9
+def get_workflows() -> WorkflowsResponse:
10
+ """
11
+ Returns a list of workflows.
12
+ """
13
+ logger.info("Getting workflows")
14
+ response = send_get_request("/api/v1/workflows")
15
+ if response is None:
16
+ return WorkflowsResponse(success=False, message="Failed to get workflows", workflows=[])
17
+ return WorkflowsResponse(success=True, message="Successfully fetched workflows", workflows=response["data"])
18
+
19
+
20
+def get_workflow_executions(exection_body: WorkflowExecutionBodyModel) -> WorkflowExecutionStatusResponseModel:
21
+ """
22
+ Returns a list of workflow executions.
23
+ """
24
+ logger.info("Getting workflow executions")
25
+ response = send_get_request(f"/api/v1/workflows/{exection_body.workflow_id}/executions")
26
+ executions = response["data"]
27
+ if executions:
28
+ status = executions[0]["status"]
29
+ if status is None:
30
+ status = "Never Ran"
31
+ else:
32
+ status = "No executions found"
33
+ return WorkflowExecutionStatusResponseModel(last_run=status)
backend/app/connectors/shuffle/utils/universal.py
new
+213
@@ -0,0 +1,213 @@
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
+
10
+
11
+def verify_shuffle_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
12
+ """
13
+ Verifies the connection to Shuffle service.
14
+
15
+ Returns:
16
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
17
+ """
18
+ logger.info(
19
+ f"Verifying the Shuffle connection to {attributes['connector_url']}",
20
+ )
21
+ try:
22
+ headers = {
23
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
24
+ }
25
+ shuffle_apps = requests.get(
26
+ f"{attributes['connector_url']}/api/v1/apps/authentication",
27
+ headers=headers,
28
+ verify=False,
29
+ )
30
+ if shuffle_apps.status_code == 200:
31
+ logger.info(
32
+ f"Connection to {attributes['connector_url']} successful",
33
+ )
34
+ return {"connectionSuccessful": True, "message": "Shuffle connection successful"}
35
+ else:
36
+ logger.error(
37
+ f"Connection to {attributes['connector_url']} failed with error: {shuffle_apps.text}",
38
+ )
39
+ return {
40
+ "connectionSuccessful": False,
41
+ "message": f"Connection to {attributes['connector_url']} failed with error: {shuffle_apps.text}",
42
+ }
43
+ except Exception as e:
44
+ logger.error(
45
+ f"Connection to {attributes['connector_url']} failed with error: {e}",
46
+ )
47
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
48
+
49
+
50
+def verify_shuffle_connection(connector_name: str) -> str:
51
+ """
52
+ Returns if connection to Shuffle service is successful.
53
+ """
54
+ logger.info("Getting Shuffle authentication token")
55
+ attributes = get_connector_info_from_db(connector_name)
56
+ if attributes is None:
57
+ logger.error("No Shuffle connector found in the database")
58
+ return None
59
+ return verify_shuffle_credentials(attributes)
60
+
61
+
62
+def send_get_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Shuffle") -> Dict[str, Any]:
63
+ """
64
+ Sends a GET request to the Shuffle service.
65
+
66
+ Args:
67
+ endpoint (str): The endpoint to send the GET request to.
68
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the GET request. Defaults to None.
69
+ connector_name (str, optional): The name of the connector to use. Defaults to "Shuffle".
70
+
71
+ Returns:
72
+ Dict[str, Any]: The response from the GET request.
73
+ """
74
+ logger.info(f"Sending GET request to {endpoint}")
75
+ attributes = get_connector_info_from_db(connector_name)
76
+ if attributes is None:
77
+ logger.error("No Graylog connector found in the database")
78
+ return None
79
+ try:
80
+ HEADERS = {
81
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
82
+ }
83
+ response = requests.get(
84
+ f"{attributes['connector_url']}{endpoint}",
85
+ headers=HEADERS,
86
+ params=params,
87
+ verify=False,
88
+ )
89
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
90
+ except Exception as e:
91
+ logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
92
+ return {"success": False, "message": f"Failed to send GET request to {endpoint} with error: {e}"}
93
+
94
+
95
+def send_post_request(endpoint: str, data: Dict[str, Any] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
96
+ """
97
+ Sends a POST request to the Graylog service.
98
+
99
+ Args:
100
+ endpoint (str): The endpoint to send the POST request to.
101
+ data (Dict[str, Any]): The data to send with the POST request.
102
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylog".
103
+
104
+ Returns:
105
+ Dict[str, Any]: The response from the POST request.
106
+ """
107
+ logger.info(f"Sending POST request to {endpoint}")
108
+ attributes = get_connector_info_from_db(connector_name)
109
+ if attributes is None:
110
+ logger.error("No Graylog connector found in the database")
111
+ return {"success": False, "message": "No Graylog connector found in the database"}
112
+
113
+ try:
114
+ HEADERS = {
115
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
116
+ }
117
+ response = requests.post(
118
+ f"{attributes['connector_url']}{endpoint}",
119
+ headers=HEADERS,
120
+ auth=(
121
+ attributes["connector_username"],
122
+ attributes["connector_password"],
123
+ ),
124
+ json=data,
125
+ verify=False,
126
+ )
127
+
128
+ if response.status_code == 204:
129
+ return {"data": None, "success": True, "message": "Successfully completed request with no content"}
130
+ else:
131
+ return {
132
+ "data": response.json(),
133
+ "success": False if response.status_code >= 400 else True,
134
+ "message": "Successfully retrieved data" if response.status_code < 400 else "Failed to retrieve data",
135
+ }
136
+ except Exception as e:
137
+ logger.debug(f"Response: {response}")
138
+ logger.error(f"Failed to send POST request to {endpoint} with error: {e}")
139
+ return {"success": False, "message": f"Failed to send POST request to {endpoint} with error: {e}"}
140
+
141
+
142
+def send_delete_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
143
+ """
144
+ Sends a DELETE request to the Graylog service.
145
+
146
+ Args:
147
+ endpoint (str): The endpoint to send the DELETE request to.
148
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the DELETE request. Defaults to None.
149
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylog".
150
+
151
+ Returns:
152
+ Dict[str, Any]: The response from the DELETE request.
153
+ """
154
+ logger.info(f"Sending DELETE request to {endpoint}")
155
+ attributes = get_connector_info_from_db(connector_name)
156
+ if attributes is None:
157
+ logger.error("No Graylog connector found in the database")
158
+ return None
159
+ try:
160
+ HEADERS = {
161
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
162
+ }
163
+ response = requests.delete(
164
+ f"{attributes['connector_url']}{endpoint}",
165
+ headers=HEADERS,
166
+ auth=(
167
+ attributes["connector_username"],
168
+ attributes["connector_password"],
169
+ ),
170
+ params=params,
171
+ verify=False,
172
+ )
173
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
174
+ except Exception as e:
175
+ logger.error(f"Failed to send DELETE request to {endpoint} with error: {e}")
176
+ return {"success": False, "message": f"Failed to send DELETE request to {endpoint} with error: {e}"}
177
+
178
+
179
+def send_put_request(endpoint: str, data: Optional[Dict[str, Any]] = None, connector_name: str = "Graylog") -> Dict[str, Any]:
180
+ """
181
+ Sends a PUT request to the Graylog service.
182
+
183
+ Args:
184
+ endpoint (str): The endpoint to send the PUT request to.
185
+ data (Optional[Dict[str, Any]]): The data to send with the PUT request.
186
+ connector_name (str, optional): The name of the connector to use. Defaults to "Graylog".
187
+
188
+ Returns:
189
+ Dict[str, Any]: The response from the PUT request.
190
+ """
191
+ logger.info(f"Sending PUT request to {endpoint}")
192
+ attributes = get_connector_info_from_db(connector_name)
193
+ if attributes is None:
194
+ logger.error("No Graylog connector found in the database")
195
+ return None
196
+ try:
197
+ HEADERS = {
198
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
199
+ }
200
+ response = requests.put(
201
+ f"{attributes['connector_url']}{endpoint}",
202
+ headers=HEADERS,
203
+ auth=(
204
+ attributes["connector_username"],
205
+ attributes["connector_password"],
206
+ ),
207
+ json=data,
208
+ verify=False,
209
+ )
210
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
211
+ except Exception as e:
212
+ logger.error(f"Failed to send PUT request to {endpoint} with error: {e}")
213
+ return {"success": False, "message": f"Failed to send PUT request to {endpoint} with error: {e}"}
backend/app/connectors/sublime/models/alerts.py
new
+79
@@ -0,0 +1,79 @@
1
+import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from sqlmodel import Field
6
+from sqlmodel import Relationship
7
+from sqlmodel import SQLModel
8
+
9
+
10
+class FlaggedRule(SQLModel, table=True):
11
+ id: Optional[int] = Field(default=None, primary_key=True)
12
+ rule_id: str
13
+ name: str
14
+ severity: Optional[str] = Field(None, description="Severity level of the flagged rule")
15
+ tags: str
16
+ sublime_alert_id: int = Field(foreign_key="sublimealerts.id")
17
+
18
+ # Relationship attribute
19
+ sublime_alert: "SublimeAlerts" = Relationship(back_populates="flagged_rules")
20
+
21
+
22
+class Mailbox(SQLModel, table=True):
23
+ id: Optional[int] = Field(default=None, primary_key=True)
24
+ external_id: Optional[str] = Field(None, description="External identifier for the mailbox")
25
+ mailbox_id: str
26
+ sublime_alert_id: int = Field(foreign_key="sublimealerts.id")
27
+
28
+ # Relationship attribute
29
+ sublime_alert: "SublimeAlerts" = Relationship(back_populates="mailbox")
30
+
31
+
32
+class TriggeredAction(SQLModel, table=True):
33
+ id: Optional[int] = Field(default=None, primary_key=True)
34
+ action_id: str
35
+ name: str
36
+ type: str
37
+ sublime_alert_id: int = Field(foreign_key="sublimealerts.id")
38
+
39
+ # Relationship attribute
40
+ sublime_alert: "SublimeAlerts" = Relationship(back_populates="triggered_actions")
41
+
42
+
43
+class Sender(SQLModel, table=True):
44
+ id: Optional[int] = Field(default=None, primary_key=True)
45
+ email: str
46
+ name: Optional[str] = Field(None, description="Name of the sender")
47
+ sublime_alert_id: int = Field(foreign_key="sublimealerts.id")
48
+
49
+ # Relationship attribute
50
+ sublime_alert: "SublimeAlerts" = Relationship(back_populates="sender")
51
+
52
+
53
+class Recipient(SQLModel, table=True):
54
+ id: Optional[int] = Field(default=None, primary_key=True)
55
+ email: str
56
+ name: Optional[str] = Field(None, description="Name of the recipient")
57
+ sublime_alert_id: int = Field(foreign_key="sublimealerts.id")
58
+
59
+ # Relationship attribute
60
+ sublime_alert: "SublimeAlerts" = Relationship(back_populates="recipients")
61
+
62
+
63
+class SublimeAlerts(SQLModel, table=True):
64
+ id: Optional[int] = Field(default=None, primary_key=True)
65
+ api_version: str
66
+ created_at: str
67
+ event_id: str
68
+ type: str
69
+ message_id: str
70
+ canonical_id: str
71
+ external_id: str
72
+ message_source_id: str
73
+ timestamp: datetime.datetime = datetime.datetime.now()
74
+
75
+ flagged_rules: List[FlaggedRule] = Relationship(back_populates="sublime_alert")
76
+ mailbox: List[Mailbox] = Relationship(back_populates="sublime_alert")
77
+ triggered_actions: List[TriggeredAction] = Relationship(back_populates="sublime_alert")
78
+ sender: List[Sender] = Relationship(back_populates="sublime_alert")
79
+ recipients: List[Recipient] = Relationship(back_populates="sublime_alert")
backend/app/connectors/sublime/routes/alerts.py
new
+35
@@ -0,0 +1,35 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.connectors.sublime.schema.alerts import AlertRequestBody
5
+from app.connectors.sublime.schema.alerts import AlertResponseBody
6
+from app.connectors.sublime.schema.alerts import SublimeAlertsResponse
7
+from app.connectors.sublime.services.alerts import collect_alerts
8
+from app.connectors.sublime.services.alerts import store_sublime_alert
9
+
10
+sublime_alerts_router = APIRouter()
11
+
12
+
13
+@sublime_alerts_router.post("/alert", description="Receive alert from Sublime and store it in the database")
14
+async def receive_sublime_alert(alert_request_body: AlertRequestBody) -> AlertResponseBody:
15
+ """
16
+ Endpoint to store alert in the `sublimealerts` table.
17
+ Invoked by the Sublime alert webhook which is configured in the Sublime UI.
18
+
19
+ Returns:
20
+ jsonify: A JSON response containing if the alert was stored successfully.
21
+ """
22
+ logger.info(f"Received alert from Sublime: {alert_request_body}")
23
+ return store_sublime_alert(alert_request_body)
24
+
25
+
26
+@sublime_alerts_router.get("/alerts", response_model=SublimeAlertsResponse, description="Get all alerts")
27
+async def get_sublime_alerts() -> SublimeAlertsResponse:
28
+ """
29
+ Endpoint to retrieve alerts from the `sublimealerts` table.
30
+
31
+ Returns:
32
+ jsonify: A JSON response containing all the alerts stored in the `sublimealerts` table.
33
+ """
34
+ logger.info("Fetching all alerts from Sublime")
35
+ return collect_alerts()
backend/app/connectors/sublime/schema/alerts.py
new
+121
@@ -0,0 +1,121 @@
1
+import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+
8
+
9
+class FlaggedRule(BaseModel):
10
+ id: str = Field(..., description="Unique identifier for the flagged rule")
11
+ name: str = Field(..., description="Name of the flagged rule")
12
+ severity: Optional[str] = Field(None, description="Severity level of the flagged rule")
13
+ tags: List[str] = Field(..., description="List of tags associated with the flagged rule")
14
+
15
+
16
+class Mailbox(BaseModel):
17
+ external_id: Optional[str] = Field(None, description="External identifier for the mailbox")
18
+ id: str = Field(..., description="Unique identifier for the mailbox")
19
+
20
+
21
+class Message(BaseModel):
22
+ canonical_id: str = Field(..., description="Canonical identifier for the message")
23
+ external_id: Optional[str] = Field(None, description="External identifier for the mailbox")
24
+ id: str = Field(..., description="Unique identifier for the message")
25
+ mailbox: Mailbox = Field(..., description="Mailbox details")
26
+ message_source_id: str = Field(..., description="Source identifier for the message")
27
+
28
+
29
+class TriggeredAction(BaseModel):
30
+ id: str = Field(..., description="Unique identifier for the triggered action")
31
+ name: str = Field(..., description="Name of the triggered action")
32
+ type: str = Field(..., description="Type of the triggered action")
33
+
34
+
35
+class Data(BaseModel):
36
+ flagged_rules: List[FlaggedRule] = Field(..., description="List of flagged rules")
37
+ message: Message = Field(..., description="Message details")
38
+ triggered_actions: List[TriggeredAction] = Field(..., description="List of triggered actions")
39
+
40
+
41
+class AlertRequestBody(BaseModel):
42
+ api_version: str = Field(..., description="API version", alias="api_version")
43
+ created_at: str = Field(..., description="Creation timestamp in ISO 8601 format", alias="created_at")
44
+ data: Data = Field(..., description="Nested data object")
45
+ id: str = Field(..., description="Unique identifier for the request body")
46
+ type: str = Field(..., description="Type of event, e.g., message.flagged")
47
+
48
+
49
+class AlertResponseBody(BaseModel):
50
+ success: bool = Field(..., description="Success status of the request")
51
+ message: str = Field(..., description="Message describing the result of the request")
52
+
53
+
54
+### SQLModel Schema
55
+class FlaggedRuleSchema(BaseModel):
56
+ rule_id: str
57
+ name: str
58
+ severity: Optional[str] = Field(None, description="Severity level of the flagged rule")
59
+ tags: str
60
+
61
+ class Config:
62
+ orm_mode = True
63
+
64
+
65
+class MailboxSchema(BaseModel):
66
+ external_id: Optional[str] = Field(None, description="External identifier for the mailbox")
67
+ mailbox_id: str
68
+
69
+ class Config:
70
+ orm_mode = True
71
+
72
+
73
+class TriggeredActionSchema(BaseModel):
74
+ action_id: str
75
+ name: str
76
+ type: str
77
+
78
+ class Config:
79
+ orm_mode = True
80
+
81
+
82
+class SenderSchema(BaseModel):
83
+ email: str
84
+ name: str
85
+
86
+ class Config:
87
+ orm_mode = True
88
+
89
+
90
+class RecipientSchema(BaseModel):
91
+ email: str
92
+ name: str
93
+
94
+ class Config:
95
+ orm_mode = True
96
+
97
+
98
+class SublimeAlertsSchema(BaseModel):
99
+ api_version: str
100
+ created_at: str
101
+ event_id: str
102
+ type: str
103
+ message_id: str
104
+ canonical_id: str
105
+ external_id: str
106
+ message_source_id: str
107
+ timestamp: datetime.datetime
108
+ flagged_rules: List[FlaggedRuleSchema]
109
+ mailbox: List[MailboxSchema]
110
+ triggered_actions: List[TriggeredActionSchema]
111
+ sender: List[SenderSchema]
112
+ recipients: List[RecipientSchema]
113
+
114
+ class Config:
115
+ orm_mode = True
116
+
117
+
118
+class SublimeAlertsResponse(BaseModel):
119
+ sublime_alerts: List[SublimeAlertsSchema]
120
+ success: bool
121
+ message: str
backend/app/connectors/sublime/services/alerts.py
new
+147
@@ -0,0 +1,147 @@
1
+import json
2
+from typing import List
3
+
4
+from fastapi import HTTPException
5
+from loguru import logger
6
+
7
+from app.connectors.sublime.models.alerts import FlaggedRule
8
+from app.connectors.sublime.models.alerts import Mailbox
9
+from app.connectors.sublime.models.alerts import Recipient
10
+from app.connectors.sublime.models.alerts import Sender
11
+from app.connectors.sublime.models.alerts import SublimeAlerts
12
+from app.connectors.sublime.models.alerts import TriggeredAction
13
+from app.connectors.sublime.schema.alerts import AlertRequestBody
14
+from app.connectors.sublime.schema.alerts import AlertResponseBody
15
+from app.connectors.sublime.schema.alerts import SublimeAlertsResponse
16
+from app.connectors.sublime.schema.alerts import SublimeAlertsSchema
17
+from app.connectors.sublime.utils.universal import send_get_request
18
+from app.db.db_session import session
19
+
20
+
21
+def create_sublime_alert(alert_request_body: AlertRequestBody) -> SublimeAlerts:
22
+ return SublimeAlerts(
23
+ api_version=alert_request_body.api_version,
24
+ created_at=alert_request_body.created_at,
25
+ event_id=alert_request_body.id,
26
+ type=alert_request_body.type,
27
+ message_id=alert_request_body.data.message.id,
28
+ canonical_id=alert_request_body.data.message.canonical_id,
29
+ external_id=alert_request_body.data.message.external_id,
30
+ message_source_id=alert_request_body.data.message.message_source_id,
31
+ )
32
+
33
+
34
+def create_flagged_rules(alert_request_body: AlertRequestBody, sublime_alert_id: int) -> List[FlaggedRule]:
35
+ flagged_rules = []
36
+ for rule in alert_request_body.data.flagged_rules:
37
+ tags_str = json.dumps(rule.tags)
38
+ flagged_rules.append(
39
+ FlaggedRule(rule_id=rule.id, name=rule.name, severity=rule.severity, tags=tags_str, sublime_alert_id=sublime_alert_id),
40
+ )
41
+ return flagged_rules
42
+
43
+
44
+def create_mailbox(alert_request_body: AlertRequestBody, sublime_alert_id: int) -> Mailbox:
45
+ return Mailbox(
46
+ external_id=alert_request_body.data.message.mailbox.external_id,
47
+ mailbox_id=alert_request_body.data.message.mailbox.id,
48
+ sublime_alert_id=sublime_alert_id,
49
+ )
50
+
51
+
52
+def create_triggered_actions(alert_request_body: AlertRequestBody, sublime_alert_id: int) -> List[TriggeredAction]:
53
+ triggered_actions = []
54
+ for action in alert_request_body.data.triggered_actions:
55
+ triggered_actions.append(
56
+ TriggeredAction(action_id=action.id, name=action.name, type=action.type, sublime_alert_id=sublime_alert_id),
57
+ )
58
+ return triggered_actions
59
+
60
+
61
+def store_sublime_alert(alert_request_body: AlertRequestBody) -> AlertResponseBody:
62
+ try:
63
+ sublime_alert = create_sublime_alert(alert_request_body)
64
+ session.add(sublime_alert)
65
+ session.flush()
66
+
67
+ flagged_rules = create_flagged_rules(alert_request_body, sublime_alert.id)
68
+ mailbox = create_mailbox(alert_request_body, sublime_alert.id)
69
+ triggered_actions = create_triggered_actions(alert_request_body, sublime_alert.id)
70
+ sender = create_sender(alert_request_body, sublime_alert.id)
71
+ recipient = create_recipient(alert_request_body, sublime_alert.id)
72
+
73
+ session.add_all(flagged_rules)
74
+ session.add(mailbox)
75
+ session.add_all(triggered_actions)
76
+ session.add(sender)
77
+ session.add(recipient)
78
+
79
+ logger.info(f"Preparing to store: {sublime_alert}")
80
+ session.commit()
81
+ logger.info(f"Alert {alert_request_body.id} stored in the database")
82
+
83
+ return AlertResponseBody(success=True, message=f"Alert {alert_request_body.id} stored in the database")
84
+ except Exception as e:
85
+ logger.error(f"Failed to store alert {alert_request_body.id} in the database: {e}")
86
+ raise HTTPException(status_code=500, detail=f"Failed to store alert {alert_request_body.id} in the database: {e}")
87
+
88
+
89
+def create_sender(alert_request_body: AlertRequestBody, sublime_alert_id: int) -> Sender:
90
+ return Sender(email=collect_sender(alert_request_body.data.message.id), name="n/a", sublime_alert_id=sublime_alert_id)
91
+
92
+
93
+def create_recipient(alert_request_body: AlertRequestBody, sublime_alert_id: int) -> Recipient:
94
+ return Recipient(email=collect_recipient(alert_request_body.data.message.id), name="n/a", sublime_alert_id=sublime_alert_id)
95
+
96
+
97
+def collect_sender(message_id: str) -> Sender:
98
+ """
99
+ Get a single Sublime Alert from the database
100
+ """
101
+ logger.info(f"Getting Sublime Alert with message_id {message_id}")
102
+ message_details = send_get_request(f"/v0/messages/{message_id}")
103
+ if not message_details["success"]:
104
+ logger.error(f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}")
105
+ raise HTTPException(
106
+ status_code=500,
107
+ detail=f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}",
108
+ )
109
+ logger.info(f"Successfully retrieved Sublime Alert with message_id {message_id}")
110
+ return message_details["data"]["sender"]["email"]
111
+
112
+
113
+def collect_recipient(message_id: str) -> Recipient:
114
+ """
115
+ Get a single Sublime Alert from the database
116
+ """
117
+ logger.info(f"Getting Sublime Alert with message_id {message_id}")
118
+ message_details = send_get_request(f"/v0/messages/{message_id}")
119
+ if not message_details["success"]:
120
+ logger.error(f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}")
121
+ raise HTTPException(
122
+ status_code=500,
123
+ detail=f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}",
124
+ )
125
+ logger.info(f"Successfully retrieved Sublime Alert with message_id {message_id}")
126
+ return message_details["data"]["recipients"][0]["email"]
127
+
128
+
129
+def collect_alerts() -> List[SublimeAlertsResponse]:
130
+ """
131
+ Get all Sublime Alerts from the database
132
+ """
133
+ logger.info("Getting all Sublime Alerts")
134
+ alerts = session.query(SublimeAlerts).all()
135
+ # Also add the relationships
136
+ for alert in alerts:
137
+ alert.flagged_rules = session.query(FlaggedRule).filter(FlaggedRule.sublime_alert_id == alert.id).all()
138
+ alert.mailbox = [session.query(Mailbox).filter(Mailbox.sublime_alert_id == alert.id).first()]
139
+ alert.triggered_actions = session.query(TriggeredAction).filter(TriggeredAction.sublime_alert_id == alert.id).all()
140
+ alert.sender = [session.query(Sender).filter(Sender.sublime_alert_id == alert.id).first()]
141
+ alert.recipients = session.query(Recipient).filter(Recipient.sublime_alert_id == alert.id).all()
142
+ logger.info("Successfully retrieved all Sublime Alerts")
143
+ return SublimeAlertsResponse(
144
+ success=True,
145
+ message="Successfully retrieved all Sublime Alerts",
146
+ sublime_alerts=[SublimeAlertsSchema.from_orm(alert) for alert in alerts],
147
+ )
backend/app/connectors/sublime/utils/universal.py
new
+98
@@ -0,0 +1,98 @@
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
+
10
+
11
+def verify_sublime_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
12
+ """
13
+ Verifies the connection to Sublime service.
14
+
15
+ Returns:
16
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
17
+ """
18
+ logger.info(
19
+ f"Verifying the Sublime connection to {attributes['connector_url']}",
20
+ )
21
+ try:
22
+ headers = {
23
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
24
+ "Content-Type": "application/json",
25
+ }
26
+ params = {
27
+ "limit": 1,
28
+ }
29
+ sublime = requests.get(
30
+ f"{attributes['connector_url']}/v0/rules",
31
+ headers=headers,
32
+ params=params,
33
+ verify=False,
34
+ )
35
+ if sublime.status_code == 200:
36
+ logger.info(
37
+ f"Connection to {attributes['connector_url']} successful",
38
+ )
39
+ return {"connectionSuccessful": True, "message": "Sublime connection successful"}
40
+ else:
41
+ logger.error(
42
+ f"Connection to {attributes['connector_url']} failed with error: {sublime.text}",
43
+ )
44
+ return {
45
+ "connectionSuccessful": False,
46
+ "message": f"Connection to {attributes['connector_url']} failed with error: {sublime.text}",
47
+ }
48
+ except Exception as e:
49
+ logger.error(
50
+ f"Connection to {attributes['connector_url']} failed with error: {e}",
51
+ )
52
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
53
+
54
+
55
+def verify_sublime_connection(connector_name: str) -> str:
56
+ """
57
+ Returns if connection to Sublime service is successful.
58
+ """
59
+ logger.info("Getting Sublime authentication token")
60
+ attributes = get_connector_info_from_db(connector_name)
61
+ if attributes is None:
62
+ logger.error("No Sublime connector found in the database")
63
+ return None
64
+ return verify_sublime_credentials(attributes)
65
+
66
+
67
+def send_get_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Sublime") -> Dict[str, Any]:
68
+ """
69
+ Sends a GET request to the Sublime service.
70
+
71
+ Args:
72
+ endpoint (str): The endpoint to send the GET request to.
73
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the GET request. Defaults to None.
74
+ connector_name (str, optional): The name of the connector to use. Defaults to "Shuffle".
75
+
76
+ Returns:
77
+ Dict[str, Any]: The response from the GET request.
78
+ """
79
+ logger.info(f"Sending GET request to {endpoint}")
80
+ attributes = get_connector_info_from_db(connector_name)
81
+ if attributes is None:
82
+ logger.error("No Sublime connector found in the database")
83
+ return None
84
+ try:
85
+ HEADERS = {
86
+ "Authorization": f"Bearer {attributes['connector_api_key']}",
87
+ "Content-Type": "application/json",
88
+ }
89
+ response = requests.get(
90
+ f"{attributes['connector_url']}{endpoint}",
91
+ headers=HEADERS,
92
+ params=params,
93
+ verify=False,
94
+ )
95
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
96
+ except Exception as e:
97
+ logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
98
+ return {"success": False, "message": f"Failed to send GET request to {endpoint} with error: {e}"}
backend/app/connectors/utils.py
new
+23
@@ -0,0 +1,23 @@
1
+from typing import Any
2
+from typing import Dict
3
+
4
+from loguru import logger
5
+from sqlmodel import Session
6
+from sqlmodel import select
7
+
8
+from app.connectors.models import Connectors
9
+from app.connectors.schema import ConnectorResponse
10
+from app.db.db_session import engine # Import the shared engine
11
+
12
+
13
+def get_connector_info_from_db(connector_name: str) -> Dict[str, Any]:
14
+ with Session(engine) as session:
15
+ query = select(Connectors).where(Connectors.connector_name == connector_name)
16
+ connector = session.exec(query).first()
17
+ if connector:
18
+ connector_pydantic = ConnectorResponse.from_orm(connector)
19
+ connector_dict = connector_pydantic.dict()
20
+ return connector_dict
21
+ else:
22
+ logger.warning("No connector found.")
23
+ return None
backend/app/connectors/velociraptor/routes/artifacts.py
new
+157
@@ -0,0 +1,157 @@
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
+
8
+from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
9
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
10
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
11
+from app.connectors.velociraptor.schema.artifacts import OSPrefixEnum
12
+from app.connectors.velociraptor.schema.artifacts import OSPrefixModel
13
+from app.connectors.velociraptor.schema.artifacts import QuarantineBody
14
+from app.connectors.velociraptor.schema.artifacts import QuarantineResponse
15
+from app.connectors.velociraptor.schema.artifacts import RunCommandBody
16
+from app.connectors.velociraptor.schema.artifacts import RunCommandResponse
17
+from app.connectors.velociraptor.services.artifacts import get_artifacts
18
+from app.connectors.velociraptor.services.artifacts import quarantine_host
19
+from app.connectors.velociraptor.services.artifacts import run_artifact_collection
20
+from app.connectors.velociraptor.services.artifacts import run_remote_command
21
+from app.db.db_session import session
22
+from app.db.universal_models import Agents
23
+
24
+# App specific imports
25
+
26
+
27
+velociraptor_artifacts_router = APIRouter()
28
+
29
+
30
+# Get all valid OS prefixes
31
+def get_valid_os_prefixes() -> List[str]:
32
+ return [prefix.name.lower() for prefix in OSPrefixEnum]
33
+
34
+
35
+# Verify the OS prefix exists and return the appropriate Enum value
36
+def verify_os_prefix_exists(os_prefix: str) -> str:
37
+ os_prefix_lower = os_prefix.lower()
38
+ os_prefix_upper = os_prefix.upper() # Convert to uppercase for Enum matching
39
+ valid_os_prefixes = get_valid_os_prefixes()
40
+
41
+ if os_prefix_lower not in valid_os_prefixes:
42
+ raise HTTPException(status_code=400, detail=f"OS prefix {os_prefix} does not exist.")
43
+
44
+ return OSPrefixEnum[os_prefix_upper].value # Use the uppercase version for Enum matching
45
+
46
+
47
+def get_os_prefix_from_os_name(os_name: str) -> str:
48
+ # Use the OSPrefixModel to get the OS prefix from the OS name
49
+ logger.info(f"Getting OS prefix from OS name {os_name}")
50
+ os_prefix_model = OSPrefixModel(os_name=os_name)
51
+ result = os_prefix_model.get_os_prefix()
52
+ logger.info(f"OS prefix for OS name {os_name} is {result}")
53
+ return result
54
+
55
+
56
+def get_velociraptor_id(hostname: str) -> str:
57
+ # Get the velociraptor_id from the hostname
58
+ logger.info(f"Getting velociraptor_id from hostname {hostname}")
59
+ agent = session.query(Agents).filter(Agents.hostname == hostname).first()
60
+ if not agent:
61
+ raise HTTPException(status_code=404, detail=f"Agent with hostname {hostname} not found")
62
+ velociraptor_id = agent.velociraptor_id
63
+ # If the velociraptor_id is `n/a`, raise an error
64
+ if velociraptor_id == "n/a":
65
+ raise HTTPException(status_code=404, detail=f"Velociraptor ID for hostname {hostname} is not available")
66
+ logger.info(f"velociraptor_id for hostname {hostname} is {velociraptor_id}")
67
+ return velociraptor_id
68
+
69
+
70
+@velociraptor_artifacts_router.get("", response_model=ArtifactsResponse, description="Get all artifacts")
71
+async def get_all_artifacts() -> ArtifactsResponse:
72
+ logger.info("Fetching all artifacts")
73
+ return get_artifacts()
74
+
75
+
76
+@velociraptor_artifacts_router.get(
77
+ "/{os_prefix}",
78
+ response_model=ArtifactsResponse,
79
+ description="Get all artifacts for a specific OS prefix",
80
+)
81
+async def get_all_artifacts_for_os_prefix(os_prefix: str = Depends(verify_os_prefix_exists)) -> ArtifactsResponse:
82
+ logger.info(f"Fetching all artifacts for OS prefix {os_prefix}")
83
+ # Get all the artifacts names that begin with the OS prefix
84
+ artifacts = get_artifacts().artifacts
85
+ artifacts_for_os_prefix = [artifact for artifact in artifacts if artifact.name.startswith(os_prefix)]
86
+ return ArtifactsResponse(success=True, message=f"All artifacts for OS prefix {os_prefix} retrieved", artifacts=artifacts_for_os_prefix)
87
+
88
+
89
+@velociraptor_artifacts_router.get(
90
+ "/hostname/{hostname}",
91
+ response_model=ArtifactsResponse,
92
+ description="Get all artifacts for a specific host's OS prefix",
93
+)
94
+async def get_all_artifacts_for_hostname(hostname: str) -> ArtifactsResponse:
95
+ logger.info(f"Fetching all artifacts for hostname {hostname}")
96
+ agent = session.query(Agents).filter(Agents.hostname == hostname).first()
97
+ if not agent:
98
+ raise HTTPException(status_code=404, detail=f"Agent with hostname {hostname} not found")
99
+ os_prefix = get_os_prefix_from_os_name(os_name=agent.os.lower())
100
+ if not os_prefix:
101
+ raise HTTPException(status_code=404, detail=f"OS prefix of {agent.os.lower()} for hostname {hostname} not found")
102
+ result = await get_all_artifacts_for_os_prefix(os_prefix)
103
+ return ArtifactsResponse(
104
+ success=True,
105
+ message=f"All available artifacts that can be ran for hostname {hostname} retrieved",
106
+ artifacts=result.artifacts,
107
+ )
108
+
109
+
110
+@velociraptor_artifacts_router.post("/collect", response_model=CollectArtifactResponse, description="Run an analyzer")
111
+async def collect_artifact(collect_artifact_body: CollectArtifactBody) -> CollectArtifactResponse:
112
+ logger.info(f"Received request to collect artifact {collect_artifact_body}")
113
+ # Check that provided artifact name applies for the provided hostname and use the `get_all_artifacts_for_hostname` function to get the list of artifacts
114
+ result = await get_all_artifacts_for_hostname(collect_artifact_body.hostname)
115
+ artifact_names = [artifact.name for artifact in result.artifacts]
116
+ if collect_artifact_body.artifact_name not in artifact_names:
117
+ raise HTTPException(
118
+ status_code=400,
119
+ detail=f"Artifact name {collect_artifact_body.artifact_name} does not apply for hostname {collect_artifact_body.hostname} or does not exist",
120
+ )
121
+ # Add the velociraptor_id to the run_analyzer_body object
122
+ collect_artifact_body.velociraptor_id = get_velociraptor_id(collect_artifact_body.hostname)
123
+ # Run the analyzer
124
+ return run_artifact_collection(collect_artifact_body)
125
+
126
+
127
+@velociraptor_artifacts_router.post("/command", response_model=RunCommandResponse, description="Run a remote command")
128
+async def run_command(run_command_body: RunCommandBody) -> RunCommandResponse:
129
+ logger.info(f"Received request to run command {run_command_body}")
130
+ result = await get_all_artifacts_for_hostname(run_command_body.hostname)
131
+ artifact_names = [artifact.name for artifact in result.artifacts]
132
+ if run_command_body.artifact_name not in artifact_names:
133
+ raise HTTPException(
134
+ status_code=400,
135
+ detail=f"Artifact name {run_command_body.artifact_name.value} does not apply for hostname {run_command_body.hostname} or does not exist",
136
+ )
137
+ # Add the velociraptor_id to the run_command_body object
138
+ run_command_body.velociraptor_id = get_velociraptor_id(run_command_body.hostname)
139
+ # Run the command
140
+ return run_remote_command(run_command_body)
141
+
142
+
143
+@velociraptor_artifacts_router.post("/quarantine", response_model=QuarantineResponse, description="Quarantine a host")
144
+async def quarantine(quarantine_body: QuarantineBody) -> QuarantineResponse:
145
+ logger.info(f"Received request to quarantine host {quarantine_body}")
146
+ result = await get_all_artifacts_for_hostname(quarantine_body.hostname)
147
+ artifact_names = [artifact.name for artifact in result.artifacts]
148
+ if quarantine_body.artifact_name not in artifact_names:
149
+ raise HTTPException(
150
+ status_code=400,
151
+ detail=f"Artifact name {quarantine_body.artifact_name.value} does not apply for hostname {quarantine_body.hostname} or does not exist",
152
+ )
153
+ # Add the velociraptor_id to the run_command_body object
154
+ # Add the velociraptor_id to the quarantine_body object
155
+ quarantine_body.velociraptor_id = get_velociraptor_id(quarantine_body.hostname)
156
+ # Quarantine the host
157
+ return quarantine_host(quarantine_body)
backend/app/connectors/velociraptor/schema/artifacts.py
new
+107
@@ -0,0 +1,107 @@
1
+from enum import Enum
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
+
10
+
11
+class Artifacts(BaseModel):
12
+ description: str = Field(..., description="Description of the artifact.")
13
+ name: str = Field(..., description="Name of the artifact.")
14
+
15
+
16
+class ArtifactsResponse(BaseModel):
17
+ message: str = Field(...)
18
+ # make artifacts optional
19
+ artifacts: Optional[List[Artifacts]]
20
+ success: str = Field(...)
21
+
22
+
23
+class OSPrefixEnum(Enum):
24
+ LINUX = "Linux."
25
+ WINDOWS = "Windows."
26
+ MACOS = "MacOS."
27
+
28
+
29
+class OSPrefixModel(BaseModel):
30
+ os_name: Optional[str]
31
+ os_prefix_mapping: Dict[str, str] = {
32
+ "windows": "Windows",
33
+ "linux": "Linux",
34
+ "mac": "MacOS",
35
+ "ubuntu": "Linux", # Add more mappings as needed
36
+ }
37
+
38
+ def get_os_prefix(self) -> Optional[str]:
39
+ if self.os_name is None:
40
+ return None
41
+ return self._map_os_name_to_prefix()
42
+
43
+ def _map_os_name_to_prefix(self) -> Optional[str]:
44
+ os_name_lower = self.os_name.lower()
45
+ for keyword, prefix in self.os_prefix_mapping.items():
46
+ if keyword in os_name_lower:
47
+ return prefix
48
+ return None
49
+
50
+
51
+class OperationEnum(str, Enum):
52
+ collect_artifact = "collect_artifact"
53
+ run_command = "run_command"
54
+ quarantine = "quarantine"
55
+
56
+
57
+class ActionEnum(str, Enum):
58
+ quarantine = "quarantine"
59
+ remove_quarantine = "remove_quarantine"
60
+
61
+
62
+class CommandArtifactsEnum(str, Enum):
63
+ windows_powershell = "Windows.System.PowerShell"
64
+ windows_cmd = "Windows.System.CmdShell"
65
+ linux_bash = "Linux.Sys.BashShell"
66
+
67
+
68
+class QuarantineArtifactsEnum(str, Enum):
69
+ windows_quarantine = "Windows.Remediation.Quarantine"
70
+ linux_quarantine = "Linux.Remediation.Quarantine"
71
+
72
+
73
+class BaseBody(BaseModel):
74
+ hostname: str = Field(..., description="Name of the client")
75
+ velociraptor_id: Optional[str] = Field(None, description="Client ID of the client")
76
+
77
+
78
+class CollectArtifactBody(BaseBody):
79
+ artifact_name: Optional[str] = Field(None, description="Name of the artifact for collection or command running")
80
+
81
+
82
+class RunCommandBody(BaseBody):
83
+ command: Optional[str] = Field(None, description="Command to run")
84
+ artifact_name: CommandArtifactsEnum = Field(None, description="Name of the artifact for command running")
85
+
86
+
87
+class QuarantineBody(BaseBody):
88
+ action: ActionEnum = Field(..., description="Action to perform")
89
+ artifact_name: QuarantineArtifactsEnum = Field(None, description="Name of the artifact for quarantine or removal of quarantine")
90
+
91
+
92
+class BaseResponse(BaseModel):
93
+ message: str = Field(...)
94
+ success: bool = Field(...) # Changed from str to bool based on your sample data
95
+ results: Optional[List[Dict[str, Any]]] = Field(None, description="Results of the operation")
96
+
97
+
98
+class CollectArtifactResponse(BaseResponse):
99
+ pass # If you have additional fields, you can define them here
100
+
101
+
102
+class RunCommandResponse(BaseResponse):
103
+ pass # If you have additional fields, you can define them here
104
+
105
+
106
+class QuarantineResponse(BaseResponse):
107
+ pass # If you have additional fields, you can define them here
backend/app/connectors/velociraptor/services/artifacts.py
new
+237
@@ -0,0 +1,237 @@
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
+universal_service = UniversalService()
15
+
16
+
17
+def create_query(query: str) -> str:
18
+ """
19
+ Create a query string.
20
+
21
+ Args:
22
+ query (str): The query to be executed.
23
+
24
+ Returns:
25
+ str: The created query string.
26
+ """
27
+ return query
28
+
29
+
30
+def get_artifact_key(analyzer_body: CollectArtifactBody) -> str:
31
+ """
32
+ Construct the artifact key.
33
+
34
+ Args:
35
+ client_id (str): The ID of the client.
36
+ artifact (str): The name of the artifact.
37
+ command (str): The command that was run, if applicable.
38
+ quarantined (bool): Whether the client is quarantined or not.
39
+
40
+ Returns:
41
+ str: The constructed artifact key.
42
+ """
43
+ action = getattr(analyzer_body, "action", None)
44
+ command = getattr(analyzer_body, "command", None)
45
+
46
+ if action == "quarantine":
47
+ return f'collect_client(client_id="{analyzer_body.velociraptor_id}", artifacts=["{analyzer_body.artifact_name}"], spec=dict(`{analyzer_body.artifact_name}`=dict()))'
48
+ elif action == "remove_quarantine":
49
+ return f'collect_client(client_id="{analyzer_body.velociraptor_id}", artifacts=["{analyzer_body.artifact_name}"], spec=dict(`{analyzer_body.artifact_name}`=dict(`RemovePolicy`="Y")))'
50
+ elif command is not None:
51
+ return f"collect_client(client_id='{analyzer_body.velociraptor_id}', urgent=true, artifacts=['{analyzer_body.artifact_name}'], env=dict(Command='{analyzer_body.command}'))"
52
+ else:
53
+ return f"collect_client(client_id='{analyzer_body.velociraptor_id}', artifacts=['{analyzer_body.artifact_name}'])"
54
+
55
+
56
+def get_artifacts() -> ArtifactsResponse:
57
+ """
58
+ Get all artifacts from Velociraptor.
59
+
60
+ Returns:
61
+ ArtifactsResponse: A dictionary containing the artifacts.
62
+ """
63
+ logger.info("Fetching artifacts from Velociraptor")
64
+ query = create_query("SELECT name,description FROM artifact_definitions()")
65
+ all_artifacts = universal_service.execute_query(query)
66
+ if all_artifacts["success"]:
67
+ artifacts = [Artifacts(**artifact) for artifact in all_artifacts["results"]]
68
+ return ArtifactsResponse(success=True, message="All artifacts retrieved", artifacts=artifacts)
69
+ else:
70
+ raise HTTPException(status_code=500, detail=f"Failed to get all artifacts: {all_artifacts['message']}")
71
+
72
+
73
+def run_artifact_collection(collect_artifact_body: CollectArtifactBody) -> CollectArtifactResponse:
74
+ """
75
+ Run an artifact collection on a client.
76
+
77
+ Args:
78
+ run_analyzer_body (RunAnalyzerBody): The body of the request.
79
+
80
+ Returns:
81
+ RunAnalyzerResponse: A dictionary containing the success status and a message.
82
+ """
83
+ try:
84
+ query = create_query(
85
+ f"SELECT collect_client(client_id='{collect_artifact_body.velociraptor_id}', artifacts=['{collect_artifact_body.artifact_name}']) FROM scope()",
86
+ )
87
+ flow = universal_service.execute_query(query)
88
+ logger.info(f"Successfully ran artifact collection on {flow}")
89
+
90
+ artifact_key = get_artifact_key(analyzer_body=collect_artifact_body)
91
+
92
+ flow_id = flow["results"][0][artifact_key]["flow_id"]
93
+ logger.info(f"Extracted flow_id: {flow_id}")
94
+
95
+ completed = universal_service.watch_flow_completion(flow_id)
96
+ logger.info(f"Successfully watched flow completion on {completed}")
97
+
98
+ results = universal_service.read_collection_results(
99
+ client_id=collect_artifact_body.velociraptor_id,
100
+ flow_id=flow_id,
101
+ artifact=collect_artifact_body.artifact_name,
102
+ )
103
+
104
+ logger.info(f"Successfully read collection results on {results}")
105
+
106
+ return CollectArtifactResponse(success=results["success"], message=results["message"], results=results["results"])
107
+ except Exception as err:
108
+ logger.error(f"Failed to run artifact collection on {collect_artifact_body}: {err}")
109
+ raise HTTPException(status_code=500, detail=f"Failed to run artifact collection on {collect_artifact_body}: {err}")
110
+
111
+
112
+def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResponse:
113
+ """
114
+ Run a remote command on a client.
115
+
116
+ Args:
117
+ run_analyzer_body (RunAnalyzerBody): The body of the request.
118
+
119
+ Returns:
120
+ RunAnalyzerResponse: A dictionary containing the success status and a message.
121
+ """
122
+ try:
123
+ run_command_body.artifact_name = run_command_body.artifact_name.value
124
+ logger.info(f"Running remote command on {run_command_body}")
125
+ query = create_query(
126
+ f"SELECT collect_client(client_id='{run_command_body.velociraptor_id}', urgent=true, artifacts=['{run_command_body.artifact_name}'], env=dict(Command='{run_command_body.command}')) "
127
+ "FROM scope()",
128
+ )
129
+ flow = universal_service.execute_query(query)
130
+ logger.info(f"Successfully ran artifact collection on {flow}")
131
+
132
+ artifact_key = get_artifact_key(analyzer_body=run_command_body)
133
+
134
+ flow_id = flow["results"][0][artifact_key]["flow_id"]
135
+ logger.info(f"Extracted flow_id: {flow_id}")
136
+
137
+ completed = universal_service.watch_flow_completion(flow_id)
138
+ logger.info(f"Successfully watched flow completion on {completed}")
139
+
140
+ results = universal_service.read_collection_results(
141
+ client_id=run_command_body.velociraptor_id,
142
+ flow_id=flow_id,
143
+ artifact=run_command_body.artifact_name,
144
+ )
145
+
146
+ logger.info(f"Successfully read collection results on {results}")
147
+
148
+ return RunCommandResponse(success=results["success"], message=results["message"], results=results["results"])
149
+ except Exception as err:
150
+ logger.error(f"Failed to run artifact collection on {run_command_body}: {err}")
151
+ raise HTTPException(status_code=500, detail=f"Failed to run artifact collection on {run_command_body}: {err}")
152
+
153
+
154
+def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse:
155
+ """
156
+ Quarantine a host.
157
+
158
+ Args:
159
+ quarantine_body (QuarantineBody): The body of the request.
160
+
161
+ Returns:
162
+ QuarantineResponse: A dictionary containing the success status and a message.
163
+ """
164
+ try:
165
+ quarantine_body.artifact_name = quarantine_body.artifact_name.value
166
+ quarantine_body.action = quarantine_body.action.value
167
+ if quarantine_body.action == "quarantine":
168
+ query = create_query(
169
+ f'SELECT collect_client(client_id="{quarantine_body.velociraptor_id}", artifacts=["{quarantine_body.artifact_name}"], spec=dict(`{quarantine_body.artifact_name}`=dict())) FROM scope()',
170
+ )
171
+ else:
172
+ query = create_query(
173
+ f'SELECT collect_client(client_id="{quarantine_body.velociraptor_id}", artifacts=["{quarantine_body.artifact_name}"], spec=dict(`{quarantine_body.artifact_name}`=dict(`RemovePolicy`="Y"))) FROM scope()',
174
+ )
175
+ flow = universal_service.execute_query(query)
176
+ logger.info(f"Successfully ran artifact collection on {flow}")
177
+
178
+ artifact_key = get_artifact_key(analyzer_body=quarantine_body)
179
+
180
+ flow_id = flow["results"][0][artifact_key]["flow_id"]
181
+ logger.info(f"Extracted flow_id: {flow_id}")
182
+
183
+ completed = universal_service.watch_flow_completion(flow_id)
184
+ logger.info(f"Successfully watched flow completion on {completed}")
185
+
186
+ results = universal_service.read_collection_results(
187
+ client_id=quarantine_body.velociraptor_id,
188
+ flow_id=flow_id,
189
+ artifact=quarantine_body.artifact_name,
190
+ )
191
+
192
+ logger.info(f"Successfully read collection results on {results}")
193
+
194
+ return QuarantineResponse(success=results["success"], message=results["message"], results=results["results"])
195
+ except Exception as err:
196
+ logger.error(f"Failed to run artifact collection on {quarantine_body}: {err}")
197
+ raise HTTPException(status_code=500, detail=f"Failed to run artifact collection on {quarantine_body}: {err}")
198
+
199
+
200
+######################## KEEP
201
+class ArtifactsService:
202
+ def delete_client(self, client_id: str) -> dict:
203
+ """
204
+ Delete a client from Velociraptor.
205
+
206
+ Args:
207
+ client_id (str): The ID of the client.
208
+
209
+ Returns:
210
+ dict: A dictionary with the success status and a message.
211
+ """
212
+ try:
213
+ query = self._create_query(
214
+ f"SELECT collect_client(client_id='server', artifacts=['Server.Utils.DeleteClient'], env=dict(ClientIdList='{client_id}',ReallyDoIt='Y')) "
215
+ "FROM scope()",
216
+ )
217
+
218
+ flow = self.universal_service.execute_query(query)
219
+ logger.info(f"Successfully ran artifact collection on {flow}")
220
+
221
+ # artifact_key = f"collect_client(client_id='server', artifacts=['Server.Utils.DeleteClient'], env=dict(ClientIdList='{client_id}',ReallyDoIt='Y'))"
222
+ flow_id = flow["results"][0][query]["flow_id"]
223
+ logger.info(f"Extracted flow_id: {flow_id}")
224
+
225
+ completed = self.universal_service.watch_flow_completion(flow_id)
226
+ logger.info(f"Successfully watched flow completion on {completed}")
227
+
228
+ return {
229
+ "message": f"Successfully deleted client {client_id}",
230
+ "success": True,
231
+ }
232
+ except Exception as err:
233
+ logger.error(f"Failed to delete client {client_id}: {err}")
234
+ return {
235
+ "message": f"Failed to delete client {client_id}",
236
+ "success": False,
237
+ }
backend/app/connectors/velociraptor/utils/universal.py
new
+283
@@ -0,0 +1,283 @@
1
+import json
2
+from datetime import datetime
3
+from typing import Any
4
+from typing import Dict
5
+
6
+import grpc
7
+import pyvelociraptor
8
+from loguru import logger
9
+from pyvelociraptor import api_pb2
10
+from pyvelociraptor import api_pb2_grpc
11
+
12
+from app.connectors.utils import get_connector_info_from_db
13
+
14
+
15
+def verify_velociraptor_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
16
+ """
17
+ Verifies the connection to Velociraptor service.
18
+
19
+ Returns:
20
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
21
+ """
22
+ try:
23
+ connector_api_key = attributes["connector_api_key"]
24
+
25
+ with open(connector_api_key, "r") as f:
26
+ f.read()
27
+
28
+ try:
29
+ config = pyvelociraptor.LoadConfigFile(connector_api_key)
30
+ creds = grpc.ssl_channel_credentials(
31
+ root_certificates=config["ca_certificate"].encode("utf8"),
32
+ private_key=config["client_private_key"].encode("utf8"),
33
+ certificate_chain=config["client_cert"].encode("utf8"),
34
+ )
35
+
36
+ options = (("grpc.ssl_target_name_override", "VelociraptorServer"),)
37
+
38
+ with grpc.secure_channel(
39
+ config["api_connection_string"],
40
+ creds,
41
+ options,
42
+ ) as channel:
43
+ stub = api_pb2_grpc.APIStub(channel)
44
+ client_query = "SELECT * FROM info()"
45
+
46
+ client_request = api_pb2.VQLCollectorArgs(
47
+ max_wait=60,
48
+ Query=[
49
+ api_pb2.VQLRequest(
50
+ Name="ClientQuery",
51
+ VQL=client_query,
52
+ ),
53
+ ],
54
+ )
55
+
56
+ r = []
57
+ for response in stub.Query(client_request):
58
+ if response.Response:
59
+ r = r + json.loads(response.Response)
60
+ return {"connectionSuccessful": True, "message": "Connection to Velociraptor successful"}
61
+ except Exception as e:
62
+ logger.error(f"Failed to verify connection to Velociraptor: {e}")
63
+ return {"connectionSuccessful": False, "message": f"Failed to verify connection to Velociraptor: {e}"}
64
+ except Exception as e:
65
+ logger.error(f"Failed to get connector_api_key from the database: {e}")
66
+ return {"connectionSuccessful": False, "message": f"Failed to get connector_api_key from the database: {e}"}
67
+
68
+
69
+def verify_velociraptor_connection(connector_name: str) -> str:
70
+ """
71
+ Verifies the connection to Velociraptor service.
72
+ """
73
+ logger.info(f"Verifying the Velociraptor connection for connector: {connector_name}")
74
+ attributes = get_connector_info_from_db(connector_name)
75
+ if attributes is None:
76
+ logger.error("No Velociraptor connector found in the database")
77
+ return None
78
+ return verify_velociraptor_credentials(attributes)
79
+
80
+
81
+class UniversalService:
82
+ """
83
+ A service class that encapsulates the logic for polling messages from Velociraptor.
84
+ """
85
+
86
+ def __init__(self) -> None:
87
+ self.setup_velociraptor_connector("Velociraptor")
88
+ self.setup_grpc_channel_and_stub()
89
+
90
+ def setup_velociraptor_connector(self, connector_name: str):
91
+ """
92
+ Collects the details of the Velociraptor connector and sets them up.
93
+
94
+ Args:
95
+ connector_name (str): The name of the Velociraptor connector.
96
+ """
97
+ attributes = get_connector_info_from_db(connector_name)
98
+ if attributes is None:
99
+ logger.error("No Velociraptor connector found in the database")
100
+ return None
101
+ self.connector_api_key = attributes["connector_api_key"]
102
+ self.config = pyvelociraptor.LoadConfigFile(self.connector_api_key)
103
+
104
+ def setup_grpc_channel_and_stub(self):
105
+ """
106
+ Sets up the gRPC channel and stub for Velociraptor.
107
+ """
108
+ creds = grpc.ssl_channel_credentials(
109
+ root_certificates=self.config["ca_certificate"].encode("utf8"),
110
+ private_key=self.config["client_private_key"].encode("utf8"),
111
+ certificate_chain=self.config["client_cert"].encode("utf8"),
112
+ )
113
+ options = (("grpc.ssl_target_name_override", "VelociraptorServer"),)
114
+ self.channel = grpc.secure_channel(
115
+ self.config["api_connection_string"],
116
+ creds,
117
+ options,
118
+ )
119
+ self.stub = api_pb2_grpc.APIStub(self.channel)
120
+
121
+ def create_vql_request(self, vql: str):
122
+ """
123
+ Creates a VQLCollectorArgs object with given VQL query.
124
+
125
+ Args:
126
+ vql (str): The VQL query.
127
+
128
+ Returns:
129
+ VQLCollectorArgs: The VQLCollectorArgs object with given VQL query.
130
+ """
131
+ return api_pb2.VQLCollectorArgs(
132
+ max_wait=1,
133
+ Query=[
134
+ api_pb2.VQLRequest(
135
+ Name="VQLRequest",
136
+ VQL=vql,
137
+ ),
138
+ ],
139
+ )
140
+
141
+ def execute_query(self, vql: str):
142
+ """
143
+ Executes a VQL query and returns the results.
144
+
145
+ Args:
146
+ vql (str): The VQL query to be executed.
147
+
148
+ Returns:
149
+ dict: A dictionary with the success status, a message, and potentially the results.
150
+ """
151
+ client_request = self.create_vql_request(vql)
152
+ try:
153
+ results = []
154
+ for response in self.stub.Query(client_request):
155
+ if response.Response:
156
+ results += json.loads(response.Response)
157
+ return {
158
+ "success": True,
159
+ "message": "Successfully executed query",
160
+ "results": results,
161
+ }
162
+ except Exception as e:
163
+ logger.error(f"Failed to execute query: {e}")
164
+ return {
165
+ "success": False,
166
+ "message": f"Failed to execute query: {e}",
167
+ }
168
+
169
+ def watch_flow_completion(self, flow_id: str):
170
+ """
171
+ Watch for the completion of a flow.
172
+
173
+ Args:
174
+ flow_id (str): The ID of the flow.
175
+
176
+ Returns:
177
+ dict: A dictionary with the success status and a message.
178
+ """
179
+ vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1"
180
+ return self.execute_query(vql)
181
+
182
+ def read_collection_results(
183
+ self,
184
+ client_id: str,
185
+ flow_id: str,
186
+ artifact: str = "Generic.Client.Info/BasicInformation",
187
+ ):
188
+ """
189
+ Read the results of a collection.
190
+
191
+ Args:
192
+ client_id (str): The client ID.
193
+ flow_id (str): The ID of the flow.
194
+ artifact (str, optional): The artifact. Defaults to 'Generic.Client.Info/BasicInformation'.
195
+
196
+ Returns:
197
+ dict: A dictionary with the success status, a message, and potentially the results.
198
+ """
199
+ vql = f"SELECT * FROM source(client_id='{client_id}', flow_id='{flow_id}', artifact='{artifact}')"
200
+ return self.execute_query(vql)
201
+
202
+ def get_client_id(self, client_name: str):
203
+ """
204
+ Get the client_id associated with a given client_name.
205
+
206
+ Args:
207
+ client_name (str): The asset name to search for.
208
+
209
+ Returns:
210
+ dict: A dictionary with the success status, a message, and potentially the client_id.
211
+ """
212
+ # Formulate queries
213
+ try:
214
+ vql_client_id = f"select client_id,os_info from clients(search='host:{client_name}')"
215
+ vql_last_seen_at = f"select last_seen_at from clients(search='host:{client_name}')"
216
+
217
+ # Get the last seen timestamp
218
+ last_seen_at = self._get_last_seen_timestamp(vql_last_seen_at)
219
+
220
+ # if last_seen_at is longer than 30 seconds from now, return False
221
+ if self._is_offline(last_seen_at):
222
+ return {
223
+ "success": False,
224
+ "message": f"{client_name} has not been seen in the last 30 seconds and "
225
+ "may not be online with the Velociraptor server.",
226
+ "results": [{"client_id": None}],
227
+ }
228
+
229
+ return self.execute_query(vql_client_id)
230
+ except Exception as e:
231
+ return {
232
+ "success": False,
233
+ "message": f"Failed to get Client ID for {client_name}: {e}",
234
+ "results": [{"client_id": None}],
235
+ }
236
+
237
+ def _get_last_seen_timestamp(self, vql: str):
238
+ """
239
+ Executes the VQL query and returns the last_seen_at timestamp.
240
+
241
+ Args:
242
+ vql (str): The VQL query.
243
+
244
+ Returns:
245
+ float: The last_seen_at timestamp.
246
+ """
247
+ return self.execute_query(vql)["results"][0]["last_seen_at"]
248
+
249
+ def _get_client_version(self, vql: str):
250
+ """
251
+ Executes the VQL query and returns the `agent_information``version` field
252
+
253
+ Args:
254
+ vql (str): The VQL query.
255
+
256
+ Returns:
257
+ str: The client version.
258
+ """
259
+ return self.execute_query(vql)["results"][0]["agent_information"]["version"]
260
+
261
+ def _get_server_version(self, vql: str):
262
+ """
263
+ Executes the VQL query and returns the velociraptor server version.
264
+
265
+ Args:
266
+ vql (str): The VQL query.
267
+
268
+ Returns:
269
+ str: The server version.
270
+ """
271
+ return self.execute_query(vql)["results"][0]["version"]["version"]
272
+
273
+ def _is_offline(self, last_seen_at: float):
274
+ """
275
+ Determines if the client is offline based on the last_seen_at timestamp.
276
+
277
+ Args:
278
+ last_seen_at (float): The last_seen_at timestamp.
279
+
280
+ Returns:
281
+ bool: True if the client is offline, False otherwise.
282
+ """
283
+ return (datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000)).total_seconds() > 30
backend/app/connectors/wazuh_indexer/models/db.py
backend/app/connectors/wazuh_indexer/routes/alerts.py
new
+97
@@ -0,0 +1,97 @@
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
+
8
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
9
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHostResponse
10
+from app.connectors.wazuh_indexer.schema.alerts import AlertsByRuleResponse
11
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchBody
12
+from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchResponse
13
+from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchBody
14
+from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchResponse
15
+from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchBody
16
+from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchResponse
17
+from app.connectors.wazuh_indexer.services.alerts import get_alerts
18
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_host
19
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule
20
+from app.connectors.wazuh_indexer.services.alerts import get_alerts_by_rule_per_host
21
+from app.connectors.wazuh_indexer.services.alerts import get_host_alerts
22
+from app.connectors.wazuh_indexer.services.alerts import get_index_alerts
23
+from app.connectors.wazuh_indexer.utils.universal import collect_indices
24
+
25
+# App specific imports
26
+
27
+
28
+wazuh_indexer_alerts_router = APIRouter()
29
+
30
+
31
+def get_index_names() -> List[str]:
32
+ indices = collect_indices()
33
+ return indices.indices_list
34
+
35
+
36
+def verify_index_name(index_alerts_search_body: IndexAlertsSearchBody) -> IndexAlertsSearchBody:
37
+ # Remove any extra spaces from index_name
38
+ index_alerts_search_body.index_name = index_alerts_search_body.index_name.strip()
39
+
40
+ managed_index_names = get_index_names()
41
+ if index_alerts_search_body.index_name not in managed_index_names:
42
+ raise HTTPException(
43
+ status_code=400,
44
+ detail=f"Index name '{index_alerts_search_body.index_name}' is not managed by Wazuh Indexer or no longer exists.",
45
+ )
46
+ return index_alerts_search_body
47
+
48
+
49
+@wazuh_indexer_alerts_router.post("", response_model=AlertsSearchResponse, description="Get all alerts")
50
+async def get_all_alerts(alerts_search_body: AlertsSearchBody) -> AlertsSearchResponse:
51
+ logger.info("Fetching all alerts")
52
+ return get_alerts(alerts_search_body)
53
+
54
+
55
+@wazuh_indexer_alerts_router.post("/host", response_model=HostAlertsSearchResponse, description="Get all alerts for a host")
56
+async def get_all_alerts_for_host(host_alerts_search_body: HostAlertsSearchBody) -> HostAlertsSearchResponse:
57
+ logger.info(f"Fetching all alerts for host {host_alerts_search_body.agent_name}")
58
+ return get_host_alerts(host_alerts_search_body)
59
+
60
+
61
+@wazuh_indexer_alerts_router.post("/index", response_model=IndexAlertsSearchResponse, description="Get all alerts for an index")
62
+async def get_all_alerts_for_index(
63
+ index_alerts_search_body: IndexAlertsSearchBody = Depends(verify_index_name),
64
+) -> IndexAlertsSearchResponse:
65
+ logger.info(f"Fetching all alerts for index {index_alerts_search_body.index_name}")
66
+ return get_index_alerts(index_alerts_search_body)
67
+
68
+
69
+@wazuh_indexer_alerts_router.post("/hosts/all", response_model=AlertsByHostResponse, description="Get number of all alerts for all hosts")
70
+async def get_all_alerts_by_host(alerts_search_body: AlertsSearchBody) -> AlertsByHostResponse:
71
+ logger.info("Fetching number of all alerts for all hosts")
72
+ return get_alerts_by_host(alerts_search_body)
73
+
74
+
75
+@wazuh_indexer_alerts_router.post("/rules/all", response_model=AlertsByRuleResponse, description="Get number of all alerts for all rules")
76
+async def get_all_alerts_by_rule(alerts_search_body: AlertsSearchBody) -> AlertsByRuleResponse:
77
+ logger.info("Fetching number of all alerts for all rules")
78
+ return get_alerts_by_rule(alerts_search_body)
79
+
80
+
81
+@wazuh_indexer_alerts_router.post(
82
+ "/rules/hosts/all",
83
+ response_model=AlertsByRulePerHostResponse,
84
+ description="Get number of all alerts for all rules per host",
85
+)
86
+async def get_all_alerts_by_rule_per_host(alerts_search_body: AlertsSearchBody) -> AlertsByRulePerHostResponse:
87
+ """
88
+ Get number of all alerts for all rules per host
89
+
90
+ Args:
91
+ alerts_search_body (AlertsSearchBody): _description_
92
+
93
+ Returns:
94
+ AlertsByRulePerHostResponse: _description_
95
+ """
96
+ logger.info("Fetching number of all alerts for all rules per host")
97
+ return get_alerts_by_rule_per_host(alerts_search_body)
backend/app/connectors/wazuh_indexer/routes/monitoring.py
new
+97
@@ -0,0 +1,97 @@
1
+from typing import Union
2
+
3
+from fastapi import APIRouter
4
+from fastapi import HTTPException
5
+
6
+from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse
7
+from app.connectors.wazuh_indexer.schema.monitoring import IndicesStatsResponse
8
+from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocationResponse
9
+from app.connectors.wazuh_indexer.schema.monitoring import ShardsResponse
10
+
11
+# from app.connectors.wazuh_indexer.schema import WazuhIndexerResponse, WazuhIndexerListResponse
12
+from app.connectors.wazuh_indexer.services.monitoring import cluster_healthcheck
13
+from app.connectors.wazuh_indexer.services.monitoring import indices_stats
14
+from app.connectors.wazuh_indexer.services.monitoring import node_allocation
15
+from app.connectors.wazuh_indexer.services.monitoring import shards
16
+
17
+wazuh_indexer_router = APIRouter()
18
+
19
+
20
+@wazuh_indexer_router.get("/health", response_model=ClusterHealthResponse, description="Fetch Wazuh Indexer cluster health")
21
+async def get_cluster_health() -> Union[ClusterHealthResponse, HTTPException]:
22
+ """
23
+ Fetch Wazuh Indexer cluster health.
24
+
25
+ This endpoint retrieves the cluster health of the Wazuh Indexer service.
26
+
27
+ Returns:
28
+ ElasticsearchResponse: A Pydantic model representing the cluster health of the Wazuh Indexer service.
29
+
30
+ Raises:
31
+ HTTPException: An exception with a 500 status code is raised if the cluster health cannot be retrieved.
32
+ """
33
+ cluster_health = cluster_healthcheck()
34
+ if cluster_health is not None:
35
+ return cluster_health
36
+ else:
37
+ raise HTTPException(status_code=500, detail="Failed to retrieve cluster health.")
38
+
39
+
40
+@wazuh_indexer_router.get("/allocation", response_model=NodeAllocationResponse, description="Fetch Wazuh Indexer node allocation")
41
+async def get_node_allocation() -> Union[NodeAllocationResponse, HTTPException]:
42
+ """
43
+ Fetch Wazuh Indexer node allocation.
44
+
45
+ This endpoint retrieves the node allocation of the Wazuh Indexer service.
46
+
47
+ Returns:
48
+ ElasticsearchResponse: A Pydantic model representing the node allocation of the Wazuh Indexer service.
49
+
50
+ Raises:
51
+ HTTPException: An exception with a 500 status code is raised if the node allocation cannot be retrieved.
52
+ """
53
+ node_allocation_response = node_allocation()
54
+ if node_allocation_response is not None:
55
+ return node_allocation_response
56
+ else:
57
+ raise HTTPException(status_code=500, detail="Failed to retrieve node allocation.")
58
+
59
+
60
+@wazuh_indexer_router.get("/indices", response_model=IndicesStatsResponse, description="Fetch Wazuh Indexer indices stats")
61
+async def get_indices_stats() -> Union[IndicesStatsResponse, HTTPException]:
62
+ """
63
+ Fetch Wazuh Indexer indices stats.
64
+
65
+ This endpoint retrieves the indices stats of the Wazuh Indexer service.
66
+
67
+ Returns:
68
+ ElasticsearchResponse: A Pydantic model representing the indices stats of the Wazuh Indexer service.
69
+
70
+ Raises:
71
+ HTTPException: An exception with a 500 status code is raised if the indices stats cannot be retrieved.
72
+ """
73
+ indices_stats_response = indices_stats()
74
+ if indices_stats_response is not None:
75
+ return indices_stats_response
76
+ else:
77
+ raise HTTPException(status_code=500, detail="Failed to retrieve indices stats.")
78
+
79
+
80
+@wazuh_indexer_router.get("/shards", response_model=ShardsResponse, description="Fetch Wazuh Indexer shards")
81
+async def get_shards() -> Union[ShardsResponse, HTTPException]:
82
+ """
83
+ Fetch Wazuh Indexer shards.
84
+
85
+ This endpoint retrieves the shards of the Wazuh Indexer service.
86
+
87
+ Returns:
88
+ ElasticsearchResponse: A Pydantic model representing the shards of the Wazuh Indexer service.
89
+
90
+ Raises:
91
+ HTTPException: An exception with a 500 status code is raised if the shards cannot be retrieved.
92
+ """
93
+ shards_response = shards()
94
+ if shards_response is not None:
95
+ return shards_response
96
+ else:
97
+ raise HTTPException(status_code=500, detail="Failed to retrieve shards.")
backend/app/connectors/wazuh_indexer/schema/alerts.py
new
+99
@@ -0,0 +1,99 @@
1
+from typing import Any
2
+from typing import Dict
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
+
11
+class Alert(BaseModel):
12
+ index_name: str
13
+ total_alerts: int
14
+ alerts: Optional[List[Dict[str, Any]]] = Field([], description="The alerts returned from the search.")
15
+
16
+
17
+class AlertsSearchBody(BaseModel):
18
+ size: int = Field(10, description="The number of alerts to return.")
19
+ timerange: str = Field("24h", description="The time range to search alerts in.")
20
+ alert_field: str = Field("syslog_level", description="The field to search alerts in.")
21
+ alert_value: str = Field("ALERT", description="The value to search alerts for.")
22
+ timestamp_field: str = Field("timestamp_utc", description="The timestamp field to search alerts in.")
23
+
24
+ @validator("timerange")
25
+ def validate_timerange(cls, value):
26
+ if value[-1] not in ("h", "d", "w", "m"):
27
+ raise ValueError("Invalid timerange format. The string should end with either 'h', 'd', 'w', or 'm'.")
28
+
29
+ # Optionally, you can check that the prefix is a number
30
+ if not value[:-1].isdigit():
31
+ raise ValueError("Invalid timerange format. The string should start with a number.")
32
+
33
+ return value
34
+
35
+
36
+class AlertsSearchResponse(BaseModel):
37
+ alerts_summary: List[Alert]
38
+ success: bool
39
+ message: str
40
+
41
+
42
+class CollectAlertsResponse(BaseModel):
43
+ alerts: List[Dict[str, Any]]
44
+ success: bool
45
+ message: str
46
+
47
+
48
+class HostAlertsSearchBody(AlertsSearchBody):
49
+ agent_name: str = Field(..., description="The name of the agent to search alerts for.")
50
+
51
+
52
+class HostAlertsSearchResponse(BaseModel):
53
+ alerts_summary: List[Alert]
54
+ success: bool
55
+ message: str
56
+
57
+
58
+class IndexAlertsSearchBody(AlertsSearchBody):
59
+ index_name: str = Field(..., description="The name of the index to search alerts for.")
60
+
61
+
62
+class IndexAlertsSearchResponse(BaseModel):
63
+ alerts_summary: List[Alert]
64
+ success: bool
65
+ message: str
66
+
67
+
68
+class AlertsByHost(BaseModel):
69
+ agent_name: str
70
+ number_of_alerts: int
71
+
72
+
73
+class AlertsByHostResponse(BaseModel):
74
+ alerts_by_host: List[AlertsByHost]
75
+ success: bool
76
+ message: str
77
+
78
+
79
+class AlertsByRule(BaseModel):
80
+ rule: str
81
+ number_of_alerts: int
82
+
83
+
84
+class AlertsByRuleResponse(BaseModel):
85
+ alerts_by_rule: List[AlertsByRule]
86
+ success: bool
87
+ message: str
88
+
89
+
90
+class AlertsByRulePerHost(BaseModel):
91
+ agent_name: str
92
+ number_of_alerts: int
93
+ rule: str
94
+
95
+
96
+class AlertsByRulePerHostResponse(BaseModel):
97
+ alerts_by_rule_per_host: List[AlertsByRulePerHost]
98
+ success: bool
99
+ message: str
backend/app/connectors/wazuh_indexer/schema/indices.py
new
+44
@@ -0,0 +1,44 @@
1
+from typing import Dict
2
+
3
+from pydantic import BaseModel
4
+from pydantic import Field
5
+
6
+
7
+class Indices(BaseModel):
8
+ indices_list: list
9
+ success: bool
10
+ message: str
11
+
12
+
13
+class IndexConfigModel(BaseModel):
14
+ SKIP_INDEX_NAMES: Dict[str, bool] = Field(
15
+ default={
16
+ "wazuh-statistics": True,
17
+ "wazuh-monitoring": True,
18
+ },
19
+ description="A dictionary containing index names to be skipped and their skip status.",
20
+ )
21
+
22
+ def is_index_skipped(self, index_name: str) -> bool:
23
+ """
24
+ Checks whether the given index name should be skipped.
25
+
26
+ Args:
27
+ index_name (str): The name of the index to check.
28
+
29
+ Returns:
30
+ bool: True if the index should be skipped, False otherwise.
31
+ """
32
+ return any(index_name.startswith(skipped) for skipped in self.SKIP_INDEX_NAMES)
33
+
34
+ def is_valid_index(self, index_name: str) -> bool:
35
+ """
36
+ Checks if the index name starts with "wazuh_" and is not in the SKIP_INDEX_NAMES list.
37
+
38
+ Args:
39
+ index_name (str): The name of the index to check.
40
+
41
+ Returns:
42
+ bool: True if the index is valid, False otherwise.
43
+ """
44
+ return index_name.startswith("wazuh") and not self.is_index_skipped(index_name)
backend/app/connectors/wazuh_indexer/schema/monitoring.py
new
+74
@@ -0,0 +1,74 @@
1
+from typing import List
2
+from typing import Optional
3
+from typing import Union
4
+
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+
8
+
9
+class ClusterHealth(BaseModel):
10
+ active_primary_shards: int
11
+ active_shards: int
12
+ active_shards_percent_as_number: Union[int, float]
13
+ cluster_name: str
14
+ delayed_unassigned_shards: int
15
+ discovered_cluster_manager: bool
16
+ discovered_master: bool
17
+ initializing_shards: int
18
+ number_of_data_nodes: int
19
+ number_of_in_flight_fetch: int
20
+ number_of_nodes: int
21
+ number_of_pending_tasks: int
22
+ relocating_shards: int
23
+ status: str
24
+ task_max_waiting_in_queue_millis: int
25
+ timed_out: bool
26
+ unassigned_shards: int
27
+
28
+
29
+class ClusterHealthResponse(BaseModel):
30
+ cluster_health: Optional[ClusterHealth]
31
+ message: str
32
+ success: bool
33
+
34
+
35
+class NodeAllocation(BaseModel):
36
+ disk_available: Optional[str] = Field(None, description="Disk available in bytes")
37
+ disk_percent: Optional[str] = Field(None, description="Disk percent")
38
+ disk_total: Optional[str] = Field(None, description="Disk total in bytes")
39
+ disk_used: Optional[str] = Field(None, description="Disk used in bytes")
40
+ node: str
41
+
42
+
43
+class NodeAllocationResponse(BaseModel):
44
+ node_allocation: Optional[List[NodeAllocation]]
45
+ message: str
46
+ success: bool
47
+
48
+
49
+class IndicesStats(BaseModel):
50
+ docs_count: str
51
+ health: str
52
+ index: str
53
+ replica_count: str
54
+ store_size: str
55
+
56
+
57
+class IndicesStatsResponse(BaseModel):
58
+ indices_stats: Optional[List[IndicesStats]]
59
+ message: str
60
+ success: bool
61
+
62
+
63
+class Shards(BaseModel):
64
+ index: str
65
+ node: Optional[str] = Field(None, description="Node name")
66
+ shard: int
67
+ state: str
68
+ size: Optional[str] = Field(None, description="Shard size in bytes")
69
+
70
+
71
+class ShardsResponse(BaseModel):
72
+ shards: Optional[List[Shards]]
73
+ message: str
74
+ success: bool
backend/app/connectors/wazuh_indexer/services/alerts.py
new
+163
@@ -0,0 +1,163 @@
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.utils.universal import AlertsQueryBuilder
23
+from app.connectors.wazuh_indexer.utils.universal import collect_indices
24
+from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
25
+
26
+# def collect_and_aggregate_alerts(field_name: str, search_body: AlertsSearchBody) -> Dict[str, int]:
27
+# indices = collect_indices()
28
+# aggregated_alerts_dict = {}
29
+
30
+# for index_name in indices.indices_list:
31
+# try:
32
+# alerts_response = collect_alerts_generic(index_name, body=search_body)
33
+# if alerts_response.success:
34
+# for alert in alerts_response.alerts:
35
+# field_value = alert["_source"][field_name]
36
+# aggregated_alerts_dict[field_value] = aggregated_alerts_dict.get(field_value, 0) + 1
37
+# except HTTPException as e:
38
+# logger.warning(f"An error occurred while processing index {index_name}: {e.detail}")
39
+
40
+# return aggregated_alerts_dict
41
+
42
+
43
+def collect_and_aggregate_alerts(field_names: List[str], search_body: AlertsSearchBody) -> Dict[str, int]:
44
+ indices = collect_indices()
45
+ aggregated_alerts_dict = {}
46
+
47
+ for index_name in indices.indices_list:
48
+ try:
49
+ alerts_response = collect_alerts_generic(index_name, body=search_body)
50
+ if alerts_response.success:
51
+ for alert in alerts_response.alerts:
52
+ composite_key = tuple(alert["_source"][field] for field in field_names)
53
+ aggregated_alerts_dict[composite_key] = aggregated_alerts_dict.get(composite_key, 0) + 1
54
+ except HTTPException as e:
55
+ logger.warning(f"An error occurred while processing index {index_name}: {e.detail}")
56
+
57
+ return aggregated_alerts_dict
58
+
59
+
60
+def collect_alerts_generic(index_name: str, body: AlertsSearchBody, is_host_specific: bool = False) -> CollectAlertsResponse:
61
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
62
+ query_builder = AlertsQueryBuilder()
63
+ query_builder.add_time_range(timerange=body.timerange, timestamp_field=body.timestamp_field)
64
+ query_builder.add_matches(matches=[(body.alert_field, body.alert_value)])
65
+ query_builder.add_sort(body.timestamp_field)
66
+
67
+ if is_host_specific:
68
+ query_builder.add_match_phrase(matches=[("agent_name", body.agent_name)])
69
+
70
+ query = query_builder.build()
71
+
72
+ try:
73
+ alerts = es_client.search(index=index_name, body=query, size=body.size)
74
+ logger.info(f"Alerts collected: {alerts}")
75
+ alerts_list = [alert for alert in alerts["hits"]["hits"]]
76
+ logger.info(f"Alerts collected: {alerts_list}")
77
+ return CollectAlertsResponse(alerts=alerts_list, success=True, message="Alerts collected successfully")
78
+ except Exception as e:
79
+ logger.debug(f"Failed to collect alerts: {e}")
80
+ return CollectAlertsResponse(alerts=[], success=False, message=f"Failed to collect alerts: {e}")
81
+
82
+
83
+def get_alerts_generic(search_body: Type[AlertsSearchBody], is_host_specific: bool = False, index_name: Optional[str] = None):
84
+ logger.info(f"Collecting Wazuh Indexer alerts for host {search_body.agent_name if is_host_specific else ''}")
85
+ alerts_summary = []
86
+ indices = collect_indices()
87
+ index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices
88
+
89
+ for index_name in index_list:
90
+ try:
91
+ alerts = collect_alerts_generic(index_name, body=search_body, is_host_specific=is_host_specific)
92
+ if alerts.success and len(alerts.alerts) > 0:
93
+ alerts_summary.append(
94
+ {
95
+ "index_name": index_name,
96
+ "total_alerts": len(alerts.alerts),
97
+ "alerts": alerts.alerts,
98
+ },
99
+ )
100
+ except HTTPException as e:
101
+ logger.warning(f"An error occurred while processing index {index_name}: {e.detail}")
102
+
103
+ if len(alerts_summary) == 0:
104
+ message = "No alerts found"
105
+ else:
106
+ message = f"Succesfully collected top {search_body.size} alerts for each index"
107
+
108
+ return {"alerts_summary": alerts_summary, "success": len(alerts_summary) > 0, "message": message}
109
+
110
+
111
+def get_alerts(search_body: AlertsSearchBody) -> AlertsSearchResponse:
112
+ result = get_alerts_generic(search_body)
113
+ return AlertsSearchResponse(**result)
114
+
115
+
116
+def get_host_alerts(search_body: HostAlertsSearchBody) -> HostAlertsSearchResponse:
117
+ result = get_alerts_generic(search_body, is_host_specific=True)
118
+ return HostAlertsSearchResponse(**result)
119
+
120
+
121
+def get_index_alerts(search_body: IndexAlertsSearchBody) -> IndexAlertsSearchResponse:
122
+ result = get_alerts_generic(search_body, index_name=search_body.index_name)
123
+ return IndexAlertsSearchResponse(**result)
124
+
125
+
126
+def get_alerts_by_host(search_body: AlertsSearchBody) -> AlertsByHostResponse:
127
+ aggregated_by_host = collect_and_aggregate_alerts(["agent_name"], search_body)
128
+ alerts_by_host_list: List[AlertsByHost] = [
129
+ AlertsByHost(agent_name=host[0], number_of_alerts=count) # host[0] because host is now a tuple
130
+ for host, count in aggregated_by_host.items()
131
+ ]
132
+ return AlertsByHostResponse(
133
+ alerts_by_host=alerts_by_host_list,
134
+ success=bool(alerts_by_host_list),
135
+ message="Successfully collected alerts by host",
136
+ )
137
+
138
+
139
+def get_alerts_by_rule(search_body: AlertsSearchBody) -> AlertsByRuleResponse:
140
+ aggregated_by_rule = collect_and_aggregate_alerts(["rule_description"], search_body)
141
+ alerts_by_rule_list: List[AlertsByRule] = [
142
+ AlertsByRule(rule=rule[0], number_of_alerts=count) # rule[0] because rule is now a tuple
143
+ for rule, count in aggregated_by_rule.items()
144
+ ]
145
+ return AlertsByRuleResponse(
146
+ alerts_by_rule=alerts_by_rule_list,
147
+ success=bool(alerts_by_rule_list),
148
+ message="Successfully collected alerts by rule",
149
+ )
150
+
151
+
152
+def get_alerts_by_rule_per_host(search_body: AlertsSearchBody) -> AlertsByRulePerHostResponse:
153
+ aggregated_by_rule_per_host = collect_and_aggregate_alerts(["agent_name", "rule_description"], search_body)
154
+ alerts_by_rule_per_host_list: List[AlertsByRulePerHost] = [
155
+ AlertsByRulePerHost(agent_name=agent_name, rule=rule, number_of_alerts=count)
156
+ for (agent_name, rule), count in aggregated_by_rule_per_host.items()
157
+ ]
158
+
159
+ return AlertsByRulePerHostResponse(
160
+ alerts_by_rule_per_host=alerts_by_rule_per_host_list,
161
+ success=bool(alerts_by_rule_per_host_list),
162
+ message="Successfully collected alerts by rule per host",
163
+ )
backend/app/connectors/wazuh_indexer/services/monitoring.py
new
+126
@@ -0,0 +1,126 @@
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
+def cluster_healthcheck() -> Union[ClusterHealthResponse, Dict[str, str]]:
21
+ """
22
+ Returns the cluster health of the Wazuh Indexer service.
23
+
24
+ Returns:
25
+ ElasticsearchResponse: A Pydantic model containing the cluster health of the Wazuh Indexer service.
26
+
27
+ Raises:
28
+ Exception: An exception is raised if the cluster health cannot be retrieved.
29
+ """
30
+ logger.info("Collecting Wazuh Indexer healthcheck")
31
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
32
+ try:
33
+ cluster_health_data = es_client.cluster.health()
34
+ cluster_health_model = ClusterHealth(**cluster_health_data)
35
+ return ClusterHealthResponse(
36
+ cluster_health=cluster_health_model,
37
+ success=True,
38
+ message="Successfully collected Wazuh Indexer cluster health",
39
+ )
40
+ except Exception as e:
41
+ logger.error(f"Cluster health check failed with error: {e}")
42
+ return {"success": False, "message": f"Cluster health check failed with error: {e}"}
43
+
44
+
45
+def node_allocation() -> Union[NodeAllocationResponse, Dict[str, bool]]:
46
+ """
47
+ Returns the node allocation of the Wazuh Indexer service.
48
+
49
+ Returns:
50
+ ElasticsearchResponse: A Pydantic model containing the node allocation of the Wazuh Indexer service.
51
+
52
+ Raises:
53
+ Exception: An exception is raised if the node allocation cannot be retrieved.
54
+ """
55
+ logger.info("Collecting Wazuh Indexer node allocation")
56
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
57
+ try:
58
+ raw_node_allocation_data = es_client.cat.allocation(format="json")
59
+ logger.info(raw_node_allocation_data)
60
+
61
+ formatted_node_allocation_data = format_node_allocation(raw_node_allocation_data)
62
+
63
+ node_allocation_models = [NodeAllocation(**node) for node in formatted_node_allocation_data]
64
+
65
+ return NodeAllocationResponse(
66
+ node_allocation=node_allocation_models,
67
+ success=True,
68
+ message="Successfully collected Wazuh Indexer node allocation",
69
+ )
70
+ except Exception as e:
71
+ logger.error(f"Node allocation check failed with error: {e}")
72
+ return {"success": False, "message": f"Node allocation check failed with error: {e}"}
73
+
74
+
75
+def indices_stats() -> Union[IndicesStatsResponse, Dict[str, str]]:
76
+ """
77
+ Returns the indices stats of the Wazuh Indexer service.
78
+
79
+ Returns:
80
+ ElasticsearchResponse: A Pydantic model containing the indices stats of the Wazuh Indexer service.
81
+
82
+ Raises:
83
+ Exception: An exception is raised if the indices stats cannot be retrieved.
84
+ """
85
+ logger.info("Collecting Wazuh Indexer indices stats")
86
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
87
+ try:
88
+ raw_indices_stats_data = es_client.cat.indices(format="json")
89
+
90
+ formatted_indices_stats_data = format_indices_stats(raw_indices_stats_data)
91
+
92
+ indices_stats_models = [IndicesStats(**index) for index in formatted_indices_stats_data]
93
+
94
+ return IndicesStatsResponse(
95
+ indices_stats=indices_stats_models,
96
+ success=True,
97
+ message="Successfully collected Wazuh Indexer indices stats",
98
+ )
99
+ except Exception as e:
100
+ logger.error(f"Indices stats check failed with error: {e}")
101
+ return {"success": False, "message": f"Indices stats check failed with error: {e}"}
102
+
103
+
104
+def shards() -> Union[ShardsResponse, Dict[str, str]]:
105
+ """
106
+ Returns the shards of the Wazuh Indexer service.
107
+
108
+ Returns:
109
+ ElasticsearchResponse: A Pydantic model containing the shards of the Wazuh Indexer service.
110
+
111
+ Raises:
112
+ Exception: An exception is raised if the shards cannot be retrieved.
113
+ """
114
+ logger.info("Collecting Wazuh Indexer shards")
115
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
116
+ try:
117
+ raw_shards_data = es_client.cat.shards(format="json")
118
+
119
+ formatted_shards_data = format_shards(raw_shards_data)
120
+
121
+ shard_models = [Shards(**shard) for shard in formatted_shards_data]
122
+
123
+ return ShardsResponse(shards=shard_models, success=True, message="Successfully collected Wazuh Indexer shards")
124
+ except Exception as e:
125
+ logger.error(f"Shards check failed with error: {e}")
126
+ return {"success": False, "message": f"Shards check failed with error: {e}"}
backend/app/connectors/wazuh_indexer/utils/universal.py
new
+288
@@ -0,0 +1,288 @@
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
+
8
+from elasticsearch7 import Elasticsearch
9
+from loguru import logger
10
+
11
+from app.connectors.utils import get_connector_info_from_db
12
+from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
13
+from app.connectors.wazuh_indexer.schema.indices import Indices
14
+
15
+
16
+def verify_wazuh_indexer_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
17
+ """
18
+ Verifies the connection to Wazuh Indexer service.
19
+
20
+ Returns:
21
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
22
+ """
23
+ logger.info(f"Verifying the wazuh-indexer connection to {attributes['connector_url']}")
24
+
25
+ try:
26
+ es = Elasticsearch(
27
+ [attributes["connector_url"]],
28
+ http_auth=(attributes["connector_username"], attributes["connector_password"]),
29
+ verify_certs=False,
30
+ timeout=15,
31
+ max_retries=10,
32
+ retry_on_timeout=False,
33
+ )
34
+ es.cluster.health()
35
+ logger.debug("Wazuh Indexer connection successful")
36
+ return {"connectionSuccessful": True, "message": "Wazuh Indexer connection successful"}
37
+ except Exception as e:
38
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")
39
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
40
+
41
+
42
+def verify_wazuh_indexer_connection(connector_name: str) -> str:
43
+ """
44
+ Returns the authentication token for the Wazuh Indexer service.
45
+
46
+ Returns:
47
+ str: Authentication token for the Wazuh Indexer service.
48
+ """
49
+ attributes = get_connector_info_from_db(connector_name)
50
+ if attributes is None:
51
+ logger.error("No Wazuh Indexer connector found in the database")
52
+ return None
53
+ return verify_wazuh_indexer_credentials(attributes)
54
+
55
+
56
+def create_wazuh_indexer_client(connector_name: str) -> Elasticsearch:
57
+ """
58
+ Returns an Elasticsearch client for the Wazuh Indexer service.
59
+
60
+ Returns:
61
+ Elasticsearch: Elasticsearch client for the Wazuh Indexer service.
62
+ """
63
+ attributes = get_connector_info_from_db(connector_name)
64
+ if attributes is None:
65
+ logger.error("No Wazuh Indexer connector found in the database")
66
+ return None
67
+ return Elasticsearch(
68
+ [attributes["connector_url"]],
69
+ http_auth=(attributes["connector_username"], attributes["connector_password"]),
70
+ verify_certs=False,
71
+ timeout=15,
72
+ max_retries=10,
73
+ retry_on_timeout=False,
74
+ )
75
+
76
+
77
+def format_node_allocation(node_allocation):
78
+ """
79
+ Format the node allocation details into a list of dictionaries. Each dictionary contains disk used, disk available, total disk, disk
80
+ usage percentage, and node name.
81
+
82
+ Args:
83
+ node_allocation: Node allocation details from Elasticsearch.
84
+
85
+ Returns:
86
+ list: A list of dictionaries containing formatted node allocation details.
87
+ """
88
+ return [
89
+ {
90
+ "disk_used": node["disk.used"],
91
+ "disk_available": node["disk.avail"],
92
+ "disk_total": node["disk.total"],
93
+ "disk_percent": node["disk.percent"],
94
+ "node": node["node"],
95
+ }
96
+ for node in node_allocation
97
+ ]
98
+
99
+
100
+def format_indices_stats(indices_stats):
101
+ """
102
+ Format the indices stats details into a list of dictionaries. Each dictionary contains the index name, the number of documents in the index,
103
+ the size of the index, and the number of shards in the index.
104
+
105
+ Args:
106
+ indices_stats: Indices stats details from Elasticsearch.
107
+
108
+ Returns:
109
+ list: A list of dictionaries containing formatted indices stats details.
110
+ """
111
+ return [
112
+ {
113
+ "index": index["index"],
114
+ "docs_count": index["docs.count"],
115
+ "store_size": index["store.size"],
116
+ "replica_count": index["rep"],
117
+ "health": index["health"],
118
+ }
119
+ for index in indices_stats
120
+ ]
121
+
122
+
123
+def format_shards(shards):
124
+ """
125
+ Format the shards details into a list of dictionaries. Each dictionary contains the index name, the shard number, the shard state, the shard
126
+ size, and the node name.
127
+
128
+ Args:
129
+ shards: Shards details from Elasticsearch.
130
+
131
+ Returns:
132
+ list: A list of dictionaries containing formatted shards details.
133
+ """
134
+ return [
135
+ {
136
+ "index": shard["index"],
137
+ "shard": shard["shard"],
138
+ "state": shard["state"],
139
+ "size": shard["store"],
140
+ "node": shard["node"],
141
+ }
142
+ for shard in shards
143
+ ]
144
+
145
+
146
+def collect_indices() -> Indices:
147
+ """
148
+ Collects the indices from Elasticsearch.
149
+
150
+ Returns:
151
+ dict: A dictionary containing the indices, shards, and indices stats.
152
+ """
153
+ logger.info("Collecting indices from Elasticsearch")
154
+ es = create_wazuh_indexer_client("Wazuh-Indexer")
155
+ try:
156
+ indices_dict = es.indices.get_alias("*")
157
+ indices_list = list(indices_dict.keys())
158
+ # Check if the index is valid
159
+ index_config = IndexConfigModel()
160
+ indices_list = [index for index in indices_list if index_config.is_valid_index(index)]
161
+ return Indices(indices_list=indices_list, success=True, message="Indices collected successfully")
162
+ except Exception as e:
163
+ logger.error(f"Failed to collect indices: {e}")
164
+ return Indices(message="Failed to collect indices", success=False)
165
+
166
+
167
+class AlertsQueryBuilder:
168
+ @staticmethod
169
+ def _get_time_range_start(timerange: str) -> str:
170
+ """
171
+ Determines the start time of the time range based on the current time and the provided timerange.
172
+
173
+ Args:
174
+ timerange (str): The time range to collect alerts from. This is a string like "24h", "1w", etc.
175
+
176
+ Returns:
177
+ str: A string representing the start time of the time range in ISO format.
178
+ """
179
+ if timerange.endswith("h"):
180
+ delta = timedelta(hours=int(timerange[:-1]))
181
+ elif timerange.endswith("d"):
182
+ delta = timedelta(days=int(timerange[:-1]))
183
+ elif timerange.endswith("w"):
184
+ delta = timedelta(weeks=int(timerange[:-1]))
185
+ else:
186
+ raise ValueError("Invalid timerange format. Expected a string like '24h', '1d', '1w', etc.")
187
+
188
+ start = datetime.utcnow() - delta
189
+ return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end
190
+
191
+ def __init__(self):
192
+ self.query = {
193
+ "query": {
194
+ "bool": {
195
+ "must": [],
196
+ },
197
+ },
198
+ "sort": [],
199
+ }
200
+
201
+ def add_time_range(self, timerange: str, timestamp_field: str):
202
+ start = self._get_time_range_start(timerange)
203
+ self.query["query"]["bool"]["must"].append({"range": {timestamp_field: {"gte": start, "lte": "now"}}})
204
+ return self
205
+
206
+ def add_matches(self, matches: Iterable[Tuple[str, str]]):
207
+ for field, value in matches:
208
+ self.query["query"]["bool"]["must"].append({"match": {field: value}})
209
+ return self
210
+
211
+ def add_match_phrase(self, matches: Iterable[Tuple[str, str]]):
212
+ for field, value in matches:
213
+ self.query["query"]["bool"]["must"].append({"match_phrase": {field: value}})
214
+ return self
215
+
216
+ def add_range(self, field: str, value: str):
217
+ self.query["query"]["bool"]["must"].append({"range": {field: {"gte": value}}})
218
+ return self
219
+
220
+ def add_sort(self, field: str, order: str = "desc"):
221
+ self.query["sort"].append({field: {"order": order}})
222
+ return self
223
+
224
+ def build(self):
225
+ return self.query
226
+
227
+
228
+class LogsQueryBuilder:
229
+ @staticmethod
230
+ def _get_time_range_start(timerange: str) -> str:
231
+ """
232
+ Determines the start time of the time range based on the current time and the provided timerange.
233
+
234
+ Args:
235
+ timerange (str): The time range to collect alerts from. This is a string like "24h", "1w", etc.
236
+
237
+ Returns:
238
+ str: A string representing the start time of the time range in ISO format.
239
+ """
240
+ if timerange.endswith("m"):
241
+ delta = timedelta(minutes=int(timerange[:-1]))
242
+ elif timerange.endswith("h"):
243
+ delta = timedelta(hours=int(timerange[:-1]))
244
+ elif timerange.endswith("d"):
245
+ delta = timedelta(days=int(timerange[:-1]))
246
+ elif timerange.endswith("w"):
247
+ delta = timedelta(weeks=int(timerange[:-1]))
248
+ else:
249
+ raise ValueError("Invalid timerange format. Expected a string like '24h', '1d', '1w', '1m', etc.")
250
+
251
+ start = datetime.utcnow() - delta
252
+ return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end
253
+
254
+ def __init__(self):
255
+ self.query = {
256
+ "query": {
257
+ "bool": {
258
+ "must": [],
259
+ },
260
+ },
261
+ "sort": [],
262
+ }
263
+
264
+ def add_time_range(self, timerange: str, timestamp_field: str):
265
+ start = self._get_time_range_start(timerange)
266
+ self.query["query"]["bool"]["must"].append({"range": {timestamp_field: {"gte": start, "lte": "now"}}})
267
+ return self
268
+
269
+ def add_matches(self, matches: Iterable[Tuple[str, str]]):
270
+ for field, value in matches:
271
+ self.query["query"]["bool"]["must"].append({"match": {field: value}})
272
+ return self
273
+
274
+ def add_match_phrase(self, matches: Iterable[Tuple[str, str]]):
275
+ for field, value in matches:
276
+ self.query["query"]["bool"]["must"].append({"match_phrase": {field: value}})
277
+ return self
278
+
279
+ def add_range(self, field: str, value: str):
280
+ self.query["query"]["bool"]["must"].append({"range": {field: {"gte": value}}})
281
+ return self
282
+
283
+ def add_sort(self, field: str, order: str = "desc"):
284
+ self.query["sort"].append({field: {"order": order}})
285
+ return self
286
+
287
+ def build(self):
288
+ return self.query
backend/app/connectors/wazuh_manager/models/rules.py
new
+16
@@ -0,0 +1,16 @@
1
+import datetime
2
+from typing import Optional
3
+
4
+from sqlmodel import Field
5
+from sqlmodel import SQLModel
6
+
7
+
8
+class DisabledRule(SQLModel, table=True):
9
+ id: Optional[int] = Field(primary_key=True)
10
+ rule_id: str = Field(index=True)
11
+ previous_level: str = Field(max_length=256)
12
+ new_level: str = Field(max_length=256)
13
+ reason_for_disabling: str = Field(max_length=256)
14
+ length_of_time: str = Field(max_length=256)
15
+ date_disabled: datetime.datetime = datetime.datetime.now()
16
+ disabled_by: str = Field(max_length=256)
backend/app/connectors/wazuh_manager/routes/rules.py
new
+100
@@ -0,0 +1,100 @@
1
+from fastapi import APIRouter
2
+from fastapi import Depends
3
+from fastapi import HTTPException
4
+from fastapi import Security
5
+
6
+# App specific imports
7
+from app.auth.routes.auth import AuthHandler
8
+from app.connectors.wazuh_manager.models.rules import DisabledRule
9
+from app.connectors.wazuh_manager.schema.rules import AllDisabledRuleResponse
10
+from app.connectors.wazuh_manager.schema.rules import RuleDisable
11
+from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse
12
+from app.connectors.wazuh_manager.schema.rules import RuleEnable
13
+from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse
14
+from app.connectors.wazuh_manager.schema.rules import RuleExclude
15
+from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse
16
+from app.connectors.wazuh_manager.services.rules import disable_rule
17
+from app.connectors.wazuh_manager.services.rules import enable_rule
18
+from app.connectors.wazuh_manager.services.rules import exclude_rule
19
+from app.db.db_session import session
20
+
21
+NEW_LEVEL = "1"
22
+wazuh_manager_router = APIRouter()
23
+auth_handler = AuthHandler()
24
+
25
+
26
+def query_disabled_rule(rule_id: str):
27
+ return session.query(DisabledRule).filter(DisabledRule.rule_id == rule_id).first()
28
+
29
+
30
+@wazuh_manager_router.get(
31
+ "/rule/disabled",
32
+ response_model=AllDisabledRuleResponse,
33
+ description="Get all disabled rules",
34
+ dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
35
+)
36
+async def get_disabled_rules() -> AllDisabledRuleResponse:
37
+ disabled_rules = session.query(DisabledRule).all()
38
+ return AllDisabledRuleResponse(disabled_rules=disabled_rules, success=True, message="Successfully fetched all disabled rules")
39
+
40
+
41
+@wazuh_manager_router.post(
42
+ "/rule/disable",
43
+ response_model=RuleDisableResponse,
44
+ description="Disable a Wazuh Rule",
45
+ dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
46
+)
47
+async def disable_wazuh_rule(rule: RuleDisable, username: str = Depends(auth_handler.get_current_user)) -> RuleDisableResponse:
48
+ if query_disabled_rule(rule.rule_id):
49
+ raise HTTPException(status_code=404, detail="Rule is already disabled")
50
+
51
+ rule_disabled = disable_rule(rule)
52
+ if rule_disabled:
53
+ new_disabled_rule = DisabledRule(
54
+ rule_id=rule.rule_id,
55
+ previous_level=rule_disabled.previous_level,
56
+ new_level=NEW_LEVEL,
57
+ reason_for_disabling=rule.reason_for_disabling,
58
+ length_of_time=rule.length_of_time,
59
+ disabled_by=username.username,
60
+ )
61
+ session.add(new_disabled_rule)
62
+ session.commit()
63
+ return rule_disabled
64
+ else:
65
+ raise HTTPException(status_code=404, detail="Was not able to disable rule")
66
+
67
+
68
+@wazuh_manager_router.post(
69
+ "/rule/enable",
70
+ response_model=RuleEnableResponse,
71
+ description="Enable a Wazuh Rule",
72
+ dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
73
+)
74
+async def enable_wazuh_rule(rule: RuleEnable) -> RuleEnableResponse:
75
+ disabled_rule = query_disabled_rule(rule.rule_id)
76
+ if not disabled_rule:
77
+ raise HTTPException(status_code=404, detail="Rule is already enabled")
78
+
79
+ previous_level = disabled_rule.previous_level
80
+ rule_enabled = enable_rule(rule, previous_level)
81
+
82
+ if rule_enabled:
83
+ session.delete(disabled_rule)
84
+ session.commit()
85
+ return rule_enabled
86
+ else:
87
+ raise HTTPException(status_code=404, detail="Was not able to enable rule")
88
+
89
+
90
+@wazuh_manager_router.post(
91
+ "/rule/exclude",
92
+ response_model=RuleExcludeResponse,
93
+ description="Retrieve recommended exclusion for a Wazuh Rule",
94
+)
95
+async def exclude_wazuh_rule(rule: RuleExclude) -> RuleExcludeResponse:
96
+ recommended_exclusion = exclude_rule(rule)
97
+ if recommended_exclusion:
98
+ return recommended_exclusion
99
+ else:
100
+ raise HTTPException(status_code=404, detail="Was not able to exclude rule")
backend/app/connectors/wazuh_manager/schema/rules.py
new
+66
@@ -0,0 +1,66 @@
1
+from typing import List
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
+
7
+
8
+class RuleDisable(BaseModel):
9
+ rule_id: str
10
+ reason_for_disabling: str
11
+ length_of_time: str
12
+
13
+
14
+class RuleDisableResponse(BaseModel):
15
+ previous_level: Optional[str]
16
+ message: str
17
+ success: bool
18
+
19
+
20
+class RuleEnable(BaseModel):
21
+ rule_id: str
22
+ reason_for_enabling: str
23
+
24
+
25
+class RuleEnableResponse(BaseModel):
26
+ new_level: Optional[str]
27
+ message: str
28
+ success: bool
29
+
30
+
31
+class AllDisabledRule(BaseModel):
32
+ rule_id: str
33
+ previous_level: str
34
+ new_level: str
35
+ reason_for_disabling: str
36
+ length_of_time: str
37
+ disabled_by: str
38
+
39
+
40
+class AllDisabledRuleResponse(BaseModel):
41
+ disabled_rules: List[AllDisabledRule]
42
+ success: bool
43
+ message: str
44
+
45
+
46
+class RuleExclude(BaseModel):
47
+ rule_value: str = Field(
48
+ ...,
49
+ description="The value of the field trying to be exclude",
50
+ example="C:\\Windows\\ServiceState\\EventLog\\Data\\lastalive1.dat",
51
+ )
52
+ input_value: str = Field(
53
+ ...,
54
+ description="The proposed value of the field trying to be exclude that would result in an exclusiong",
55
+ example="C:\\\\Windows\\\\ServiceState\\\\EventLog\\\\Data\\\\lastalive1\.dat",
56
+ )
57
+
58
+
59
+class RuleExcludeResponse(BaseModel):
60
+ success: bool
61
+ message: str
62
+ recommended_exclusion: str = Field(
63
+ ...,
64
+ description="The recommended exclusion for the rule",
65
+ example="C:\\\\Windows\\\\ServiceState\\\\EventLog\\\\Data\\\\lastalive1\.dat",
66
+ )
backend/app/connectors/wazuh_manager/services/rules.py
new
+142
@@ -0,0 +1,142 @@
1
+from typing import Any
2
+from typing import Dict
3
+from typing import List
4
+from typing import Tuple
5
+from typing import Union
6
+
7
+import pcre2
8
+import xmltodict
9
+from loguru import logger
10
+
11
+from app.connectors.wazuh_manager.schema.rules import RuleDisable
12
+from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse
13
+from app.connectors.wazuh_manager.schema.rules import RuleEnable
14
+from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse
15
+from app.connectors.wazuh_manager.schema.rules import RuleExclude
16
+from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse
17
+from app.connectors.wazuh_manager.utils.universal import restart_service
18
+from app.connectors.wazuh_manager.utils.universal import send_get_request
19
+from app.connectors.wazuh_manager.utils.universal import send_put_request
20
+
21
+
22
+def fetch_filename(rule_id: str) -> str:
23
+ endpoint = "rules"
24
+ params = {"rule_ids": rule_id}
25
+ filename_data = send_get_request(endpoint=endpoint, params=params)
26
+ if not filename_data["success"]:
27
+ raise ValueError(filename_data["message"])
28
+ return filename_data["data"]["data"]["affected_items"][0]["filename"]
29
+
30
+
31
+def fetch_file_content(filename: str) -> str:
32
+ endpoint = f"rules/files/{filename}"
33
+ file_content_data = send_get_request(endpoint=endpoint)
34
+ if not file_content_data["success"]:
35
+ raise ValueError(file_content_data["message"])
36
+ return file_content_data["data"]["data"]["affected_items"][0]["group"]
37
+
38
+
39
+def set_rule_level(file_content: Any, rule_id: str, new_level: str) -> Tuple[str, Any]:
40
+ previous_level = None
41
+ if isinstance(file_content, dict):
42
+ file_content = [file_content]
43
+ for group_block in file_content:
44
+ rule_block = group_block.get("rule", None)
45
+ if rule_block:
46
+ if isinstance(rule_block, dict):
47
+ rule_block = [rule_block]
48
+ for rule in rule_block:
49
+ if rule["@id"] == rule_id:
50
+ previous_level = rule["@level"]
51
+ rule["@level"] = new_level
52
+ break
53
+ return previous_level, file_content
54
+
55
+
56
+def convert_to_xml(updated_file_content: Union[Dict[str, str], List[Dict[str, str]]]) -> str:
57
+ xml_content_list = []
58
+ for group in updated_file_content:
59
+ xml_dict = {"group": group}
60
+ xml_content = xmltodict.unparse(xml_dict, pretty=True)
61
+ xml_content = xml_content.replace('<?xml version="1.0" encoding="utf-8"?>', "")
62
+ xml_content_list.append(xml_content)
63
+ xml_content = "\n".join(xml_content_list)
64
+ xml_content = xml_content.strip()
65
+ return xml_content
66
+
67
+
68
+def upload_updated_rule(filename: str, xml_content: str):
69
+ response = send_put_request(
70
+ endpoint=f"rules/files/{filename}",
71
+ data=xml_content,
72
+ params={"overwrite": "true"},
73
+ )
74
+ if not response["success"]:
75
+ raise ValueError(response["message"])
76
+
77
+
78
+def process_rule(rule, rule_action_func, ResponseModel):
79
+ filename, file_content = fetch_filename_and_content(rule.rule_id)
80
+ previous_level, updated_file_content = rule_action_func(file_content, rule.rule_id)
81
+ xml_content = convert_to_xml(updated_file_content)
82
+ upload_updated_rule(filename, xml_content)
83
+ restart_service()
84
+ return ResponseModel(
85
+ previous_level=previous_level,
86
+ success=True,
87
+ message=f"Rule {rule.rule_id} successfully processed in file {filename}.",
88
+ )
89
+
90
+
91
+def fetch_filename_and_content(rule_id: str) -> Tuple[str, str]:
92
+ filename = fetch_filename(rule_id)
93
+ file_content = fetch_file_content(filename)
94
+ return filename, file_content
95
+
96
+
97
+def disable_rule(rule: RuleDisable) -> RuleDisableResponse:
98
+ return process_rule(rule, lambda fc, rid: set_rule_level(fc, rid, "1"), RuleDisableResponse)
99
+
100
+
101
+def enable_rule(rule: RuleEnable, previous_level: str) -> RuleEnableResponse:
102
+ return process_rule(rule, lambda fc, rid: set_rule_level(fc, rid, previous_level), RuleEnableResponse)
103
+
104
+
105
+################# ! EXCLUDE RULE ! #################
106
+
107
+
108
+def make_pcre2_compatible(input_string: str) -> str:
109
+ """
110
+ Convert the input string to a PCRE2 compatible regex pattern.
111
+
112
+ Parameters:
113
+ - input_string (str): The input string to convert.
114
+
115
+ Returns:
116
+ - str: The PCRE2 compatible regex pattern.
117
+ """
118
+ # PCRE2 uses \\ to escape a backslash
119
+ return input_string.replace("\\", "\\\\")
120
+
121
+
122
+def exclude_rule(rule: RuleExclude) -> RuleExcludeResponse:
123
+ try:
124
+ # Convert rule_value to a PCRE2 compatible regex pattern
125
+ pcre2_pattern = make_pcre2_compatible(rule.rule_value)
126
+
127
+ compiled_pattern = pcre2.compile(pcre2_pattern)
128
+ print(f"Compiled Pattern: {compiled_pattern}") # Debugging line
129
+
130
+ print(f"Input Value: {rule.input_value}") # Debugging line
131
+
132
+ match_data = compiled_pattern.match(rule.input_value)
133
+
134
+ if match_data:
135
+ return RuleExcludeResponse(success=True, message="Successfully excluded rule", recommended_exclusion=rule.input_value)
136
+ else:
137
+ return RuleExcludeResponse(success=False, message="Failed to exclude rule", recommended_exclusion="")
138
+
139
+ except Exception as e:
140
+ print(f"Exception: {e}") # Debugging line
141
+ logger.error(f"Failed to exclude rule: {e}")
142
+ return RuleExcludeResponse(success=False, message=f"Failed to exclude rule: {e}", recommended_exclusion="")
backend/app/connectors/wazuh_manager/utils/universal.py
new
+253
@@ -0,0 +1,253 @@
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
+
10
+
11
+def verify_wazuh_manager_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
12
+ """
13
+ Verifies the connection to Wazuh manager service.
14
+
15
+ Returns:
16
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
17
+ """
18
+ logger.info(f"Verifying the wazuh-manager connection to {attributes['connector_url']}")
19
+
20
+ try:
21
+ wazuh_auth_token = requests.get(
22
+ f"{attributes['connector_url']}/security/user/authenticate",
23
+ auth=(
24
+ attributes["connector_username"],
25
+ attributes["connector_password"],
26
+ ),
27
+ verify=False,
28
+ )
29
+
30
+ if wazuh_auth_token.status_code == 200:
31
+ logger.debug("Wazuh Authentication Token successful")
32
+ return {"connectionSuccessful": True, "message": "Wazuh Manager authentication successful"}
33
+ else:
34
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {wazuh_auth_token.text}")
35
+
36
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed"}
37
+ except Exception as e:
38
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")
39
+
40
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error."}
41
+
42
+
43
+def verify_wazuh_manager_connection(connector_name: str) -> str:
44
+ """
45
+ Returns the authentication token for the Wazuh manager service.
46
+
47
+ Returns:
48
+ str: Authentication token for the Wazuh manager service.
49
+ """
50
+ logger.info("Getting Wazuh Manager authentication token")
51
+ attributes = get_connector_info_from_db(connector_name)
52
+ if attributes is None:
53
+ logger.error("No Wazuh Manager connector found in the database")
54
+ return None
55
+ return verify_wazuh_manager_credentials(attributes)
56
+
57
+
58
+def create_wazuh_manager_client(connector_name: str) -> str:
59
+ """
60
+ Returns the authentication token for the Wazuh manager service.
61
+
62
+ Returns:
63
+ str: Authentication token for the Wazuh manager service.
64
+ """
65
+ logger.info("Getting Wazuh Manager authentication token")
66
+ attributes = get_connector_info_from_db(connector_name)
67
+ if attributes is None:
68
+ logger.error("No Wazuh Manager connector found in the database")
69
+ return None
70
+ try:
71
+ wazuh_auth_token = requests.get(
72
+ f"{attributes['connector_url']}/security/user/authenticate",
73
+ auth=(
74
+ attributes["connector_username"],
75
+ attributes["connector_password"],
76
+ ),
77
+ verify=False,
78
+ )
79
+
80
+ if wazuh_auth_token.status_code == 200:
81
+ logger.debug("Wazuh Authentication Token successful")
82
+ wazuh_auth_token = wazuh_auth_token.json()
83
+ wazuh_auth_token = wazuh_auth_token["data"]["token"]
84
+
85
+ return {"Authorization": f"Bearer {wazuh_auth_token}"}
86
+ else:
87
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {wazuh_auth_token.text}")
88
+
89
+ return None
90
+ except Exception as e:
91
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")
92
+
93
+ return None
94
+
95
+
96
+def send_get_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Wazuh-Manager") -> Dict[str, Any]:
97
+ """
98
+ Sends a GET request to the Wazuh Manager service.
99
+
100
+ Args:
101
+ endpoint (str): The endpoint to send the GET request to.
102
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the GET request. Defaults to None.
103
+ connector_name (str, optional): The name of the connector to use. Defaults to "Wazuh-Manager".
104
+
105
+ Returns:
106
+ Dict[str, Any]: The response from the GET request.
107
+ """
108
+ logger.info(f"Sending GET request to {endpoint}")
109
+ wazuh_manager_client = create_wazuh_manager_client(connector_name)
110
+ attributes = get_connector_info_from_db(connector_name)
111
+ if attributes is None:
112
+ logger.error("No Wazuh Manager connector found in the database")
113
+ return None
114
+ try:
115
+ response = requests.get(
116
+ f"{attributes['connector_url']}/{endpoint}",
117
+ headers=wazuh_manager_client,
118
+ params=params,
119
+ verify=False,
120
+ )
121
+ response.raise_for_status()
122
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
123
+ except Exception as e:
124
+ logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
125
+ return {"success": False, "message": f"Failed to send GET request to {endpoint} with error: {e}"}
126
+
127
+
128
+def send_post_request(endpoint: str, data: Dict[str, Any], connector_name: str = "Wazuh-Manager") -> Dict[str, Any]:
129
+ """
130
+ Sends a POST request to the Wazuh Manager service.
131
+
132
+ Args:
133
+ endpoint (str): The endpoint to send the POST request to.
134
+ data (Dict[str, Any]): The data to send with the POST request.
135
+ connector_name (str, optional): The name of the connector to use. Defaults to "Wazuh-Manager".
136
+
137
+ Returns:
138
+ Dict[str, Any]: The response from the POST request.
139
+ """
140
+ logger.info(f"Sending POST request to {endpoint}")
141
+ wazuh_manager_client = create_wazuh_manager_client(connector_name)
142
+ attributes = get_connector_info_from_db(connector_name)
143
+ if attributes is None:
144
+ logger.error("No Wazuh Manager connector found in the database")
145
+ return None
146
+ try:
147
+ response = requests.post(
148
+ f"{attributes['connector_url']}/{endpoint}",
149
+ headers=wazuh_manager_client,
150
+ json=data,
151
+ verify=False,
152
+ )
153
+ response.raise_for_status()
154
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
155
+ except Exception as e:
156
+ logger.error(f"Failed to send POST request to {endpoint} with error: {e}")
157
+ return {"success": False, "message": f"Failed to send POST request to {endpoint} with error: {e}"}
158
+
159
+
160
+def send_put_request(
161
+ endpoint: str,
162
+ data: Optional[Dict[str, Any]],
163
+ params: Optional[Dict[str, str]] = None,
164
+ connector_name: str = "Wazuh-Manager",
165
+) -> Dict[str, Any]:
166
+ """
167
+ Sends a PUT request to the Wazuh Manager service.
168
+
169
+ Args:
170
+ endpoint (str): The endpoint to send the PUT request to.
171
+ data (Dict[str, Any]): The data to send with the PUT request.
172
+ connector_name (str, optional): The name of the connector to use. Defaults to "Wazuh-Manager".
173
+
174
+ Returns:
175
+ Dict[str, Any]: The response from the PUT request.
176
+ """
177
+ logger.info(f"Sending PUT request to {endpoint}")
178
+ wazuh_manager_client = create_wazuh_manager_client(connector_name)
179
+ attributes = get_connector_info_from_db(connector_name)
180
+ if attributes is None:
181
+ logger.error("No Wazuh Manager connector found in the database")
182
+ return None
183
+ try:
184
+ response = requests.put(
185
+ f"{attributes['connector_url']}/{endpoint}",
186
+ headers=wazuh_manager_client,
187
+ params=params,
188
+ data=data,
189
+ verify=False,
190
+ )
191
+ response.raise_for_status()
192
+ return {"data": response.json(), "success": True, "message": "Successfully retrieved data"}
193
+ except Exception as e:
194
+ logger.error(f"Failed to send PUT request to {endpoint} with error: {e}")
195
+ return {"success": False, "message": f"Failed to send PUT request to {endpoint} with error: {e}"}
196
+
197
+
198
+def send_delete_request(endpoint: str, params: Optional[Dict[str, Any]] = None, connector_name: str = "Wazuh-Manager") -> Dict[str, Any]:
199
+ """
200
+ Sends a DELETE request to the Wazuh Manager service.
201
+
202
+ Args:
203
+ endpoint (str): The endpoint to send the DELETE request to.
204
+ params (Optional[Dict[str, Any]], optional): The parameters to send with the DELETE request. Defaults to None.
205
+ connector_name (str, optional): The name of the connector to use. Defaults to "Wazuh-Manager".
206
+
207
+ Returns:
208
+ Dict[str, Any]: The response from the DELETE request.
209
+ """
210
+ logger.info(f"Sending DELETE request to {endpoint}")
211
+ wazuh_manager_client = create_wazuh_manager_client(connector_name)
212
+ attributes = get_connector_info_from_db(connector_name)
213
+ if attributes is None:
214
+ logger.error("No Wazuh Manager connector found in the database")
215
+ return None
216
+ try:
217
+ response = requests.delete(
218
+ f"{attributes['connector_url']}/{endpoint}",
219
+ headers=wazuh_manager_client,
220
+ params=params,
221
+ verify=False,
222
+ )
223
+ response.raise_for_status()
224
+ return {"data": response.json(), "success": True, "message": "Successfully deleted data"}
225
+ except Exception as e:
226
+ logger.error(f"Failed to send DELETE request to {endpoint} with error: {e}")
227
+ return {"success": False, "message": f"Failed to send DELETE request to {endpoint} with error: {e}"}
228
+
229
+
230
+def restart_service(connector_name: str = "Wazuh-Manager") -> Dict[str, Any]:
231
+ """
232
+ Restarts the Wazuh Manager service.
233
+
234
+ Returns:
235
+ Dict[str, Any]: The response from the DELETE request.
236
+ """
237
+ logger.info("Restarting Wazuh Manager service")
238
+ wazuh_manager_client = create_wazuh_manager_client(connector_name)
239
+ attributes = get_connector_info_from_db(connector_name)
240
+ if attributes is None:
241
+ logger.error("No Wazuh Manager connector found in the database")
242
+ return None
243
+ try:
244
+ response = requests.put(
245
+ f"{attributes['connector_url']}/manager/restart",
246
+ headers=wazuh_manager_client,
247
+ verify=False,
248
+ )
249
+ response.raise_for_status()
250
+ return {"data": response.json(), "success": True, "message": "Successfully restarted service"}
251
+ except Exception as e:
252
+ logger.error(f"Failed to restart Wazuh Manager service with error: {e}")
253
+ return {"success": False, "message": f"Failed to restart Wazuh Manager service with error: {e}"}
backend/app/customers/routes/customers.py
new
+239
@@ -0,0 +1,239 @@
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from fastapi import Query
4
+from loguru import logger
5
+from starlette.status import HTTP_401_UNAUTHORIZED
6
+
7
+# App specific imports
8
+from app.customers.schema.customers import AgentModel
9
+from app.customers.schema.customers import AgentsResponse
10
+from app.customers.schema.customers import CustomerFullResponse
11
+from app.customers.schema.customers import CustomerMetaRequestBody
12
+from app.customers.schema.customers import CustomerMetaResponse
13
+from app.customers.schema.customers import CustomerRequestBody
14
+from app.customers.schema.customers import CustomerResponse
15
+from app.customers.schema.customers import CustomersResponse
16
+from app.db.db_session import session
17
+from app.db.universal_models import Agents
18
+from app.db.universal_models import Customers
19
+from app.db.universal_models import CustomersMeta
20
+
21
+# from app.healthchecks.agents.schema.agents import AgentModel
22
+from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
23
+from app.healthchecks.agents.schema.agents import TimeCriteriaModel
24
+from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck
25
+from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck
26
+
27
+customers_router = APIRouter()
28
+
29
+
30
+def verify_admin(user):
31
+ if not user.is_admin:
32
+ raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Unauthorized")
33
+
34
+
35
+def verify_unique_customer_code(customer: CustomerRequestBody):
36
+ existing_customer = session.query(Customers).filter(Customers.customer_code == customer.customer_code).first()
37
+ if existing_customer:
38
+ raise HTTPException(status_code=400, detail="Customer with this customer_code already exists")
39
+
40
+
41
+@customers_router.post("", response_model=CustomerResponse, description="Create a new customer")
42
+async def create_customer(customer: CustomerRequestBody) -> CustomerResponse:
43
+ verify_unique_customer_code(customer)
44
+ logger.info(f"Creating new customer: {customer}")
45
+ new_customer = Customers(**customer.dict())
46
+ session.add(new_customer)
47
+ session.commit()
48
+ return CustomerResponse(customer=customer, success=True, message="Customer created successfully")
49
+
50
+
51
+@customers_router.get("", response_model=CustomersResponse, description="Get all customers")
52
+async def get_customers() -> CustomersResponse:
53
+ logger.info("Fetching all customers")
54
+ customers = session.query(Customers).all()
55
+ # Explode the customers list into a list of Customer objects
56
+ customers = [CustomerRequestBody.parse_obj(customer.__dict__) for customer in customers]
57
+ return CustomersResponse(customers=customers, success=True, message="Customers fetched successfully")
58
+
59
+
60
+@customers_router.get("/{customer_code}", response_model=CustomerResponse, description="Get customer by customer_code")
61
+async def get_customer(customer_code: str) -> CustomerResponse:
62
+ logger.info(f"Fetching customer with customer_code: {customer_code}")
63
+ customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
64
+ if not customer:
65
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
66
+ return CustomerResponse(
67
+ customer=CustomerRequestBody.parse_obj(customer.__dict__),
68
+ success=True,
69
+ message="Customer fetched successfully",
70
+ )
71
+
72
+
73
+@customers_router.put("/{customer_code}", response_model=CustomerResponse, description="Update customer by customer_code")
74
+async def update_customer(customer_code: str, customer: CustomerRequestBody) -> CustomerResponse:
75
+ logger.info(f"Updating customer with customer_code: {customer_code}")
76
+ existing_customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
77
+ if not existing_customer:
78
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
79
+ existing_customer.update_from_model(customer)
80
+ session.commit()
81
+ return CustomerResponse(
82
+ customer=CustomerRequestBody.parse_obj(customer.__dict__),
83
+ success=True,
84
+ message="Customer updated successfully",
85
+ )
86
+
87
+
88
+# ! TODO: Fix delete customer
89
+# @customers_router.delete("/{customer_code}", response_model=CustomerResponse, description="Delete customer by customer_code")
90
+# async def delete_customer(customer_code: str) -> CustomerResponse:
91
+# logger.info(f"Deleting customer with customer_code: {customer_code}")
92
+# existing_customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
93
+# if not existing_customer:
94
+# raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
95
+# session.delete(existing_customer)
96
+# session.commit()
97
+# return CustomerResponse(customer=CustomerRequestBody.parse_obj(existing_customer.__dict__), success=True, message="Customer deleted successfully")
98
+
99
+
100
+@customers_router.post("/{customer_code}/meta", response_model=CustomerMetaResponse, description="Add new customer meta")
101
+async def add_customer_meta(customer_code: str, customer_meta: CustomerMetaRequestBody) -> CustomerMetaResponse:
102
+ logger.info(f"Adding new customer meta: {customer_meta}")
103
+ existing_customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
104
+ if not existing_customer:
105
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
106
+ # Get the customer_code and customer_name from the existing customer and add it to the customer_meta object
107
+ logger.info(f"Got existing customer: {existing_customer}")
108
+ new_customer_meta = CustomersMeta(**customer_meta.dict())
109
+ new_customer_meta.customer_code = existing_customer.customer_code
110
+ new_customer_meta.customer_name = existing_customer.customer_name
111
+ session.add(new_customer_meta)
112
+ session.commit()
113
+ return CustomerMetaResponse(customer_meta=customer_meta, success=True, message="Customer meta added successfully")
114
+
115
+
116
+@customers_router.get("/{customer_code}/meta", response_model=CustomerMetaResponse, description="Get customer meta by customer_code")
117
+async def get_customer_meta(customer_code: str) -> CustomerMetaResponse:
118
+ logger.info(f"Fetching customer meta with customer_code: {customer_code}")
119
+ customer_meta = session.query(CustomersMeta).filter(CustomersMeta.customer_code == customer_code).first()
120
+ if not customer_meta:
121
+ raise HTTPException(status_code=404, detail=f"Customer meta with customer_code {customer_code} not found")
122
+ return CustomerMetaResponse(
123
+ customer_meta=CustomerMetaRequestBody.parse_obj(customer_meta.__dict__),
124
+ success=True,
125
+ message="Customer meta fetched successfully",
126
+ )
127
+
128
+
129
+@customers_router.put("/{customer_code}/meta", response_model=CustomerMetaResponse, description="Update customer meta by customer_code")
130
+async def update_customer_meta(customer_code: str, customer_meta: CustomerMetaRequestBody) -> CustomerMetaResponse:
131
+ logger.info(f"Updating customer meta with customer_code: {customer_code}")
132
+ existing_customer_meta = session.query(CustomersMeta).filter(CustomersMeta.customer_code == customer_code).first()
133
+ if not existing_customer_meta:
134
+ raise HTTPException(status_code=404, detail=f"Customer meta with customer_code {customer_code} not found")
135
+
136
+ # Update the existing record with new values
137
+ existing_customer_meta.update_from_model(customer_meta)
138
+ logger.info(f"Updated existing customer meta: {existing_customer_meta}")
139
+
140
+ # Commit the changes to the database
141
+ session.commit()
142
+
143
+ return CustomerMetaResponse(
144
+ customer_meta=CustomerMetaRequestBody.parse_obj(customer_meta.__dict__),
145
+ success=True,
146
+ message="Customer meta updated successfully",
147
+ )
148
+
149
+
150
+# ! TODO: Fix delete customer meta
151
+# @customers_router.delete("/{customer_code}/meta", response_model=CustomerMetaResponse, description="Delete customer meta by customer_code")
152
+# async def delete_customer_meta(customer_code: str) -> CustomerMetaResponse:
153
+# logger.info(f"Deleting customer meta with customer_code: {customer_code}")
154
+# existing_customer_meta = session.query(CustomersMeta).filter(CustomersMeta.customer_code == customer_code).first()
155
+# if not existing_customer_meta:
156
+# raise HTTPException(status_code=404, detail=f"Customer meta with customer_code {customer_code} not found")
157
+# session.delete(existing_customer_meta)
158
+# session.commit()
159
+# return CustomerMetaResponse(customer_meta=CustomerMetaRequestBody.parse_obj(existing_customer_meta.__dict__), success=True, message="Customer meta deleted successfully")
160
+
161
+
162
+@customers_router.get(
163
+ "/{customer_code}/full",
164
+ response_model=CustomerFullResponse,
165
+ description="Get customer and customer meta by customer_code",
166
+)
167
+async def get_customer_full(customer_code: str) -> CustomerFullResponse:
168
+ logger.info(f"Fetching customer and customer meta with customer_code: {customer_code}")
169
+ customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
170
+ if not customer:
171
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
172
+ customer_meta = session.query(CustomersMeta).filter(CustomersMeta.customer_code == customer_code).first()
173
+ if not customer_meta:
174
+ raise HTTPException(status_code=404, detail=f"Customer meta with customer_code {customer_code} not found")
175
+ return CustomerFullResponse(
176
+ customer=CustomerRequestBody.parse_obj(customer.__dict__),
177
+ customer_meta=CustomerMetaRequestBody.parse_obj(customer_meta.__dict__),
178
+ success=True,
179
+ message="Customer and customer meta fetched successfully",
180
+ )
181
+
182
+
183
+# Get Agents for the given customer_code
184
+@customers_router.get("/{customer_code}/agents", response_model=AgentsResponse, description="Get agents for the given customer_code")
185
+async def get_agents(customer_code: str) -> AgentsResponse:
186
+ logger.info(f"Fetching agents for customer_code: {customer_code}")
187
+ customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
188
+ if not customer:
189
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
190
+ agents = session.query(Agents).filter(Agents.customer_code == customer_code).all()
191
+ # Explode the agents list into a list of Agent objects
192
+ agents = [AgentModel.parse_obj(agent.__dict__) for agent in agents]
193
+ return AgentsResponse(agents=agents, success=True, message="Agents fetched successfully")
194
+
195
+
196
+# Retrieve the agents for the given customer_code then perform a healthcheck on them
197
+@customers_router.get(
198
+ "/{customer_code}/agents/healthcheck/wazuh",
199
+ response_model=AgentHealthCheckResponse,
200
+ description="Get agents healthcheck for the given customer_code",
201
+)
202
+async def get_wazuh_agents_healthcheck(
203
+ customer_code: str,
204
+ minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),
205
+ hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),
206
+ days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),
207
+) -> AgentHealthCheckResponse:
208
+ logger.info(f"Fetching agents for customer_code: {customer_code}")
209
+ customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
210
+ if not customer:
211
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
212
+ agents = session.query(Agents).filter(Agents.customer_code == customer_code).all()
213
+ # Explode the agents list into a list of Agent objects
214
+ agents = [AgentModel.parse_obj(agent.__dict__) for agent in agents]
215
+ time_criteria = TimeCriteriaModel(minutes=minutes, hours=hours, days=days)
216
+ return wazuh_agents_healthcheck(agents, time_criteria)
217
+
218
+
219
+# Retrieve the agents for the given customer_code then perform a healthcheck on them
220
+@customers_router.get(
221
+ "/{customer_code}/agents/healthcheck/velociraptor",
222
+ response_model=AgentHealthCheckResponse,
223
+ description="Get agents healthcheck for the given customer_code",
224
+)
225
+async def get_velociraptor_agents_healthcheck(
226
+ customer_code: str,
227
+ minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),
228
+ hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),
229
+ days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),
230
+) -> AgentHealthCheckResponse:
231
+ logger.info(f"Fetching agents for customer_code: {customer_code}")
232
+ customer = session.query(Customers).filter(Customers.customer_code == customer_code).first()
233
+ if not customer:
234
+ raise HTTPException(status_code=404, detail=f"Customer with customer_code {customer_code} not found")
235
+ agents = session.query(Agents).filter(Agents.customer_code == customer_code).all()
236
+ # Explode the agents list into a list of Agent objects
237
+ agents = [AgentModel.parse_obj(agent.__dict__) for agent in agents]
238
+ time_criteria = TimeCriteriaModel(minutes=minutes, hours=hours, days=days)
239
+ return velociraptor_agents_healthcheck(agents, time_criteria)
backend/app/customers/schema/customers.py
new
+121
@@ -0,0 +1,121 @@
1
+from datetime import datetime
2
+from typing import List
3
+from typing import Optional
4
+
5
+from pydantic import BaseModel
6
+from pydantic import Field
7
+
8
+
9
+class CustomerRequestBody(BaseModel):
10
+ customer_code: str = Field(..., description="Unique code for the customer")
11
+ customer_name: str = Field(..., description="Name of the customer")
12
+ contact_last_name: str = Field(..., description="Last name of the contact person")
13
+ contact_first_name: str = Field(..., description="First name of the contact person")
14
+
15
+ parent_customer_code: Optional[str] = Field(None, description="Code for the parent customer")
16
+ phone: Optional[str] = Field(None, description="Phone number")
17
+ address_line1: Optional[str] = Field(None, description="First line of the address")
18
+ address_line2: Optional[str] = Field(None, description="Second line of the address")
19
+ city: Optional[str] = Field(None, description="City")
20
+ state: Optional[str] = Field(None, description="State")
21
+ postal_code: Optional[str] = Field(None, description="Postal Code")
22
+ country: Optional[str] = Field(None, description="Country")
23
+ customer_type: Optional[str] = Field(None, description="Type of the customer")
24
+ logo_file: Optional[str] = Field(None, description="Logo file for the customer")
25
+
26
+ class Config:
27
+ schema_extra = {
28
+ "example": {
29
+ "customer_code": "CUST123",
30
+ "customer_name": "Sample Customer",
31
+ "contact_last_name": "Doe",
32
+ "contact_first_name": "John",
33
+ "phone": "123-456-7890",
34
+ "address_line1": "123 Main St",
35
+ "address_line2": "Apt 4",
36
+ "city": "Anytown",
37
+ "state": "CA",
38
+ "postal_code": "12345",
39
+ "country": "USA",
40
+ "customer_type": "Enterprise",
41
+ "logo_file": "logo.png",
42
+ },
43
+ }
44
+
45
+
46
+class CustomerResponse(BaseModel):
47
+ customer: Optional[CustomerRequestBody]
48
+ success: bool
49
+ message: str
50
+
51
+
52
+class CustomersResponse(BaseModel):
53
+ customers: list[CustomerRequestBody]
54
+ success: bool
55
+ message: str
56
+
57
+
58
+############# Customer Meta
59
+class CustomerMetaRequestBody(BaseModel):
60
+ customer_meta_graylog_index: str = Field(..., description="Graylog index for the customer")
61
+ customer_meta_graylog_stream: str = Field(..., description="Graylog stream for the customer")
62
+ customer_meta_influx_org: str = Field(..., description="InfluxDB organization for the customer")
63
+ customer_meta_grafana_org: str = Field(..., description="Grafana organization for the customer")
64
+ customer_meta_wazuh_group: str = Field(..., description="Wazuh group for the customer")
65
+ index_retention: int = Field(..., description="Index retention for the customer")
66
+ wazuh_registration_port: int = Field(..., description="Wazuh registration port for the customer")
67
+ wazuh_log_ingestion_port: int = Field(..., description="Wazuh log ingestion port for the customer")
68
+
69
+ class Config:
70
+ schema_extra = {
71
+ "example": {
72
+ "customer_meta_graylog_index": "graylog_index",
73
+ "customer_meta_graylog_stream": "graylog_stream",
74
+ "customer_meta_influx_org": "influx_org",
75
+ "customer_meta_grafana_org": "grafana_org",
76
+ "customer_meta_wazuh_group": "wazuh_group",
77
+ "index_retention": 30,
78
+ "wazuh_registration_port": 1514,
79
+ "wazuh_log_ingestion_port": 1515,
80
+ },
81
+ }
82
+
83
+
84
+class CustomerMetaResponse(BaseModel):
85
+ customer_meta: Optional[CustomerMetaRequestBody]
86
+ success: bool
87
+ message: str
88
+
89
+
90
+############# Customer Full Response
91
+class CustomerFullResponse(BaseModel):
92
+ customer: Optional[CustomerRequestBody]
93
+ customer_meta: Optional[CustomerMetaRequestBody]
94
+ success: bool
95
+ message: str
96
+
97
+
98
+############# Agent Model #############
99
+class AgentModel(BaseModel):
100
+ id: Optional[int]
101
+ os: Optional[str]
102
+ label: Optional[str]
103
+ wazuh_last_seen: Optional[datetime]
104
+ velociraptor_last_seen: Optional[datetime]
105
+ velociraptor_agent_version: Optional[str]
106
+ ip_address: Optional[str]
107
+ agent_id: Optional[str]
108
+ hostname: Optional[str]
109
+ critical_asset: Optional[bool]
110
+ velociraptor_id: Optional[str]
111
+ wazuh_agent_version: Optional[str]
112
+ customer_code: Optional[str]
113
+
114
+ class Config:
115
+ orm_mode = True
116
+
117
+
118
+class AgentsResponse(BaseModel):
119
+ agents: Optional[List[AgentModel]] = Field([], description="List of agents")
120
+ success: bool
121
+ message: str
backend/app/db/all_models.py
new
+8
@@ -0,0 +1,8 @@
1
+# all_models.py
2
+from app.auth.models.users import User
3
+from app.connectors.models import Connectors
4
+from app.connectors.sublime.models.alerts import SublimeAlerts
5
+from app.connectors.wazuh_manager.models.rules import DisabledRule
6
+from app.db.universal_models import Agents
7
+from app.db.universal_models import Customers
8
+from app.db.universal_models import CustomersMeta
backend/app/db/db_populate.py
new
+171
@@ -0,0 +1,171 @@
1
+from loguru import logger
2
+from sqlmodel import Session
3
+
4
+from app.auth.models.users import Role
5
+from app.connectors.models import Connectors
6
+
7
+
8
+def add_connectors_if_not_exist(session: Session):
9
+ # List of connectors to add
10
+ connector_list = [
11
+ {
12
+ "connector_name": "Wazuh-Indexer",
13
+ "connector_type": "4.4.1",
14
+ "connector_url": "https://ashwix01.socfortress.local:9200",
15
+ "connector_username": "admin",
16
+ "connector_password": "hmx7KPy15XPhJkgjlFrVgrWZ+Aid6QNm",
17
+ "connector_api_key": None,
18
+ "connector_configured": True,
19
+ "connector_accepts_username_password": True,
20
+ },
21
+ {
22
+ "connector_name": "Wazuh-Manager",
23
+ "connector_type": "4.4.1",
24
+ "connector_url": "https://ashwzhma.socfortress.local:55000",
25
+ "connector_username": "wazuh-wui",
26
+ "connector_password": "wazuh-wui",
27
+ "connector_api_key": None,
28
+ "connector_configured": True,
29
+ "connector_accepts_username_password": True,
30
+ },
31
+ {
32
+ "connector_name": "Graylog",
33
+ "connector_type": "5.0.7",
34
+ "connector_url": "http://ashgrl02.socfortress.local:9000",
35
+ "connector_username": "socfortress_graylog_manager",
36
+ "connector_password": "R{2PvE5TQkU7[xS$pX>fw>`y",
37
+ "connector_api_key": None,
38
+ "connector_configured": True,
39
+ "connector_accepts_username_password": True,
40
+ },
41
+ {
42
+ "connector_name": "Shuffle",
43
+ "connector_type": "1.1.0",
44
+ "connector_url": "https://ASHDKR02.socfortress.local:3443",
45
+ "connector_username": "sting",
46
+ "connector_password": "string",
47
+ "connector_api_key": "bc5d1e18-6230-40f0-b032-6ed898c307c5",
48
+ "connector_configured": True,
49
+ "connector_accepts_api_key": True,
50
+ },
51
+ {
52
+ "connector_name": "DFIR-IRIS",
53
+ "connector_type": "2.0",
54
+ "connector_url": "https://ashirs01.socfortress.local",
55
+ "connector_username": None,
56
+ "connector_password": None,
57
+ "connector_api_key": "I3Hwvkpvdk8Z0XRFlyGm4WXGw8jksnEzvKNoD9BobtSQ2AgWmdo_p-pfmJCg_ev2cm8I-zgWzAfya3jLBWZ6qw",
58
+ "connector_configured": True,
59
+ "connector_accepts_api_key": True,
60
+ },
61
+ {
62
+ "connector_name": "Velociraptor",
63
+ "connector_type": "0.6.8",
64
+ "connector_url": "https://ashvlo01.socfortress.local:8001",
65
+ "connector_username": None,
66
+ "connector_password": None,
67
+ "connector_api_key": "C:\\Users\\walto\\Desktop\\GitHub\\CoPilot\\backend\\file-store\\api.config.yaml",
68
+ "connector_configured": True,
69
+ "connector_accepts_file": True,
70
+ },
71
+ {
72
+ "connector_name": "RabbitMQ",
73
+ "connector_type": "3",
74
+ "connector_url": "ashdkr02.socfortress.local:5672",
75
+ "connector_username": "guest",
76
+ "connector_password": "guest",
77
+ "connector_api_key": None,
78
+ "connector_configured": True,
79
+ "connector_accepts_username_password": True,
80
+ },
81
+ {
82
+ "connector_name": "Sublime",
83
+ "connector_type": "3",
84
+ "connector_url": "http://ashdkr02.socfortress.local:8000",
85
+ "connector_username": None,
86
+ "connector_password": None,
87
+ "connector_api_key": "7653trxhakxn4wxdh8bbatbvu97hm8fopos7wztzjrwfd12gf5i2kyebhvke9rt4",
88
+ "connector_configured": True,
89
+ "connector_accepts_api_key": True,
90
+ },
91
+ {
92
+ "connector_name": "InfluxDB",
93
+ "connector_type": "3",
94
+ "connector_url": "http://ashdkr02.socfortress.local:8086",
95
+ "connector_username": "SOCFortress",
96
+ "connector_password": None,
97
+ "connector_api_key": "gOLoFKucQXXd5d1rDx59YYktIz6OfrHIe4jRowJKZ8iB4IcZES8rOhRPaDEejEkahch8Ze2FiMzZxbQ9ZV8K6g==",
98
+ "connector_configured": True,
99
+ "connector_accepts_api_key": True,
100
+ },
101
+ {
102
+ "connector_name": "AskSocfortress",
103
+ "connector_type": "3",
104
+ "connector_url": "https://api.socfortress.co/rule",
105
+ "connector_username": None,
106
+ "connector_password": None,
107
+ "connector_api_key": "CkKmw1B9NM1hG669tC4sTazLm1HlRfSXVvMZkxa9",
108
+ "connector_configured": True,
109
+ "connector_accepts_api_key": True,
110
+ },
111
+ {
112
+ "connector_name": "SocfortressThreatIntel",
113
+ "connector_type": "3",
114
+ "connector_url": "https://intel.socfortress.co/search",
115
+ "connector_username": None,
116
+ "connector_password": None,
117
+ "connector_api_key": "ozH1jHp1zmacCePYrAZmxarJCGptcMth93a86Jq8",
118
+ "connector_configured": True,
119
+ "connector_accepts_api_key": True,
120
+ },
121
+ {
122
+ "connector_name": "Cortex",
123
+ "connector_type": "3",
124
+ "connector_url": "http://ashvlo01.socfortress.local:9001",
125
+ "connector_username": None,
126
+ "connector_password": None,
127
+ "connector_api_key": "+k/DvVYMEYURbc8sUdXA5/hW9VhJZV3v",
128
+ "connector_configured": True,
129
+ "connector_accepts_api_key": True,
130
+ },
131
+ ]
132
+
133
+ for connector_data in connector_list:
134
+ # Check if connector already exists in the database
135
+ existing_connector = session.query(Connectors).filter_by(connector_name=connector_data["connector_name"]).first()
136
+
137
+ if existing_connector is None:
138
+ # If connector does not exist, create new connector entry
139
+ new_connector = Connectors(**connector_data)
140
+ session.add(new_connector)
141
+ logger.info(f"Added new connector: {connector_data['connector_name']}")
142
+
143
+ # Commit the changes if any new connectors were added
144
+ session.commit()
145
+
146
+
147
+def add_roles_if_not_exist(session: Session):
148
+ # List of roles to add
149
+ role_list = [
150
+ {
151
+ "name": "admin",
152
+ "description": "Administrator",
153
+ },
154
+ {
155
+ "name": "analyst",
156
+ "description": "SOC Analyst",
157
+ },
158
+ ]
159
+
160
+ for role_data in role_list:
161
+ # Check if role already exists in the database
162
+ existing_role = session.query(Role).filter_by(name=role_data["name"]).first()
163
+
164
+ if existing_role is None:
165
+ # If role does not exist, create new role entry
166
+ new_role = Role(**role_data)
167
+ session.add(new_role)
168
+ logger.info(f"Added new role: {role_data['name']}")
169
+
170
+ # Commit the changes if any new roles were added
171
+ session.commit()
backend/app/db/db_session.py
new
+7
@@ -0,0 +1,7 @@
1
+from sqlmodel import Session
2
+from sqlmodel import create_engine
3
+
4
+from settings import SQLALCHEMY_DATABASE_URI
5
+
6
+engine = create_engine(SQLALCHEMY_DATABASE_URI)
7
+session = Session(bind=engine)
backend/app/db/db_setup.py
new
+31
@@ -0,0 +1,31 @@
1
+from loguru import logger
2
+from sqlalchemy import inspect
3
+from sqlmodel import Session
4
+from sqlmodel import SQLModel
5
+
6
+# from app.db.all_models import *
7
+from app.db.db_populate import add_connectors_if_not_exist
8
+from app.db.db_populate import add_roles_if_not_exist
9
+
10
+
11
+def create_tables(engine):
12
+ logger.info("Creating tables")
13
+
14
+ # Create an inspector object based on the engine
15
+ inspector = inspect(engine)
16
+
17
+ # Get the names of all tables in the database
18
+ existing_tables = inspector.get_table_names()
19
+
20
+ # Loop through all your models (tables)
21
+ for table in SQLModel.metadata.sorted_tables:
22
+ if table.name not in existing_tables:
23
+ # Only create the table if it doesn't exist
24
+ table.create(bind=engine)
25
+ logger.info(f"Table {table.name} created.")
26
+
27
+ # After creating all tables, add connectors if they don't exist
28
+ with Session(engine) as session:
29
+ add_connectors_if_not_exist(session)
30
+ add_roles_if_not_exist(session)
31
+ session.commit()
backend/app/db/universal_models.py
new
+122
@@ -0,0 +1,122 @@
1
+from datetime import datetime
2
+from typing import Optional
3
+
4
+from sqlmodel import Field
5
+from sqlmodel import Relationship
6
+from sqlmodel import SQLModel
7
+
8
+
9
+class Customers(SQLModel, table=True):
10
+ id: Optional[int] = Field(primary_key=True)
11
+ customer_code: str = Field(max_length=11, nullable=False)
12
+ parent_customer_code: Optional[str] = Field(max_length=11)
13
+ customer_name: str = Field(max_length=50, nullable=False)
14
+ contact_last_name: Optional[str] = Field(max_length=50)
15
+ contact_first_name: Optional[str] = Field(max_length=50)
16
+ phone: Optional[str] = Field(max_length=50)
17
+ address_line1: Optional[str] = Field(max_length=1024)
18
+ address_line2: Optional[str] = Field(max_length=1024)
19
+ city: Optional[str] = Field(max_length=50)
20
+ state: Optional[str] = Field(max_length=50)
21
+ postal_code: Optional[str] = Field(max_length=15)
22
+ country: Optional[str] = Field(max_length=50)
23
+ customer_type: Optional[str] = Field(max_length=50)
24
+ logo_file: Optional[str] = Field(max_length=64)
25
+ created_at: datetime = Field(default=datetime.utcnow())
26
+
27
+ agents: list["Agents"] = Relationship(back_populates="customer")
28
+ meta: Optional["CustomersMeta"] = Relationship(back_populates="customer")
29
+
30
+ def update_from_model(self, customer):
31
+ self.customer_code = customer.customer_code
32
+ self.parent_customer_code = customer.parent_customer_code
33
+ self.customer_name = customer.customer_name
34
+ self.contact_last_name = customer.contact_last_name
35
+ self.contact_first_name = customer.contact_first_name
36
+ self.phone = customer.phone
37
+ self.address_line1 = customer.address_line1
38
+ self.address_line2 = customer.address_line2
39
+ self.city = customer.city
40
+ self.state = customer.state
41
+ self.postal_code = customer.postal_code
42
+ self.country = customer.country
43
+ self.customer_type = customer.customer_type
44
+ self.logo_file = customer.logo_file
45
+
46
+
47
+class CustomersMeta(SQLModel, table=True):
48
+ id: Optional[int] = Field(primary_key=True)
49
+ customer_code: str = Field(foreign_key="customers.customer_code", nullable=False)
50
+ customer_name: str = Field(max_length=255)
51
+ customer_meta_graylog_index: str = Field(max_length=1024)
52
+ customer_meta_graylog_stream: str = Field(max_length=1024)
53
+ customer_meta_influx_org: str = Field(max_length=1024)
54
+ customer_meta_grafana_org: str = Field(max_length=1024)
55
+ customer_meta_wazuh_group: str = Field(max_length=1024)
56
+ index_retention: Optional[int] = Field()
57
+ wazuh_registration_port: Optional[int] = Field()
58
+ wazuh_log_ingestion_port: Optional[int] = Field()
59
+
60
+ # Link back to Customers
61
+ customer: Optional["Customers"] = Relationship(back_populates="meta")
62
+
63
+ def update_from_model(self, customer_meta):
64
+ if hasattr(customer_meta, "customer_code"):
65
+ self.customer_code = customer_meta.customer_code
66
+ if hasattr(customer_meta, "customer_name"):
67
+ self.customer_name = customer_meta.customer_name
68
+ self.customer_meta_graylog_index = customer_meta.customer_meta_graylog_index
69
+ self.customer_meta_graylog_stream = customer_meta.customer_meta_graylog_stream
70
+ self.customer_meta_influx_org = customer_meta.customer_meta_influx_org
71
+ self.customer_meta_grafana_org = customer_meta.customer_meta_grafana_org
72
+ self.customer_meta_wazuh_group = customer_meta.customer_meta_wazuh_group
73
+ self.index_retention = customer_meta.index_retention
74
+ self.wazuh_registration_port = customer_meta.wazuh_registration_port
75
+ self.wazuh_log_ingestion_port = customer_meta.wazuh_log_ingestion_port
76
+
77
+
78
+class Agents(SQLModel, table=True):
79
+ id: Optional[int] = Field(primary_key=True)
80
+ agent_id: str = Field(index=True)
81
+ ip_address: str = Field(max_length=256)
82
+ os: str = Field(max_length=256)
83
+ hostname: str = Field(max_length=256)
84
+ label: str = Field(max_length=256)
85
+ critical_asset: bool = Field(default=False)
86
+ wazuh_last_seen: datetime
87
+ velociraptor_id: str = Field(max_length=256)
88
+ velociraptor_last_seen: datetime
89
+ wazuh_agent_version: str = Field(max_length=256)
90
+ velociraptor_agent_version: str = Field(max_length=256)
91
+ customer_code: Optional[str] = Field(foreign_key="customers.customer_code")
92
+
93
+ customer: Optional[Customers] = Relationship(back_populates="agents")
94
+
95
+ @classmethod
96
+ def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code):
97
+ return cls(
98
+ agent_id=wazuh_agent.agent_id,
99
+ hostname=wazuh_agent.agent_name,
100
+ ip_address=wazuh_agent.agent_ip,
101
+ os=wazuh_agent.agent_os,
102
+ label=wazuh_agent.agent_label,
103
+ wazuh_last_seen=wazuh_agent.agent_last_seen_as_datetime,
104
+ wazuh_agent_version=wazuh_agent.wazuh_agent_version,
105
+ velociraptor_id=velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a",
106
+ velociraptor_last_seen=velociraptor_agent.client_last_seen_as_datetime,
107
+ velociraptor_agent_version=velociraptor_agent.client_version,
108
+ customer_code=customer_code,
109
+ )
110
+
111
+ def update_from_model(self, wazuh_agent, velociraptor_agent, customer_code):
112
+ self.agent_id = wazuh_agent.agent_id
113
+ self.hostname = wazuh_agent.agent_name
114
+ self.ip_address = wazuh_agent.agent_ip
115
+ self.os = wazuh_agent.agent_os
116
+ self.label = wazuh_agent.agent_label
117
+ self.wazuh_last_seen = wazuh_agent.agent_last_seen_as_datetime
118
+ self.wazuh_agent_version = wazuh_agent.wazuh_agent_version
119
+ self.velociraptor_id = velociraptor_agent.client_id if velociraptor_agent.client_id else "n/a"
120
+ self.velociraptor_last_seen = velociraptor_agent.client_last_seen_as_datetime
121
+ self.velociraptor_agent_version = velociraptor_agent.client_version
122
+ self.customer_code = customer_code
backend/app/healthchecks/agents/routes/agents.py
new
+94
@@ -0,0 +1,94 @@
1
+from fastapi import APIRouter
2
+from fastapi import HTTPException
3
+from fastapi import Query
4
+from loguru import logger
5
+from starlette.status import HTTP_401_UNAUTHORIZED
6
+
7
+from app.db.db_session import session
8
+from app.db.universal_models import Agents
9
+from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse
10
+from app.healthchecks.agents.schema.agents import HostLogsSearchBody
11
+from app.healthchecks.agents.schema.agents import HostLogsSearchResponse
12
+from app.healthchecks.agents.schema.agents import TimeCriteriaModel
13
+from app.healthchecks.agents.services.agents import host_logs
14
+from app.healthchecks.agents.services.agents import velociraptor_agent_healthcheck
15
+from app.healthchecks.agents.services.agents import velociraptor_agents_healthcheck
16
+from app.healthchecks.agents.services.agents import wazuh_agent_healthcheck
17
+from app.healthchecks.agents.services.agents import wazuh_agents_healthcheck
18
+
19
+healtcheck_agents_router = APIRouter()
20
+
21
+
22
+def verify_admin(user):
23
+ if not user.is_admin:
24
+ raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Unauthorized")
25
+
26
+
27
+@healtcheck_agents_router.get("/wazuh", response_model=AgentHealthCheckResponse, description="Get Wazuh agents healthcheck")
28
+async def get_wazuh_agent_healthcheck(
29
+ minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),
30
+ hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),
31
+ days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),
32
+) -> AgentHealthCheckResponse:
33
+ time_criteria = TimeCriteriaModel(minutes=minutes, hours=hours, days=days)
34
+ agents = session.query(Agents).all()
35
+ return wazuh_agents_healthcheck(agents, time_criteria)
36
+
37
+
38
+# Get single agent by agent_id
39
+@healtcheck_agents_router.get(
40
+ "/wazuh/{agent_id}",
41
+ response_model=AgentHealthCheckResponse,
42
+ description="Get Wazuh agent healthcheck by agent_id",
43
+)
44
+async def get_wazuh_agent_healthcheck_by_agent_id(
45
+ agent_id: str,
46
+ minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),
47
+ hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),
48
+ days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),
49
+) -> AgentHealthCheckResponse:
50
+ time_criteria = TimeCriteriaModel(minutes=minutes, hours=hours, days=days)
51
+ agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()
52
+ if not agent:
53
+ raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
54
+ return wazuh_agent_healthcheck(agent, time_criteria)
55
+
56
+
57
+@healtcheck_agents_router.get("/velociraptor", response_model=AgentHealthCheckResponse, description="Get Velociraptor agents healthcheck")
58
+async def get_velociraptor_agent_healthcheck(
59
+ minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),
60
+ hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),
61
+ days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),
62
+) -> AgentHealthCheckResponse:
63
+ time_criteria = TimeCriteriaModel(minutes=minutes, hours=hours, days=days)
64
+ agents = session.query(Agents).all()
65
+ return velociraptor_agents_healthcheck(agents, time_criteria)
66
+
67
+
68
+# Get single agent by agent_id
69
+@healtcheck_agents_router.get(
70
+ "/velociraptor/{agent_id}",
71
+ response_model=AgentHealthCheckResponse,
72
+ description="Get Velociraptor agent healthcheck by agent_id",
73
+)
74
+async def get_velociraptor_agent_healthcheck_by_agent_id(
75
+ agent_id: str,
76
+ minutes: int = Query(60, description="Number of minutes within which the agent should have been last seen to be considered healthy."),
77
+ hours: int = Query(0, description="Number of hours within which the agent should have been last seen to be considered healthy."),
78
+ days: int = Query(0, description="Number of days within which the agent should have been last seen to be considered healthy."),
79
+) -> AgentHealthCheckResponse:
80
+ time_criteria = TimeCriteriaModel(minutes=minutes, hours=hours, days=days)
81
+ agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()
82
+ if not agent:
83
+ raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")
84
+ return velociraptor_agent_healthcheck(agent, time_criteria)
85
+
86
+
87
+@healtcheck_agents_router.post("/logs", response_model=HostLogsSearchResponse, description="Get host logs")
88
+async def get_host_logs(body: HostLogsSearchBody) -> HostLogsSearchResponse:
89
+ logger.info(f"Received request to get host logs for {body.agent_name}")
90
+ # Verify the agent exists
91
+ agent = session.query(Agents).filter(Agents.hostname == body.agent_name).first()
92
+ if not agent:
93
+ raise HTTPException(status_code=404, detail=f"Agent with hostname {body.agent_name} not found")
94
+ return host_logs(body)
backend/app/healthchecks/agents/schema/agents.py
new
+102
@@ -0,0 +1,102 @@
1
+from datetime import datetime
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):
13
+ id: Optional[int]
14
+ os: Optional[str]
15
+ label: Optional[str]
16
+ wazuh_last_seen: Optional[datetime]
17
+ velociraptor_last_seen: Optional[datetime]
18
+ velociraptor_agent_version: Optional[str]
19
+ ip_address: Optional[str]
20
+ agent_id: Optional[str]
21
+ hostname: Optional[str]
22
+ critical_asset: Optional[bool]
23
+ velociraptor_id: Optional[str]
24
+ wazuh_agent_version: Optional[str]
25
+ customer_code: Optional[str]
26
+
27
+ class Config:
28
+ orm_mode = True
29
+
30
+
31
+class ExtendedAgentModel(AgentModel):
32
+ unhealthy_wazuh_agent: Optional[bool] = Field(None, description="Whether the agent is unhealthy in Wazuh")
33
+ unhealthy_velociraptor_agent: Optional[bool] = Field(None, description="Whether the agent is unhealthy in Velociraptor")
34
+ unhealthy_recent_logs_collected: Optional[bool] = Field(None, description="Whether the agent has not collected logs recently")
35
+
36
+
37
+class AgentHealthCheckResponse(BaseModel):
38
+ healthy_wazuh_agents: Optional[List[ExtendedAgentModel]]
39
+ unhealthy_wazuh_agents: Optional[List[ExtendedAgentModel]]
40
+ healthy_velociraptor_agents: Optional[List[ExtendedAgentModel]]
41
+ unhealthy_velociraptor_agents: Optional[List[ExtendedAgentModel]]
42
+ healthy_recent_logs_collected: Optional[List[ExtendedAgentModel]]
43
+ unhealthy_recent_logs_collected: Optional[List[ExtendedAgentModel]]
44
+ message: str
45
+ success: bool
46
+
47
+
48
+class TimeCriteriaModel(BaseModel):
49
+ minutes: int = Field(60, description="Number of minutes within which the agent should have been last seen to be considered healthy.")
50
+ hours: int = Field(0, description="Number of hours within which the agent should have been last seen to be considered healthy.")
51
+ days: int = Field(0, description="Number of days within which the agent should have been last seen to be considered healthy.")
52
+
53
+
54
+########## Logs Schemas ##########
55
+
56
+
57
+class Log(BaseModel):
58
+ index_name: str
59
+ total_logs: int
60
+ logs: Optional[List[Dict[str, Any]]] = Field([], description="The logs returned from the search.")
61
+
62
+
63
+class LogsSearchBody(BaseModel):
64
+ size: int = Field(1, description="The number of logs to return.")
65
+ timerange: str = Field("24h", description="The time range to search logs in.")
66
+ log_field: str = Field("syslog_level", description="The field to search logs in.")
67
+ log_value: str = Field("INFO", description="The value to search logs for.")
68
+ timestamp_field: str = Field("timestamp_utc", description="The timestamp field to search logs in.")
69
+
70
+ @validator("timerange")
71
+ def validate_timerange(cls, value):
72
+ if value[-1] not in ("h", "d", "w", "m"):
73
+ raise ValueError("Invalid timerange format. The string should end with either 'h', 'd', 'w', or 'm'.")
74
+
75
+ # Optionally, you can check that the prefix is a number
76
+ if not value[:-1].isdigit():
77
+ raise ValueError("Invalid timerange format. The string should start with a number.")
78
+
79
+ return value
80
+
81
+
82
+class LogsSearchResponse(BaseModel):
83
+ logs_summary: List[Log]
84
+ success: bool
85
+ message: str
86
+
87
+
88
+class CollectLogsResponse(BaseModel):
89
+ logs: List[Dict[str, Any]]
90
+ success: bool
91
+ message: str
92
+
93
+
94
+class HostLogsSearchBody(LogsSearchBody):
95
+ agent_name: str = Field(..., description="The name of the agent to search logs for.")
96
+
97
+
98
+class HostLogsSearchResponse(BaseModel):
99
+ logs_summary: Optional[List[Log]] = Field([], description="The logs summary returned from the search.")
100
+ healthy: bool = Field(False, description="Whether the host is healthy or not.")
101
+ success: bool
102
+ message: str
backend/app/healthchecks/agents/services/agents.py
new
+211
@@ -0,0 +1,211 @@
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(agent: AgentModel, time_criteria: TimeCriteriaModel) -> ExtendedAgentModel:
23
+ current_time = datetime.now()
24
+ wazuh_last_seen = agent.wazuh_last_seen
25
+
26
+ if wazuh_last_seen > current_time:
27
+ logger.info(f"Agent {agent} has a wazuh_last_seen time in the future: {wazuh_last_seen}")
28
+ return ExtendedAgentModel(**agent.dict(), unhealthy_wazuh_agent=True)
29
+
30
+ # Calculate the total time delta based on the criteria
31
+ total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60
32
+ time_delta = timedelta(minutes=total_minutes)
33
+
34
+ is_unhealthy = (current_time - wazuh_last_seen) > time_delta
35
+ return ExtendedAgentModel(**agent.dict(), unhealthy_wazuh_agent=is_unhealthy)
36
+
37
+
38
+def is_velociraptor_agent_unhealthy(agent: AgentModel, time_criteria: TimeCriteriaModel) -> ExtendedAgentModel:
39
+ current_time = datetime.now()
40
+ velociraptor_last_seen = agent.velociraptor_last_seen
41
+
42
+ if velociraptor_last_seen > current_time:
43
+ logger.info(f"Agent {agent} has a velociraptor_last_seen time in the future: {velociraptor_last_seen}")
44
+ return ExtendedAgentModel(**agent.dict(), unhealthy_velociraptor_agent=True)
45
+
46
+ # Calculate the total time delta based on the criteria
47
+ total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60
48
+ time_delta = timedelta(minutes=total_minutes)
49
+
50
+ is_unhealthy = (current_time - velociraptor_last_seen) > time_delta
51
+ return ExtendedAgentModel(**agent.dict(), unhealthy_velociraptor_agent=is_unhealthy)
52
+
53
+
54
+def wazuh_agents_healthcheck(agents: list, time_criteria: TimeCriteriaModel) -> AgentHealthCheckResponse:
55
+ healthy_wazuh_agents = []
56
+ unhealthy_wazuh_agents = []
57
+ for agent in agents:
58
+ # If agent_id is `000` skip it because this is the Wazuh manager
59
+ if agent.agent_id == "000":
60
+ continue
61
+ logger.info(f"Checking agent {agent} for health")
62
+ extended_agent = is_wazuh_agent_unhealthy(agent, time_criteria)
63
+ logger.info(f"Extended agent: {extended_agent}")
64
+ if extended_agent.unhealthy_wazuh_agent:
65
+ unhealthy_wazuh_agents.append(extended_agent)
66
+ else:
67
+ healthy_wazuh_agents.append(extended_agent)
68
+
69
+ return AgentHealthCheckResponse(
70
+ healthy_wazuh_agents=healthy_wazuh_agents,
71
+ unhealthy_wazuh_agents=unhealthy_wazuh_agents,
72
+ success=True,
73
+ message="Wazuh agent healthcheck fetched successfully",
74
+ )
75
+
76
+
77
+def wazuh_agent_healthcheck(agent: AgentModel, time_criteria: TimeCriteriaModel) -> AgentHealthCheckResponse:
78
+ extended_agent = is_wazuh_agent_unhealthy(agent, time_criteria)
79
+ if extended_agent.unhealthy_wazuh_agent:
80
+ return AgentHealthCheckResponse(
81
+ healthy_wazuh_agents=[],
82
+ unhealthy_wazuh_agents=[extended_agent],
83
+ success=True,
84
+ message="Wazuh agent healthcheck fetched successfully",
85
+ )
86
+ else:
87
+ return AgentHealthCheckResponse(
88
+ healthy_wazuh_agents=[extended_agent],
89
+ unhealthy_wazuh_agents=[],
90
+ success=True,
91
+ message="Wazuh agent healthcheck fetched successfully",
92
+ )
93
+
94
+
95
+def velociraptor_agents_healthcheck(agents: list, time_criteria: TimeCriteriaModel) -> AgentHealthCheckResponse:
96
+ healthy_velociraptor_agents = []
97
+ unhealthy_velociraptor_agents = []
98
+ for agent in agents:
99
+ # If agent_id is `000` skip it because this is the Wazuh manager
100
+ if agent.agent_id == "000":
101
+ continue
102
+ logger.info(f"Checking agent {agent} for health")
103
+ extended_agent = is_velociraptor_agent_unhealthy(agent, time_criteria)
104
+ logger.info(f"Extended agent: {extended_agent}")
105
+ if extended_agent.unhealthy_velociraptor_agent:
106
+ unhealthy_velociraptor_agents.append(extended_agent)
107
+ else:
108
+ healthy_velociraptor_agents.append(extended_agent)
109
+
110
+ return AgentHealthCheckResponse(
111
+ healthy_velociraptor_agents=healthy_velociraptor_agents,
112
+ unhealthy_velociraptor_agents=unhealthy_velociraptor_agents,
113
+ success=True,
114
+ message="Velociraptor agent healthcheck fetched successfully",
115
+ )
116
+
117
+
118
+def velociraptor_agent_healthcheck(agent: AgentModel, time_criteria: TimeCriteriaModel) -> AgentHealthCheckResponse:
119
+ extended_agent = is_velociraptor_agent_unhealthy(agent, time_criteria)
120
+ if extended_agent.unhealthy_velociraptor_agent:
121
+ return AgentHealthCheckResponse(
122
+ healthy_velociraptor_agents=[],
123
+ unhealthy_velociraptor_agents=[extended_agent],
124
+ success=True,
125
+ message="Velociraptor agent healthcheck fetched successfully",
126
+ )
127
+ else:
128
+ return AgentHealthCheckResponse(
129
+ healthy_velociraptor_agents=[extended_agent],
130
+ unhealthy_velociraptor_agents=[],
131
+ success=True,
132
+ message="Velociraptor agent healthcheck fetched successfully",
133
+ )
134
+
135
+
136
+def host_logs(search_body: HostLogsSearchBody) -> HostLogsSearchResponse:
137
+ result = get_logs_generic(search_body, is_host_specific=True)
138
+ logger.info(f"Host logs search result: {result}")
139
+
140
+ # Initialize variable to keep track of total logs
141
+ total_logs = 0
142
+
143
+ # Loop through each item in logs_summary to count total logs
144
+ for log_summary in result["logs_summary"]:
145
+ total_logs += log_summary["total_logs"]
146
+
147
+ # Check if there are any logs
148
+ if total_logs > 0:
149
+ return HostLogsSearchResponse(
150
+ success=True,
151
+ healthy=True,
152
+ message=f"Host is healthy. At least one log was found within the specified time range of {search_body.timerange}",
153
+ )
154
+ else:
155
+ return HostLogsSearchResponse(
156
+ success=True,
157
+ healthy=False,
158
+ message=f"Host is unhealthy. No logs were found within the specified time range of {search_body.timerange}",
159
+ )
160
+
161
+
162
+def get_logs_generic(search_body: Type[LogsSearchBody], is_host_specific: bool = False, index_name: Optional[str] = None):
163
+ logger.info(f"Collecting Wazuh Indexer alerts for host {search_body.agent_name if is_host_specific else ''}")
164
+ logs_summary = []
165
+ indices = collect_indices()
166
+ index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices
167
+
168
+ for index_name in index_list:
169
+ try:
170
+ logs = collect_logs_generic(index_name, body=search_body, is_host_specific=is_host_specific)
171
+ if logs.success and len(logs.logs) > 0:
172
+ logs_summary.append(
173
+ {
174
+ "index_name": index_name,
175
+ "total_logs": len(logs.logs),
176
+ "logs": logs.logs,
177
+ },
178
+ )
179
+ break # Only collect logs from the first index that has logs
180
+ except HTTPException as e:
181
+ logger.warning(f"An error occurred while processing index {index_name}: {e.detail}")
182
+
183
+ if len(logs_summary) == 0:
184
+ message = "No logs found"
185
+ else:
186
+ message = f"Succesfully collected top {search_body.size} logs for each index"
187
+
188
+ return {"logs_summary": logs_summary, "success": len(logs_summary) > 0, "message": message}
189
+
190
+
191
+def collect_logs_generic(index_name: str, body: LogsSearchBody, is_host_specific: bool = False) -> CollectLogsResponse:
192
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
193
+ query_builder = LogsQueryBuilder()
194
+ query_builder.add_time_range(timerange=body.timerange, timestamp_field=body.timestamp_field)
195
+ query_builder.add_matches(matches=[(body.log_field, body.log_value)])
196
+ query_builder.add_sort(body.timestamp_field)
197
+
198
+ if is_host_specific:
199
+ query_builder.add_match_phrase(matches=[("agent_name", body.agent_name)])
200
+
201
+ query = query_builder.build()
202
+
203
+ try:
204
+ logs = es_client.search(index=index_name, body=query, size=body.size)
205
+ logger.info(f"logs collected: {logs}")
206
+ logs_list = [log for log in logs["hits"]["hits"]]
207
+ logger.info(f"logs collected: {logs_list}")
208
+ return CollectLogsResponse(logs=logs_list, success=True, message="logs collected successfully")
209
+ except Exception as e:
210
+ logger.debug(f"Failed to collect logs: {e}")
211
+ return CollectLogsResponse(logs=[], success=False, message=f"Failed to collect logs: {e}")
backend/app/integrations/alert_escalation/routes/general_alert.py
new
+14
@@ -0,0 +1,14 @@
1
+from fastapi import APIRouter
2
+from loguru import logger
3
+
4
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
5
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
6
+from app.integrations.alert_escalation.services.general_alert import create_alert
7
+
8
+integration_general_alerts_router = APIRouter()
9
+
10
+
11
+@integration_general_alerts_router.post("/create", response_model=CreateAlertResponse, description="Create an alert in IRIS")
12
+async def create_alert_route(create_alert_request: CreateAlertRequest) -> CreateAlertResponse:
13
+ logger.info(f"Creating alert {create_alert_request.alert_id} in IRIS")
14
+ return create_alert(create_alert_request)
backend/app/integrations/alert_escalation/schema/general_alert.py
new
+115
@@ -0,0 +1,115 @@
1
+from enum import Enum
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):
13
+ MISP_VALUE = "misp_value"
14
+ OPENCTI_VALUE = "opencti_value"
15
+ THREAT_INTEL_VALUE = "threat_intel_value"
16
+
17
+
18
+class CreateAlertRequest(BaseModel):
19
+ index_name: str = Field(..., description="The name of the index to search alerts for.")
20
+ alert_id: str = Field(..., description="The alert id to create.")
21
+
22
+
23
+class CreateAlertResponse(BaseModel):
24
+ success: bool
25
+ message: str
26
+ alert_id: int = Field(..., description="The alert id as created in IRIS.")
27
+
28
+
29
+class GenericSourceModel(BaseModel):
30
+ agent_name: str = Field(..., description="The name of the agent.")
31
+ agent_id: str = Field(..., description="The id of the agent.")
32
+ agent_labels_customer: str = Field(..., description="The customer of the agent.")
33
+ rule_id: str = Field(..., description="The id of the rule.")
34
+ rule_level: int = Field(..., description="The level of the rule.")
35
+ rule_description: str = Field(..., description="The description of the rule.")
36
+ timestamp: str = Field(..., description="The timestamp of the alert.")
37
+ timestamp_utc: Optional[str] = Field(..., description="The UTC timestamp of the alert.")
38
+
39
+ class Config:
40
+ extra = Extra.allow
41
+
42
+
43
+class GenericAlertModel(BaseModel):
44
+ _index: str
45
+ _id: str
46
+ _version: int
47
+ _source: GenericSourceModel # Nested model
48
+ asset_type_id: Optional[int] = Field(
49
+ None,
50
+ description="The asset type id of the alert which is needed for when we add the asset to IRIS.",
51
+ )
52
+ ioc_value: Optional[str] = Field(None, description="The IoC value of the alert which is needed for when we add the IoC to IRIS.")
53
+ ioc_type: Optional[str] = Field(None, description="The IoC type of the alert which is needed for when we add the IoC to IRIS.")
54
+
55
+ class Config:
56
+ extra = Extra.allow
57
+
58
+
59
+# Sample data from `get_single_alert_details`
60
+sample_data = {
61
+ "_index": "some_index",
62
+ "_id": "some_id",
63
+ "_version": 1,
64
+ "_source": {
65
+ "agent_name": "some_agent_name",
66
+ "agent_id": "some_agent_id",
67
+ # ... other fields
68
+ },
69
+ # ... other fields
70
+}
71
+
72
+
73
+########### Create Alerts Schemas ###########
74
+class IrisAsset(BaseModel):
75
+ asset_name: str = Field(..., description="Name of the asset", example="Server01")
76
+ asset_ip: str = Field(..., description="IP address of the asset", example="192.168.1.1")
77
+ asset_description: str = Field(..., description="Description of the asset", example="Windows Server")
78
+ asset_type_id: int = Field(..., description="Type ID of the asset", example=1)
79
+
80
+
81
+class IrisIoc(BaseModel):
82
+ ioc_value: str = Field(..., description="Value of the IoC", example="www.google.com")
83
+ ioc_description: str = Field(..., description="Description of the IoC", example="Google")
84
+ ioc_tlp_id: int = Field(1, description="TLP ID of the IoC", example=1)
85
+ ioc_type_id: int = Field(20, description="Type ID of the IoC", example=20)
86
+
87
+
88
+class IrisAlertContext(BaseModel):
89
+ alert_id: str = Field(..., description="ID of the alert", example="123")
90
+ alert_name: str = Field(..., description="Name of the alert", example="Intrusion Detected")
91
+ alert_level: int = Field(..., description="Severity level of the alert", example=3)
92
+ rule_id: str = Field(..., description="ID of the rule that triggered the alert", example="2001")
93
+ asset_name: str = Field(..., description="Name of the affected asset", example="Server01")
94
+ asset_ip: str = Field(..., description="IP address of the affected asset", example="192.168.1.1")
95
+ asset_type: int = Field(..., description="Type ID of the affected asset", example=1)
96
+ process_id: Optional[str] = Field("No process ID found", description="Process ID involved in the alert", example="4567")
97
+ rule_mitre_id: Optional[str] = Field("n/a", description="MITRE ATT&CK ID of the rule", example="T1234")
98
+ rule_mitre_tactic: Optional[str] = Field("n/a", description="MITRE ATT&CK Tactic", example="Execution")
99
+ rule_mitre_technique: Optional[str] = Field("n/a", description="MITRE ATT&CK Technique", example="Scripting")
100
+
101
+
102
+class IrisAlertPayload(BaseModel):
103
+ alert_title: str = Field(..., description="Title of the alert", example="Intrusion Detected")
104
+ alert_description: str = Field(..., description="Description of the alert", example="Intrusion Detected by Firewall")
105
+ alert_source: str = Field(..., description="Source of the alert", example="Wazuh")
106
+ assets: List[IrisAsset] = Field(..., description="List of affected assets")
107
+ alert_status_id: int = Field(..., description="Status ID of the alert", example=3)
108
+ alert_severity_id: int = Field(..., description="Severity ID of the alert", example=5)
109
+ alert_customer_id: int = Field(..., description="Customer ID related to the alert", example=1)
110
+ alert_source_content: Dict[str, Any] = Field(..., description="Original content from the alert source")
111
+ alert_context: IrisAlertContext = Field(..., description="Contextual information about the alert")
112
+ alert_iocs: Optional[List[IrisIoc]] = Field(None, description="List of IoCs related to the alert")
113
+
114
+ def to_dict(self):
115
+ return self.dict(exclude_none=True)
backend/app/integrations/alert_escalation/services/general_alert.py
new
+123
@@ -0,0 +1,123 @@
1
+from typing import Optional
2
+from typing import Set
3
+
4
+from fastapi import HTTPException
5
+from loguru import logger
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.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
10
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
11
+from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
12
+from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel
13
+from app.integrations.alert_escalation.schema.general_alert import GenericSourceModel
14
+from app.integrations.alert_escalation.schema.general_alert import IrisAlertContext
15
+from app.integrations.alert_escalation.schema.general_alert import IrisAlertPayload
16
+from app.integrations.alert_escalation.schema.general_alert import IrisAsset
17
+from app.integrations.alert_escalation.schema.general_alert import IrisIoc
18
+from app.integrations.alert_escalation.schema.general_alert import ValidIocFields
19
+from app.integrations.alert_escalation.utils.universal import get_agent_data
20
+from app.integrations.alert_escalation.utils.universal import get_asset_type_id
21
+from app.integrations.alert_escalation.utils.universal import validate_ioc_type
22
+
23
+
24
+def valid_ioc_fields() -> Set[str]:
25
+ """
26
+ Getter for the set of valid IoC fields.
27
+ Returns
28
+ -------
29
+ Set[str]
30
+ The set of valid IoC fields.
31
+ """
32
+ return {field.value for field in ValidIocFields}
33
+
34
+
35
+def get_single_alert_details(alert_details: CreateAlertRequest) -> GenericAlertModel:
36
+ logger.info(f"Fetching alert details for alert {alert_details.alert_id} in index {alert_details.index_name}")
37
+ es_client = create_wazuh_indexer_client("Wazuh-Indexer")
38
+ try:
39
+ alert = es_client.get(index=alert_details.index_name, id=alert_details.alert_id)
40
+ source_model = GenericSourceModel(**alert["_source"])
41
+ return GenericAlertModel(_source=source_model, _id=alert["_id"], _index=alert["_index"], _version=alert["_version"])
42
+ except Exception as e:
43
+ logger.debug(f"Failed to collect alert details: {e}")
44
+ raise HTTPException(status_code=400, detail=f"Failed to collect alert details: {e}")
45
+
46
+
47
+def build_ioc_payload(alert_details: GenericAlertModel) -> Optional[IrisIoc]:
48
+ for field in valid_ioc_fields():
49
+ if hasattr(alert_details._source, field):
50
+ ioc_value = getattr(alert_details._source, field)
51
+ ioc_type = validate_ioc_type(ioc_value=ioc_value)
52
+ return IrisIoc(ioc_value=ioc_value, ioc_description="IoC found in alert", ioc_tlp_id=1, ioc_type_id=ioc_type)
53
+ return None
54
+
55
+
56
+def build_asset_payload(agent_data, alert_details) -> IrisAsset:
57
+ return IrisAsset(
58
+ asset_name=agent_data.hostname,
59
+ asset_ip=agent_data.ip_address,
60
+ asset_description=agent_data.os,
61
+ asset_type_id=alert_details.asset_type_id,
62
+ )
63
+
64
+
65
+def build_alert_context_payload(alert_details: GenericAlertModel, agent_data) -> IrisAlertContext:
66
+ return IrisAlertContext(
67
+ alert_id=alert_details._id,
68
+ alert_name=alert_details._source.rule_description,
69
+ alert_level=alert_details._source.rule_level,
70
+ rule_id=alert_details._source.rule_id,
71
+ asset_name=agent_data.hostname,
72
+ asset_ip=agent_data.ip_address,
73
+ asset_type=alert_details.asset_type_id,
74
+ process_id=getattr(alert_details._source, "process_id", "No process id found"),
75
+ rule_mitre_id=getattr(alert_details._source, "rule_mitre_id", "No rule mitre id found"),
76
+ rule_mitre_tactic=getattr(alert_details._source, "rule_mitre_tactic", "No rule mitre tactic found"),
77
+ rule_mitre_technique=getattr(alert_details._source, "rule_mitre_technique", "No rule mitre technique found"),
78
+ )
79
+
80
+
81
+def build_alert_payload(alert_details: GenericAlertModel, agent_data, ioc_payload: Optional[IrisIoc]) -> IrisAlertPayload:
82
+ asset_payload = build_asset_payload(agent_data, alert_details)
83
+ context_payload = build_alert_context_payload(alert_details, agent_data)
84
+ if ioc_payload:
85
+ logger.info(f"Alert has IoC: {ioc_payload}")
86
+ return IrisAlertPayload(
87
+ alert_title=alert_details._source.rule_description,
88
+ alert_description=alert_details._source.rule_description,
89
+ alert_source="CoPilot",
90
+ assets=[asset_payload],
91
+ alert_status_id=3,
92
+ alert_severity_id=5,
93
+ alert_customer_id=1,
94
+ alert_source_content=alert_details._source,
95
+ alert_context=context_payload,
96
+ alert_iocs=[ioc_payload],
97
+ )
98
+ else:
99
+ logger.info("Alert does not have IoC")
100
+ return IrisAlertPayload(
101
+ alert_title=alert_details._source.rule_description,
102
+ alert_description=alert_details._source.rule_description,
103
+ alert_source="CoPilot",
104
+ assets=[asset_payload],
105
+ alert_status_id=3,
106
+ alert_severity_id=5,
107
+ alert_customer_id=1,
108
+ alert_source_content=alert_details._source,
109
+ alert_context=context_payload,
110
+ )
111
+
112
+
113
+def create_alert(alert: CreateAlertRequest) -> CreateAlertResponse:
114
+ logger.info(f"Creating alert {alert.alert_id} in IRIS")
115
+ alert_details = get_single_alert_details(alert_details=alert)
116
+ agent_data = get_agent_data(agent_id=alert_details._source.agent_id)
117
+ alert_details.asset_type_id = get_asset_type_id(os=agent_data.os)
118
+ ioc_payload = build_ioc_payload(alert_details)
119
+ iris_alert_payload = build_alert_payload(alert_details, agent_data, ioc_payload)
120
+ client, alert = initialize_client_and_alert("DFIR-IRIS")
121
+ result = fetch_and_validate_data(client, alert.add_alert, iris_alert_payload.to_dict())
122
+ alert_id = result["data"]["alert_id"]
123
+ return CreateAlertResponse(alert_id=alert_id, success=True, message=f"Alert {alert_id} created successfully")
backend/app/integrations/alert_escalation/utils/universal.py
new
+406
@@ -0,0 +1,406 @@
1
+import ipaddress
2
+import re
3
+from abc import ABC
4
+from typing import Any
5
+from typing import Dict
6
+from typing import Optional
7
+from typing import Union
8
+
9
+import regex
10
+from elasticsearch7 import Elasticsearch
11
+from fastapi import HTTPException
12
+from loguru import logger
13
+
14
+from app.connectors.utils import get_connector_info_from_db
15
+from app.db.all_models import Agents
16
+from app.db.db_session import session
17
+from app.healthchecks.agents.schema.agents import AgentModel
18
+
19
+
20
+#################### ! DFIR IRIS ASSET VALIDATOR ! ####################
21
+class AssetValidator(ABC):
22
+ """
23
+ Base class for asset validators.
24
+
25
+ Attributes:
26
+ os (str): The OS to be validated.
27
+ """
28
+
29
+ ASSET_TYPE_ID: int = 1
30
+
31
+ def __init__(self, os: str) -> None:
32
+ """
33
+ Initialize a Validator.
34
+
35
+ Args:
36
+ os (str): The OS to be validated.
37
+ """
38
+ self.os = os.lower()
39
+
40
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
41
+ """
42
+ Validate the OS.
43
+
44
+ If the OS matches the type of this validator,
45
+ the method returns a dictionary indicating success, the matching message, and the asset type id.
46
+
47
+ Returns:
48
+ Dict[str, Union[bool, str, int]]: The validation result.
49
+ """
50
+ raise NotImplementedError
51
+
52
+
53
+class WindowsAssetValidator(AssetValidator):
54
+ """
55
+ Class to check if an OS is Windows.
56
+ """
57
+
58
+ ASSET_TYPE_ID = 9
59
+
60
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
61
+ if "windows" in self.os:
62
+ return {
63
+ "success": True,
64
+ "message": f"{self.os} is a valid Windows OS.",
65
+ "asset_type_id": self.ASSET_TYPE_ID,
66
+ }
67
+ else:
68
+ return {
69
+ "success": False,
70
+ "message": f"{self.os} is not a Windows OS.",
71
+ "asset_type_id": self.ASSET_TYPE_ID,
72
+ }
73
+
74
+
75
+class LinuxAssetValidator(AssetValidator):
76
+ """
77
+ Class to check if an OS is Linux.
78
+ """
79
+
80
+ ASSET_TYPE_ID = 4
81
+
82
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
83
+ if "linux" in self.os:
84
+ return {
85
+ "success": True,
86
+ "message": f"{self.os} is a valid Linux OS.",
87
+ "asset_type_id": self.ASSET_TYPE_ID,
88
+ }
89
+ else:
90
+ return {
91
+ "success": False,
92
+ "message": f"{self.os} is not a Linux OS.",
93
+ "asset_type_id": self.ASSET_TYPE_ID,
94
+ }
95
+
96
+
97
+class FirewallAssetValidator(AssetValidator):
98
+ """
99
+ Class to check if an OS is Firewall.
100
+ """
101
+
102
+ ASSET_TYPE_ID = 2
103
+
104
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
105
+ if "firewall" in self.os:
106
+ return {
107
+ "success": True,
108
+ "message": f"{self.os} is a valid Firewall OS.",
109
+ "asset_type_id": self.ASSET_TYPE_ID,
110
+ }
111
+ else:
112
+ return {
113
+ "success": False,
114
+ "message": f"{self.os} is not a Firewall OS.",
115
+ "asset_type_id": self.ASSET_TYPE_ID,
116
+ }
117
+
118
+
119
+class UbuntuAssetValidator(AssetValidator):
120
+ """
121
+ Class to check if an OS is Ubuntu.
122
+ """
123
+
124
+ ASSET_TYPE_ID = 4
125
+
126
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
127
+ if "ubuntu" in self.os:
128
+ return {
129
+ "success": True,
130
+ "message": f"{self.os} is a valid Ubuntu OS.",
131
+ "asset_type_id": self.ASSET_TYPE_ID,
132
+ }
133
+ else:
134
+ return {
135
+ "success": False,
136
+ "message": f"{self.os} is not an Ubuntu OS.",
137
+ "asset_type_id": self.ASSET_TYPE_ID,
138
+ }
139
+
140
+
141
+class AssetTypeResolver:
142
+ """
143
+ Class to iterate over asset validators and return the successful validator's asset type id.
144
+ """
145
+
146
+ def __init__(self, os: str):
147
+ """
148
+ Initialize AssetTypeResolver.
149
+
150
+ Args:
151
+ os (str): The OS to be validated.
152
+ """
153
+ self.os = os
154
+ self.validators = [
155
+ WindowsAssetValidator,
156
+ LinuxAssetValidator,
157
+ FirewallAssetValidator,
158
+ UbuntuAssetValidator,
159
+ ]
160
+
161
+ def get_asset_type_id(self) -> int:
162
+ """
163
+ Iterate over validators and return the successful validator's asset type id.
164
+
165
+ Returns:
166
+ int: The asset type id.
167
+ """
168
+ for Validator in self.validators:
169
+ validator = Validator(self.os)
170
+ result = validator.validate()
171
+ if result["success"] is True:
172
+ return result["asset_type_id"]
173
+
174
+ # Return default asset type id (1) if no validators succeed
175
+ return 1
176
+
177
+
178
+#################### ! DFIR IRIS ASSET VALIDATOR END ! ####################
179
+
180
+
181
+#################### ! DFIR IRIS IOC VALIDATOR ! ##########################
182
+
183
+
184
+class IoCValidator(ABC):
185
+ """
186
+ Base class for validators.
187
+
188
+ Attributes:
189
+ value (str): The value to be validated.
190
+ """
191
+
192
+ PATTERN: Optional[str] = None # type: ignore
193
+ IOC_TYPE: Optional[int] = None # type: ignore
194
+
195
+ def __init__(self, value: str) -> None:
196
+ """
197
+ Initialize a Validator.
198
+
199
+ Args:
200
+ value (str): The value to be validated.
201
+ """
202
+ self.value = value
203
+
204
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
205
+ """
206
+ Validate the value.
207
+
208
+ If the value matches the pattern,
209
+ the method returns a dictionary indicating success, the matching message, and the IOC type.
210
+
211
+ Returns:
212
+ Dict[str, Union[bool, str, int]]: The validation result.
213
+ """
214
+ logger.info(f"Validating {self.value} against {self.PATTERN}.")
215
+ if self.PATTERN and regex.match(self.PATTERN, self.value, re.IGNORECASE):
216
+ return {
217
+ "success": True,
218
+ "message": f"{self.value} matches the pattern.",
219
+ "ioc_type": self.IOC_TYPE,
220
+ }
221
+ else:
222
+ return {
223
+ "success": False,
224
+ "message": f"{self.value} does not match the pattern.",
225
+ "ioc_type": self.IOC_TYPE,
226
+ }
227
+
228
+
229
+class IPv4AddressValidator(IoCValidator):
230
+ """
231
+ Class to check if a string is a valid IPv4 address.
232
+ """
233
+
234
+ IOC_TYPE = 76
235
+
236
+ def validate(self) -> Dict[str, Union[bool, str, int]]:
237
+ """
238
+ Validate if the given value is a valid IPv4 address.
239
+
240
+ Returns:
241
+ dict: A dictionary containing success status, message, and the associated IoC type.
242
+ """
243
+ try:
244
+ # if the value is like this `162.159.133.233|443` strip the port
245
+ if "|" in self.value:
246
+ self.value = self.value.split("|")[0]
247
+ logger.info(f"Validating {self.value} as an IPv4 address.")
248
+ ipaddress.IPv4Address(self.value)
249
+ return {
250
+ "success": True,
251
+ "message": f"{self.value} is a valid IPv4 address.",
252
+ "ioc_type": self.IOC_TYPE,
253
+ }
254
+ except ValueError:
255
+ return {
256
+ "success": False,
257
+ "message": f"{self.value} is not a valid IPv4 address.",
258
+ "ioc_type": self.IOC_TYPE,
259
+ }
260
+
261
+
262
+class HashValidator(IoCValidator):
263
+ """
264
+ Class to check if a string is a valid SHA256 hash.
265
+ """
266
+
267
+ PATTERN = r"^[a-fA-F\d]{64}$"
268
+ IOC_TYPE = 113
269
+
270
+
271
+class DomainValidator(IoCValidator):
272
+ """
273
+ Class to check if a string is a valid domain name.
274
+ """
275
+
276
+ PATTERN = r"^([a-z0-9]+(-[a-z0-9]+)*\.)+[a-z]{2,}$"
277
+ IOC_TYPE = 20
278
+
279
+
280
+#################### ! DFIR IRIS IOC VALIDATOR END ! ##########################
281
+
282
+
283
+def verify_wazuh_indexer_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:
284
+ """
285
+ Verifies the connection to Wazuh Indexer service.
286
+
287
+ Returns:
288
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
289
+ """
290
+ logger.info(f"Verifying the wazuh-indexer connection to {attributes['connector_url']}")
291
+
292
+ try:
293
+ es = Elasticsearch(
294
+ [attributes["connector_url"]],
295
+ http_auth=(attributes["connector_username"], attributes["connector_password"]),
296
+ verify_certs=False,
297
+ timeout=15,
298
+ max_retries=10,
299
+ retry_on_timeout=False,
300
+ )
301
+ es.cluster.health()
302
+ logger.debug("Wazuh Indexer connection successful")
303
+ return {"connectionSuccessful": True, "message": "Wazuh Indexer connection successful"}
304
+ except Exception as e:
305
+ logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")
306
+ return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}
307
+
308
+
309
+def verify_wazuh_indexer_connection(connector_name: str) -> str:
310
+ """
311
+ Returns the authentication token for the Wazuh Indexer service.
312
+
313
+ Returns:
314
+ str: Authentication token for the Wazuh Indexer service.
315
+ """
316
+ attributes = get_connector_info_from_db(connector_name)
317
+ if attributes is None:
318
+ logger.error("No Wazuh Indexer connector found in the database")
319
+ return None
320
+ return verify_wazuh_indexer_credentials(attributes)
321
+
322
+
323
+def create_wazuh_indexer_client(connector_name: str) -> Elasticsearch:
324
+ """
325
+ Returns an Elasticsearch client for the Wazuh Indexer service.
326
+
327
+ Returns:
328
+ Elasticsearch: Elasticsearch client for the Wazuh Indexer service.
329
+ """
330
+ attributes = get_connector_info_from_db(connector_name)
331
+ if attributes is None:
332
+ logger.error("No Wazuh Indexer connector found in the database")
333
+ return None
334
+ return Elasticsearch(
335
+ [attributes["connector_url"]],
336
+ http_auth=(attributes["connector_username"], attributes["connector_password"]),
337
+ verify_certs=False,
338
+ timeout=15,
339
+ max_retries=10,
340
+ retry_on_timeout=False,
341
+ )
342
+
343
+
344
+def get_agent_data(agent_id: str) -> AgentModel:
345
+ """
346
+ Get agent data based on the agent id from the agents table.
347
+
348
+ Args:
349
+ agent_id (str): Agent id.
350
+
351
+ Returns:
352
+ Dict[str, Any]: Agent data.
353
+ """
354
+ agent_details = session.query(Agents).filter(Agents.agent_id == agent_id).first()
355
+ if agent_details is not None:
356
+ return agent_details
357
+ else:
358
+ raise HTTPException(status_code=404, detail=f"Agent with id {agent_id} not found in agents table")
359
+
360
+
361
+def get_asset_type_id(os: str) -> int:
362
+ """
363
+ Use AssetTypeResolver to determine the asset type ID to set within DFIR-IRIS.
364
+
365
+ Parameters
366
+ ----------
367
+ os : str
368
+ The operating system (OS) string used to resolve the asset type ID.
369
+
370
+ Returns
371
+ -------
372
+ int
373
+ The ID corresponding to the asset type.
374
+ """
375
+ asset_resolver = AssetTypeResolver(os)
376
+ return asset_resolver.get_asset_type_id()
377
+
378
+
379
+def validate_ioc_type(ioc_value: str) -> str:
380
+ """
381
+ Validate IoC type using validators.
382
+
383
+ Parameters
384
+ ----------
385
+ ioc_value : str
386
+ The value to validate the IoC type.
387
+
388
+ Returns
389
+ -------
390
+ str
391
+ The type of the IoC. Returns None if validation fails.
392
+ """
393
+ validators = [IPv4AddressValidator, HashValidator, DomainValidator]
394
+ ioc_type = None
395
+
396
+ for Validator in validators:
397
+ validator = Validator(ioc_value)
398
+ result = validator.validate()
399
+
400
+ if result["success"]:
401
+ ioc_type = result["ioc_type"]
402
+ break
403
+
404
+ if ioc_type is None:
405
+ logger.error("Failed to validate IoC value.")
406
+ return ioc_type
backend/app/integrations/dnstwist/routes/analyze.py
new
+41
@@ -0,0 +1,41 @@
1
+import regex
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
+
14
+def is_domain(domain: str) -> DomainRequestBody:
15
+ """
16
+ Check if the provided domain is valid.
17
+
18
+ Args:
19
+ domain (str): The domain to check.
20
+
21
+ Returns:
22
+ bool: True if the domain is valid, False otherwise.
23
+ """
24
+ logger.info(f"Checking if domain {domain} is valid.")
25
+ pattern = regex.compile(
26
+ r"^(?:[a-zA-Z0-9]+([-._]?[a-zA-Z0-9]+)*\.)+[a-zA-Z]{2,}$",
27
+ )
28
+ if not pattern.match(domain):
29
+ raise HTTPException(status_code=400, detail="Invalid domain")
30
+ return DomainRequestBody(domain=domain)
31
+
32
+
33
+@dnstwist_router.post("/analyze", response_model=DomainAnalysisResponse, status_code=200, description="Analyze domain with DNS Twist")
34
+async def analyze(body: DomainRequestBody = Depends(is_domain)):
35
+ return analyze_domain(body.domain)
36
+
37
+
38
+# ! TODO: Add phishing analysis - Need more clarification on this
39
+# @dnstwist_router.post('/analyze/phishing', response_model=DomainAnalysisResponse, status_code=200, description='Analyze domain with DNS Twist')
40
+# async def analyze_phishing(body: DomainRequestBody = Depends(is_domain)):
41
+# return analyze_domain_phishing(body.domain)
backend/app/integrations/dnstwist/schema/analyze.py
new
+23
@@ -0,0 +1,23 @@
1
+from typing import List
2
+from typing import Optional
3
+
4
+from pydantic import BaseModel
5
+from pydantic import Field
6
+
7
+
8
+class DomainData(BaseModel):
9
+ dns_a: Optional[List[str]]
10
+ dns_mx: Optional[List[str]]
11
+ dns_ns: Optional[List[str]]
12
+ domain: str
13
+ fuzzer: str
14
+
15
+
16
+class DomainAnalysisResponse(BaseModel):
17
+ data: List[DomainData]
18
+ message: str
19
+ success: bool
20
+
21
+
22
+class DomainRequestBody(BaseModel):
23
+ domain: str = Field("socfortress.co", description="The domain to analyze.")
backend/app/integrations/dnstwist/services/analyze.py
new
+42
@@ -0,0 +1,42 @@
1
+import dnstwist
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
+ """
10
+ Analyze the domain using dnstwist and return the results for registered domains.
11
+
12
+ Args:
13
+ domain (DomainRequestBody): The domain to analyze.
14
+
15
+ Returns:
16
+ DomainAnalysisResponse: The response from DNS Twist.
17
+ """
18
+ logger.info(f"Analyzing domain {domain} with DNS Twist.")
19
+ logger.info("Analyzing domain for registered domains.")
20
+ data = dnstwist.run(domain=domain, registered=True, format="json")
21
+ return DomainAnalysisResponse(data=data, message="Domain analysis completed.", success=True)
22
+
23
+
24
+def analyze_domain_phishing(domain: DomainRequestBody) -> DomainAnalysisResponse:
25
+ """
26
+ Analyze the domain using dnstwist and return the results for registered domains.
27
+
28
+ Args:
29
+ domain (DomainRequestBody): The domain to analyze.
30
+
31
+ Returns:
32
+ DomainAnalysisResponse: The response from DNS Twist.
33
+ """
34
+ logger.info(f"Analyzing domain {domain} with DNS Twist.")
35
+ logger.info("Analyzing domain for registered domains.")
36
+ data = dnstwist.run(
37
+ domain=domain,
38
+ registered=True,
39
+ format="json",
40
+ lsh=True,
41
+ )
42
+ return DomainAnalysisResponse(data=data, message="Domain analysis completed.", success=True)
backend/app/integrations/dnstwist/utils/universal.py
backend/app/smtp/routes/configure.py
new
+70
@@ -0,0 +1,70 @@
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
11
+
12
+smtp_router = APIRouter()
13
+auth_handler = AuthHandler()
14
+
15
+
16
+@smtp_router.post("/{user_id}/register", response_model=SMTPResponse, status_code=200, description="Register new SMTP for user")
17
+async def register(user_id: int, smtp: SMTPInput):
18
+ users = select_all_users()
19
+ logger.info(users)
20
+ if not any(x.id == user_id for x in users):
21
+ raise HTTPException(status_code=400, detail="User not found")
22
+ # Check if SMTP already exists for user
23
+ smtp_found = session.query(SMTP).filter(SMTP.user_id == user_id).first()
24
+ if smtp_found:
25
+ raise HTTPException(status_code=400, detail="SMTP already exists for user")
26
+ hashed_pwd = auth_handler.get_password_hash(smtp.smtp_password)
27
+ u = SMTP(email=smtp.email, smtp_password=hashed_pwd, smtp_server=smtp.smtp_server, smtp_port=smtp.smtp_port, user_id=user_id)
28
+ session.add(u)
29
+ session.commit()
30
+ return {"message": "SMTP created successfully", "success": True}
31
+
32
+
33
+@smtp_router.get("/{user_id}", response_model=SMTP, status_code=200, description="Get SMTP for user")
34
+async def get_smtp(user_id: int):
35
+ users = select_all_users()
36
+ if not any(x.id == user_id for x in users):
37
+ raise HTTPException(status_code=400, detail="User not found")
38
+ smtp_found = session.query(SMTP).filter(SMTP.user_id == user_id).first()
39
+ if not smtp_found:
40
+ raise HTTPException(status_code=400, detail="SMTP not found for user")
41
+ return smtp_found
42
+
43
+
44
+@smtp_router.put("/{user_id}", response_model=SMTPResponse, status_code=200, description="Update SMTP for user")
45
+async def update_smtp(user_id: int, smtp: SMTPInput):
46
+ users = select_all_users()
47
+ if not any(x.id == user_id for x in users):
48
+ raise HTTPException(status_code=400, detail="User not found")
49
+ smtp_found = session.query(SMTP).filter(SMTP.user_id == user_id).first()
50
+ if not smtp_found:
51
+ raise HTTPException(status_code=400, detail="SMTP not found for user")
52
+ smtp_found.email = smtp.email
53
+ smtp_found.smtp_server = smtp.smtp_server
54
+ smtp_found.smtp_port = smtp.smtp_port
55
+ smtp_found.smtp_password = auth_handler.get_password_hash(smtp.smtp_password)
56
+ session.commit()
57
+ return {"message": "SMTP updated successfully", "success": True}
58
+
59
+
60
+@smtp_router.delete("/{user_id}", response_model=SMTPResponse, status_code=200, description="Delete SMTP for user")
61
+async def delete_smtp(user_id: int):
62
+ users = select_all_users()
63
+ if not any(x.id == user_id for x in users):
64
+ raise HTTPException(status_code=400, detail="User not found")
65
+ smtp_found = session.query(SMTP).filter(SMTP.user_id == user_id).first()
66
+ if not smtp_found:
67
+ raise HTTPException(status_code=400, detail="SMTP not found for user")
68
+ session.delete(smtp_found)
69
+ session.commit()
70
+ return {"message": "SMTP deleted successfully", "success": True}
backend/app/smtp/routes/reports.py
new
+31
@@ -0,0 +1,31 @@
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
11
+
12
+smtp_reports_router = APIRouter()
13
+auth_handler = AuthHandler()
14
+
15
+
16
+# ! TODO: Add SMTP reporting all things. Example is in the services/reports.py and services/create_report.py file
17
+@smtp_reports_router.post("/{user_id}/register", response_model=SMTPResponse, status_code=200, description="Register new SMTP for user")
18
+async def register(user_id: int, smtp: SMTPInput):
19
+ users = select_all_users()
20
+ logger.info(users)
21
+ if not any(x.id == user_id for x in users):
22
+ raise HTTPException(status_code=400, detail="User not found")
23
+ # Check if SMTP already exists for user
24
+ smtp_found = session.query(SMTP).filter(SMTP.user_id == user_id).first()
25
+ if smtp_found:
26
+ raise HTTPException(status_code=400, detail="SMTP already exists for user")
27
+ hashed_pwd = auth_handler.get_password_hash(smtp.smtp_password)
28
+ u = SMTP(email=smtp.email, smtp_password=hashed_pwd, smtp_server=smtp.smtp_server, smtp_port=smtp.smtp_port, user_id=user_id)
29
+ session.add(u)
30
+ session.commit()
31
+ return {"message": "SMTP created successfully", "success": True}
backend/app/smtp/schema/configure.py
new
+6
@@ -0,0 +1,6 @@
1
+from pydantic import BaseModel
2
+
3
+
4
+class SMTPResponse(BaseModel):
5
+ message: str
6
+ success: bool
backend/app/smtp/services/create_report.py
new
+143
@@ -0,0 +1,143 @@
1
+import urllib.request
2
+from typing import List
3
+
4
+import matplotlib
5
+from loguru import logger
6
+from reportlab.lib.pagesizes import letter
7
+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
12
+from reportlab.platypus import Paragraph
13
+from reportlab.platypus import SimpleDocTemplate
14
+from reportlab.platypus import Spacer
15
+
16
+matplotlib.use(
17
+ "Agg",
18
+) # set the backend to Agg which is a non-interactive backend suitable
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
26
+
27
+
28
+def fetch_alert_data(service, fetch_func):
29
+ """
30
+ Fetches alert data using the provided function.
31
+
32
+ Args:
33
+ service: An instance of the service to use for fetching data.
34
+ fetch_func (function): The function to use to fetch the data.
35
+
36
+ Returns:
37
+ Returns the result of the fetch function.
38
+ """
39
+ alerts = fetch_func()
40
+ logger.info(alerts)
41
+ return alerts
42
+
43
+
44
+def create_bar_chart(alerts: dict, title: str, output_filename: str) -> None:
45
+ """
46
+ Creates a horizontal bar chart of alerts and saves it to a file.
47
+
48
+ Args:
49
+ alerts (dict): A dictionary containing alert data.
50
+ title (str): The title for the chart.
51
+ output_filename (str): The filename to save the chart to.
52
+
53
+ Returns:
54
+ None
55
+ """
56
+ entities = [alert["hostname"] for alert in alerts["alerts_by_host"]]
57
+ num_alerts = [alert["number_of_alerts"] for alert in alerts["alerts_by_host"]]
58
+
59
+ plt.figure(figsize=(10, 10))
60
+ plt.barh(entities, num_alerts, color="blue")
61
+ plt.xlabel("Number of Alerts")
62
+ plt.ylabel("Hostnames")
63
+ plt.title(title)
64
+ plt.tight_layout()
65
+ plt.savefig(output_filename)
66
+
67
+
68
+def create_pie_chart(alerts: dict, title: str, output_filename: str) -> None:
69
+ """
70
+ Creates a pie chart of alerts and saves it to a file.
71
+
72
+ Args:
73
+ alerts (dict): A dictionary containing alert data.
74
+ title (str): The title for the chart.
75
+ output_filename (str): The filename to save the chart to.
76
+
77
+ Returns:
78
+ None
79
+ """
80
+ entities = [alert["rule"] for alert in alerts["alerts_by_rule"]]
81
+ num_alerts = [alert["number_of_alerts"] for alert in alerts["alerts_by_rule"]]
82
+
83
+ plt.figure(figsize=(10, 6))
84
+ plt.pie(num_alerts, labels=entities, autopct="%1.1f%%")
85
+ plt.legend(entities, loc="lower right", bbox_to_anchor=(1.0, 1.0)) # Add this line to include a legend
86
+ plt.title(title)
87
+ plt.tight_layout()
88
+ plt.savefig(output_filename)
89
+
90
+
91
+def create_pdf(title: str, image_filenames: List[str], pdf_filename: str) -> None:
92
+ """
93
+ Creates a PDF containing images.
94
+
95
+ Args:
96
+ title (str): The title for the PDF.
97
+ image_filenames (List[str]): A list of image filenames to include in the PDF.
98
+ pdf_filename (str): The filename to save the PDF to.
99
+
100
+ Returns:
101
+ None
102
+ """
103
+ # Download the SOC Fortress logo
104
+ logo_url = "https://socfortress-images.s3.amazonaws.com/socfortress_logo_orange.png"
105
+ logo_filename = "socfortress_logo_orange.png"
106
+ urllib.request.urlretrieve(logo_url, logo_filename)
107
+
108
+ doc = SimpleDocTemplate(pdf_filename, pagesize=letter)
109
+ styles = getSampleStyleSheet()
110
+ Story = []
111
+
112
+ # Add a cover page
113
+ Story.append(Spacer(1, 2 * inch))
114
+ Story.append(Image(logo_filename, 5 * inch, 5 * inch)) # Adjust size as needed
115
+ Story.append(Spacer(1, 1 * inch))
116
+ style = styles["Title"]
117
+ Story.append(Paragraph(title, style))
118
+ Story.append(Spacer(1, 2 * inch))
119
+
120
+ # Add the images
121
+ for i, image_filename in enumerate(image_filenames):
122
+ Story.append(Image(image_filename, 6 * inch, 4 * inch)) # Adjust size as needed
123
+ Story.append(Spacer(1, 0.2 * inch))
124
+
125
+ doc.build(Story)
126
+
127
+
128
+def create_alerts_report_pdf() -> None:
129
+ """
130
+ Creates a PDF report of alerts including a bar chart and a pie chart.
131
+
132
+ Returns:
133
+ None
134
+ """
135
+ service = AlertsService()
136
+
137
+ alerts_by_host = fetch_alert_data(service, service.collect_alerts_by_host)
138
+ create_bar_chart(alerts_by_host, "Number of Alerts by Host", "alerts_by_host.png")
139
+
140
+ alerts_by_rules = fetch_alert_data(service, service.collect_alerts_by_rule)
141
+ create_pie_chart(alerts_by_rules, "Number of Alerts by Rule", "alerts_by_rule.png")
142
+
143
+ create_pdf("Test", ["alerts_by_host.png", "alerts_by_rule.png"], "alerts_report.pdf")
backend/app/smtp/services/reports.py
new
+116
@@ -0,0 +1,116 @@
1
+import smtplib
2
+from email import encoders
3
+from email.mime.base import MIMEBase
4
+from email.mime.multipart import MIMEMultipart
5
+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
10
+from app.services.smtp.universal import UniversalEmailCredentials
11
+
12
+# ! SEND REPORT
13
+
14
+
15
+class EmailReportSender:
16
+ """
17
+ Class for sending an email report with PDF attachments.
18
+ """
19
+
20
+ def __init__(self, to_email: str):
21
+ """
22
+ Constructor for the EmailReportSender class.
23
+
24
+ Args:
25
+ to_email (str): The email address to send the report to.
26
+ """
27
+ self.to_email = to_email
28
+
29
+ def _get_credentials(self) -> dict:
30
+ """
31
+ Fetches the email credentials.
32
+
33
+ Returns:
34
+ dict: A dictionary containing the email credentials. If no credentials are found,
35
+ the dictionary contains an "error" key.
36
+ """
37
+ try:
38
+ return UniversalEmailCredentials.read_all()["emails_configured"][0]
39
+ except IndexError:
40
+ return {"error": "No email credentials found"}
41
+
42
+ def create_email_message(self, subject: str, body: str) -> MIMEMultipart:
43
+ """
44
+ Creates an email message with the provided subject and body.
45
+
46
+ Args:
47
+ subject (str): The subject of the email.
48
+ body (str): The body of the email.
49
+
50
+ Returns:
51
+ MIMEMultipart: An email message object. If an error occurs while fetching credentials,
52
+ the return value is a dictionary containing an "error" key.
53
+ """
54
+ msg = MIMEMultipart()
55
+ credentials = self._get_credentials()
56
+ if "error" in credentials:
57
+ return credentials
58
+ msg["From"] = credentials["email"]
59
+ msg["To"] = self.to_email
60
+ msg["Subject"] = subject
61
+ msg.attach(MIMEText(body, "html"))
62
+ return msg
63
+
64
+ def attach_pdfs(self, msg: MIMEMultipart, filenames: List[str]) -> MIMEMultipart:
65
+ """
66
+ Attaches PDF files to an email message.
67
+
68
+ Args:
69
+ msg (MIMEMultipart): The email message to attach the PDFs to.
70
+ filenames (List[str]): A list of filenames of the PDFs to attach.
71
+
72
+ Returns:
73
+ MIMEMultipart: The email message with the attached PDFs.
74
+ """
75
+ for filename in filenames:
76
+ with open(filename, "rb") as attachment_file:
77
+ part = MIMEBase("application", "octet-stream")
78
+ part.set_payload(attachment_file.read())
79
+ encoders.encode_base64(part)
80
+ part.add_header("Content-Disposition", f"attachment; filename= {filename}")
81
+ msg.attach(part)
82
+ return msg
83
+
84
+ def send_email_with_pdf(self):
85
+ """
86
+ Sends an email with a PDF report.
87
+
88
+ Returns:
89
+ dict: A dictionary containing a "message" key describing the result of the operation
90
+ and a "success" key indicating whether the operation was successful.
91
+ """
92
+ # Generate the PDF report
93
+ create_alerts_report_pdf()
94
+
95
+ # Render the email body
96
+ template = EmailTemplate("email_template")
97
+ body = template.render_html_body(template_name="email_template")
98
+
99
+ # Create the email message and attach the PDF
100
+ msg = self.create_email_message("Test Report", body)
101
+ if isinstance(msg, dict) and "error" in msg:
102
+ return {"message": msg["error"], "success": False}
103
+ msg = self.attach_pdfs(msg, ["alerts_report.pdf"])
104
+
105
+ credentials = self._get_credentials()
106
+ if "error" in credentials:
107
+ return {"message": credentials["error"], "success": False}
108
+
109
+ # Send the email
110
+ with smtplib.SMTP(credentials["smtp_server"], credentials["smtp_port"]) as server:
111
+ server.starttls()
112
+ server.login(credentials["email"], credentials["password"])
113
+ text = msg.as_string()
114
+ server.sendmail(credentials["email"], self.to_email, text)
115
+
116
+ return {"message": "Report sent successfully", "success": True}
backend/app/utils.py
new
+3
@@ -0,0 +1,3 @@
1
+def allowed_file(filename):
2
+ ALLOWED_EXTENSIONS = {"yaml", "txt"}
3
+ return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
backend/copilot.py
new
+103
@@ -0,0 +1,103 @@
1
+import uvicorn
2
+from fastapi import FastAPI
3
+from fastapi import HTTPException
4
+from fastapi import Request
5
+from fastapi.middleware.cors import CORSMiddleware
6
+from fastapi.responses import JSONResponse
7
+
8
+from app.agents.routes.agents import agents_router
9
+from app.auth.routes.auth import user_router
10
+from app.connectors.cortex.routes.analyzers import cortex_analyzer_router
11
+from app.connectors.dfir_iris.routes.alerts import dfir_iris_alerts_router
12
+from app.connectors.dfir_iris.routes.assets import assets_router
13
+from app.connectors.dfir_iris.routes.cases import cases_router
14
+from app.connectors.dfir_iris.routes.notes import notes_router
15
+from app.connectors.dfir_iris.routes.users import dfir_iris_users_router
16
+from app.connectors.graylog.routes.collector import graylog_collector_router
17
+from app.connectors.graylog.routes.events import graylog_events_router
18
+from app.connectors.graylog.routes.management import graylog_management_router
19
+from app.connectors.graylog.routes.monitoring import graylog_monitoring_router
20
+from app.connectors.graylog.routes.pipelines import graylog_pipelines_router
21
+from app.connectors.graylog.routes.streams import graylog_streams_router
22
+from app.connectors.routes import connector_router
23
+from app.connectors.shuffle.routes.workflows import shuffle_workflows_router
24
+from app.connectors.sublime.routes.alerts import sublime_alerts_router
25
+from app.connectors.velociraptor.routes.artifacts import velociraptor_artifacts_router
26
+from app.connectors.wazuh_indexer.routes.alerts import wazuh_indexer_alerts_router
27
+
28
+# from app.connectors.wazuh_indexer.routes.routes import wazuh_indexer_router
29
+from app.connectors.wazuh_indexer.routes.monitoring import wazuh_indexer_router
30
+from app.connectors.wazuh_manager.routes.rules import wazuh_manager_router
31
+from app.customers.routes.customers import customers_router
32
+from app.db.db_session import engine
33
+from app.db.db_setup import create_tables
34
+from app.healthchecks.agents.routes.agents import healtcheck_agents_router
35
+from app.integrations.alert_escalation.routes.general_alert import (
36
+ integration_general_alerts_router,
37
+)
38
+from app.integrations.dnstwist.routes.analyze import dnstwist_router
39
+from app.smtp.routes.configure import smtp_router
40
+
41
+app = FastAPI(description="CoPilot API", version="0.1.0", title="CoPilot API")
42
+
43
+# Allow all origins, methods and headers
44
+app.add_middleware(
45
+ CORSMiddleware,
46
+ allow_origins=["*"],
47
+ allow_credentials=True,
48
+ allow_methods=["*"],
49
+ allow_headers=["*"],
50
+)
51
+
52
+
53
+@app.exception_handler(HTTPException)
54
+async def custom_http_exception_handler(request: Request, exc: HTTPException):
55
+ return JSONResponse(
56
+ status_code=exc.status_code,
57
+ content={
58
+ "success": False,
59
+ "message": exc.detail,
60
+ },
61
+ )
62
+
63
+
64
+app.include_router(connector_router, prefix="/connectors", tags=["connectors"])
65
+app.include_router(wazuh_indexer_router, prefix="/wazuh_indexer", tags=["wazuh-indexer"])
66
+app.include_router(user_router, prefix="/auth", tags=["auth"])
67
+app.include_router(wazuh_manager_router, prefix="/wazuh_manager", tags=["wazuh-manager"])
68
+app.include_router(agents_router, prefix="/agents", tags=["agents"])
69
+app.include_router(graylog_monitoring_router, prefix="/graylog", tags=["graylog"])
70
+app.include_router(graylog_collector_router, prefix="/graylog", tags=["graylog"])
71
+app.include_router(graylog_events_router, prefix="/graylog", tags=["graylog"])
72
+app.include_router(graylog_pipelines_router, prefix="/graylog", tags=["graylog"])
73
+app.include_router(graylog_streams_router, prefix="/graylog", tags=["graylog"])
74
+app.include_router(graylog_management_router, prefix="/graylog", tags=["graylog"])
75
+app.include_router(wazuh_indexer_alerts_router, prefix="/alerts", tags=["alerts"])
76
+app.include_router(cases_router, prefix="/cases", tags=["cases"])
77
+app.include_router(notes_router, prefix="/notes", tags=["notes"])
78
+app.include_router(assets_router, prefix="/assets", tags=["assets"])
79
+app.include_router(dfir_iris_alerts_router, prefix="/alerts", tags=["soc-alerts"])
80
+app.include_router(dfir_iris_users_router, prefix="/users", tags=["dfir_iris-users"])
81
+app.include_router(cortex_analyzer_router, prefix="/analyzers", tags=["cortex-analyzers"])
82
+app.include_router(velociraptor_artifacts_router, prefix="/artifacts", tags=["velociraptor"])
83
+app.include_router(shuffle_workflows_router, prefix="/workflows", tags=["shuffle"])
84
+app.include_router(sublime_alerts_router, prefix="/sublime", tags=["sublime"])
85
+app.include_router(customers_router, prefix="/customers", tags=["customers"])
86
+app.include_router(healtcheck_agents_router, prefix="/healthcheck", tags=["healthcheck"])
87
+app.include_router(smtp_router, prefix="/smtp", tags=["smtp"])
88
+app.include_router(dnstwist_router, prefix="/dnstwist", tags=["dnstwist"])
89
+app.include_router(integration_general_alerts_router, prefix="/alerts", tags=["alerts"])
90
+
91
+
92
+@app.on_event("startup")
93
+async def init_db():
94
+ create_tables(engine)
95
+
96
+
97
+@app.get("/")
98
+def hello():
99
+ return {"message": "Hello World"}
100
+
101
+
102
+if __name__ == "__main__":
103
+ uvicorn.run(app, host="localhost", port=5000)
backend/requirements.in
new
+33
@@ -0,0 +1,33 @@
1
+bcrypt
2
+blueprint
3
+cortex4py
4
+dfir_iris_client
5
+dnstwist
6
+elasticsearch7==7.10.1
7
+environs
8
+fastapi
9
+libmagic
10
+loguru
11
+marshmallow-sqlalchemy
12
+matplotlib
13
+mitreattack-python
14
+openai
15
+passlib[bcrypt]
16
+pcre2
17
+pika
18
+psycopg2-binary
19
+python-multipart
20
+pytest
21
+python-magic
22
+passlib
23
+PyJWT
24
+pydantic[email]
25
+python-jose[cryptography]
26
+pyvelociraptor~=0.1
27
+regex
28
+uvicorn
29
+sqlmodel
30
+reportlab
31
+requests
32
+xmltodict
33
+werkzeug
backend/settings.py
new
+26
@@ -0,0 +1,26 @@
1
+"""Application configuration.
2
+
3
+Most configuration is set via environment variables.
4
+
5
+For local development, use a .env file to set
6
+environment variables.
7
+"""
8
+from pathlib import Path
9
+
10
+from environs import Env
11
+
12
+env = Env()
13
+env.read_env()
14
+
15
+basedir = Path().absolute()
16
+db_path = str(basedir / "copilot.db")
17
+
18
+ENV = env.str("SECRET_KEY", default="production")
19
+DEBUG = env.bool("FLASK_DEBUG", default=False)
20
+SECRET_KEY = env.str("SECRET_KEY", "not-a-secret")
21
+SQLALCHEMY_DATABASE_URI = env.str("SQLALCHEMY_DATABASE_URI", f"sqlite:///{db_path}")
22
+SQLALCHEMY_TRACK_MODIFICATIONS = env.bool(
23
+ "SQLALCHEMY_TRACK_MODIFICATIONS",
24
+ default=False,
25
+)
26
+UPLOAD_FOLDER = env.str("UPLOAD_FOLDER", str(Path.home() / "Desktop/copilot_uploads"))