def deserialize(self):
"""Extract the numpy array"""
if self.raw_buffer is not None:
# Compression was not used.
numpy_array = np.frombuffer(self.raw_buffer,
dtype=self.dtype).reshape(self.shape)
elif self.compressed_label_blocks is not None:
# label compression was used.
numpy_array = deserialize_uint64_blocks(
self.compressed_label_blocks, self.shape)
elif self.compressed_mask_array is not None:
numpy_array, _, _ = decode_mask_array(
lz4.uncompress(self.compressed_mask_array), self.shape)
numpy_array = np.asarray(numpy_array, order='C')
else:
numpy_array = np.ndarray(shape=self.shape, dtype=self.dtype)
# See serialization of 2D, 1D and 0D arrays, above.
if numpy_array.ndim <= 2:
buf = lz4.uncompress(self.serialized_subarrays[0])
numpy_array[:] = np.frombuffer(buf, self.dtype).reshape(
numpy_array.shape)
else:
for subarray, serialized_subarray in zip(
numpy_array, self.serialized_subarrays):
buf = lz4.uncompress(serialized_subarray)
subarray[:] = np.frombuffer(buf, self.dtype).reshape(
subarray.shape)
if self.layout == 'F':
numpy_array = numpy_array.transpose()
return numpy_array
def get(self):
buf = self.sock.recv(8)
header, length = _HEADER.unpack_from(buf)
version = header >> 28 & 0xf
msg_id = header >> 16 & 0xff
msg_type = header >> 8 & 0xff
compression = header & 1 == 1
if version != 0:
return None
subcls = _MESSAGE_TYPES.get(msg_type)
if not subcls:
return None
buf = None
if length > 0:
buf = self.sock.recv(length)
if compression:
buf = lz4.uncompress(buf)
self.last_recv = datetime.datetime.now()
return subcls.unpack(msg_id, buf)
def unpack(stream):
while True:
data = stream.read(4)
if len(data) != 4:
break
(pktSize,) = struct.unpack('<I', data)
data = stream.read(8+2)
if len(data) != 8+2:
sys.stderr.write('short read')
break
(checksum, lenfname) = struct.unpack('<QH', data)
fname = stream.read(lenfname)
if len(fname) != lenfname:
sys.stderr.write('short read')
break
data = stream.read(4)
if len(data) != 4:
sys.stderr.write('short read')
break
(fsize,) = struct.unpack('<I', data)
compressedSize = pktSize - 4 - 8 - 2 - 4 - lenfname
data = stream.read(compressedSize)
if len(data) != compressedSize:
sys.stderr.write('short read')
break
data = lz4.uncompress(data)
got = siphashc.siphash('\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0', data)
if got == checksum:
sys.stderr.write('%s: %d -> %d\n' % (fname, compressedSize, fsize))
else:
sys.stderr.write('%s: checksum fail: got %d, want %d\n' % (fname, got, checksum))
def delete(self, eventId):
with self._sem:
with self.lmdb.begin(write=True) as txn:
index_db = self.lmdb.open_db(key=self.index_name,
txn=txn,
dupsort=True)
mtimes_db = self.lmdb.open_db(key=self.mtimes_name, txn=txn)
deltas_db = self.lmdb.open_db(key=self.deltas_name, txn=txn)
self.logger.debug(("Removing Event {ev} from the "
"database").format(ev=eventId))
with txn.cursor(db=index_db) as cursor:
if not cursor.set_key(eventId.encode('utf-8')):
raise KeyNotFoundError(
"Key {ev} not found in {path}".format(
ev=eventId, path=self.db_path))
self.logger.debug(
"Found {n} deltas".format(n=cursor.count()))
for delta_key in cursor.iternext_dup(keys=False):
if self.logger.isEnabledFor(self.logger.TRACE):
data = txn.get(delta_key, db=deltas_db)
self.logger.trace("Deleting delta {id}\n".format(
id=int.from_bytes(
delta_key, byteorder='big', signed=False))
+ prettify(uncompress(data)))
txn.delete(delta_key, db=deltas_db)
self.logger.debug("Deleting MTIME")
txn.delete(eventId.encode('utf-8'), db=mtimes_db)
self.logger.debug("Deleting index entry")
txn.delete(eventId.encode('utf-8'), db=index_db)
def get(self):
buf = self.sock.recv(8)
header, length = _HEADER.unpack_from(buf)
version = header >> 28 & 0xf
msg_id = header >> 16 & 0xff
msg_type = header >> 8 & 0xff
compression = header & 1 == 1
if version != 0:
return None
subcls = _MESSAGE_TYPES.get(msg_type)
if not subcls:
return None
buf = None
if length > 0:
buf = self.sock.recv(length)
if compression:
buf = lz4.uncompress(buf)
self.last_recv = datetime.datetime.now()
return subcls.unpack(msg_id, buf)
def unpack(stream):
while True:
data = stream.read(4)
if len(data) != 4:
break
(pktSize, ) = struct.unpack('<I', data)
data = stream.read(8 + 2)
if len(data) != 8 + 2:
sys.stderr.write('short read')
break
(checksum, lenfname) = struct.unpack('<QH', data)
fname = stream.read(lenfname)
if len(fname) != lenfname:
sys.stderr.write('short read')
break
data = stream.read(4)
if len(data) != 4:
sys.stderr.write('short read')
break
(fsize, ) = struct.unpack('<I', data)
compressedSize = pktSize - 4 - 8 - 2 - 4 - lenfname
data = stream.read(compressedSize)
if len(data) != compressedSize:
sys.stderr.write('short read')
break
data = lz4.uncompress(data)
got = siphashc.siphash('\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0', data)
if got == checksum:
sys.stderr.write('%s: %d -> %d\n' % (fname, compressedSize, fsize))
else:
sys.stderr.write('%s: checksum fail: got %d, want %d\n' %
(fname, got, checksum))
def _peek(self):
with self._lmdb.begin(write=False, buffers=True) as txn:
with txn.cursor(db=self._queue_db) as cursor:
if cursor.first(): hash_key = cursor.value()
else: raise Empty()
buf = txn.get(hash_key, db=self._hashes_db)
return pickle.loads(
uncompress(buf)) if self.compression else pickle.loads(buf)
def _get(self):
with self._sem:
with self._lmdb.begin(write=True, buffers=True) as txn:
with txn.cursor(db=self._queue_db) as cursor:
if cursor.first():
buf = cursor.pop(cursor.key())
return pickle.loads(uncompress(
buf)) if self.compression else pickle.loads(buf)
else:
raise Empty()
def decrypt(inputPath, keyPath, outputPath, force=False, mode="raw"):
"""
This function performs vernam decryption using an input file and a key,
result is stored in output file.
Args:
* inputPath : path to the file used as input
* keyPath : path to the file used as key
* outputPath : path to the file used as output
* mode : working mode, Modes are:
* raw: traditional way
* lz4: all input data will be decompressed via lz4
* human: all data will be converted to ownBase32 beforehand
Returns:
None
"""
if not os.path.exists(inputPath):
print("Could not find input file {}".format(inputPath),
file=sys.stderr)
sys.exit(10)
if os.path.exists(outputPath):
if force is False:
print("output file exists, won't overwrite {}".format(outputPath),
file=sys.stderr)
sys.exit(11)
else:
sys.stderr.write("Output file will be overwritten as requested.\n")
if mode == "human":
seek, inputCryp = message.readHumanMessage(inputPath)
offset = seek / 3
size = int(math.ceil(((len(inputCryp) + 1 / 3) / 2.)) * 2)
l2r = True
else:
offset, l2r, inputCryp = message.readMessage(keyPath, inputPath)
offset = offset[0]
size = len(inputCryp)
key = keymanagement.getKeyBytes(keyPath, size, offset=offset, l2r=l2r)
key = bytearray(key[0])
if mode == "human":
key = keymanagement.ba2humankeyba(key)
if len(key) != len(inputCryp):
key = key[:len(inputCryp) - len(key)]
inputCryp = ownbase32.string2ob32ba(inputCryp)
clear = ownbase32.ba2ob32string(vernam(inputCryp, key))
print(clear)
else:
clear = vernam(bytearray(inputCryp), key)
if mode == "lz4":
clear = uncompress(str(clear))
open(outputPath, 'wb').write(clear)
def read(self, byte_count=None):
while self.decompressed_buffer.tell() < byte_count:
try:
fname = self.file_iter.next()
except StopIteration:
break
compressed = open(fname).read()
new_uncompressed_bytes = lz4.uncompress(compressed) if compressed else ""
self.decompressed_buffer.write(new_uncompressed_bytes)
decompressed = self.decompressed_buffer.getvalue()
self.decompressed_buffer.close()
self.decompressed_buffer = StringIO()
self.decompressed_buffer.write(decompressed[byte_count:])
return decompressed[:byte_count]
def pump(self):
"""Receive raw data from socket and decode it as a dict"""
data = {} # fancy data
# Extract raw JSON data from socket and
# convert it into a dict (if possible)
try:
data_raw, addr = self.sock.recvfrom(NET_MAX_BYTES)
if self._use_lz4:
data_str = lz4.uncompress(data_raw)
else:
data_str = data_raw
data = self._codec.decode(data_str)
except Exception as e:
pass
# on_data_received is only called if data is not empty
if (isinstance(data, dict) or isinstance(data, list)) and len(data):
self.on_data_received(data, addr[0], addr[1])
def get_merged(self, eventId):
self.logger.debug("Merging event: " + eventId)
with self.lmdb.begin() as txn:
eventId = eventId.encode('utf-8')
index_db = self.lmdb.open_db(key=self.index_name,
txn=txn,
dupsort=True)
deltas_db = self.lmdb.open_db(key=self.deltas_name, txn=txn)
with txn.cursor(db=index_db) as cursor:
result = {}
if cursor.set_key(eventId):
for delta in cursor.iternext_dup():
result = self.merger.merge(
result,
json.loads(uncompress(txn.get(delta,
db=deltas_db))),
meta={'timestamp': str(int(time.time() * 1000))})
self.logger.trace("Merged event data for {id}:\n".format(
id=eventId.decode('utf-8')) + prettify(result))
return result
def lz4_pickle_load(filename):
path = pathlib.Path(filename)
with path.open('rb') as f:
return pickle.loads(lz4.uncompress(f.read()))
def _unpack(self, buf):
if self.compression:
return pickle.loads(uncompress(buf))
else:
return pickle.loads(buf)
def uncompressLz4(compressedData, size):
header = struct.pack("<I", size)
compressedDataWithHeader = header + compressedData
uncompressedData = lz4.uncompress(compressedDataWithHeader)
return uncompressedData
def test_datafeed_update_ip(self, glet_mock, SR_mock):
config = {}
chassis = mock.Mock()
chassis.request_sub_channel.return_value = None
ochannel = mock.Mock()
chassis.request_pub_channel.return_value = ochannel
chassis.request_rpc_channel.return_value = None
rpcmock = mock.Mock()
rpcmock.get.return_value = {'error': None, 'result': 'OK'}
chassis.send_rpc.return_value = rpcmock
b = minemeld.ft.taxii.DataFeed(FTNAME, chassis, config)
inputs = ['a']
output = False
b.connect(inputs, output)
b.mgmtbus_initialize()
b.start()
# __init__ + get chkp + delete chkp
self.assertEqual(len(SR_mock.mock_calls), 6)
SR_mock.reset_mock()
SR_mock.return_value.zcard.return_value = 1
# unicast
b.filtered_update('a',
indicator='1.1.1.1',
value={
'type': 'IPv4',
'confidence': 100,
'share_level': 'green',
'sources': ['test.1']
})
for call in SR_mock.mock_calls:
name, args, kwargs = call
if name == '().pipeline().__enter__().hset':
break
else:
self.fail(msg='hset not found')
self.assertEqual(args[2].startswith('lz4'), True)
stixdict = json.loads(lz4.uncompress(args[2][3:]))
indicator = stixdict['indicators'][0]
cyboxprops = indicator['observable']['object']['properties']
self.assertEqual(cyboxprops['address_value'], '1.1.1.1')
self.assertEqual(cyboxprops['xsi:type'], 'AddressObjectType')
SR_mock.reset_mock()
# CIDR
b.filtered_update('a',
indicator='1.1.1.0/24',
value={
'type': 'IPv4',
'confidence': 100,
'share_level': 'green',
'sources': ['test.1']
})
for call in SR_mock.mock_calls:
name, args, kwargs = call
if name == '().pipeline().__enter__().hset':
break
else:
self.fail(msg='hset not found')
self.assertEqual(args[2].startswith('lz4'), True)
stixdict = json.loads(lz4.uncompress(args[2][3:]))
indicator = stixdict['indicators'][0]
cyboxprops = indicator['observable']['object']['properties']
self.assertEqual(cyboxprops['address_value'], '1.1.1.0/24')
self.assertEqual(cyboxprops['xsi:type'], 'AddressObjectType')
SR_mock.reset_mock()
# fake range
b.filtered_update('a',
indicator='1.1.1.1-1.1.1.1',
value={
'type': 'IPv4',
'confidence': 100,
'share_level': 'green',
'sources': ['test.1']
})
for call in SR_mock.mock_calls:
name, args, kwargs = call
if name == '().pipeline().__enter__().hset':
break
else:
self.fail(msg='hset not found')
self.assertEqual(args[2].startswith('lz4'), True)
stixdict = json.loads(lz4.uncompress(args[2][3:]))
indicator = stixdict['indicators'][0]
cyboxprops = indicator['observable']['object']['properties']
self.assertEqual(cyboxprops['address_value'], '1.1.1.1')
self.assertEqual(cyboxprops['xsi:type'], 'AddressObjectType')
SR_mock.reset_mock()
# fake range 2
b.filtered_update('a',
indicator='1.1.1.0-1.1.1.31',
value={
'type': 'IPv4',
'confidence': 100,
'share_level': 'green',
'sources': ['test.1']
})
for call in SR_mock.mock_calls:
name, args, kwargs = call
if name == '().pipeline().__enter__().hset':
break
else:
self.fail(msg='hset not found')
self.assertEqual(args[2].startswith('lz4'), True)
stixdict = json.loads(lz4.uncompress(args[2][3:]))
indicator = stixdict['indicators'][0]
cyboxprops = indicator['observable']['object']['properties']
self.assertEqual(cyboxprops['address_value'], '1.1.1.0/27')
self.assertEqual(cyboxprops['xsi:type'], 'AddressObjectType')
SR_mock.reset_mock()
SR_mock.return_value.zcard.return_value = 1
# real range
b.filtered_update('a',
indicator='1.1.1.0-1.1.1.33',
value={
'type': 'IPv4',
'confidence': 100,
'share_level': 'green',
'sources': ['test.1']
})
for call in SR_mock.mock_calls:
name, args, kwargs = call
if name == '().pipeline().__enter__().hset':
break
else:
self.fail(msg='hset not found')
self.assertEqual(args[2].startswith('lz4'), True)
stixdict = json.loads(lz4.uncompress(args[2][3:]))
indicator = stixdict['indicators']
cyboxprops = indicator[0]['observable']['object']['properties']
self.assertEqual(cyboxprops['address_value'], '1.1.1.0/27')
self.assertEqual(cyboxprops['xsi:type'], 'AddressObjectType')
cyboxprops = indicator[1]['observable']['object']['properties']
self.assertEqual(cyboxprops['address_value'], '1.1.1.32/31')
self.assertEqual(cyboxprops['xsi:type'], 'AddressObjectType')
SR_mock.reset_mock()
b.stop()
def lz4_pickle_load(filename):
path = pathlib.Path(filename)
with path.open('rb') as f:
return pickle.loads(lz4.uncompress(f.read()))
import lz4
import sys
DATA = open("/dev/urandom", "rb").read(128 * 1024) # Read 128kb
sys.exit(DATA != lz4.uncompress(lz4.compress(DATA)) and 1 or 0)
