[connector] Add the connector helper class (#8)

Author: Samuel Hassine

examples/config.yml.sample

@@ -1,7 +1,3 @@
 opencti:
   api_url: 'http://opencti.example.com'
-  api_key : 'XXxxXxXXxxX'
-  verbose: true
-
-mitre:
-  repository_path_cti: '/home/user/git/cti'
+  token : 'ChangeMe'

examples/stix2/export.py

@@ -10,10 +10,10 @@
 config = yaml.load(open(os.path.dirname(__file__) + '/../config.yml'))
 
 # Export file
-export_file = './exports/report.json'
+export_file = './report.json'
 
 # OpenCTI initialization
-opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['api_key'], config['opencti']['log_file'], config['opencti']['verbose'])
+opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['token'])
 
 # Import the bundle
 bundle = opencti.stix2_export_entity('report', '{REPORT_ID}', 'full')

examples/stix2/import.py

@@ -9,10 +9,10 @@
 config = yaml.load(open(os.path.dirname(__file__) + '/../config.yml'))
 
 # File to import
-file_to_import = config['mitre']['repository_path_cti'] + '/apt1.json'
+file_to_import = './enterprise-attack.json'
 
 # OpenCTI initialization
-opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['api_key'], config['opencti']['log_file'], config['opencti']['verbose'])
+opencti = OpenCTI(config['opencti']['api_url'], config['opencti']['token'])
 
 # Import the bundle
 opencti.stix2_import_bundle_from_file(file_to_import, False)

pycti/__init__.py

@@ -1 +1,2 @@
-from .opencti import OpenCTI
+from .opencti import OpenCTI
+from .opencti_connector import OpenCTIConnector

pycti/opencti.py

@@ -7,6 +7,7 @@
 import json
 import uuid
 import base64
+import logging
 
 from pycti.opencti_stix2 import OpenCTIStix2
 
@@ -15,28 +16,18 @@ class OpenCTI:
     """
         Python API for OpenCTI
         :param url: OpenCTI URL
-        :param key: The API key
-        :param verbose: Log all requests. Defaults to None
-        :param stdout: Display log to stdout. Defaults to None
+        :param token: The API key
     """
 
-    def __init__(self, url, key, log_file='', verbose=True, stdout=True):
+    def __init__(self, url, token):
         self.api_url = url + '/graphql'
-        self.log_file = log_file
-        self.verbose = verbose
-        self.stdout = stdout
         self.request_headers = {
-            'Authorization': 'Bearer ' + key,
+            'Authorization': 'Bearer ' + token,
             'Content-Type': 'application/json'
         }
 
     def log(self, message):
-        if self.stdout:
-            print(message)
-        if self.verbose and len(self.log_file) > 0:
-            file = open(self.log_file, 'a')
-            file.write('[' + datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S') + '] ' + message + "\n")
-            file.close()
+        logging.info('[' + datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S') + '] ' + message + "\n")
 
     def query(self, query, variables={}):
         r = requests.post(self.api_url, json={'query': query, 'variables': variables}, headers=self.request_headers)
@@ -409,7 +400,8 @@ def get_stix_relation_by_id(self, id):
         result = self.query(query, {'id': id})
         return result['data']['stixRelation']
 
-    def get_stix_relations(self, from_id=None, to_id=None, type='stix_relation', first_seen=None, last_seen=None, inferred=False):
+    def get_stix_relations(self, from_id=None, to_id=None, type='stix_relation', first_seen=None, last_seen=None,
+                           inferred=False):
         self.log('Getting relations, from: ' + from_id + ', to: ' + to_id + '...')
         if type == 'revoked-by':
             return []
@@ -1717,8 +1709,10 @@ def create_incident_if_not_exists(self,
         object_result = self.check_existing_stix_domain_entity(stix_id, name, 'Incident')
         if object_result is not None:
             self.update_stix_domain_entity_field(object_result['id'], 'name', name)
-            description is not None and self.update_stix_domain_entity_field(object_result['id'], 'description', description)
-            first_seen is not None and self.update_stix_domain_entity_field(object_result['id'], 'first_seen', first_seen)
+            description is not None and self.update_stix_domain_entity_field(object_result['id'], 'description',
+                                                                             description)
+            first_seen is not None and self.update_stix_domain_entity_field(object_result['id'], 'first_seen',
+                                                                            first_seen)
             last_seen is not None and self.update_stix_domain_entity_field(object_result['id'], 'last_seen', last_seen)
             return object_result
         else:
@@ -2509,7 +2503,8 @@ def create_course_of_action(self, name, description, id=None, stix_id=None, crea
         })
         return result['data']['courseOfActionAdd']
 
