Made changes to data_svc store

app/api/rest_api.py

@@ -70,6 +70,10 @@ async def landing(self, request):
 
     @check_authorization
     async def rest_core(self, request):
+        print(" ")
+        print("rest_core")
+        print("request: %s"%request)
+        print(" ")
         try:
             access = dict(access=tuple(await self.auth_svc.get_permissions(request)))
             data = dict(await request.json())

app/api/v2/handlers/ability_api.py

@@ -60,7 +60,9 @@ async def get_ability_names_from_database(self):
             return ability_names
         except Exception as e:
             # Handle the error
+            print("ERRRO IN get_ability_names_from_database")
             print(f"An error occurred while connecting to the database: {e}")
+
         finally:
             if 'session' in locals():
                 session.close()

app/api/v2/handlers/agent_api.py

@@ -67,6 +67,7 @@ async def get_agent_names_from_database(self):
             return agent_names
         except Exception as e:
             # Handle the error
+            print(" ERROR in get_agent_names_from_database")
             print(f"An error occurred while connecting to the database: {e}")
         finally:
             if 'session' in locals():

app/service/data_svc.py

@@ -24,6 +24,10 @@
 from app.utility.base_service import BaseService
 
 from neo4j import GraphDatabase
+from enum import Enum
+import enum
+import json
+
 
 MIN_MODULE_LEN = 1
 
@@ -50,99 +54,153 @@ class DataService(DataServiceInterface, BaseService):
 
     def __init__(self):
         self.log = self.add_service('data_svc', self)
-        # self.schema = dict(agents=[], planners=[], adversaries=[], abilities=[], sources=[], operations=[],
-        #                    schedules=[], plugins=[], obfuscators=[], objectives=[], data_encoders=[])
-        # self.ram = copy.deepcopy(self.schema)
+        self.schema = dict(agents=[], planners=[], adversaries=[], abilities=[], sources=[], operations=[],
+                           schedules=[], plugins=[], obfuscators=[], objectives=[], data_encoders=[])
+        self.ram = copy.deepcopy(self.schema)
         # Connect to Neo4j Database
         neo4j_uri = "bolt://localhost:7687"
         neo4j_user = "neo4j"
         neo4j_password = "calderaadmin"
         self.driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))
 
-    async def store(self, c_object):
-        """
-        Store the object in Neo4j database.
-        Assumes c_object has properties like unique, trait, name, value, etc.
-        """
 
-        # The query to create a fact node in Neo4j
-        query = """
-        CREATE (f:Fact {
-            unique: $unique,
-            trait: $trait,
-            name: $name,
-            value: $value,
-            created: $created,
-            score: $score,
-            source: $source,
-            origin_type: $origin_type,
-            links: $links,
-            relationships: $relationships,
-            limit_count: $limit_count,
-            collected_by: $collected_by,
-            technique_id: $technique_id
-        })
-        """
+    # async def store(self, c_object):
+    #     try:
+    #         print(" ")
+    #         # print(dir(c_object))
+    #         print(self.ram)
+    #         return c_object.store(self.ram)
+    #     except Exception as e:
+    #         self.log.error('[!] can only store first-class objects: %s' % e)
 
-        # Parameters to be passed to the query
-        params = {
-            'unique': getattr(c_object, 'unique', None),
-            'trait': getattr(c_object, 'trait', None),
-            'name': getattr(c_object, 'name', None),
-            'value': getattr(c_object, 'value', None),
-            'created': getattr(c_object, 'created', None),
-            'score': getattr(c_object, 'score', None),
-            'source': getattr(c_object, 'source', None),
-            'origin_type': getattr(c_object, 'origin_type', None),
-            'links': getattr(c_object, 'links', []),
-            'relationships': getattr(c_object, 'relationships', []),
-            'limit_count': getattr(c_object, 'limit_count', None),
-            'collected_by': getattr(c_object, 'collected_by', []),
-            'technique_id': getattr(c_object, 'technique_id', None)
-        }
 
