class UnitTestSuite(object):
def setup(self):
current_date = datetime.now()
self.event_index_name = current_date.strftime("events-%Y%m%d")
self.previous_event_index_name = (
current_date - timedelta(days=1)).strftime("events-%Y%m%d")
self.alert_index_name = current_date.strftime("alerts-%Y%m")
self.parse_config()
# Elasticsearch
self.es_client = ElasticsearchClient(
list('{0}'.format(s) for s in self.options.esservers))
# RabbitMQ
mqConnString = 'amqp://{0}:{1}@{2}:{3}//'.format(
self.options.mquser, self.options.mqpassword,
self.options.mqalertserver, self.options.mqport)
mqAlertConn = Connection(mqConnString)
alertExchange = Exchange(name=self.options.alertExchange,
type='topic',
durable=True,
delivery_mode=1)
alertExchange(mqAlertConn).declare()
alertQueue = Queue(self.options.queueName,
exchange=alertExchange,
routing_key=self.options.alerttopic,
durable=False,
no_ack=(not self.options.mqack))
alertQueue(mqAlertConn).declare()
self.rabbitmq_alerts_consumer = mqAlertConn.Consumer(alertQueue,
accept=['json'])
if pytest.config.option.delete_indexes:
self.reset_elasticsearch()
self.setup_elasticsearch()
if pytest.config.option.delete_queues:
self.reset_rabbitmq()
def parse_config(self):
default_config = os.path.join(os.path.dirname(__file__), "config.conf")
options = DotDict()
options.configfile = default_config
options.esservers = list(
getConfig('esservers', 'http://localhost:9200',
options.configfile).split(','))
options.alertExchange = getConfig('alertexchange', 'alerts',
options.configfile)
options.queueName = getConfig('alertqueuename', 'alertBot',
options.configfile)
options.alerttopic = getConfig('alerttopic', 'mozdef.*',
options.configfile)
options.mquser = getConfig('mquser', 'guest', options.configfile)
options.mqalertserver = getConfig('mqalertserver', 'localhost',
options.configfile)
options.mqpassword = getConfig('mqpassword', 'guest',
options.configfile)
options.mqport = getConfig('mqport', 5672, options.configfile)
options.mqack = getConfig('mqack', True, options.configfile)
self.options = options
def reset_rabbitmq(self):
self.rabbitmq_alerts_consumer.channel.queue_purge()
def teardown(self):
if pytest.config.option.delete_indexes:
self.reset_elasticsearch()
if pytest.config.option.delete_queues:
self.reset_rabbitmq()
self.rabbitmq_alerts_consumer.connection.close()
self.rabbitmq_alerts_consumer.close()
def populate_test_event(self, event, event_type='event'):
self.es_client.save_event(body=event, doc_type=event_type)
def populate_test_object(self, event, event_type='event'):
self.es_client.save_object(index='events',
body=event,
doc_type=event_type)
def setup_elasticsearch(self):
default_mapping_file = os.path.join(
os.path.dirname(__file__), "../config/defaultMappingTemplate.json")
mapping_str = ''
with open(default_mapping_file) as data_file:
mapping_str = data_file.read()
self.es_client.create_index(self.event_index_name, mapping=mapping_str)
self.es_client.create_alias('events', self.event_index_name)
self.es_client.create_index(self.previous_event_index_name,
mapping=mapping_str)
self.es_client.create_alias('events-previous',
self.previous_event_index_name)
self.es_client.create_index(self.alert_index_name, mapping=mapping_str)
self.es_client.create_alias('alerts', self.alert_index_name)
def reset_elasticsearch(self):
self.es_client.delete_index(self.event_index_name, True)
self.es_client.delete_index('events', True)
self.es_client.delete_index(self.previous_event_index_name, True)
self.es_client.delete_index('events-previous', True)
self.es_client.delete_index(self.alert_index_name, True)
self.es_client.delete_index('alerts', True)
def flush(self, index_name):
self.es_client.flush(index_name)
def random_ip(self):
return str(random.randint(1, 255)) + "." + str(random.randint(
1, 255)) + "." + str(random.randint(1, 255)) + "." + str(
random.randint(1, 255))
def generate_default_event(self):
current_timestamp = UnitTestSuite.current_timestamp_lambda()
source_ip = self.random_ip()
event = {
"_index": "events",
"_type": "event",
"_source": {
"category": "excategory",
"utctimestamp": current_timestamp,
"receivedtimestamp": current_timestamp,
"mozdefhostname": "mozdefhost",
"hostname": "exhostname",
"severity": "NOTICE",
"source": "exsource",
"summary": "Example summary",
"tags": ['tag1', 'tag2'],
"details": {
"sourceipaddress": source_ip,
"hostname": "exhostname"
}
}
}
return event
def verify_event(self, event, expected_event):
assert sorted(event.keys()) == sorted(expected_event.keys())
for key, value in expected_event.iteritems():
if key == 'receivedtimestamp':
assert type(event[key]) == unicode
else:
assert event[
key] == value, 'Incorrect match for {0}, expected: {1}'.format(
key, value)
@staticmethod
def current_timestamp():
return toUTC(datetime.now()).isoformat()
@staticmethod
def subtract_from_timestamp(date_timedelta, timestamp=None):
if timestamp is None:
timestamp = UnitTestSuite.current_timestamp()
utc_time = parse(timestamp)
custom_date = utc_time - timedelta(**date_timedelta)
return custom_date.isoformat()
@staticmethod
def create_timestamp_from_now(hour, minute, second):
return toUTC(datetime.now().replace(hour=hour,
minute=minute,
second=second).isoformat())
@staticmethod
def current_timestamp_lambda():
return lambda: UnitTestSuite.current_timestamp()
@staticmethod
def subtract_from_timestamp_lambda(date_timedelta, timestamp=None):
return lambda: UnitTestSuite.subtract_from_timestamp(
date_timedelta, timestamp)
@staticmethod
def create_timestamp_from_now_lambda(hour, minute, second):
return lambda: UnitTestSuite.create_timestamp_from_now(
hour, minute, second)
def esRotateIndexes():
if options.output == 'syslog':
logger.addHandler(SysLogHandler(address=(options.sysloghostname, options.syslogport)))
else:
sh = logging.StreamHandler(sys.stderr)
sh.setFormatter(formatter)
logger.addHandler(sh)
logger.debug('started')
try:
es = ElasticsearchClient((list('{0}'.format(s) for s in options.esservers)))
indices = es.get_indices()
# calc dates for use in index names events-YYYYMMDD, alerts-YYYYMM, etc.
odate_day = date.strftime(toUTC(datetime.now()) - timedelta(days=1), '%Y%m%d')
odate_month = date.strftime(toUTC(datetime.now()) - timedelta(days=1), '%Y%m')
ndate_day = date.strftime(toUTC(datetime.now()), '%Y%m%d')
ndate_month = date.strftime(toUTC(datetime.now()), '%Y%m')
# examine each index in the .conf file
# for rotation settings
for (index, dobackup, rotation, pruning) in zip(options.indices, options.dobackup, options.rotation, options.pruning):
try:
if rotation != 'none':
oldindex = index
newindex = index
if rotation == 'daily':
oldindex += '-%s' % odate_day
newindex += '-%s' % ndate_day
elif rotation == 'monthly':
oldindex += '-%s' % odate_month
newindex += '-%s' % ndate_month
# do not rotate before the month ends
if oldindex == newindex:
logger.debug('do not rotate %s index, month has not changed yet' % index)
continue
if newindex not in indices:
logger.debug('Creating %s index' % newindex)
es.create_index(newindex)
# set aliases: events to events-YYYYMMDD
# and events-previous to events-YYYYMMDD-1
logger.debug('Setting {0} alias to index: {1}'.format(index, newindex))
es.create_alias(index, newindex)
if oldindex in indices:
logger.debug('Setting {0}-previous alias to index: {1}'.format(index, oldindex))
es.create_alias('%s-previous' % index, oldindex)
else:
logger.debug('Old index %s is missing, do not change %s-previous alias' % (oldindex, index))
except Exception as e:
logger.error("Unhandled exception while rotating %s, terminating: %r" % (index, e))
indices = es.get_indices()
# Create weekly aliases for certain indices
week_ago_date = toUTC(datetime.now()) - timedelta(weeks=1)
week_ago_str = week_ago_date.strftime('%Y%m%d')
current_date = toUTC(datetime.now())
for index in options.weekly_rotation_indices:
weekly_index_alias = '%s-weekly' % index
logger.debug('Trying to re-alias {0} to indices since {1}'.format(weekly_index_alias, week_ago_str))
existing_weekly_indices = []
for day_obj in daterange(week_ago_date, current_date):
day_str = day_obj.strftime('%Y%m%d')
day_index = index + '-' + str(day_str)
if day_index in indices:
existing_weekly_indices.append(day_index)
else:
logger.debug('%s not found, so cant assign weekly alias' % day_index)
if existing_weekly_indices:
logger.debug('Creating {0} alias for {1}'.format(weekly_index_alias, existing_weekly_indices))
es.create_alias_multiple_indices(weekly_index_alias, existing_weekly_indices)
else:
logger.warning('No indices within the past week to assign events-weekly to')
except Exception as e:
logger.error("Unhandled exception, terminating: %r" % e)
class UnitTestSuite(object):
def setup(self):
self.event_index_name = datetime.now().strftime("events-%Y%m%d")
self.previous_event_index_name = (
datetime.now() - timedelta(days=1)).strftime("events-%Y%m%d")
self.alert_index_name = datetime.now().strftime("alerts-%Y%m")
self.es_client = ElasticsearchClient(ES['servers'])
if pytest.config.option.delete_indexes:
self.reset_elasticsearch()
self.setup_elasticsearch()
def teardown(self):
if pytest.config.option.delete_indexes:
self.reset_elasticsearch()
def populate_test_event(self, event, event_type='event'):
self.es_client.save_event(body=event, doc_type=event_type)
self.es_client.flush(self.event_index_name)
def setup_elasticsearch(self):
self.es_client.create_index(self.event_index_name)
self.es_client.create_alias('events', self.event_index_name)
self.es_client.create_index(self.previous_event_index_name)
self.es_client.create_alias('events-previous',
self.previous_event_index_name)
self.es_client.create_index(self.alert_index_name)
self.es_client.create_alias('alerts', self.alert_index_name)
def reset_elasticsearch(self):
self.es_client.delete_index(self.event_index_name, True)
self.es_client.delete_index('events', True)
self.es_client.delete_index(self.previous_event_index_name, True)
self.es_client.delete_index('events-previous', True)
self.es_client.delete_index(self.alert_index_name, True)
self.es_client.delete_index('alerts', True)
def random_ip(self):
return str(random.randint(1, 255)) + "." + str(random.randint(
1, 255)) + "." + str(random.randint(1, 255)) + "." + str(
random.randint(1, 255))
def generate_default_event(self):
current_timestamp = UnitTestSuite.current_timestamp_lambda()
source_ip = self.random_ip()
event = {
"_index": "events",
"_type": "event",
"_source": {
"category": "excategory",
"utctimestamp": current_timestamp,
"hostname": "exhostname",
"severity": "NOTICE",
"source": "exsource",
"summary": "Example summary",
"tags": ['tag1', 'tag2'],
"details": {
"sourceipaddress": source_ip,
"hostname": "exhostname"
}
}
}
return event
def verify_event(self, event, expected_event):
assert sorted(event.keys()) == sorted(expected_event.keys())
for key, value in expected_event.iteritems():
if key == 'receivedtimestamp':
assert type(event[key]) == unicode
else:
assert event[
key] == value, 'Incorrect match for {0}, expected: {1}'.format(
key, value)
@staticmethod
def current_timestamp():
return toUTC(datetime.now()).isoformat()
@staticmethod
def subtract_from_timestamp(date_timedelta, timestamp=None):
if timestamp is None:
timestamp = UnitTestSuite.current_timestamp()
utc_time = parse(timestamp)
custom_date = utc_time - timedelta(**date_timedelta)
return custom_date.isoformat()
@staticmethod
def current_timestamp_lambda():
return lambda: UnitTestSuite.current_timestamp()
@staticmethod
def subtract_from_timestamp_lambda(date_timedelta, timestamp=None):
return lambda: UnitTestSuite.subtract_from_timestamp(
date_timedelta, timestamp)
index_settings['settings'] = {
"index": {
"refresh_interval": refresh_interval,
"number_of_shards": number_of_shards,
"number_of_replicas": number_of_replicas,
"search.slowlog.threshold.query.warn": slowlog_threshold_query_warn,
"search.slowlog.threshold.fetch.warn": slowlog_threshold_fetch_warn,
"mapping.total_fields.limit": mapping_total_fields_limit
}
}
if event_index_name not in all_indices:
print "Creating " + event_index_name
client.create_index(event_index_name, index_config=index_settings)
client.create_alias('events', event_index_name)
if previous_event_index_name not in all_indices:
print "Creating " + previous_event_index_name
client.create_index(previous_event_index_name, index_config=index_settings)
client.create_alias('events-previous', previous_event_index_name)
if alert_index_name not in all_indices:
print "Creating " + alert_index_name
client.create_index(alert_index_name)
client.create_alias('alerts', alert_index_name)
if weekly_index_alias not in all_indices:
print "Creating " + weekly_index_alias
client.create_alias_multiple_indices(weekly_index_alias, [event_index_name, previous_event_index_name])
