| 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 fastapi import HTTPException |
| 9 | from loguru import logger |
| 10 | from pyvelociraptor import api_pb2 |
| 11 | from pyvelociraptor import api_pb2_grpc |
| 12 | |
| 13 | from app.connectors.utils import get_connector_info_from_db |
| 14 | from app.db.db_session import AsyncSessionLocal |
| 15 | from app.db.db_session import get_db_session |
| 16 | |
| 17 | |
| 18 | async def verify_velociraptor_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: |
| 19 | """ |
| 20 | Verifies the connection to Velociraptor service. |
| 21 | |
| 22 | Returns: |
| 23 | dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful. |
| 24 | """ |
| 25 | try: |
| 26 | connector_api_key = attributes["connector_api_key"] |
| 27 | |
| 28 | with open(connector_api_key, "r") as f: |
| 29 | f.read() |
| 30 | |
| 31 | try: |
| 32 | config = pyvelociraptor.LoadConfigFile(connector_api_key) |
| 33 | creds = grpc.ssl_channel_credentials( |
| 34 | root_certificates=config["ca_certificate"].encode("utf8"), |
| 35 | private_key=config["client_private_key"].encode("utf8"), |
| 36 | certificate_chain=config["client_cert"].encode("utf8"), |
| 37 | ) |
| 38 | |
| 39 | options = (("grpc.ssl_target_name_override", "VelociraptorServer"),) |
| 40 | |
| 41 | with grpc.secure_channel( |
| 42 | config["api_connection_string"], |
| 43 | creds, |
| 44 | options, |
| 45 | ) as channel: |
| 46 | stub = api_pb2_grpc.APIStub(channel) |
| 47 | client_query = "SELECT * FROM info()" |
| 48 | |
| 49 | client_request = api_pb2.VQLCollectorArgs( |
| 50 | max_wait=60, |
| 51 | Query=[ |
| 52 | api_pb2.VQLRequest( |
| 53 | Name="ClientQuery", |
| 54 | VQL=client_query, |
| 55 | ), |
| 56 | ], |
| 57 | ) |
| 58 | |
| 59 | r = [] |
| 60 | for response in stub.Query(client_request): |
| 61 | if response.Response: |
| 62 | r = r + json.loads(response.Response) |
| 63 | return { |
| 64 | "connectionSuccessful": True, |
| 65 | "message": "Connection to Velociraptor successful", |
| 66 | } |
| 67 | except Exception as e: |
| 68 | logger.error(f"Failed to verify connection to Velociraptor: {e}") |
| 69 | return { |
| 70 | "connectionSuccessful": False, |
| 71 | "message": f"Failed to verify connection to Velociraptor: {e}", |
| 72 | } |
| 73 | except Exception as e: |
| 74 | logger.error(f"Failed to get connector_api_key from the database: {e}") |
| 75 | return { |
| 76 | "connectionSuccessful": False, |
| 77 | "message": f"Failed to get connector_api_key from the database: {e}", |
| 78 | } |
| 79 | |
| 80 | |
| 81 | async def verify_velociraptor_connection(connector_name: str) -> str: |
| 82 | """ |
| 83 | Verifies the connection to Velociraptor service. |
| 84 | """ |
| 85 | logger.info( |
| 86 | f"Verifying the Velociraptor connection for connector: {connector_name}", |
| 87 | ) |
| 88 | async with get_db_session() as session: # This will correctly enter the context manager |
| 89 | attributes = await get_connector_info_from_db(connector_name, session) |
| 90 | if attributes is None: |
| 91 | logger.error("No Velociraptor connector found in the database") |
| 92 | return None |
| 93 | return await verify_velociraptor_credentials(attributes) |
| 94 | |
| 95 | |
| 96 | class UniversalService: |
| 97 | """ |
| 98 | A service class that encapsulates the logic for polling messages from Velociraptor. |
| 99 | """ |
| 100 | |
| 101 | # ! Modify this to use AsyncSessionLocal Begin - ALSO SEE BELOW CLASS METHOD |
| 102 | def __init__(self) -> None: |
| 103 | self.connector_api_key = None |
| 104 | self.config = None |
| 105 | |
| 106 | async def setup_velociraptor_connector(self, connector_name: str): |
| 107 | async with AsyncSessionLocal() as session: |
| 108 | attributes = await get_connector_info_from_db(connector_name, session) |
| 109 | if attributes is None: |
| 110 | logger.error("No Velociraptor connector found in the database") |
| 111 | return None |
| 112 | self.connector_api_key = attributes["connector_api_key"] |
| 113 | self.config = pyvelociraptor.LoadConfigFile(self.connector_api_key) |
| 114 | |
| 115 | # ! Modify this to use AsyncSessionLocal End |
| 116 | |
| 117 | async def setup_grpc_channel_and_stub(self): |
| 118 | """ |
| 119 | Sets up the gRPC channel and stub for Velociraptor. |
| 120 | """ |
| 121 | creds = grpc.ssl_channel_credentials( |
| 122 | root_certificates=self.config["ca_certificate"].encode("utf8"), |
| 123 | private_key=self.config["client_private_key"].encode("utf8"), |
| 124 | certificate_chain=self.config["client_cert"].encode("utf8"), |
| 125 | ) |
| 126 | options = (("grpc.ssl_target_name_override", "VelociraptorServer"),) |
| 127 | self.channel = grpc.secure_channel( |
| 128 | self.config["api_connection_string"], |
| 129 | creds, |
| 130 | options, |
| 131 | ) |
| 132 | self.stub = api_pb2_grpc.APIStub(self.channel) |
| 133 | |
| 134 | # ! Modify this to use AsyncSessionLocal Begin |
| 135 | @classmethod |
| 136 | async def create(cls, connector_name: str): |
| 137 | instance = cls() |
| 138 | await instance.setup_velociraptor_connector(connector_name) |
| 139 | await instance.setup_grpc_channel_and_stub() |
| 140 | return instance |
| 141 | |
| 142 | # ! Modify this to use AsyncSessionLocal End |
| 143 | |
| 144 | def create_vql_request(self, vql: str, org_id: str = "root"): |
| 145 | """ |
| 146 | Creates a VQLCollectorArgs object with given VQL query. |
| 147 | |
| 148 | Args: |
| 149 | vql (str): The VQL query. |
| 150 | |
| 151 | Returns: |
| 152 | VQLCollectorArgs: The VQLCollectorArgs object with given VQL query. |
| 153 | """ |
| 154 | return api_pb2.VQLCollectorArgs( |
| 155 | max_wait=1, |
| 156 | org_id=org_id, |
| 157 | Query=[ |
| 158 | api_pb2.VQLRequest( |
| 159 | Name="VQLRequest", |
| 160 | VQL=vql, |
| 161 | ), |
| 162 | ], |
| 163 | ) |
| 164 | |
| 165 | def execute_query(self, vql: str, org_id: str = "root"): |
| 166 | """ |
| 167 | Executes a VQL query and returns the results. |
| 168 | |
| 169 | Args: |
| 170 | vql (str): The VQL query to be executed. |
| 171 | |
| 172 | Returns: |
| 173 | dict: A dictionary with the success status, a message, and potentially the results. |
| 174 | """ |
| 175 | logger.info(f"Executing query: {vql}") |
| 176 | |
| 177 | client_request = self.create_vql_request(vql, org_id) |
| 178 | |
| 179 | try: |
| 180 | results = [] |
| 181 | for response in self.stub.Query(client_request, timeout=30): |
| 182 | if response.Response: |
| 183 | results += json.loads(response.Response) |
| 184 | elif response.log: |
| 185 | logger.info(f"Log: {response.log}") |
| 186 | |
| 187 | return { |
| 188 | "success": True, |
| 189 | "message": "Successfully executed query", |
| 190 | "results": results, |
| 191 | } |
| 192 | except grpc.RpcError as e: # Catch gRPC-specific errors |
| 193 | if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED: |
| 194 | logger.error("Failed to execute query due to timeout.") |
| 195 | raise HTTPException( |
| 196 | status_code=500, |
| 197 | detail="Failed to execute query due to timeout. Make sure the Velocraptor server has stopped this artifact collection.", |
| 198 | ) |
| 199 | else: |
| 200 | logger.error(f"Failed to execute query: {e}") |
| 201 | raise HTTPException( |
| 202 | status_code=500, |
| 203 | detail=f"Failed to execute query: {e.details()}", |
| 204 | ) |
| 205 | except Exception as e: |
| 206 | logger.error(f"Failed to execute query: {e}") |
| 207 | raise HTTPException(status_code=500, detail=f"Failed to execute query: {e}") |
| 208 | |
| 209 | def watch_flow_completion(self, flow_id: str, org_id: str = "root"): |
| 210 | """ |
| 211 | Watch for the completion of a flow. |
| 212 | |
| 213 | Args: |
| 214 | flow_id (str): The ID of the flow. |
| 215 | |
| 216 | Returns: |
| 217 | dict: A dictionary with the success status and a message. |
| 218 | """ |
| 219 | vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1" |
| 220 | logger.info(f"Watching flow {flow_id} for completion") |
| 221 | return self.execute_query(vql, org_id) |
| 222 | |
| 223 | def read_collection_results( |
| 224 | self, |
| 225 | client_id: str, |
| 226 | flow_id: str, |
| 227 | org_id: str = "root", |
| 228 | artifact: str = "Generic.Client.Info/BasicInformation", |
| 229 | ): |
| 230 | """ |
| 231 | Read the results of a collection. |
| 232 | |
| 233 | Args: |
| 234 | client_id (str): The client ID. |
| 235 | flow_id (str): The ID of the flow. |
| 236 | artifact (str, optional): The artifact. Defaults to 'Generic.Client.Info/BasicInformation'. |
| 237 | |
| 238 | Returns: |
| 239 | dict: A dictionary with the success status, a message, and potentially the results. |
| 240 | """ |
| 241 | vql = f"SELECT * FROM source(client_id='{client_id}', flow_id='{flow_id}', artifact='{artifact}')" |
| 242 | return self.execute_query(vql, org_id) |
| 243 | |
| 244 | async def get_client_id(self, client_name: str): |
| 245 | """ |
| 246 | Get the client_id associated with a given client_name. |
| 247 | |
| 248 | Args: |
| 249 | client_name (str): The asset name to search for. |
| 250 | |
| 251 | Returns: |
| 252 | dict: A dictionary with the success status, a message, and potentially the client_id. |
| 253 | """ |
| 254 | # Formulate queries |
| 255 | try: |
| 256 | vql_client_id = f"select client_id,os_info from clients(search='host:{client_name}')" |
| 257 | vql_last_seen_at = f"select last_seen_at from clients(search='host:{client_name}')" |
| 258 | |
| 259 | # Get the last seen timestamp |
| 260 | logger.info(f"Getting last seen at timestamp for {client_name}") |
| 261 | |
| 262 | last_seen_at = await self._get_last_seen_timestamp(vql_last_seen_at) |
| 263 | |
| 264 | logger.info(f"Last seen at timestamp for {client_name}: {last_seen_at}") |
| 265 | |
| 266 | # if last_seen_at is longer than 30 seconds from now, return False |
| 267 | if await self._is_offline(last_seen_at): |
| 268 | return self.execute_query(vql_client_id) |
| 269 | |
| 270 | return self.execute_query(vql_client_id) |
| 271 | except Exception as e: |
| 272 | return { |
| 273 | "success": False, |
| 274 | "message": f"Failed to get Client ID for {client_name}: {e}", |
| 275 | "results": [{"client_id": None}], |
| 276 | } |
| 277 | |
| 278 | async def get_client_id_via_client_id(self, client_id: str): |
| 279 | """ |
| 280 | Get the client_id associated with a given client_id. |
| 281 | |
| 282 | Args: |
| 283 | client_id (str): The client_id to search for. |
| 284 | |
| 285 | Returns: |
| 286 | dict: A dictionary with the success status, a message, and potentially the client_id. |
| 287 | """ |
| 288 | # Formulate queries |
| 289 | try: |
| 290 | vql_client_id = f"select client_id,os_info from clients(search='client_id:{client_id}')" |
| 291 | vql_last_seen_at = f"select last_seen_at from clients(search='client_id:{client_id}')" |
| 292 | |
| 293 | # Get the last seen timestamp |
| 294 | logger.info(f"Getting last seen at timestamp for {client_id}") |
| 295 | |
| 296 | last_seen_at = await self._get_last_seen_timestamp(vql_last_seen_at) |
| 297 | |
| 298 | logger.info(f"Last seen at timestamp for {client_id}: {last_seen_at}") |
| 299 | |
| 300 | # if last_seen_at is longer than 30 seconds from now, return False |
| 301 | if await self._is_offline(last_seen_at): |
| 302 | return self.execute_query(vql_client_id) |
| 303 | |
| 304 | return self.execute_query(vql_client_id) |
| 305 | except Exception as e: |
| 306 | return { |
| 307 | "success": False, |
| 308 | "message": f"Failed to get Client ID for {client_id}: {e}", |
| 309 | "results": [{"client_id": None}], |
| 310 | } |
| 311 | |
| 312 | async def _get_last_seen_timestamp(self, vql: str): |
| 313 | """ |
| 314 | Executes the VQL query and returns the last_seen_at timestamp. |
| 315 | |
| 316 | Args: |
| 317 | vql (str): The VQL query. |
| 318 | |
| 319 | Returns: |
| 320 | float: The last_seen_at timestamp. |
| 321 | """ |
| 322 | return self.execute_query(vql)["results"][0]["last_seen_at"] |
| 323 | |
| 324 | async def _get_client_version(self, vql: str): |
| 325 | """ |
| 326 | Executes the VQL query and returns the `agent_information``version` field |
| 327 | |
| 328 | Args: |
| 329 | vql (str): The VQL query. |
| 330 | |
| 331 | Returns: |
| 332 | str: The client version. |
| 333 | """ |
| 334 | return self.execute_query(vql)["results"][0]["agent_information"]["version"] |
| 335 | |
| 336 | async def _get_server_version(self, vql: str): |
| 337 | """ |
| 338 | Executes the VQL query and returns the velociraptor server version. |
| 339 | |
| 340 | Args: |
| 341 | vql (str): The VQL query. |
| 342 | |
| 343 | Returns: |
| 344 | str: The server version. |
| 345 | """ |
| 346 | try: |
| 347 | return self.execute_query(vql)["results"][0]["version"]["version"] |
| 348 | except IndexError as e: |
| 349 | raise HTTPException( |
| 350 | status_code=500, |
| 351 | detail=f"Failed to get server version: {e}", |
| 352 | ) |
| 353 | |
| 354 | async def _is_offline(self, last_seen_at: float): |
| 355 | """ |
| 356 | Determines if the client is offline based on the last_seen_at timestamp. |
| 357 | |
| 358 | Args: |
| 359 | last_seen_at (float): The last_seen_at timestamp. |
| 360 | |
| 361 | Returns: |
| 362 | bool: True if the client is offline, False otherwise. |
| 363 | """ |
| 364 | return (datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000)).total_seconds() > 30 |