blacked

.pre-commit-config.yaml

@@ -25,7 +25,7 @@ repos:
       rev: 23.3.0
       hooks:
           - id: black
-            language_version: python3.9
+            language_version: python3.11
 
     - repo: https://github.com/asottile/setup-cfg-fmt
       rev: v2.3.0

README.md

@@ -1,2 +1,3 @@
 # CoPilot
+
 SOCFortress CoPilot

backend/README.MD

@@ -79,7 +79,6 @@ classDiagram
     }
 ```
 
-
 # Connector Classes
 
 ```mermaid
@@ -132,7 +131,7 @@ classDiagram
     }
 ```
 
-# Routes 
+# Routes
 
 ```mermaid
 graph TD;
@@ -165,4 +164,3 @@ graph TD;
     O --> V[Return Data]
 
 ```
-

backend/app/__init__.py

@@ -37,15 +37,15 @@
 migrate = Migrate(app, db)
 ma = Marshmallow(app)
 
-from app.routes.connectors import bp as connectors_bp  # Import the blueprint
 from app.routes.agents import bp as agents_bp  # Import the blueprint
-from app.routes.rules import bp as rules_bp  # Import the blueprint
-from app.routes.graylog import bp as graylog_bp  # Import the blueprint
 from app.routes.alerts import bp as alerts_bp  # Import the blueprint
-from app.routes.wazuhindexer import bp as wazuhindexer_bp  # Import the blueprint
+from app.routes.connectors import bp as connectors_bp  # Import the blueprint
+from app.routes.dfir_iris import bp as dfir_iris_bp  # Import the blueprint
+from app.routes.graylog import bp as graylog_bp  # Import the blueprint
+from app.routes.rules import bp as rules_bp  # Import the blueprint
 from app.routes.shuffle import bp as shuffle_bp  # Import the blueprint
 from app.routes.velociraptor import bp as velociraptor_bp  # Import the blueprint
-from app.routes.dfir_iris import bp as dfir_iris_bp  # Import the blueprint
+from app.routes.wazuhindexer import bp as wazuhindexer_bp  # Import the blueprint
 
 app.register_blueprint(connectors_bp)  # Register the connectors blueprint
 app.register_blueprint(agents_bp)  # Register the agents blueprint

backend/app/models/connectors.py

@@ -1,23 +1,26 @@
 import importlib
 import json
 import os
-import pika
+from abc import ABC
+from abc import abstractmethod
 from dataclasses import dataclass
+
+import grpc
+import pika
+import pyvelociraptor
 import requests
-from abc import ABC, abstractmethod
 from elasticsearch7 import Elasticsearch
+from flask import current_app
 from loguru import logger
-from sqlalchemy.orm.exc import NoResultFound
-import pyvelociraptor
 from pyvelociraptor import api_pb2
 from pyvelociraptor import api_pb2_grpc
-from werkzeug.utils import secure_filename
-import grpc
-
 from sqlalchemy.exc import SQLAlchemyError
-from flask import current_app
+from sqlalchemy.orm.exc import NoResultFound
+from werkzeug.utils import secure_filename
 
-from app.models.models import Connectors, connectors_schema, ConnectorsAvailable
+from app.models.models import Connectors
+from app.models.models import ConnectorsAvailable
+from app.models.models import connectors_schema
 
 
 def dynamic_import(module_name, class_name):
@@ -98,7 +101,7 @@ def verify_connection(self):
         :return: A dictionary containing the status of the connection attempt and information about the cluster's health.
         """
         logger.info(
-            f"Verifying the wazuh-indexer connection to {self.attributes['connector_url']}"
+            f"Verifying the wazuh-indexer connection to {self.attributes['connector_url']}",
         )
         try:
             es = Elasticsearch(
@@ -117,7 +120,7 @@ def verify_connection(self):
             return {"connectionSuccessful": True}
         except Exception as e:
             logger.error(
-                f"Connection to {self.attributes['connector_url']} failed with error: {e}"
+                f"Connection to {self.attributes['connector_url']} failed with error: {e}",
             )
             return {"connectionSuccessful": False, "clusterHealth": None}
 
@@ -140,7 +143,7 @@ def verify_connection(self):
             dict: A dictionary containing 'connectionSuccessful' status and 'roles' if the connection is successful.
         """
         logger.info(
-            f"Verifying the graylog connection to {self.attributes['connector_url']}"
+            f"Verifying the graylog connection to {self.attributes['connector_url']}",
         )
         try:
             graylog_roles = requests.get(
@@ -153,17 +156,17 @@ def verify_connection(self):
             )
             if graylog_roles.status_code == 200:
                 logger.info(
-                    f"Connection to {self.attributes['connector_url']} successful"
+                    f"Connection to {self.attributes['connector_url']} successful",
                 )
                 return {"connectionSuccessful": True}
             else:
                 logger.error(
-                    f"Connection to {self.attributes['connector_url']} failed with error: {graylog_roles.text}"
+                    f"Connection to {self.attributes['connector_url']} failed with error: {graylog_roles.text}",
                 )
                 return {"connectionSuccessful": False, "roles": None}
         except Exception as e:
             logger.error(
-                f"Connection to {self.attributes['connector_url']} failed with error: {e}"
+                f"Connection to {self.attributes['connector_url']} failed with error: {e}",
             )
             return {"connectionSuccessful": False, "roles": None}
 
@@ -186,7 +189,7 @@ def verify_connection(self):
             dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
         """
         logger.info(
-            f"Verifying the wazuh-manager connection to {self.attributes['connector_url']}"
+            f"Verifying the wazuh-manager connection to {self.attributes['connector_url']}",
         )
         try:
             wazuh_auth_token = requests.get(
@@ -204,12 +207,12 @@ def verify_connection(self):
                 return {"connectionSuccessful": True, "authToken": wazuh_auth_token}
             else:
                 logger.error(
-                    f"Connection to {self.attributes['connector_url']} failed with error: {wazuh_auth_token.text}"
+                    f"Connection to {self.attributes['connector_url']} failed with error: {wazuh_auth_token.text}",
                 )
                 return {"connectionSuccessful": False, "authToken": None}
         except Exception as e:
             logger.error(
-                f"Connection to {self.attributes['connector_url']} failed with error: {e}"
+                f"Connection to {self.attributes['connector_url']} failed with error: {e}",
             )
             return {"connectionSuccessful": False, "authToken": None}
 
@@ -241,11 +244,11 @@ def verify_connection(self):
             dict: A dictionary containing 'connectionSuccessful' status and 'apps' if the connection is successful.
         """
         logger.info(
-            f"Verifying the shuffle connection to {self.attributes['connector_url']}"
+            f"Verifying the shuffle connection to {self.attributes['connector_url']}",
         )
         try:
             headers = {
-                "Authorization": f"Bearer {self.attributes['connector_api_key']}"
+                "Authorization": f"Bearer {self.attributes['connector_api_key']}",
             }
             shuffle_apps = requests.get(
                 f"{self.attributes['connector_url']}/api/v1/apps",
@@ -254,17 +257,17 @@ def verify_connection(self):
             )
             if shuffle_apps.status_code == 200:
                 logger.info(
-                    f"Connection to {self.attributes['connector_url']} successful"
+                    f"Connection to {self.attributes['connector_url']} successful",
                 )
                 return {"connectionSuccessful": True}
             else:
                 logger.error(
-                    f"Connection to {self.attributes['connector_url']} failed with error: {shuffle_apps.text}"
+                    f"Connection to {self.attributes['connector_url']} failed with error: {shuffle_apps.text}",
                 )
                 return {"connectionSuccessful": False}
         except Exception as e:
             logger.error(
-                f"Connection to {self.attributes['connector_url']} failed with error: {e}"
+                f"Connection to {self.attributes['connector_url']} failed with error: {e}",
             )
             return {"connectionSuccessful": False}
 
@@ -287,11 +290,11 @@ def verify_connection(self):
             dict: A dictionary containing 'connectionSuccessful' status and 'response' if the connection is successful.
         """
         logger.info(
-            f"Verifying the dfir-iris connection to {self.attributes['connector_url']}"
+            f"Verifying the dfir-iris connection to {self.attributes['connector_url']}",
         )
         try:
             headers = {
-                "Authorization": f"Bearer {self.attributes['connector_api_key']}"
+                "Authorization": f"Bearer {self.attributes['connector_api_key']}",
             }
             dfir_iris = requests.get(
                 f"{self.attributes['connector_url']}/api/ping",
@@ -301,17 +304,17 @@ def verify_connection(self):
             # See if 200 is returned
             if dfir_iris.status_code == 200:
                 logger.info(
-                    f"Connection to {self.attributes['connector_url']} successful"
+                    f"Connection to {self.attributes['connector_url']} successful",
                 )
                 return {"connectionSuccessful": True}
             else:
                 logger.error(
-                    f"Connection to {self.attributes['connector_url']} failed with error: {dfir_iris.text}"
+                    f"Connection to {self.attributes['connector_url']} failed with error: {dfir_iris.text}",
                 )
                 return {"connectionSuccessful": False, "response": None}
         except Exception as e:
             logger.error(
-                f"Connection to {self.attributes['connector_url']} failed with error: {e}"
+                f"Connection to {self.attributes['connector_url']} failed with error: {e}",
             )
             return {"connectionSuccessful": False, "response": None}
 
@@ -351,7 +354,9 @@ def verify_connection(self):
                 options = (("grpc.ssl_target_name_override", "VelociraptorServer"),)
 
                 with grpc.secure_channel(
-                    config["api_connection_string"], creds, options
+                    config["api_connection_string"],
+                    creds,
+                    options,
                 ) as channel:
                     stub = api_pb2_grpc.APIStub(channel)
                     client_query = "SELECT * FROM info()"
@@ -395,7 +400,7 @@ def verify_connection(self):
         Verifies the connection to RabbitMQ service.
         """
         logger.info(
-            f"Verifying the rabbitmq connection to {self.attributes['connector_url']}"
+            f"Verifying the rabbitmq connection to {self.attributes['connector_url']}",
         )
         try:
             # For the connector_url, strip out the host and port and use that for the connection
@@ -415,15 +420,15 @@ def verify_connection(self):
             connection = pika.BlockingConnection(parameters)
             if connection.is_open:
                 logger.info(
-                    f"Connection to {self.attributes['connector_url']} successful"
+                    f"Connection to {self.attributes['connector_url']} successful",
                 )
                 return {"connectionSuccessful": True}
             else:
                 logger.error(f"Connection to {self.attributes['connector_url']} failed")
                 return {"connectionSuccessful": False, "response": None}
         except Exception as e:
             logger.error(
-                f"Connection to {self.attributes['connector_url']} failed with error: {e}"
+                f"Connection to {self.attributes['connector_url']} failed with error: {e}",
             )
             return {"connectionSuccessful": False, "response": None}
 

backend/app/routes/agents.py

@@ -1,15 +1,16 @@
-from flask import Blueprint, jsonify, request
+from flask import Blueprint
+from flask import jsonify
+from flask import request
 from loguru import logger
-from app.models.connectors import Connector, WazuhManagerConnector
 
-from app.services.agents.agents import AgentService, AgentSyncService
-
-from app.services.WazuhManager.universal import UniversalService
+from app.models.connectors import Connector
+from app.models.connectors import WazuhManagerConnector
+from app.services.agents.agents import AgentService
+from app.services.agents.agents import AgentSyncService
 from app.services.WazuhManager.agent import WazuhManagerAgentService
+from app.services.WazuhManager.universal import UniversalService
 from app.services.WazuhManager.vulnerability import VulnerabilityService
 
-
-
 bp = Blueprint("agents", __name__)
 
 

backend/app/routes/alerts.py

@@ -1,8 +1,12 @@
-from flask import Blueprint, jsonify, request
+from flask import Blueprint
+from flask import jsonify
+from flask import request
 from loguru import logger
-from app.models.connectors import Connector, WazuhManagerConnector
 
-from app.services.agents.agents import AgentService, AgentSyncService
+from app.models.connectors import Connector
+from app.models.connectors import WazuhManagerConnector
+from app.services.agents.agents import AgentService
+from app.services.agents.agents import AgentSyncService
 from app.services.WazuhIndexer.alerts import AlertsService
 
 bp = Blueprint("alerts", __name__)

backend/app/routes/connectors.py

@@ -1,13 +1,13 @@
-from flask import Blueprint, jsonify, request
+from flask import Blueprint
+from flask import jsonify
+from flask import request
 from loguru import logger
-from app.models.models import (
-    ConnectorsAvailable,
-    Connectors,
-    connectors_available_schema,
-)
 
-from app.services.connectors.connectors import ConnectorService
 from app import db
+from app.models.models import Connectors
+from app.models.models import ConnectorsAvailable
+from app.models.models import connectors_available_schema
+from app.services.connectors.connectors import ConnectorService
 
 bp = Blueprint("connectors", __name__)
 
@@ -48,7 +48,7 @@ def get_connector_details(id):
     # Call service function instead of direct function call
     service = ConnectorService(db)
     connector_validated = service.validate_connector_exists(
-        int(id)
+        int(id),
     )  # convert id to integer
     logger.info(connector_validated)
     if connector_validated["success"] == False:
@@ -70,14 +70,16 @@ def update_connector_route(id):
         id (str): The id of the connector to be updated.
 
     Returns:
-        json: A JSON response containing the success status of the update operation and a message indicating the status. If the update operation was successful, it returns the connector name and the status of the connection verification.
+        json: A JSON response containing the success status of the update operation and
+        a message indicating the status. If the update operation was successful,
+        it returns the connector name and the status of the connection verification.
     """
     api_key_connector = ["Shuffle", "DFIR-IRIS", "Velociraptor"]
 
     request_data = request.get_json()
     service = ConnectorService(db)
     connector_validated = service.validate_connector_exists(
-        int(id)
+        int(id),
     )  # convert id to integer
     logger.info(connector_validated)
     if connector_validated["success"] == False: