[connector] Enhance the connector helper class, rename others (#8)

Author: Samuel Hassine

examples/observables/get_observables.py

@@ -3,7 +3,7 @@
 import os
 import yaml
 
-from pycti.opencti import OpenCTI
+from pycti.opencti_api_client import OpenCTI
 
 # Load configuration
 config = yaml.load(open(os.path.dirname(__file__) + '/../config.yml'))

examples/stix2/export.py

@@ -4,7 +4,7 @@
 import yaml
 import json
 
-from opencti import OpenCTI
+from pycti import OpenCTIApiClient
 
 # Load configuration
 config = yaml.load(open(os.path.dirname(__file__) + '/../config.yml'))
@@ -13,7 +13,7 @@
 export_file = './report.json'
 
 # OpenCTI initialization
-opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['token'])
+opencti = OpenCTIApiClient(config['opencti']['api_url'], config['opencti']['token'])
 
 # Import the bundle
 bundle = opencti.stix2_export_entity('report', '{REPORT_ID}', 'full')

examples/stix2/import.py

@@ -3,7 +3,7 @@
 import os
 import yaml
 
-from opencti import OpenCTI
+from pycti import OpenCTIApiClient
 
 # Load configuration
 config = yaml.load(open(os.path.dirname(__file__) + '/../config.yml'))
@@ -12,7 +12,7 @@
 file_to_import = './enterprise-attack.json'
 
 # OpenCTI initialization
-opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['token'])
+opencti = OpenCTIApiClient(config['opencti']['api_url'], config['opencti']['token'])
 
 # Import the bundle
-opencti.stix2_import_bundle_from_file(file_to_import, False)
+opencti.stix2_import_bundle_from_file(file_to_import, False)

pycti/__init__.py

@@ -1,2 +1,2 @@
-from .opencti import OpenCTI
-from .opencti_connector import OpenCTIConnector
+from .opencti_api_client import OpenCTIApiClient
+from .opencti_connector_helper import OpenCTIConnectorHelper

pycti/opencti.py renamed to pycti/opencti_api_client.py

@@ -12,7 +12,7 @@
 from pycti.opencti_stix2 import OpenCTIStix2
 
 
-class OpenCTI:
+class OpenCTIApiClient:
     """
         Python API for OpenCTI
         :param url: OpenCTI URL

pycti/opencti_connector.py renamed to pycti/opencti_connector_helper.py

@@ -9,7 +9,7 @@
 EXCHANGE_NAME = 'amqp.opencti'
 
 
-class OpenCTIConnector:
+class OpenCTIConnectorHelper:
     """
         Python API for OpenCTI connector
         :param identifier: Connector identifier
@@ -35,40 +35,48 @@ def __init__(self, identifier, config, rabbitmq_hostname, rabbitmq_port, rabbitm
         # Encode the configuration
         config_encoded = base64.b64encode(json.dumps(self.config).encode('utf-8')).decode('utf-8')
 
+        # Connect to RabbitMQ
+        self.connection = self._connect()
+        self.channel = self._create_channel()
+        logging.info('Successfully connected to RabbitMQ')
+
         # Declare the queue for the connector
-        channel = self._connect()
-        channel.queue_delete(self.queue_name)
-        channel.queue_declare(self.queue_name, durable=True, arguments={'config': config_encoded})
-        channel.queue_bind(queue=self.queue_name, exchange=EXCHANGE_NAME, routing_key=self.routing_key)
-        channel.close()
-        self._disconnect()
+        self.channel.queue_delete(self.queue_name)
+        self.channel.queue_declare(self.queue_name, durable=True, arguments={'config': config_encoded})
+        self.channel.queue_bind(queue=self.queue_name, exchange=EXCHANGE_NAME, routing_key=self.routing_key)
 
     def _connect(self):
         try:
             credentials = pika.PlainCredentials(self.rabbitmq_username, self.rabbitmq_password)
             parameters = pika.ConnectionParameters(self.rabbitmq_hostname, self.rabbitmq_port, '/', credentials)
-            self.connection = pika.BlockingConnection(parameters)
+            return pika.BlockingConnection(parameters)
+        except:
+            logging.error('Unable to connect to RabbitMQ with the given parameters')
+
+    def _create_channel(self):
+        try:
             channel = self.connection.channel()
             channel.exchange_declare(exchange=EXCHANGE_NAME, exchange_type='direct', durable=True)
             return channel
         except:
-            logging.error('Unable to connect to RabbitMQ with the given parameters')
+            logging.error('Unable to open channel to RabbitMQ with the given parameters')
 
-    def _disconnect(self):
-        if self.connection is not None and not self.connection.is_closed:
-            self.connection.close()
+    def _reconnect(self):
+        self.connection = self._connect()
+        self.channel = self._create_channel()
 
     def send_stix2_bundle(self, bundle):
         """
             This method send a STIX2 bundle to RabbitMQ to be consumed by workers
             :param bundle: A valid STIX2 bundle
         """
-        logging.info('Sending a STI2 bundle to RabbitMQ...')
+        if not self.channel.is_open:
+            self._reconnect()
 
         # Validate the STIX 2 bundle
         validation = validate_string(bundle)
         if not validation.is_valid:
-            raise ValueError('The bundle is not a valid STIX2 JSON')
+            raise ValueError('The bundle is not a valid STIX2 JSON:' + bundle)
 
         # Prepare the message
         message = {
@@ -77,8 +85,5 @@ def send_stix2_bundle(self, bundle):
         }
 
         # Send the message
-        channel = self._connect()
-        channel.basic_publish(EXCHANGE_NAME, self.routing_key, json.dumps(message))
-        channel.close()
-        self._disconnect()
+        self.channel.basic_publish(EXCHANGE_NAME, self.routing_key, json.dumps(message))
         logging.info('STIX2 bundle has been sent')