| 1 | from datetime import datetime |
| 2 | |
| 3 | from fastapi import HTTPException |
| 4 | from loguru import logger |
| 5 | |
| 6 | from app.agents.schema.agents import AgentModifyResponse |
| 7 | from app.agents.velociraptor.schema.agents import VelociraptorAgent |
| 8 | from app.connectors.velociraptor.utils.universal import UniversalService |
| 9 | |
| 10 | |
| 11 | def create_query(query: str) -> str: |
| 12 | """ |
| 13 | Create a query string. |
| 14 | |
| 15 | Args: |
| 16 | query (str): The query to be executed. |
| 17 | |
| 18 | Returns: |
| 19 | str: The created query string. |
| 20 | """ |
| 21 | return query |
| 22 | |
| 23 | |
| 24 | async def collect_velociraptor_clients(org_id: str) -> list: |
| 25 | """ |
| 26 | Collects all clients from Velociraptor. |
| 27 | |
| 28 | Returns: |
| 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( |
| 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 | |
| 59 | async def collect_velociraptor_agent(agent_name: str) -> VelociraptorAgent: |
| 60 | """ |
| 61 | Retrieves the client ID, last_seen_at and client version based on the agent name from Velociraptor. |
| 62 | |
| 63 | Args: |
| 64 | agent_name (str): The name of the agent. |
| 65 | |
| 66 | Returns: |
| 67 | str: The client ID if found, None otherwise. |
| 68 | str: The last seen at timestamp if found, Default timsetamp otherwise. |
| 69 | """ |
| 70 | logger.info(f"Collecting agent {agent_name} from Velociraptor") |
| 71 | velociraptor_service = await UniversalService.create("Velociraptor") |
| 72 | try: |
| 73 | client_id = await velociraptor_service.get_client_id(agent_name) |
| 74 | client_id = client_id["results"][0]["client_id"] |
| 75 | except (KeyError, IndexError, TypeError) as e: |
| 76 | logger.error(f"Failed to get client ID for {agent_name}. Error: {e}") |
| 77 | return VelociraptorAgent( |
| 78 | client_id="Unknown", |
| 79 | client_last_seen="Unknown", |
| 80 | client_version="Unknown", |
| 81 | ) |
| 82 | |
| 83 | try: |
| 84 | vql_last_seen_at = f"select last_seen_at from clients(search='host:{agent_name}')" |
| 85 | last_seen_at = await velociraptor_service._get_last_seen_timestamp( |
| 86 | vql_last_seen_at, |
| 87 | ) |
| 88 | client_last_seen = datetime.fromtimestamp( |
| 89 | int(last_seen_at) / 1000000, |
| 90 | ).strftime( |
| 91 | "%Y-%m-%dT%H:%M:%S+00:00", |
| 92 | ) # Converting to string format |
| 93 | except Exception as e: |
| 94 | logger.error( |
| 95 | f"Failed to get or convert last seen at for {agent_name}. Error: {e}", |
| 96 | ) |
| 97 | client_last_seen = "1970-01-01T00:00:00+00:00" |
| 98 | |
| 99 | try: |
| 100 | vql_client_version = f"select * from clients(search='host:{agent_name}')" |
| 101 | # client_version = UniversalService()._get_client_version(vql_client_version) |
| 102 | client_version = await velociraptor_service._get_client_version( |
| 103 | vql_client_version, |
| 104 | ) |
| 105 | except Exception as e: |
| 106 | logger.error(f"Failed to get client version for {agent_name}. Error: {e}") |
| 107 | client_version = "Unknown" |
| 108 | |
| 109 | return VelociraptorAgent( |
| 110 | client_id=client_id, |
| 111 | client_last_seen=client_last_seen, |
| 112 | client_version=client_version, |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | async def collect_velociraptor_agent_via_client_id(client_id: str) -> VelociraptorAgent: |
| 117 | """ |
| 118 | Retrieves the client ID, last_seen_at and client version based on the agent name from Velociraptor. |
| 119 | |
| 120 | Args: |
| 121 | agent_name (str): The name of the agent. |
| 122 | |
| 123 | Returns: |
| 124 | str: The client ID if found, None otherwise. |
| 125 | str: The last seen at timestamp if found, Default timsetamp otherwise. |
| 126 | """ |
| 127 | logger.info(f"Collecting agent {client_id} from Velociraptor") |
| 128 | velociraptor_service = await UniversalService.create("Velociraptor") |
| 129 | try: |
| 130 | client_id = await velociraptor_service.get_client_id_via_client_id(client_id) |
| 131 | client_id = client_id["results"][0]["client_id"] |
| 132 | except (KeyError, IndexError, TypeError) as e: |
| 133 | logger.error(f"Failed to get client ID for {client_id}. Error: {e}") |
| 134 | return VelociraptorAgent( |
| 135 | client_id="Unknown", |
| 136 | client_last_seen="Unknown", |
| 137 | client_version="Unknown", |
| 138 | ) |
| 139 | |
| 140 | try: |
| 141 | vql_last_seen_at = f"select last_seen_at from clients(search='host:{client_id}')" |
| 142 | last_seen_at = await velociraptor_service._get_last_seen_timestamp( |
| 143 | vql_last_seen_at, |
| 144 | ) |
| 145 | client_last_seen = datetime.fromtimestamp( |
| 146 | int(last_seen_at) / 1000000, |
| 147 | ).strftime( |
| 148 | "%Y-%m-%dT%H:%M:%S+00:00", |
| 149 | ) # Converting to string format |
| 150 | except Exception as e: |
| 151 | logger.error( |
| 152 | f"Failed to get or convert last seen at for {client_id}. Error: {e}", |
| 153 | ) |
| 154 | client_last_seen = "1970-01-01T00:00:00+00:00" |
| 155 | |
| 156 | try: |
| 157 | vql_client_version = f"select * from clients(search='host:{client_id}')" |
| 158 | # client_version = UniversalService()._get_client_version(vql_client_version) |
| 159 | client_version = await velociraptor_service._get_client_version( |
| 160 | vql_client_version, |
| 161 | ) |
| 162 | except Exception as e: |
| 163 | logger.error(f"Failed to get client version for {client_id}. Error: {e}") |
| 164 | client_version = "Unknown" |
| 165 | |
| 166 | return VelociraptorAgent( |
| 167 | client_id=client_id, |
| 168 | client_last_seen=client_last_seen, |
| 169 | client_version=client_version, |
| 170 | ) |
| 171 | |
| 172 | |
| 173 | def execute_query(universal_service, query: str) -> dict: |
| 174 | """ |
| 175 | Executes a query using the provided universal service. |
| 176 | |
| 177 | Args: |
| 178 | universal_service: The universal service to use for executing the query. |
| 179 | query: The query to execute. |
| 180 | |
| 181 | Returns: |
| 182 | A dictionary containing the result of the query execution. |
| 183 | """ |
| 184 | flow = universal_service.execute_query(query) |
| 185 | logger.info(f"Successfully ran artifact collection on {flow}") |
| 186 | return flow |
| 187 | |
| 188 | |
| 189 | def check_flow_success(flow: dict, client_id: str) -> dict: |
| 190 | """ |
| 191 | Checks the success status of a flow and returns a dictionary with a message and success status. |
| 192 | |
| 193 | Args: |
| 194 | flow (dict): The flow dictionary containing the success status. |
| 195 | client_id (str): The ID of the velociraptor client. |
| 196 | |
| 197 | Returns: |
| 198 | dict: A dictionary with a message and success status. |
| 199 | |
| 200 | Raises: |
| 201 | Exception: If there is an error while deleting the velociraptor client. |
| 202 | """ |
| 203 | if flow["success"]: |
| 204 | logger.info(f"Successfully deleted velociraptor client {client_id}") |
| 205 | return { |
| 206 | "message": f"Successfully deleted velociraptor client {client_id}", |
| 207 | "success": True, |
| 208 | } |
| 209 | else: |
| 210 | logger.error(f"Failed to delete velociraptor client {client_id}") |
| 211 | return handle_exception( |
| 212 | e="Failed to delete velociraptor client", |
| 213 | client_id=client_id, |
| 214 | ) |
| 215 | |
| 216 | |
| 217 | def handle_exception(e: Exception, client_id: str) -> dict: |
| 218 | """ |
| 219 | Handles exceptions that occur during the deletion of a Velociraptor client. |
| 220 | |
| 221 | Args: |
| 222 | e (Exception): The exception that occurred. |
| 223 | client_id (str): The ID of the client being deleted. |
| 224 | |
| 225 | Raises: |
| 226 | HTTPException: An HTTP exception with a status code of 500 and a detailed error message. |
| 227 | |
| 228 | Returns: |
| 229 | dict: An empty dictionary. |
| 230 | """ |
| 231 | logger.error(f"Failed to delete client {client_id}: {e}") |
| 232 | raise HTTPException( |
| 233 | status_code=500, |
| 234 | detail=f"Failed to delete Velociraptor client {client_id}: {e}", |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | async def delete_agent_velociraptor(client_id: str) -> AgentModifyResponse: |
| 239 | """ |
| 240 | Deletes an agent with the specified client ID. |
| 241 | |
| 242 | Args: |
| 243 | client_id (str): The ID of the client to delete. |
| 244 | |
| 245 | Returns: |
| 246 | AgentModifyResponse: An object representing the result of the agent deletion operation. |
| 247 | """ |
| 248 | try: |
| 249 | await delete_client(client_id=client_id) |
| 250 | return AgentModifyResponse(success=True, message="Agent deleted successfully") |
| 251 | except Exception as e: |
| 252 | return handle_exception(e, client_id) |
| 253 | |
| 254 | |
| 255 | async def delete_client(client_id: str) -> dict: |
| 256 | """ |
| 257 | Deletes a client with the specified client ID. |
| 258 | |
| 259 | Args: |
| 260 | client_id (str): The ID of the client to be deleted. |
| 261 | |
| 262 | Returns: |
| 263 | dict: A dictionary containing the result of the deletion operation. |
| 264 | """ |
| 265 | universal_service = await UniversalService.create("Velociraptor") |
| 266 | try: |
| 267 | query = create_query( |
| 268 | f"SELECT collect_client(client_id='server', artifacts=['Server.Utils.DeleteClient'], env=dict(ClientIdList='{client_id}',ReallyDoIt='Y')) FROM scope()", |
| 269 | ) |
| 270 | flow = execute_query(universal_service, query) |
| 271 | return check_flow_success(flow, client_id) |
| 272 | except Exception as e: |
| 273 | return handle_exception(e, client_id) |