def _process_message(self, message):
m = Msg().from_frame(message)
logger.debug(f"handling: {len(m.data)}")
try:
m.data = [Indicator(**i) for i in m.data]
except Exception as e:
logger.error(e)
logger.debug(traceback.print_exc())
try:
[p.process(m.data) for p in self.plugins]
except (KeyboardInterrupt, SystemExit):
return
except Exception as e:
logger.error(e)
logger.debug(traceback.print_exc())
try:
m.data = [i.__dict__() for i in m.data]
except Exception as e:
logger.error(e)
m.data = []
self.push_s.send_msg(m)
logger.debug('done...')
def test_indicators_create_sha1(store, indicator):
indicator['indicator'] = 'd52380918a07322c50f1bfa2b43af3bb54cb33db'
indicator['group'] = 'everyone'
indicator['itype'] = 'sha1'
m = Msg(data=[indicator])
x = store.indicators.indicators_create(m)
def _send_fireball(self, data, f_size):
if len(data) == 0:
logger.error('no data to send')
return []
self.loop = ioloop.IOLoop()
self.socket.close()
self.socket = self.context.socket(zmq.DEALER)
self.socket.connect(self.remote)
self.stream = ZMQStream(self.socket, io_loop=self.loop)
self.stream.on_recv(self._handle_message_fireball)
self.stream.io_loop.call_later(SNDTIMEO, self._fireball_timeout)
self.response = []
if not isinstance(data, list):
data = [data]
self.num_responses = int((len(data) / f_size))
if (len(data) % f_size) != 0:
self.num_responses += 1
logger.debug('responses expected: %i' % self.num_responses)
batch = []
for d in data:
batch.append(d)
if len(batch) == f_size:
self.socket.send_msg(
Msg(mtype=Msg.INDICATORS_CREATE, data=batch)
)
batch = []
if len(batch):
self.socket.send_msg(Msg(mtype=Msg.INDICATORS_CREATE, data=batch))
logger.debug("starting loop to receive")
self.loop.start()
# clean up FDs
self.loop.close()
self.stream.close()
self.socket.close()
return self.response
def recv_msg(self, relay=False):
m = self.recv_multipart()
if relay:
assert isinstance(relay, zmq.Socket) or isinstance(relay, Socket)
return relay.send_msg(m)
return Msg().from_frame(m)
def test_indicators_create(store, indicator):
m = Msg()
m.data = indicator
x = store.indicators.indicators_create(m)
# TODO- fix this based on store handle
assert x == 1
indicator['last_at'] = arrow.utcnow().strftime('%Y-%m-%dT%H:%M:%S.%fZ')
indicator['tags'] = ['malware']
m.data = indicator
x = store.indicators.indicators_create(m)
assert x == 1
def test_indicators_search_ipv4(store, indicator):
indicator['indicator'] = '192.168.1.1'
indicator['itype'] = 'ipv4'
indicator['tags'] = 'botnet'
m = Msg()
m.data = indicator
x = store.indicators.indicators_create(m)
assert x == 1
for i in ['192.168.1.1', '192.168.1.0/24']:
m.data = [{
'indicator': i,
}]
x = store.indicators.indicators_search(m)
assert len(list(x)) > 0
def test_indicators_search_bulk(store, indicator):
m = Msg(data=indicator)
x = store.indicators.indicators_create(m)
assert x == 1
m.data = [{
'indicator': 'example.com',
'tags': 'botnet',
'confidence': 1
}, {
'indicator': 'example2.com',
}]
x = list(store.indicators.indicators_search(m))
assert len(x) == 1
assert x[0]['indicator'] == 'example.com'
def test_msgs():
def _send_multipart(m):
assert msgpack.unpackb(m[0]) == Msg.INDICATORS_SEARCH
assert unpack(m[1]) == []
m = Msg(mtype=Msg.INDICATORS_SEARCH, data=[])
ctx = Context()
s = ctx.socket(zmq.REQ)
s.send_multipart = _send_multipart
s.send_msg(m)
def _handle_message_fireball(self, m):
logger.debug('message received')
m = self._check_recv(loads(Msg.from_frame(m).data))
self.response.append(m)
self.num_responses -= 1
logger.debug('num responses remaining: %i' % self.num_responses)
if self.num_responses == 0:
logger.debug('finishing up...')
self.loop.stop()
def _process_message(self, message):
try:
[
self._process_indicator(i)
for i in Msg().from_frame(message).data
]
except (KeyboardInterrupt, SystemExit):
return
except Exception as e:
logger.error(e)
if logger.getEffectiveLevel() == logging.DEBUG:
traceback.print_exc()
def test_indicators_search_ipv6(store, indicator):
indicator['indicator'] = '2001:4860:4860::8888'
indicator['itype'] = 'ipv6'
indicator['tags'] = 'botnet'
m = Msg(data=indicator)
x = store.indicators.indicators_create(m)
assert x == 1
m.data = [{
'indicator': '2001:4860:4860::8888',
}]
x = store.indicators.indicators_search(m)
assert len(list(x)) > 0
m.data = [{
'indicator': '2001:4860::/32',
}]
x = store.indicators.indicators_search(m)
assert len(list(x)) > 0
def _send(self, mtype, data='[]', nowait=False, decode=True):
if not self.is_connected:
self.socket.connect(self.remote)
self.is_connected = True
if isinstance(data, str):
data = data.encode('utf-8')
self.socket.send_msg(Msg(mtype=mtype, data=data))
if self.nowait or nowait:
if self.autoclose:
self.socket.close()
return
return self._recv(decode=decode, close=self.autoclose)
def test_indicators_search_fqdn(store, indicator):
m = Msg()
m.data = [{
'indicator': 'example.com',
}]
x = store.indicators.indicators_search(m)
assert len(list(x)) == 0
x = store.indicators.indicators_search(m)
assert len(list(x)) == 1
indicator['tags'] = 'botnet'
indicator['indicator'] = 'example2.com'
m.data = indicator
x = store.indicators.indicators_create(m)
assert x == 1
m.data = [{
'indicator': 'example2.com',
}]
x = store.indicators.indicators_search(m)
assert len(list(x)) == 1
x = store.indicators.indicators_search(m)
assert len(list(x)) > 0
m.data = [{'indicator': 'example2.com', 'tags': 'malware'}]
x = store.indicators.indicators_search(m)
assert len(list(x)) == 0
def _recv_multipart():
return Msg(id=msgpack.packb(1234),
mtype=Msg.INDICATORS_SEARCH,
data=[]).to_frame()