-    def create_course_of_action_if_not_exists(self, name, description, id=None, stix_id=None, created=None, modified=None):
+    def create_course_of_action_if_not_exists(self, name, description, id=None, stix_id=None, created=None,
+                                              modified=None):
         object_result = self.check_existing_stix_domain_entity(stix_id, name, 'Course-Of-Action')
         if object_result is not None:
             return object_result

pycti/opencti_connector.py

@@ -0,0 +1,84 @@
+# coding: utf-8
+
+import pika
+import logging
+import json
+import base64
+from stix2validator import validate_string
+
+EXCHANGE_NAME = 'amqp.opencti'
+
+
+class OpenCTIConnector:
+    """
+        Python API for OpenCTI connector
+        :param identifier: Connector identifier
+        :param config: Connector configuration
+        :param rabbitmq_hostname: RabbitMQ hostname
+        :param rabbitmq_port: RabbitMQ hostname
+        :param rabbitmq_username: RabbitMQ hostname
+        :param rabbitmq_password: RabbitMQ password
+    """
+
+    def __init__(self, identifier, config, rabbitmq_hostname, rabbitmq_port, rabbitmq_username, rabbitmq_password):
+        # Initialize configuration
+        self.connection = None
+        self.identifier = identifier
+        self.config = config
+        self.rabbitmq_hostname = rabbitmq_hostname
+        self.rabbitmq_port = rabbitmq_port
+        self.rabbitmq_username = rabbitmq_username
+        self.rabbitmq_password = rabbitmq_password
+        self.queue_name = 'import-connectors-' + self.identifier
+        self.routing_key = 'import.connectors.' + self.identifier
+
+        # Encode the configuration
+        config_encoded = base64.b64encode(json.dumps(self.config).encode('utf-8')).decode('utf-8')
+
+        # 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()
+
+    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)
+            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')
+
+    def _disconnect(self):
+        if self.connection is not None and not self.connection.is_closed:
+            self.connection.close()
+
+    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...')
+
+        # Validate the STIX 2 bundle
+        validation = validate_string(bundle)
+        if not validation.is_valid:
+            raise ValueError('The bundle is not a valid STIX2 JSON')
+
+        # Prepare the message
+        message = {
+            'type': 'stix2-bundle',
+            'content': base64.b64encode(bundle.encode('utf-8')).decode('utf-8')
+        }
+
+        # Send the message
+        channel = self._connect()
+        channel.basic_publish(EXCHANGE_NAME, self.routing_key, json.dumps(message))
+        channel.close()
+        self._disconnect()
+        logging.info('STIX2 bundle has been sent')

requirements.txt

@@ -5,4 +5,6 @@ pypandoc
 python-dateutil
 datefinder
 stix2
-pytz
+pytz
+pika
+stix2-validator

setup.py

@@ -7,12 +7,14 @@
 
 try:
     from pypandoc import convert
+
     read_md = lambda f: convert(f, 'rst')
 except ImportError:
     print("warning: pypandoc module not found, could not convert Markdown to RST")
     read_md = lambda f: open(f, 'r').read()
 
-VERSION = "1.2.7"
+VERSION = "1.2.8"
+
 
 class VerifyVersionCommand(install):
     description = 'verify that the git tag matches our version'
@@ -25,6 +27,7 @@ def run(self):
             )
             sys.exit(info)
 
+
 setup(
     name='pycti',
     version=VERSION,
@@ -50,7 +53,8 @@ def run(self):
         'Topic :: Software Development :: Libraries :: Python Modules'
     ],
     include_package_data=True,
-    install_requires=['requests', 'PyYAML', 'python-dateutil', 'datefinder', 'stix2', 'pytz'],
+    install_requires=['requests', 'PyYAML', 'python-dateutil', 'datefinder', 'stix2', 'stix2-validator', 'pytz',
+                      'pika'],
     cmdclass={
         'verify': VerifyVersionCommand,
     }