From 2cc5711aef2c0d730722e8dcc1615b09cf61aa40 Mon Sep 17 00:00:00 2001 From: Taylor Date: Wed, 26 Jun 2024 17:42:42 -0500 Subject: [PATCH 1/7] collect velo org per agent testing --- backend/app/agents/services/sync.py | 26 ++++++++-- .../app/agents/velociraptor/schema/agents.py | 51 +++++++++++++++++++ .../agents/velociraptor/services/agents.py | 23 ++++++++- .../velociraptor/services/artifacts.py | 1 + 4 files changed, 96 insertions(+), 5 deletions(-) diff --git a/backend/app/agents/services/sync.py b/backend/app/agents/services/sync.py index 05d1e62b..ea27852a 100644 --- a/backend/app/agents/services/sync.py +++ b/backend/app/agents/services/sync.py @@ -12,7 +12,7 @@ from app.agents.schema.agents import SyncedAgentsResponse from app.agents.schema.agents import SyncedWazuhAgent from app.agents.velociraptor.schema.agents import VelociraptorAgent -from app.agents.velociraptor.schema.agents import VelociraptorClients +from app.agents.velociraptor.schema.agents import VelociraptorClients, VelociraptorOrganizations from app.agents.wazuh.schema.agents import WazuhAgent from app.agents.wazuh.schema.agents import WazuhAgentsList from app.connectors.models import Connectors @@ -40,7 +40,7 @@ async def fetch_wazuh_agents() -> WazuhAgentsList: ) -async def fetch_velociraptor_clients() -> VelociraptorClients: +async def fetch_velociraptor_clients(org_id: str) -> VelociraptorClients: """ Fetches clients from Velociraptor service. @@ -50,11 +50,27 @@ async def fetch_velociraptor_clients() -> VelociraptorClients: Returns: VelociraptorClientsList: The fetched clients. """ - collected_velociraptor_agents = await velociraptor_services.collect_velociraptor_clients() + collected_velociraptor_agents = await velociraptor_services.collect_velociraptor_clients(org_id=org_id) return VelociraptorClients( clients=collected_velociraptor_agents, ) +async def fetch_velociraptor_organizations() -> VelociraptorOrganizations: + """ + Fetches organizations from Velociraptor service. + + Args: + None + + Returns: + VelociraptorOrgsList: The fetched orgs. + """ + collected_velociraptor_orgs = await velociraptor_services.collect_velociraptor_organizations() + logger.info(f"Collected Velociraptor Orgs: {collected_velociraptor_orgs}") + return VelociraptorOrganizations( + organizations=collected_velociraptor_orgs, + ) + async def fetch_velociraptor_agent(agent_name: str) -> VelociraptorAgent: """ @@ -298,8 +314,12 @@ async def sync_agents_velociraptor() -> SyncedAgentsResponse: :rtype: SyncedAgentsResponse """ agents_added_list: List[VelociraptorAgent] = [] + velo_orgs = await fetch_velociraptor_organizations() + logger.info(f"Collected Velociraptor Orgs: {velo_orgs}") + return None velociraptor_clients = await fetch_velociraptor_clients() + logger.info(f"Collected Velociraptor Clients: {velociraptor_clients}") velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else [] async with get_db_session() as session: # Create a new session here diff --git a/backend/app/agents/velociraptor/schema/agents.py b/backend/app/agents/velociraptor/schema/agents.py index 10027db3..b67d309c 100644 --- a/backend/app/agents/velociraptor/schema/agents.py +++ b/backend/app/agents/velociraptor/schema/agents.py @@ -53,3 +53,54 @@ class VelociraptorClient(BaseModel): class VelociraptorClients(BaseModel): clients: List[VelociraptorClient] + + +class Version(BaseModel): + name: str + version: str + commit: str + build_time: str + ci_build_url: str + compiler: str + + +class Installer(BaseModel): + service_name: str + install_path: str + service_description: Optional[str] = None + + +class LocalBuffer(BaseModel): + memory_size: int + disk_size: int + filename_linux: str + filename_windows: str + filename_darwin: str + + +class ClientConfig(BaseModel): + server_urls: List[str] + ca_certificate: str + nonce: str + writeback_darwin: str + writeback_linux: str + writeback_windows: str + tempdir_windows: str + max_poll: int + nanny_max_connection_delay: int + windows_installer: Installer + darwin_installer: Installer + version: Version + use_self_signed_ssl: bool + pinned_server_name: str + max_upload_size: int + local_buffer: LocalBuffer + + +class Organization(BaseModel): + Name: str + OrgId: str + _client_config: ClientConfig + +class VelociraptorOrganizations(BaseModel): + organizations: List[Organization] diff --git a/backend/app/agents/velociraptor/services/agents.py b/backend/app/agents/velociraptor/services/agents.py index c81c9943..d423bccd 100644 --- a/backend/app/agents/velociraptor/services/agents.py +++ b/backend/app/agents/velociraptor/services/agents.py @@ -21,7 +21,7 @@ def create_query(query: str) -> str: return query -async def collect_velociraptor_clients() -> list: +async def collect_velociraptor_clients(org_id: str) -> list: """ Collects all clients from Velociraptor. @@ -29,10 +29,29 @@ async def collect_velociraptor_clients() -> list: list: A list of all clients. """ velociraptor_service = await UniversalService.create("Velociraptor") + # query = create_query( + # "SELECT * FROM clients()", + # ) query = create_query( - "SELECT * FROM clients()", + f"SELECT * FROM query(org_id='{org_id}', query='SELECT * FROM clients()')", ) flow = velociraptor_service.execute_query(query) + logger.info(f"Successfully ran artifact collection on {flow}") + return flow["results"] + +async def collect_velociraptor_organizations() -> list: + """ + Collects all organizations from Velociraptor. + + Returns: + list: A list of all organizations. + """ + velociraptor_service = await UniversalService.create("Velociraptor") + query = create_query( + "SELECT * FROM orgs()", + ) + flow = velociraptor_service.execute_query(query) + logger.info(f"Successfully ran artifact collection on {flow}") return flow["results"] diff --git a/backend/app/connectors/velociraptor/services/artifacts.py b/backend/app/connectors/velociraptor/services/artifacts.py index ddab29f6..a0d1bf36 100644 --- a/backend/app/connectors/velociraptor/services/artifacts.py +++ b/backend/app/connectors/velociraptor/services/artifacts.py @@ -112,6 +112,7 @@ async def run_artifact_collection( """ velociraptor_service = await UniversalService.create("Velociraptor") try: + # ! Can specify org_id with org_id='OL680' ! # query = create_query( f"SELECT collect_client(client_id='{collect_artifact_body.velociraptor_id}', artifacts=['{collect_artifact_body.artifact_name}']) FROM scope()", ) From 8f620e0bca58e0613e6f1fb6d01a73c742748a5d Mon Sep 17 00:00:00 2001 From: Taylor Date: Thu, 27 Jun 2024 09:14:52 -0500 Subject: [PATCH 2/7] 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. --- backend/app/agents/services/sync.py | 98 ++++++++++++++--------------- 1 file changed, 49 insertions(+), 49 deletions(-) diff --git a/backend/app/agents/services/sync.py b/backend/app/agents/services/sync.py index ea27852a..810be721 100644 --- a/backend/app/agents/services/sync.py +++ b/backend/app/agents/services/sync.py @@ -316,55 +316,55 @@ async def sync_agents_velociraptor() -> SyncedAgentsResponse: agents_added_list: List[VelociraptorAgent] = [] velo_orgs = await fetch_velociraptor_organizations() logger.info(f"Collected Velociraptor Orgs: {velo_orgs}") - return None - - velociraptor_clients = await fetch_velociraptor_clients() - logger.info(f"Collected Velociraptor Clients: {velociraptor_clients}") - velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else [] - - async with get_db_session() as session: # Create a new session here - existing_agents_query = select(Agents) - result = await session.execute(existing_agents_query) - existing_agents = result.scalars().all() - - for agent in existing_agents: - logger.info(f"Collecting Velociraptor Agent for {agent.hostname}") - - try: - # Build the velociraptor_agent where the hostname or `client_id` is that equal to the `agents` - velociraptor_agent = next( - ( - client - for client in velociraptor_clients - if client.os_info.hostname == agent.hostname or client.client_id == agent.velociraptor_id - ), - None, - ) - # Convert Unix epoch timestamp to datetime - last_seen_at = datetime.fromtimestamp( - int(velociraptor_agent.last_seen_at) / 1e6, - ) # Divide by 1e6 to convert from microseconds to seconds - # Convert datetime to ISO 8601 format without fractional seconds - last_seen_at_iso = last_seen_at.replace(tzinfo=timezone.utc).isoformat(timespec="seconds") - velociraptor_agent = VelociraptorAgent( - velociraptor_id=velociraptor_agent.client_id, - velociraptor_last_seen=last_seen_at_iso, - velociraptor_agent_version=velociraptor_agent.agent_information.version, - ) - - except Exception as e: - logger.error( - f"Failed to collect Velociraptor Agent for {agent.hostname}: {e}", - ) - continue - - if velociraptor_agent: - # Update the agent with the Velociraptor client's details - await update_agent_with_velociraptor_in_db(session, agent, velociraptor_agent) - agents_added_list.append(velociraptor_agent) - - # Close the session - await session.close() + for org in velo_orgs.organizations: + + velociraptor_clients = await fetch_velociraptor_clients(org_id=org.OrgId) + logger.info(f"Collected Velociraptor Clients: {velociraptor_clients}") + velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else [] + + async with get_db_session() as session: # Create a new session here + existing_agents_query = select(Agents) + result = await session.execute(existing_agents_query) + existing_agents = result.scalars().all() + + for agent in existing_agents: + logger.info(f"Collecting Velociraptor Agent for {agent.hostname}") + + try: + # Build the velociraptor_agent where the hostname or `client_id` is that equal to the `agents` + velociraptor_agent = next( + ( + client + for client in velociraptor_clients + if client.os_info.hostname == agent.hostname or client.client_id == agent.velociraptor_id + ), + None, + ) + # Convert Unix epoch timestamp to datetime + last_seen_at = datetime.fromtimestamp( + int(velociraptor_agent.last_seen_at) / 1e6, + ) # Divide by 1e6 to convert from microseconds to seconds + # Convert datetime to ISO 8601 format without fractional seconds + last_seen_at_iso = last_seen_at.replace(tzinfo=timezone.utc).isoformat(timespec="seconds") + velociraptor_agent = VelociraptorAgent( + velociraptor_id=velociraptor_agent.client_id, + velociraptor_last_seen=last_seen_at_iso, + velociraptor_agent_version=velociraptor_agent.agent_information.version, + ) + + except Exception as e: + logger.error( + f"Failed to collect Velociraptor Agent for {agent.hostname}: {e}", + ) + continue + + if velociraptor_agent: + # Update the agent with the Velociraptor client's details + await update_agent_with_velociraptor_in_db(session, agent, velociraptor_agent) + agents_added_list.append(velociraptor_agent) + + # Close the session + await session.close() logger.info(f"Agents Added List: {agents_added_list}") return SyncedAgentsResponse( From 8d97c1a56259276510412bab12b8248681f5c4f2 Mon Sep 17 00:00:00 2001 From: Taylor Date: Thu, 27 Jun 2024 09:25:36 -0500 Subject: [PATCH 3/7] refactor: add velo org to agents table --- .../fed7739bd07c_add_velo_org_to_agents.py | 31 +++++++++++++++++++ backend/app/db/universal_models.py | 5 +++ 2 files changed, 36 insertions(+) create mode 100644 backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py diff --git a/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py b/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py new file mode 100644 index 00000000..2d6ae2b5 --- /dev/null +++ b/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py @@ -0,0 +1,31 @@ +"""Add Velo Org to Agents + +Revision ID: fed7739bd07c +Revises: 39c3aaec0084 +Create Date: 2024-06-27 09:22:59.354696 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql +import sqlmodel.sql.sqltypes + +# revision identifiers, used by Alembic. +revision: str = 'fed7739bd07c' +down_revision: Union[str, None] = '39c3aaec0084' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('agents', sa.Column('velociraptor_org', sqlmodel.sql.sqltypes.AutoString(length=256), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('agents', 'velociraptor_org') + # ### end Alembic commands ### diff --git a/backend/app/db/universal_models.py b/backend/app/db/universal_models.py index 5a3dea3c..f4046516 100644 --- a/backend/app/db/universal_models.py +++ b/backend/app/db/universal_models.py @@ -101,6 +101,8 @@ class Agents(SQLModel, table=True): velociraptor_agent_version: Optional[str] = Field(max_length=256) customer_code: Optional[str] = Field(foreign_key="customers.customer_code", max_length=256) quarantined: bool = Field(default=False) + velociraptor_org: Optional[str] = Field(max_length=256) + customer: Optional[Customers] = Relationship(back_populates="agents") @@ -129,6 +131,7 @@ def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code): if velociraptor_agent and velociraptor_agent.client_version else None, customer_code=customer_code, + velociraptor_org=velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None, ) @classmethod @@ -177,6 +180,7 @@ def update_from_model(self, wazuh_agent, velociraptor_agent, customer_code): velociraptor_agent.client_version if velociraptor_agent and velociraptor_agent.client_version else None ) self.customer_code = customer_code + self.velociraptor_org = velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None def update_wazuh_agent_from_model(self, wazuh_agent, customer_code): if wazuh_agent.agent_last_seen == "Unknown" or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00": @@ -209,6 +213,7 @@ def update_velociraptor_details(self, velociraptor_agent): velociraptor_agent.client_version if velociraptor_agent and velociraptor_agent.client_version else None ) logger.info(f"Updated with Velociraptor details: {self}") + self.velociraptor_org = velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None class LogEntry(SQLModel, table=True): From 222ce98a0febfeef737339c07a1f7ff145c4d36e Mon Sep 17 00:00:00 2001 From: Taylor Date: Thu, 27 Jun 2024 09:29:42 -0500 Subject: [PATCH 4/7] add velo org when collecting agents --- backend/app/agents/services/sync.py | 1 + backend/app/agents/velociraptor/schema/agents.py | 1 + 2 files changed, 2 insertions(+) diff --git a/backend/app/agents/services/sync.py b/backend/app/agents/services/sync.py index 810be721..c65ae546 100644 --- a/backend/app/agents/services/sync.py +++ b/backend/app/agents/services/sync.py @@ -350,6 +350,7 @@ async def sync_agents_velociraptor() -> SyncedAgentsResponse: velociraptor_id=velociraptor_agent.client_id, velociraptor_last_seen=last_seen_at_iso, velociraptor_agent_version=velociraptor_agent.agent_information.version, + velociraptor_org=org.OrgId, ) except Exception as e: diff --git a/backend/app/agents/velociraptor/schema/agents.py b/backend/app/agents/velociraptor/schema/agents.py index b67d309c..0fe14c6a 100644 --- a/backend/app/agents/velociraptor/schema/agents.py +++ b/backend/app/agents/velociraptor/schema/agents.py @@ -10,6 +10,7 @@ class VelociraptorAgent(BaseModel): client_id: Optional[str] = Field("n/a", alias="velociraptor_id") client_last_seen: str = Field(..., alias="velociraptor_last_seen") client_version: str = Field(..., alias="velociraptor_agent_version") + client_org: str = Field(..., alias="velociraptor_org") @property def client_last_seen_as_datetime(self): From 818b4ff25148c9fdf862b7d92faaf705ccc4e330 Mon Sep 17 00:00:00 2001 From: Taylor Date: Thu, 27 Jun 2024 10:44:05 -0500 Subject: [PATCH 5/7] 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. --- .../velociraptor/routes/artifacts.py | 49 +++++++++++++++++++ .../velociraptor/schema/artifacts.py | 1 + .../velociraptor/services/artifacts.py | 16 +++--- .../velociraptor/utils/universal.py | 1 + 4 files changed, 59 insertions(+), 8 deletions(-) diff --git a/backend/app/connectors/velociraptor/routes/artifacts.py b/backend/app/connectors/velociraptor/routes/artifacts.py index e1d112cd..cd287d67 100644 --- a/backend/app/connectors/velociraptor/routes/artifacts.py +++ b/backend/app/connectors/velociraptor/routes/artifacts.py @@ -125,6 +125,39 @@ async def get_velociraptor_id(session: AsyncSession, hostname: str) -> str: logger.info(f"velociraptor_id for hostname {hostname} is {agent.velociraptor_id}") return agent.velociraptor_id +async def get_velociraptor_org(session: AsyncSession, hostname: str) -> str: + """ + Retrieves the velociraptor_org associated with the given hostname. + + Args: + session (AsyncSession): The database session. + hostname (str): The hostname of the agent. + + Returns: + str: The velociraptor_org associated with the hostname. + + Raises: + HTTPException: If the agent with the given hostname is not found or if the velociraptor_org is not available. + """ + logger.info(f"Getting velociraptor_org from hostname {hostname}") + result = await session.execute(select(Agents).filter(Agents.hostname == hostname)) + agent = result.scalars().first() + + if not agent: + raise HTTPException( + status_code=404, + detail=f"Agent with hostname {hostname} not found", + ) + + if agent.velociraptor_org is None: + raise HTTPException( + status_code=404, + detail=f"Velociraptor ORG for hostname {hostname} is not available", + ) + + logger.info(f"velociraptor_org for hostname {hostname} is {agent.velociraptor_org}") + return agent.velociraptor_org + async def update_agent_quarantine_status( session: AsyncSession, @@ -321,6 +354,11 @@ async def collect_artifact( collect_artifact_body.hostname, ) + collect_artifact_body.velociraptor_org = await get_velociraptor_org( + session, + collect_artifact_body.hostname, + ) + # Assuming run_artifact_collection is an async function and takes a session as a parameter return await run_artifact_collection(collect_artifact_body) @@ -358,6 +396,11 @@ async def run_command( session, run_command_body.hostname, ) + + run_command_body.velociraptor_org = await get_velociraptor_org( + session, + run_command_body.hostname, + ) # Run the command return await run_remote_command(run_command_body) @@ -396,6 +439,12 @@ async def quarantine( session, quarantine_body.hostname, ) + + quarantine_body.velociraptor_org = await get_velociraptor_org( + session, + quarantine_body.hostname, + ) + # Quarantine the host quarantine_response = await quarantine_host(quarantine_body) diff --git a/backend/app/connectors/velociraptor/schema/artifacts.py b/backend/app/connectors/velociraptor/schema/artifacts.py index 0e4dda86..9c36439d 100644 --- a/backend/app/connectors/velociraptor/schema/artifacts.py +++ b/backend/app/connectors/velociraptor/schema/artifacts.py @@ -73,6 +73,7 @@ class QuarantineArtifactsEnum(str, Enum): class BaseBody(BaseModel): hostname: str = Field(..., description="Name of the client") velociraptor_id: Optional[str] = Field(None, description="Client ID of the client") + velociraptor_org: Optional[str] = Field(None, description="Organization of the client") class CollectArtifactBody(BaseBody): diff --git a/backend/app/connectors/velociraptor/services/artifacts.py b/backend/app/connectors/velociraptor/services/artifacts.py index a0d1bf36..c3f1102c 100644 --- a/backend/app/connectors/velociraptor/services/artifacts.py +++ b/backend/app/connectors/velociraptor/services/artifacts.py @@ -46,24 +46,24 @@ def get_artifact_key(analyzer_body: CollectArtifactBody) -> str: if action == "quarantine": return ( - f'collect_client(client_id="{analyzer_body.velociraptor_id}", ' + f'collect_client(org_id="{analyzer_body.velociraptor_org}", client_id="{analyzer_body.velociraptor_id}", ' f'artifacts=["{analyzer_body.artifact_name}"], ' f"spec=dict(`{analyzer_body.artifact_name}`=dict()))" ) elif action == "remove_quarantine": return ( - f'collect_client(client_id="{analyzer_body.velociraptor_id}", ' + f'collect_client(org_id="{analyzer_body.velociraptor_org}", client_id="{analyzer_body.velociraptor_id}", ' f'artifacts=["{analyzer_body.artifact_name}"], ' f'spec=dict(`{analyzer_body.artifact_name}`=dict(`RemovePolicy`="Y")))' ) elif command is not None: return ( - f"collect_client(client_id='{analyzer_body.velociraptor_id}', " + f"collect_client(org_id='{analyzer_body.velociraptor_org}', client_id='{analyzer_body.velociraptor_id}', " f"urgent=true, artifacts=['{analyzer_body.artifact_name}'], " f"env=dict(Command='{analyzer_body.command}'))" ) else: - return f"collect_client(client_id='{analyzer_body.velociraptor_id}', " f"artifacts=['{analyzer_body.artifact_name}'])" + return f"collect_client(org_id='{analyzer_body.velociraptor_org}', client_id='{analyzer_body.velociraptor_id}', " f"artifacts=['{analyzer_body.artifact_name}'])" async def get_artifacts() -> ArtifactsResponse: @@ -114,7 +114,7 @@ async def run_artifact_collection( try: # ! Can specify org_id with org_id='OL680' ! # query = create_query( - f"SELECT collect_client(client_id='{collect_artifact_body.velociraptor_id}', artifacts=['{collect_artifact_body.artifact_name}']) FROM scope()", + f"SELECT collect_client(org_id='{collect_artifact_body.velociraptor_org}', client_id='{collect_artifact_body.velociraptor_id}', artifacts=['{collect_artifact_body.artifact_name}']) FROM scope()", ) flow = velociraptor_service.execute_query(query) logger.info(f"Successfully ran artifact collection on {flow}") @@ -171,7 +171,7 @@ async def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResp logger.info(f"Running remote command on {run_command_body}") query = create_query( ( - f"SELECT collect_client(client_id='{run_command_body.velociraptor_id}', " + f"SELECT collect_client(org_id='{run_command_body.velociraptor_org}', client_id='{run_command_body.velociraptor_id}', " f"urgent=true, artifacts=['{run_command_body.artifact_name}'], " f"env=dict(Command='{run_command_body.command}')) " "FROM scope()" @@ -226,7 +226,7 @@ async def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse if quarantine_body.action == "quarantine": query = create_query( ( - f'SELECT collect_client(client_id="{quarantine_body.velociraptor_id}", ' + f'SELECT collect_client(org_id="{quarantine_body.velociraptor_org}", client_id="{quarantine_body.velociraptor_id}", ' f'artifacts=["{quarantine_body.artifact_name}"], ' f"spec=dict(`{quarantine_body.artifact_name}`=dict())) " "FROM scope()" @@ -235,7 +235,7 @@ async def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse else: query = create_query( ( - f'SELECT collect_client(client_id="{quarantine_body.velociraptor_id}", ' + f'SELECT collect_client(org_id="{quarantine_body.velociraptor_org}", client_id="{quarantine_body.velociraptor_id}", ' f'artifacts=["{quarantine_body.artifact_name}"], ' f'spec=dict(`{quarantine_body.artifact_name}`=dict(`RemovePolicy`="Y"))) ' "FROM scope()" diff --git a/backend/app/connectors/velociraptor/utils/universal.py b/backend/app/connectors/velociraptor/utils/universal.py index b13df0f5..b1393b33 100644 --- a/backend/app/connectors/velociraptor/utils/universal.py +++ b/backend/app/connectors/velociraptor/utils/universal.py @@ -213,6 +213,7 @@ def watch_flow_completion(self, flow_id: str): dict: A dictionary with the success status and a message. """ vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1" + #vql = f"SELECT * FROM query(org_id='OL680', query='SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1')" logger.info(f"Watching flow {flow_id} for completion") return self.execute_query(vql) From 1468a825c40b1aeee3f8e37913925e0fc736fcb3 Mon Sep 17 00:00:00 2001 From: Taylor Date: Thu, 27 Jun 2024 11:20:13 -0500 Subject: [PATCH 6/7] 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. --- backend/app/threat_intel/schema/socfortress.py | 7 ++++++- backend/app/threat_intel/services/socfortress.py | 1 + 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/backend/app/threat_intel/schema/socfortress.py b/backend/app/threat_intel/schema/socfortress.py index be642216..db4bfa40 100644 --- a/backend/app/threat_intel/schema/socfortress.py +++ b/backend/app/threat_intel/schema/socfortress.py @@ -1,8 +1,9 @@ from typing import List from typing import Optional +import os from pydantic import BaseModel -from pydantic import Field +from pydantic import Field, validator class SocfortressThreatIntelRequest(BaseModel): @@ -54,6 +55,10 @@ class SocfortressProcessNameAnalysisRequest(BaseModel): description="The process name to evaluate.", ) + @validator('process_name', pre=True) + def extract_filename(cls, v): + return os.path.basename(v) + class Path(BaseModel): directory: str diff --git a/backend/app/threat_intel/services/socfortress.py b/backend/app/threat_intel/services/socfortress.py index dd6cf085..b000db12 100644 --- a/backend/app/threat_intel/services/socfortress.py +++ b/backend/app/threat_intel/services/socfortress.py @@ -227,6 +227,7 @@ async def get_process_analysis_response( # Using .get() with default values data = response_data.get("data", {}) + logger.info(f"Data {data}") success = response_data.get("success", False) message = response_data.get("message", "No message provided") From 716b6ed043d4fd090389744cee5c97f721acc51e Mon Sep 17 00:00:00 2001 From: Taylor Date: Thu, 27 Jun 2024 11:29:14 -0500 Subject: [PATCH 7/7] precommit fixes --- .../fed7739bd07c_add_velo_org_to_agents.py | 15 ++++++++------- backend/app/agents/services/sync.py | 5 +++-- backend/app/agents/velociraptor/schema/agents.py | 1 + .../app/agents/velociraptor/services/agents.py | 1 + .../connectors/velociraptor/routes/artifacts.py | 1 + .../connectors/velociraptor/services/artifacts.py | 13 +++++++++++-- .../connectors/velociraptor/utils/universal.py | 2 +- backend/app/db/universal_models.py | 1 - backend/app/threat_intel/schema/socfortress.py | 7 ++++--- 9 files changed, 30 insertions(+), 16 deletions(-) diff --git a/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py b/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py index 2d6ae2b5..f65ce5fc 100644 --- a/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py +++ b/backend/alembic/versions/fed7739bd07c_add_velo_org_to_agents.py @@ -5,27 +5,28 @@ Create Date: 2024-06-27 09:22:59.354696 """ -from typing import Sequence, Union +from typing import Sequence +from typing import Union -from alembic import op import sqlalchemy as sa -from sqlalchemy.dialects import mysql import sqlmodel.sql.sqltypes +from alembic import op + # revision identifiers, used by Alembic. -revision: str = 'fed7739bd07c' -down_revision: Union[str, None] = '39c3aaec0084' +revision: str = "fed7739bd07c" +down_revision: Union[str, None] = "39c3aaec0084" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - op.add_column('agents', sa.Column('velociraptor_org', sqlmodel.sql.sqltypes.AutoString(length=256), nullable=True)) + op.add_column("agents", sa.Column("velociraptor_org", sqlmodel.sql.sqltypes.AutoString(length=256), nullable=True)) # ### end Alembic commands ### def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - op.drop_column('agents', 'velociraptor_org') + op.drop_column("agents", "velociraptor_org") # ### end Alembic commands ### diff --git a/backend/app/agents/services/sync.py b/backend/app/agents/services/sync.py index c65ae546..410d3f80 100644 --- a/backend/app/agents/services/sync.py +++ b/backend/app/agents/services/sync.py @@ -12,7 +12,8 @@ from app.agents.schema.agents import SyncedAgentsResponse from app.agents.schema.agents import SyncedWazuhAgent from app.agents.velociraptor.schema.agents import VelociraptorAgent -from app.agents.velociraptor.schema.agents import VelociraptorClients, VelociraptorOrganizations +from app.agents.velociraptor.schema.agents import VelociraptorClients +from app.agents.velociraptor.schema.agents import VelociraptorOrganizations from app.agents.wazuh.schema.agents import WazuhAgent from app.agents.wazuh.schema.agents import WazuhAgentsList from app.connectors.models import Connectors @@ -55,6 +56,7 @@ async def fetch_velociraptor_clients(org_id: str) -> VelociraptorClients: clients=collected_velociraptor_agents, ) + async def fetch_velociraptor_organizations() -> VelociraptorOrganizations: """ Fetches organizations from Velociraptor service. @@ -317,7 +319,6 @@ async def sync_agents_velociraptor() -> SyncedAgentsResponse: velo_orgs = await fetch_velociraptor_organizations() logger.info(f"Collected Velociraptor Orgs: {velo_orgs}") for org in velo_orgs.organizations: - velociraptor_clients = await fetch_velociraptor_clients(org_id=org.OrgId) logger.info(f"Collected Velociraptor Clients: {velociraptor_clients}") velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else [] diff --git a/backend/app/agents/velociraptor/schema/agents.py b/backend/app/agents/velociraptor/schema/agents.py index 0fe14c6a..92ecdd29 100644 --- a/backend/app/agents/velociraptor/schema/agents.py +++ b/backend/app/agents/velociraptor/schema/agents.py @@ -103,5 +103,6 @@ class Organization(BaseModel): OrgId: str _client_config: ClientConfig + class VelociraptorOrganizations(BaseModel): organizations: List[Organization] diff --git a/backend/app/agents/velociraptor/services/agents.py b/backend/app/agents/velociraptor/services/agents.py index d423bccd..87afdd1d 100644 --- a/backend/app/agents/velociraptor/services/agents.py +++ b/backend/app/agents/velociraptor/services/agents.py @@ -39,6 +39,7 @@ async def collect_velociraptor_clients(org_id: str) -> list: logger.info(f"Successfully ran artifact collection on {flow}") return flow["results"] + async def collect_velociraptor_organizations() -> list: """ Collects all organizations from Velociraptor. diff --git a/backend/app/connectors/velociraptor/routes/artifacts.py b/backend/app/connectors/velociraptor/routes/artifacts.py index cd287d67..c0692d3c 100644 --- a/backend/app/connectors/velociraptor/routes/artifacts.py +++ b/backend/app/connectors/velociraptor/routes/artifacts.py @@ -125,6 +125,7 @@ async def get_velociraptor_id(session: AsyncSession, hostname: str) -> str: logger.info(f"velociraptor_id for hostname {hostname} is {agent.velociraptor_id}") return agent.velociraptor_id + async def get_velociraptor_org(session: AsyncSession, hostname: str) -> str: """ Retrieves the velociraptor_org associated with the given hostname. diff --git a/backend/app/connectors/velociraptor/services/artifacts.py b/backend/app/connectors/velociraptor/services/artifacts.py index c3f1102c..842b2221 100644 --- a/backend/app/connectors/velociraptor/services/artifacts.py +++ b/backend/app/connectors/velociraptor/services/artifacts.py @@ -63,7 +63,10 @@ def get_artifact_key(analyzer_body: CollectArtifactBody) -> str: f"env=dict(Command='{analyzer_body.command}'))" ) else: - return f"collect_client(org_id='{analyzer_body.velociraptor_org}', client_id='{analyzer_body.velociraptor_id}', " f"artifacts=['{analyzer_body.artifact_name}'])" + return ( + f"collect_client(org_id='{analyzer_body.velociraptor_org}', client_id='{analyzer_body.velociraptor_id}', " + f"artifacts=['{analyzer_body.artifact_name}'])" + ) async def get_artifacts() -> ArtifactsResponse: @@ -114,7 +117,13 @@ async def run_artifact_collection( try: # ! Can specify org_id with org_id='OL680' ! # query = create_query( - f"SELECT collect_client(org_id='{collect_artifact_body.velociraptor_org}', client_id='{collect_artifact_body.velociraptor_id}', artifacts=['{collect_artifact_body.artifact_name}']) FROM scope()", + ( + f"SELECT collect_client(" + f"org_id='{collect_artifact_body.velociraptor_org}', " + f"client_id='{collect_artifact_body.velociraptor_id}', " + f"artifacts=['{collect_artifact_body.artifact_name}']) " + f"FROM scope()" + ), ) flow = velociraptor_service.execute_query(query) logger.info(f"Successfully ran artifact collection on {flow}") diff --git a/backend/app/connectors/velociraptor/utils/universal.py b/backend/app/connectors/velociraptor/utils/universal.py index b1393b33..95c032f4 100644 --- a/backend/app/connectors/velociraptor/utils/universal.py +++ b/backend/app/connectors/velociraptor/utils/universal.py @@ -213,7 +213,7 @@ def watch_flow_completion(self, flow_id: str): dict: A dictionary with the success status and a message. """ vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1" - #vql = f"SELECT * FROM query(org_id='OL680', query='SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1')" + # vql = f"SELECT * FROM query(org_id='OL680', query='SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1')" logger.info(f"Watching flow {flow_id} for completion") return self.execute_query(vql) diff --git a/backend/app/db/universal_models.py b/backend/app/db/universal_models.py index f4046516..749dc9de 100644 --- a/backend/app/db/universal_models.py +++ b/backend/app/db/universal_models.py @@ -103,7 +103,6 @@ class Agents(SQLModel, table=True): quarantined: bool = Field(default=False) velociraptor_org: Optional[str] = Field(max_length=256) - customer: Optional[Customers] = Relationship(back_populates="agents") @classmethod diff --git a/backend/app/threat_intel/schema/socfortress.py b/backend/app/threat_intel/schema/socfortress.py index db4bfa40..9e8bad12 100644 --- a/backend/app/threat_intel/schema/socfortress.py +++ b/backend/app/threat_intel/schema/socfortress.py @@ -1,9 +1,10 @@ +import os from typing import List from typing import Optional -import os from pydantic import BaseModel -from pydantic import Field, validator +from pydantic import Field +from pydantic import validator class SocfortressThreatIntelRequest(BaseModel): @@ -55,7 +56,7 @@ class SocfortressProcessNameAnalysisRequest(BaseModel): description="The process name to evaluate.", ) - @validator('process_name', pre=True) + @validator("process_name", pre=True) def extract_filename(cls, v): return os.path.basename(v)