-        # Execute the query
-        with self.driver.session() as session:
-            session.run(query, params)
+#  # The following is a test function to see if I can get the neo4j database to work   
+    async def store(self, c_object):
+        def convert_values(obj):
+            """
+            Convert object values to Neo4j-compatible format.
+            """
+            converted_props = {}
+            for key, value in obj.__dict__.items():
+                print(key, value)
+                if isinstance(value, dict):
+                    # Convert dictionary to JSON string
+                    converted_props[key] = json.dumps(value)
+                elif isinstance(value, enum.Enum):
+                    # Convert Enum to its string representation
+                    converted_props[key] = str(value)
+                elif isinstance(value, set):
+                    # Convert set to list
+                    converted_props[key] = list(value)
+                elif key == 'name' and isinstance(value, str):
+                    # Convert name attribute to string
+                    converted_props[key] = value
+                elif key == 'goals':
+                    # Handle goals attribute separately
+                    converted_props[key] = [goal.name for goal in value]
+                else:
+                    # No conversion needed for other types
+                    converted_props[key] = value
+            return converted_props
 
+        try:
+            with self.driver.session() as session:
+                # Determine the node label and name from the c_object
+                label = c_object.__class__.__name__
+                name = getattr(c_object, 'name', None)
+
+                if name:
+                    # Check if a node with the same name already exists
+                    query = f"MATCH (:{label} {{name: $name}}) RETURN COUNT(*) AS count"
+                    result = session.run(query, name=name)
+                    count = result.single()['count']
+
+                    if count > 0:
+                        # Append the object to the existing node
+                        query = f"MATCH (n:{label} {{name: $name}}) SET n += $props"
+                        session.run(query, name=name, props=convert_values(c_object))
+                    else:
+                        # Create a new node with the object
+                        query = f"CREATE (:{label} $props)"
+                        session.run(query, props=convert_values(c_object))
+                else:
+                    self.log.error("[!] 'name' attribute not found in the c_object")
+
+            # Return the stored object
+            return c_object
+        except Exception as e:
+            self.log.error('[!] Error storing object: %s' % e)
+        finally:
+            # Perform any necessary cleanup or closing operations
+            print("Closing Neo4j session")
+            if 'session' in locals():
+                session.close()
+
+
+    # @staticmethod
+    # def _iter_data_files(self):
+    #     """Yield paths to data files managed by caldera.
+
+    #     The files paths are relative to the root caldera folder, so they
+    #     will begin with "data/".
+
+    #     Note:
+    #         This will skip any files starting with '.' (e.g., '.gitkeep').
+    #     """
+    #     for data_glob in DATA_FILE_GLOBS:
+    #         for f in glob.glob(data_glob):
+    #             yield f
+
+# The following is a test function to see if I can get the neo4j database to work
+# NEED TO LOK AT THIS FUNCTION FURTHER NOT QUITE RIGHT
     @staticmethod
-    def _iter_data_files(self):
-        """Yield paths to data files managed by caldera.
+    def _iter_data_files():
+        """Yield paths to data files managed by caldera and store them in the Neo4j database.
 
         The files paths are relative to the root caldera folder, so they
         will begin with "data/".
 
         Note:
-            This will skip any files starting with '.' (e.g., '.gitkeep').
-        """
-        for data_glob in DATA_FILE_GLOBS:
-            for f in glob.glob(data_glob):
-                yield f
+            This will skip any files starting with '.' (e.g., '.gitkeep).
 
-        with self.driver.session() as session:
-            # Create empty abilities node
-            session.run("CREATE (:Abilities)")
+        Args:
+            driver (GraphDatabase.driver): Neo4j driver instance for database connection.
 
-            # Create empty adversaries node
-            session.run("CREATE (:Adversaries)")
-
-            # Create empty facts node
-            session.run("CREATE (:Facts)")
-
-            # Create empty objectives node
-            session.run("CREATE (:Objectives)")
-
-            # Create empty payloads node
-            session.run("CREATE (:Payloads)")
+        """
+        # NEED TO RETHINK THIS
+        # Connect to Neo4j Database
+        neo4j_uri = "bolt://localhost:7687"
+        neo4j_user = "neo4j"
+        neo4j_password = "calderaadmin"
+        driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))
 
