Create universal.py
taylor_socfortress committed
Jul 10, 2023 at 16:42 UTC
22b3c84ff5f97db7eb5859f6e051a00fd4c91d5f
1 file changed
+104
backend/app/services/DFIR_IRIS/universal.py
new
+104
@@ -0,0 +1,104 @@
1
+from app.models.agents import (
2
+ AgentMetadata,
3
+ agent_metadata_schema,
4
+ agent_metadatas_schema,
5
+)
6
+from typing import Dict, List
7
+from app import db
8
+from datetime import datetime
9
+import requests
10
+from loguru import logger
11
+from elasticsearch7 import Elasticsearch
12
+from app.models.connectors import connector_factory, Connector
13
+import dfir_iris_client
14
+from dfir_iris_client.session import ClientSession
15
+from typing import Any
16
+from typing import Dict
17
+from typing import Optional
18
+from typing import Set
19
+from typing import Tuple
20
+from dfir_iris_client.case import Case
21
+from dfir_iris_client.helper.utils import assert_api_resp
22
+from dfir_iris_client.helper.utils import get_data_from_resp
23
+
24
+
25
+class UniversalService:
26
+ """
27
+ A service class that encapsulates the logic for polling messages from DFIR-IRIS.
28
+ """
29
+
30
+ def __init__(self, connector_name: str) -> None:
31
+ self.connector_url, self.connector_api_key = self.collect_iris_details(connector_name)
32
+
33
+ def collect_iris_details(self, connector_name: str):
34
+ """
35
+ Collects the details of the DFIR-IRIS connector.
36
+
37
+ Args:
38
+ connector_name (str): The name of the DFIR-IRIS connector.
39
+
40
+ Returns:
41
+ tuple: A tuple containing the connection URL, and api key.
42
+ """
43
+ connector_instance = connector_factory.create(connector_name, connector_name)
44
+ connection_successful = connector_instance.verify_connection()
45
+ if connection_successful:
46
+ connection_details = Connector.get_connector_info_from_db(connector_name)
47
+ return (
48
+ connection_details.get("connector_url"),
49
+ connection_details.get("connector_api_key"),
50
+ )
51
+ else:
52
+ return None, None
53
+
54
+ def create_session(self) -> Optional[ClientSession]:
55
+ """
56
+ Create a session with DFIR-IRIS.
57
+
58
+ This method creates a session with DFIR-IRIS and returns the session object.
59
+ If a session cannot be established, an error is logged and None is returned.
60
+
61
+ Returns:
62
+ session: A session object for DFIR-IRIS.
63
+ """
64
+ try:
65
+ logger.info("Creating session with DFIR-IRIS.")
66
+ session = ClientSession(
67
+ host=self.connector_url,
68
+ apikey=self.connector_api_key,
69
+ agent="iris-client",
70
+ ssl_verify=False,
71
+ timeout=120,
72
+ proxy=None,
73
+ )
74
+ logger.info("Session created.")
75
+ return {"success": True, "session": session}
76
+ except Exception as e:
77
+ logger.error(f"Error creating session with DFIR-IRIS: {e}")
78
+ return {
79
+ "success": False,
80
+ "message": "Connection to DFIR-IRIS unsuccessful.",
81
+ }
82
+
83
+ def fetch_and_parse_data(self, session, action, *args):
84
+ """
85
+ General method to fetch and parse data from DFIR-IRIS.
86
+
87
+ Args:
88
+ session: ClientSession object.
89
+ action: callable, the action to be performed (e.g., list_cases or get_case)
90
+ args: arguments for the action callable
91
+
92
+ Returns:
93
+ dict: A dictionary containing the data and a success status.
94
+ """
95
+ try:
96
+ logger.info(f"Executing {action.__name__}... on args: {args}")
97
+ status = action(*args)
98
+ assert_api_resp(status, soft_fail=False)
99
+ data = get_data_from_resp(status)
100
+ logger.info(f"Successfully executed {action.__name__}")
101
+ return {"success": True, "data": data}
102
+ except Exception as err:
103
+ logger.error(f"Failed to execute {action.__name__}: {err}")
104
+ return {"success": False}