/
worker_import.py
63 lines (54 loc) · 2.26 KB
/
worker_import.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
# coding: utf-8
import os
import yaml
import pika
import json
import base64
import time
from logger import Logger
from pycti import OpenCTI
class WorkerImport:
def __init__(self, verbose=True):
# Initialize logger
self.logger = Logger(os.path.dirname(os.path.abspath(__file__)) + '/logs/worker.log')
# Load configuration
self.config = yaml.load(open(os.path.dirname(os.path.abspath(__file__)) + '/config.yml'))
# Initialize OpenCTI client
self.opencti = OpenCTI(
self.config['opencti']['api_url'],
self.config['opencti']['api_key'],
os.path.dirname(os.path.abspath(__file__)) + '/logs/worker.log',
self.config['opencti']['verbose']
)
def import_action(self, ch, method, properties, body):
try:
data = json.loads(body)
self.logger.log('Receiving new action of type: { ' + data['type'] + ' }')
if data['type'] == 'import.stix2.bundle':
self.opencti.stix2_import_bundle(base64.b64decode(data['content']).decode('utf-8'))
except Exception as e:
self.logger.log('An unexpected error occurred: { ' + str(e) + ' }')
return False
def consume(self):
# Initialize the RabbitMQ connection
credentials = pika.PlainCredentials(self.config['rabbitmq']['username'], self.config['rabbitmq']['password'])
connection = pika.BlockingConnection(pika.ConnectionParameters(
host=self.config['rabbitmq']['hostname'],
port=self.config['rabbitmq']['port'],
virtual_host='/',
credentials=credentials
))
channel = connection.channel()
channel.exchange_declare(exchange='opencti', exchange_type='topic', durable=True)
channel.queue_declare('opencti-import', durable=True)
channel.queue_bind(exchange='opencti', queue='opencti-import', routing_key='import.*.*')
channel.basic_consume(queue='opencti-import', on_message_callback=self.import_action, auto_ack=True)
channel.start_consuming()
if __name__ == '__main__':
worker_import = WorkerImport()
while True:
try:
worker_import.consume()
except Exception as e:
print(e)
time.sleep(5)