Velo collect agent org (#256)
* collect velo org per agent testing * refactor: Collect Velociraptor agents per organization This code change modifies the `sync_agents_velociraptor` function in the `sync.py` file. It now collects Velociraptor agents per organization, allowing for more efficient synchronization. This refactor improves the organization and management of Velociraptor agents in the system. * refactor: add velo org to agents table * add velo org when collecting agents * refactor: Add velociraptor_org field to BaseBody model This code change adds the `velociraptor_org` field to the `BaseBody` model in the `artifacts.py` file. The `velociraptor_org` field allows for specifying the organization of the client when collecting artifacts. This enhancement improves the flexibility and accuracy of artifact collection in the system. * refactor: Extract filename from process_name in SocfortressThreatIntelRequest This code change adds a validator to the `process_name` field in the `SocfortressThreatIntelRequest` model. The validator extracts the filename from the `process_name` using the `os.path.basename` function. This enhancement improves the consistency and reliability of the filename extraction process in the Socfortress threat intelligence module. * precommit fixes
taylor_socfortress committed
Jun 27, 2024 at 11:30 UTC
17a0656586d32a55689eef9169c670214f8ea251
11 files changed
+259
-59
backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py
new
+32
@@ -0,0 +1,32 @@
1
+"""Add Velo Org to Agents
2
+
3
+Revision ID: fed7739bd07c
4
+Revises: 39c3aaec0084
5
+Create Date: 2024-06-27 09:22:59.354696
6
+
7
+"""
8
+from typing import Sequence
9
+from typing import Union
10
+
11
+import sqlalchemy as sa
12
+import sqlmodel.sql.sqltypes
13
+
14
+from alembic import op
15
+
16
+# revision identifiers, used by Alembic.
17
+revision: str = "fed7739bd07c"
18
+down_revision: Union[str, None] = "39c3aaec0084"
19
+branch_labels: Union[str, Sequence[str], None] = None
20
+depends_on: Union[str, Sequence[str], None] = None
21
+
22
+
23
+def upgrade() -> None:
24
+ # ### commands auto generated by Alembic - please adjust! ###
25
+ op.add_column("agents", sa.Column("velociraptor_org", sqlmodel.sql.sqltypes.AutoString(length=256), nullable=True))
26
+ # ### end Alembic commands ###
27
+
28
+
29
+def downgrade() -> None:
30
+ # ### commands auto generated by Alembic - please adjust! ###
31
+ op.drop_column("agents", "velociraptor_org")
32
+ # ### end Alembic commands ###
backend/app/agents/services/sync.py
+71
-49
@@ -13,6 +13,7 @@ from app.agents.schema.agents import SyncedAgentsResponse
13
from app.agents.schema.agents import SyncedWazuhAgent
14
from app.agents.velociraptor.schema.agents import VelociraptorAgent
15
from app.agents.velociraptor.schema.agents import VelociraptorClients
16
+from app.agents.velociraptor.schema.agents import VelociraptorOrganizations
17
from app.agents.wazuh.schema.agents import WazuhAgent
18
from app.agents.wazuh.schema.agents import WazuhAgentsList
19
from app.connectors.models import Connectors
@@ -40,7 +41,7 @@ async def fetch_wazuh_agents() -> WazuhAgentsList:
41
)
42
43
43
-async def fetch_velociraptor_clients() -> VelociraptorClients:
44
+async def fetch_velociraptor_clients(org_id: str) -> VelociraptorClients:
45
"""
46
Fetches clients from Velociraptor service.
47
@@ -50,12 +51,29 @@ async def fetch_velociraptor_clients() -> VelociraptorClients:
51
Returns:
52
VelociraptorClientsList: The fetched clients.
53
"""
53
- collected_velociraptor_agents = await velociraptor_services.collect_velociraptor_clients()
54
+ collected_velociraptor_agents = await velociraptor_services.collect_velociraptor_clients(org_id=org_id)
55
return VelociraptorClients(
56
clients=collected_velociraptor_agents,
57
)
58
59
60
+async def fetch_velociraptor_organizations() -> VelociraptorOrganizations:
61
+ """
62
+ Fetches organizations from Velociraptor service.
63
+
64
+ Args:
65
+ None
66
+
67
+ Returns:
68
+ VelociraptorOrgsList: The fetched orgs.
69
+ """
70
+ collected_velociraptor_orgs = await velociraptor_services.collect_velociraptor_organizations()
71
+ logger.info(f"Collected Velociraptor Orgs: {collected_velociraptor_orgs}")
72
+ return VelociraptorOrganizations(
73
+ organizations=collected_velociraptor_orgs,
74
+ )
75
+
76
+
77
async def fetch_velociraptor_agent(agent_name: str) -> VelociraptorAgent:
78
"""
79
Fetches agent details from Velociraptor service.
@@ -298,53 +316,57 @@ async def sync_agents_velociraptor() -> SyncedAgentsResponse:
316
:rtype: SyncedAgentsResponse
317
"""
318
agents_added_list: List[VelociraptorAgent] = []
301
-
302
- velociraptor_clients = await fetch_velociraptor_clients()
303
- velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else []
304
-
305
- async with get_db_session() as session: # Create a new session here
306
- existing_agents_query = select(Agents)
307
- result = await session.execute(existing_agents_query)
308
- existing_agents = result.scalars().all()
309
-
310
- for agent in existing_agents:
311
- logger.info(f"Collecting Velociraptor Agent for {agent.hostname}")
312
-
313
- try:
314
- # Build the velociraptor_agent where the hostname or `client_id` is that equal to the `agents`
315
- velociraptor_agent = next(
316
- (
317
- client
318
- for client in velociraptor_clients
319
- if client.os_info.hostname == agent.hostname or client.client_id == agent.velociraptor_id
320
- ),
321
- None,
322
- )
323
- # Convert Unix epoch timestamp to datetime
324
- last_seen_at = datetime.fromtimestamp(
325
- int(velociraptor_agent.last_seen_at) / 1e6,
326
- ) # Divide by 1e6 to convert from microseconds to seconds
327
- # Convert datetime to ISO 8601 format without fractional seconds
328
- last_seen_at_iso = last_seen_at.replace(tzinfo=timezone.utc).isoformat(timespec="seconds")
329
- velociraptor_agent = VelociraptorAgent(
330
- velociraptor_id=velociraptor_agent.client_id,
331
- velociraptor_last_seen=last_seen_at_iso,
332
- velociraptor_agent_version=velociraptor_agent.agent_information.version,
333
- )
334
-
335
- except Exception as e:
336
- logger.error(
337
- f"Failed to collect Velociraptor Agent for {agent.hostname}: {e}",
338
- )
339
- continue
340
-
341
- if velociraptor_agent:
342
- # Update the agent with the Velociraptor client's details
343
- await update_agent_with_velociraptor_in_db(session, agent, velociraptor_agent)
344
- agents_added_list.append(velociraptor_agent)
345
-
346
- # Close the session
347
- await session.close()
319
+ velo_orgs = await fetch_velociraptor_organizations()
320
+ logger.info(f"Collected Velociraptor Orgs: {velo_orgs}")
321
+ for org in velo_orgs.organizations:
322
+ velociraptor_clients = await fetch_velociraptor_clients(org_id=org.OrgId)
323
+ logger.info(f"Collected Velociraptor Clients: {velociraptor_clients}")
324
+ velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else []
325
+
326
+ async with get_db_session() as session: # Create a new session here
327
+ existing_agents_query = select(Agents)
328
+ result = await session.execute(existing_agents_query)
329
+ existing_agents = result.scalars().all()
330
+
331
+ for agent in existing_agents:
332
+ logger.info(f"Collecting Velociraptor Agent for {agent.hostname}")
333
+
334
+ try:
335
+ # Build the velociraptor_agent where the hostname or `client_id` is that equal to the `agents`
336
+ velociraptor_agent = next(
337
+ (
338
+ client
339
+ for client in velociraptor_clients
340
+ if client.os_info.hostname == agent.hostname or client.client_id == agent.velociraptor_id
341
+ ),
342
+ None,
343
+ )
344
+ # Convert Unix epoch timestamp to datetime
345
+ last_seen_at = datetime.fromtimestamp(
346
+ int(velociraptor_agent.last_seen_at) / 1e6,
347
+ ) # Divide by 1e6 to convert from microseconds to seconds
348
+ # Convert datetime to ISO 8601 format without fractional seconds
349
+ last_seen_at_iso = last_seen_at.replace(tzinfo=timezone.utc).isoformat(timespec="seconds")
350
+ velociraptor_agent = VelociraptorAgent(
351
+ velociraptor_id=velociraptor_agent.client_id,
352
+ velociraptor_last_seen=last_seen_at_iso,
353
+ velociraptor_agent_version=velociraptor_agent.agent_information.version,
354
+ velociraptor_org=org.OrgId,
355
+ )
356
+
357
+ except Exception as e:
358
+ logger.error(
359
+ f"Failed to collect Velociraptor Agent for {agent.hostname}: {e}",
360
+ )
361
+ continue
362
+
363
+ if velociraptor_agent:
364
+ # Update the agent with the Velociraptor client's details
365
+ await update_agent_with_velociraptor_in_db(session, agent, velociraptor_agent)
366
+ agents_added_list.append(velociraptor_agent)
367
+
368
+ # Close the session
369
+ await session.close()
370
371
logger.info(f"Agents Added List: {agents_added_list}")
372
return SyncedAgentsResponse(
backend/app/agents/velociraptor/schema/agents.py
+53
@@ -10,6 +10,7 @@ class VelociraptorAgent(BaseModel):
10
client_id: Optional[str] = Field("n/a", alias="velociraptor_id")
11
client_last_seen: str = Field(..., alias="velociraptor_last_seen")
12
client_version: str = Field(..., alias="velociraptor_agent_version")
13
+ client_org: str = Field(..., alias="velociraptor_org")
14
15
@property
16
def client_last_seen_as_datetime(self):
@@ -53,3 +54,55 @@ class VelociraptorClient(BaseModel):
54
55
class VelociraptorClients(BaseModel):
56
clients: List[VelociraptorClient]
57
+
58
+
59
+class Version(BaseModel):
60
+ name: str
61
+ version: str
62
+ commit: str
63
+ build_time: str
64
+ ci_build_url: str
65
+ compiler: str
66
+
67
+
68
+class Installer(BaseModel):
69
+ service_name: str
70
+ install_path: str
71
+ service_description: Optional[str] = None
72
+
73
+
74
+class LocalBuffer(BaseModel):
75
+ memory_size: int
76
+ disk_size: int
77
+ filename_linux: str
78
+ filename_windows: str
79
+ filename_darwin: str
80
+
81
+
82
+class ClientConfig(BaseModel):
83
+ server_urls: List[str]
84
+ ca_certificate: str
85
+ nonce: str
86
+ writeback_darwin: str
87
+ writeback_linux: str
88
+ writeback_windows: str
89
+ tempdir_windows: str
90
+ max_poll: int
91
+ nanny_max_connection_delay: int
92
+ windows_installer: Installer
93
+ darwin_installer: Installer
94
+ version: Version
95
+ use_self_signed_ssl: bool
96
+ pinned_server_name: str
97
+ max_upload_size: int
98
+ local_buffer: LocalBuffer
99
+
100
+
101
+class Organization(BaseModel):
102
+ Name: str
103
+ OrgId: str
104
+ _client_config: ClientConfig
105
+
106
+
107
+class VelociraptorOrganizations(BaseModel):
108
+ organizations: List[Organization]
backend/app/agents/velociraptor/services/agents.py
+22
-2
@@ -21,7 +21,7 @@ def create_query(query: str) -> str:
21
return query
22
23
24
-async def collect_velociraptor_clients() -> list:
24
+async def collect_velociraptor_clients(org_id: str) -> list:
25
"""
26
Collects all clients from Velociraptor.
27
@@ -29,10 +29,30 @@ async def collect_velociraptor_clients() -> list:
29
list: A list of all clients.
30
"""
31
velociraptor_service = await UniversalService.create("Velociraptor")
32
+ # query = create_query(
33
+ # "SELECT * FROM clients()",
34
+ # )
35
query = create_query(
33
- "SELECT * FROM clients()",
36
+ f"SELECT * FROM query(org_id='{org_id}', query='SELECT * FROM clients()')",
37
)
38
flow = velociraptor_service.execute_query(query)
39
+ logger.info(f"Successfully ran artifact collection on {flow}")
40
+ return flow["results"]
41
+
42
+
43
+async def collect_velociraptor_organizations() -> list:
44
+ """
45
+ Collects all organizations from Velociraptor.
46
+
47
+ Returns:
48
+ list: A list of all organizations.
49
+ """
50
+ velociraptor_service = await UniversalService.create("Velociraptor")
51
+ query = create_query(
52
+ "SELECT * FROM orgs()",
53
+ )
54
+ flow = velociraptor_service.execute_query(query)
55
+ logger.info(f"Successfully ran artifact collection on {flow}")
56
return flow["results"]
57
58
backend/app/connectors/velociraptor/routes/artifacts.py
+50
@@ -126,6 +126,40 @@ async def get_velociraptor_id(session: AsyncSession, hostname: str) -> str:
126
return agent.velociraptor_id
127
128
129
+async def get_velociraptor_org(session: AsyncSession, hostname: str) -> str:
130
+ """
131
+ Retrieves the velociraptor_org associated with the given hostname.
132
+
133
+ Args:
134
+ session (AsyncSession): The database session.
135
+ hostname (str): The hostname of the agent.
136
+
137
+ Returns:
138
+ str: The velociraptor_org associated with the hostname.
139
+
140
+ Raises:
141
+ HTTPException: If the agent with the given hostname is not found or if the velociraptor_org is not available.
142
+ """
143
+ logger.info(f"Getting velociraptor_org from hostname {hostname}")
144
+ result = await session.execute(select(Agents).filter(Agents.hostname == hostname))
145
+ agent = result.scalars().first()
146
+
147
+ if not agent:
148
+ raise HTTPException(
149
+ status_code=404,
150
+ detail=f"Agent with hostname {hostname} not found",
151
+ )
152
+
153
+ if agent.velociraptor_org is None:
154
+ raise HTTPException(
155
+ status_code=404,
156
+ detail=f"Velociraptor ORG for hostname {hostname} is not available",
157
+ )
158
+
159
+ logger.info(f"velociraptor_org for hostname {hostname} is {agent.velociraptor_org}")
160
+ return agent.velociraptor_org
161
+
162
+
163
async def update_agent_quarantine_status(
164
session: AsyncSession,
165
quarantine_body: QuarantineBody,
@@ -321,6 +355,11 @@ async def collect_artifact(
355
collect_artifact_body.hostname,
356
)
357
358
+ collect_artifact_body.velociraptor_org = await get_velociraptor_org(
359
+ session,
360
+ collect_artifact_body.hostname,
361
+ )
362
+
363
# Assuming run_artifact_collection is an async function and takes a session as a parameter
364
return await run_artifact_collection(collect_artifact_body)
365
@@ -358,6 +397,11 @@ async def run_command(
397
session,
398
run_command_body.hostname,
399
)
400
+
401
+ run_command_body.velociraptor_org = await get_velociraptor_org(
402
+ session,
403
+ run_command_body.hostname,
404
+ )
405
# Run the command
406
return await run_remote_command(run_command_body)
407
@@ -396,6 +440,12 @@ async def quarantine(
440
session,
441
quarantine_body.hostname,
442
)
443
+
444
+ quarantine_body.velociraptor_org = await get_velociraptor_org(
445
+ session,
446
+ quarantine_body.hostname,
447
+ )
448
+
449
# Quarantine the host
450
quarantine_response = await quarantine_host(quarantine_body)
451
backend/app/connectors/velociraptor/schema/artifacts.py
+1
@@ -73,6 +73,7 @@ class QuarantineArtifactsEnum(str, Enum):
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
+ velociraptor_org: Optional[str] = Field(None, description="Organization of the client")
77
78
79
class CollectArtifactBody(BaseBody):
backend/app/connectors/velociraptor/services/artifacts.py
+18
-8
@@ -46,24 +46,27 @@ def get_artifact_key(analyzer_body: CollectArtifactBody) -> str:
46
47
if action == "quarantine":
48
return (
49
- f'collect_client(client_id="{analyzer_body.velociraptor_id}", '
49
+ f'collect_client(org_id="{analyzer_body.velociraptor_org}", client_id="{analyzer_body.velociraptor_id}", '
50
f'artifacts=["{analyzer_body.artifact_name}"], '
51
f"spec=dict(`{analyzer_body.artifact_name}`=dict()))"
52
)
53
elif action == "remove_quarantine":
54
return (
55
- f'collect_client(client_id="{analyzer_body.velociraptor_id}", '
55
+ f'collect_client(org_id="{analyzer_body.velociraptor_org}", client_id="{analyzer_body.velociraptor_id}", '
56
f'artifacts=["{analyzer_body.artifact_name}"], '
57
f'spec=dict(`{analyzer_body.artifact_name}`=dict(`RemovePolicy`="Y")))'
58
)
59
elif command is not None:
60
return (
61
- f"collect_client(client_id='{analyzer_body.velociraptor_id}', "
61
+ f"collect_client(org_id='{analyzer_body.velociraptor_org}', client_id='{analyzer_body.velociraptor_id}', "
62
f"urgent=true, artifacts=['{analyzer_body.artifact_name}'], "
63
f"env=dict(Command='{analyzer_body.command}'))"
64
)
65
else:
66
- return f"collect_client(client_id='{analyzer_body.velociraptor_id}', " f"artifacts=['{analyzer_body.artifact_name}'])"
66
+ return (
67
+ f"collect_client(org_id='{analyzer_body.velociraptor_org}', client_id='{analyzer_body.velociraptor_id}', "
68
+ f"artifacts=['{analyzer_body.artifact_name}'])"
69
+ )
70
71
72
async def get_artifacts() -> ArtifactsResponse:
@@ -112,8 +115,15 @@ async def run_artifact_collection(
115
"""
116
velociraptor_service = await UniversalService.create("Velociraptor")
117
try:
118
+ # ! Can specify org_id with org_id='OL680' ! #
119
query = create_query(
116
- f"SELECT collect_client(client_id='{collect_artifact_body.velociraptor_id}', artifacts=['{collect_artifact_body.artifact_name}']) FROM scope()",
120
+ (
121
+ f"SELECT collect_client("
122
+ f"org_id='{collect_artifact_body.velociraptor_org}', "
123
+ f"client_id='{collect_artifact_body.velociraptor_id}', "
124
+ f"artifacts=['{collect_artifact_body.artifact_name}']) "
125
+ f"FROM scope()"
126
+ ),
127
)
128
flow = velociraptor_service.execute_query(query)
129
logger.info(f"Successfully ran artifact collection on {flow}")
@@ -170,7 +180,7 @@ async def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResp
180
logger.info(f"Running remote command on {run_command_body}")
181
query = create_query(
182
(
173
- f"SELECT collect_client(client_id='{run_command_body.velociraptor_id}', "
183
+ f"SELECT collect_client(org_id='{run_command_body.velociraptor_org}', client_id='{run_command_body.velociraptor_id}', "
184
f"urgent=true, artifacts=['{run_command_body.artifact_name}'], "
185
f"env=dict(Command='{run_command_body.command}')) "
186
"FROM scope()"
@@ -225,7 +235,7 @@ async def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse
235
if quarantine_body.action == "quarantine":
236
query = create_query(
237
(
228
- f'SELECT collect_client(client_id="{quarantine_body.velociraptor_id}", '
238
+ f'SELECT collect_client(org_id="{quarantine_body.velociraptor_org}", client_id="{quarantine_body.velociraptor_id}", '
239
f'artifacts=["{quarantine_body.artifact_name}"], '
240
f"spec=dict(`{quarantine_body.artifact_name}`=dict())) "
241
"FROM scope()"
@@ -234,7 +244,7 @@ async def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse
244
else:
245
query = create_query(
246
(
237
- f'SELECT collect_client(client_id="{quarantine_body.velociraptor_id}", '
247
+ f'SELECT collect_client(org_id="{quarantine_body.velociraptor_org}", client_id="{quarantine_body.velociraptor_id}", '
248
f'artifacts=["{quarantine_body.artifact_name}"], '
249
f'spec=dict(`{quarantine_body.artifact_name}`=dict(`RemovePolicy`="Y"))) '
250
"FROM scope()"
backend/app/connectors/velociraptor/utils/universal.py
+1
@@ -213,6 +213,7 @@ class UniversalService:
213
dict: A dictionary with the success status and a message.
214
"""
215
vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1"
216
+ # vql = f"SELECT * FROM query(org_id='OL680', query='SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1')"
217
logger.info(f"Watching flow {flow_id} for completion")
218
return self.execute_query(vql)
219
backend/app/db/universal_models.py
+4
@@ -101,6 +101,7 @@ class Agents(SQLModel, table=True):
101
velociraptor_agent_version: Optional[str] = Field(max_length=256)
102
customer_code: Optional[str] = Field(foreign_key="customers.customer_code", max_length=256)
103
quarantined: bool = Field(default=False)
104
+ velociraptor_org: Optional[str] = Field(max_length=256)
105
106
customer: Optional[Customers] = Relationship(back_populates="agents")
107
@@ -129,6 +130,7 @@ class Agents(SQLModel, table=True):
130
if velociraptor_agent and velociraptor_agent.client_version
131
else None,
132
customer_code=customer_code,
133
+ velociraptor_org=velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None,
134
)
135
136
@classmethod
@@ -177,6 +179,7 @@ class Agents(SQLModel, table=True):
179
velociraptor_agent.client_version if velociraptor_agent and velociraptor_agent.client_version else None
180
)
181
self.customer_code = customer_code
182
+ self.velociraptor_org = velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None
183
184
def update_wazuh_agent_from_model(self, wazuh_agent, customer_code):
185
if wazuh_agent.agent_last_seen == "Unknown" or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00":
@@ -209,6 +212,7 @@ class Agents(SQLModel, table=True):
212
velociraptor_agent.client_version if velociraptor_agent and velociraptor_agent.client_version else None
213
)
214
logger.info(f"Updated with Velociraptor details: {self}")
215
+ self.velociraptor_org = velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None
216
217
218
class LogEntry(SQLModel, table=True):
backend/app/threat_intel/schema/socfortress.py
+6
@@ -1,8 +1,10 @@
1
+import os
2
from typing import List
3
from typing import Optional
4
5
from pydantic import BaseModel
6
from pydantic import Field
7
+from pydantic import validator
8
9
10
class SocfortressThreatIntelRequest(BaseModel):
@@ -54,6 +56,10 @@ class SocfortressProcessNameAnalysisRequest(BaseModel):
56
description="The process name to evaluate.",
57
)
58
59
+ @validator("process_name", pre=True)
60
+ def extract_filename(cls, v):
61
+ return os.path.basename(v)
62
+
63
64
class Path(BaseModel):
65
directory: str
backend/app/threat_intel/services/socfortress.py
+1
@@ -227,6 +227,7 @@ async def get_process_analysis_response(
227
228
# Using .get() with default values
229
data = response_data.get("data", {})
230
+ logger.info(f"Data {data}")
231
success = response_data.get("success", False)
232
message = response_data.get("message", "No message provided")
233