-            # Create empty results node
-            session.run("CREATE (:Results)")
 
-            # Create empty sources node
-            session.run("CREATE (:Sources)")
+        try:
+            with driver.session() as session:
+                for data_glob in DATA_FILE_GLOBS:
+                    for file_path in glob.glob(data_glob):
+                        print(" adding filepath: %s"%file_path)
+                        try:
+                            # Create a node for each file path in the database
+                            session.run("CREATE (:File {path: $path})", path=file_path)
+
+                            # Yield the file path for further processing
+                            yield file_path
+
+                        except Exception as e:
+                            # Handle any specific exceptions raised during node creation
+                            print(f"Error creating node for file path: {file_path}. Error: {str(e)}")
+
+        except Exception as e:
+            # Handle any specific exceptions raised during the session creation or execution
+            print(f"Error in Neo4j session: {str(e)}")
 
-            # Create empty object_store node
-            session.run("CREATE (:ObjectStore)")
+        finally:
+            # Perform any necessary cleanup or closing operations
+            print("Closing Neo4j session")
+            if 'session' in locals():
+                session.close()
+# End of test function
 
     @staticmethod
     def _delete_file(path):
@@ -194,6 +252,7 @@ async def save_state(self):
         #     self.log.debug('Restored data from persistent storage')
         # self.log.debug('There are %s jobs in the scheduler' % len(self.ram['schedules']))
 
+# The following is a test function to see if I can get the neo4j database to work
     async def restore_state(self, profile_id=None):
         """
         Restore the object database from the Neo4j database for a particular profile
@@ -226,6 +285,7 @@ async def restore_state(self, profile_id=None):
         except Exception as e:
             self.log.error(f'[!] RESTORE_STATE: {e}')
 
+# End of test function
 
     async def apply(self, collection):
         if collection not in self.ram:
@@ -238,11 +298,7 @@ async def load_data(self, plugins=()):
     async def reload_data(self, plugins=()):
         await self._load(plugins)
 
-    # async def store(self, c_object):
-    #     try:
-    #         return c_object.store(self.ram)
-    #     except Exception as e:
-    #         self.log.error('[!] can only store first-class objects: %s' % e)
+
 
     # async def locate(self, object_name, match=None):
     #     try:
@@ -251,6 +307,7 @@ async def reload_data(self, plugins=()):
     #         self.log.error('[!] LOCATE: %s' % e)
 
     async def locate(self, object_name, match=None):
+        print("in locate")
         print("object_name: %s"%object_name)
         print("object type: %s"%type(object_name))
 

app/service/file_svc.py

@@ -17,6 +17,8 @@
 from app.utility.base_service import BaseService
 from app.utility.payload_encoder import xor_file, xor_bytes
 
+from neo4j import GraphDatabase
+
 FILE_ENCRYPTION_FLAG = '%encrypted%'
 
 
@@ -29,6 +31,11 @@ def __init__(self):
         self.encryptor = self._get_encryptor()
         self.encrypt_output = False if self.get_config('encrypt_files') is False else True
         self.packers = dict()
+        # Connect to Neo4j Database
+        neo4j_uri = "bolt://localhost:7687"
+        neo4j_user = "neo4j"
+        neo4j_password = "calderaadmin"
+        self.driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))
 
     async def get_file(self, headers):
         headers = CIMultiDict(headers)
@@ -159,15 +166,16 @@ async def add_special_payload(self, name, func):
             self.special_payloads[name] = func
 
             try:
-                async with self.driver.session() as session:
-                    await session.write_transaction(self._create_special_payload, name)
+                session = self.driver.session()
+                session.write_transaction(self._create_special_payload, name)
             except Exception as e:
                 # Handle database connection error
+                print(" ERROR in add_special_payload")
                 print(f"An error occurred while connecting to the database: {e}")
             finally:
                 self.driver.close()
 
-    async def _create_special_payload(self, tx, name):
+    def _create_special_payload(self, tx, name):
         """
         Create a special payload node in the Neo4j database