def test_watcher(redis_connection):
redis_connection.time = RedisTime()
rds = redis_connection
queue_name = get_random_id()
out_queue = NamedQueue(queue_name, rds)
try:
# Create a server and hijack its running flag and the current time in 'redis'
client = WatcherClient(rds)
server = WatcherServer(rds, rds)
server.running = ToggleTrue()
rds.time.current = 0
assert out_queue.length() == 0
# Send a simple event to occur soon
client.touch(10, 'one-second', queue_name, {'first': 'one'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
rds.time.current = 12 # Jump forward 12 seconds
server.try_run()
assert out_queue.length() == 1
assert out_queue.pop() == {'first': 'one'}
# Send a simple event to occur soon, then change our mind
client.touch(10, 'one-second', queue_name, {'first': 'one'})
client.touch(20, 'one-second', queue_name, {'first': 'one'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
# Set events to occur, in inverse order, reuse a key, overwrite content and timeout
client.touch(200, 'one-second', queue_name, {'first': 'last'})
client.touch(100, '100-second', queue_name, {'first': '100'})
client.touch(50, '50-second', queue_name, {'first': '50'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
for _ in range(15):
rds.time.current += 20
server.try_run()
assert out_queue.length() == 3
assert out_queue.pop() == {'first': '50'}
assert out_queue.pop() == {'first': '100'}
assert out_queue.pop() == {'first': 'last'}
# Send a simple event to occur soon, then stop it
rds.time.current = 0
client.touch(10, 'one-second', queue_name, {'first': 'one'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
client.clear('one-second')
rds.time.current = 12 # Jump forward 12 seconds
server.try_run()
assert out_queue.length() == 0 # still nothing because it was cleared
finally:
out_queue.delete()
def default_authenticator(auth, req, ses, storage):
# This is assemblyline authentication procedure
# It will try to authenticate the user in the following order until a method is successful
# apikey
# username/password
# PKI DN
#
# During the authentication procedure the user/pass and DN methods will be subject to OTP challenge
# if OTP is allowed on the server and has been turned on by the user
#
# Apikey authentication procedure is not subject to OTP challenge but has limited functionality
apikey = auth.get('apikey', None)
otp = auth.get('otp', 0)
webauthn_auth_resp = auth.get('webauthn_auth_resp', None)
state = ses.pop('state', None)
password = auth.get('password', None)
uname = auth.get('username', None)
oauth_token = auth.get('oauth_token', None)
if not uname:
raise AuthenticationException('No user specified for authentication')
# Bruteforce protection
auth_fail_queue = NamedQueue("ui-failed-%s" % uname, **nonpersistent_config)
if auth_fail_queue.length() >= config.auth.internal.max_failures:
# Failed 'max_failures' times, stop trying... This will timeout in 'failure_ttl' seconds
raise AuthenticationException("Maximum password retry of {retry} was reached. "
"This account is locked for the next {ttl} "
"seconds...".format(retry=config.auth.internal.max_failures,
ttl=config.auth.internal.failure_ttl))
try:
validated_user, priv = validate_apikey(uname, apikey, storage)
if validated_user:
return validated_user, priv
validated_user, priv = validate_oauth(uname, oauth_token)
if not validated_user:
validated_user, priv = validate_ldapuser(uname, password, storage)
if not validated_user:
validated_user, priv = validate_userpass(uname, password, storage)
if validated_user:
validate_2fa(validated_user, otp, state, webauthn_auth_resp, storage)
return validated_user, priv
except AuthenticationException:
# Failure appended, push failure parameters
auth_fail_queue.push({
'remote_addr': req.remote_addr,
'host': req.host,
'full_path': req.full_path
})
raise
raise AuthenticationException("None of the authentication methods succeeded")
class HeartbeatFormatter(object):
def __init__(self, sender, log, config=None, redis=None):
self.sender = sender
self.log = log
self.config = config or forge.get_config()
self.datastore = forge.get_datastore(self.config)
self.redis = redis or get_client(
host=self.config.core.redis.nonpersistent.host,
port=self.config.core.redis.nonpersistent.port,
private=False,
)
self.redis_persist = get_client(
host=self.config.core.redis.persistent.host,
port=self.config.core.redis.persistent.port,
private=False,
)
self.status_queue = CommsQueue(STATUS_QUEUE, self.redis)
self.dispatch_active_hash = Hash(DISPATCH_TASK_HASH,
self.redis_persist)
self.dispatcher_submission_queue = NamedQueue(SUBMISSION_QUEUE,
self.redis)
self.ingest_scanning = Hash('m-scanning-table', self.redis_persist)
self.ingest_unique_queue = PriorityQueue('m-unique',
self.redis_persist)
self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.redis_persist)
self.ingest_complete_queue = NamedQueue(COMPLETE_QUEUE_NAME,
self.redis)
self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.redis_persist)
constants = forge.get_constants(self.config)
self.c_rng = constants.PRIORITY_RANGES['critical']
self.h_rng = constants.PRIORITY_RANGES['high']
self.m_rng = constants.PRIORITY_RANGES['medium']
self.l_rng = constants.PRIORITY_RANGES['low']
self.c_s_at = self.config.core.ingester.sampling_at['critical']
self.h_s_at = self.config.core.ingester.sampling_at['high']
self.m_s_at = self.config.core.ingester.sampling_at['medium']
self.l_s_at = self.config.core.ingester.sampling_at['low']
self.to_expire = {k: 0 for k in metrics.EXPIRY_METRICS}
if self.config.core.expiry.batch_delete:
self.delete_query = f"expiry_ts:[* TO {self.datastore.ds.now}-{self.config.core.expiry.delay}" \
f"{self.datastore.ds.hour}/DAY]"
else:
self.delete_query = f"expiry_ts:[* TO {self.datastore.ds.now}-{self.config.core.expiry.delay}" \
f"{self.datastore.ds.hour}]"
self.scheduler = BackgroundScheduler(daemon=True)
self.scheduler.add_job(
self._reload_expiry_queues,
'interval',
seconds=self.config.core.metrics.export_interval * 4)
self.scheduler.start()
def _reload_expiry_queues(self):
try:
self.log.info("Refreshing expiry queues...")
for collection_name in metrics.EXPIRY_METRICS:
try:
collection = getattr(self.datastore, collection_name)
self.to_expire[collection_name] = collection.search(
self.delete_query,
rows=0,
fl='id',
track_total_hits="true")['total']
except SearchException:
self.to_expire[collection_name] = 0
except Exception:
self.log.exception(
"Unknown exception occurred while reloading expiry queues:")
def send_heartbeat(self, m_type, m_name, m_data, instances):
if m_type == "dispatcher":
try:
instances = sorted(Dispatcher.all_instances(
self.redis_persist))
inflight = {
_i: Dispatcher.instance_assignment_size(
self.redis_persist, _i)
for _i in instances
}
queues = {
_i: Dispatcher.all_queue_lengths(self.redis, _i)
for _i in instances
}
msg = {
"sender": self.sender,
"msg": {
"inflight": {
"max": self.config.core.dispatcher.max_inflight,
"outstanding": self.dispatch_active_hash.length(),
"per_instance": [inflight[_i] for _i in instances]
},
"instances": len(instances),
"metrics": m_data,
"queues": {
"ingest":
self.dispatcher_submission_queue.length(),
"start": [queues[_i]['start'] for _i in instances],
"result":
[queues[_i]['result'] for _i in instances],
"command":
[queues[_i]['command'] for _i in instances]
},
"component": m_name,
}
}
self.status_queue.publish(
DispatcherMessage(msg).as_primitives())
self.log.info(f"Sent dispatcher heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating DispatcherMessage")
elif m_type == "ingester":
try:
c_q_len = self.ingest_unique_queue.count(*self.c_rng)
h_q_len = self.ingest_unique_queue.count(*self.h_rng)
m_q_len = self.ingest_unique_queue.count(*self.m_rng)
l_q_len = self.ingest_unique_queue.count(*self.l_rng)
msg = {
"sender": self.sender,
"msg": {
"instances": instances,
"metrics": m_data,
"processing": {
"inflight": self.ingest_scanning.length()
},
"processing_chance": {
"critical": 1 - drop_chance(c_q_len, self.c_s_at),
"high": 1 - drop_chance(h_q_len, self.h_s_at),
"low": 1 - drop_chance(l_q_len, self.l_s_at),
"medium": 1 - drop_chance(m_q_len, self.m_s_at)
},
"queues": {
"critical": c_q_len,
"high": h_q_len,
"ingest": self.ingest_queue.length(),
"complete": self.ingest_complete_queue.length(),
"low": l_q_len,
"medium": m_q_len
}
}
}
self.status_queue.publish(IngestMessage(msg).as_primitives())
self.log.info(f"Sent ingester heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating IngestMessage")
elif m_type == "alerter":
try:
msg = {
"sender": self.sender,
"msg": {
"instances": instances,
"metrics": m_data,
"queues": {
"alert": self.alert_queue.length()
}
}
}
self.status_queue.publish(AlerterMessage(msg).as_primitives())
self.log.info(f"Sent alerter heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating AlerterMessage")
elif m_type == "expiry":
try:
msg = {
"sender": self.sender,
"msg": {
"instances": instances,
"metrics": m_data,
"queues": self.to_expire
}
}
self.status_queue.publish(ExpiryMessage(msg).as_primitives())
self.log.info(f"Sent expiry heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating ExpiryMessage")
elif m_type == "archive":
try:
msg = {
"sender": self.sender,
"msg": {
"instances": instances,
"metrics": m_data
}
}
self.status_queue.publish(ArchiveMessage(msg).as_primitives())
self.log.info(f"Sent archive heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating ArchiveMessage")
elif m_type == "scaler":
try:
msg = {
"sender": self.sender,
"msg": {
"instances": instances,
"metrics": m_data,
}
}
self.status_queue.publish(ScalerMessage(msg).as_primitives())
self.log.info(f"Sent scaler heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating WatcherMessage")
elif m_type == "scaler_status":
try:
msg = {
"sender": self.sender,
"msg": {
"service_name": m_name,
"metrics": m_data,
}
}
self.status_queue.publish(
ScalerStatusMessage(msg).as_primitives())
self.log.info(f"Sent scaler status heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating WatcherMessage")
elif m_type == "service":
try:
busy, idle = get_working_and_idle(self.redis, m_name)
msg = {
"sender": self.sender,
"msg": {
"instances": len(busy) + len(idle),
"metrics": m_data,
"activity": {
'busy': len(busy),
'idle': len(idle)
},
"queue": get_service_queue(m_name,
self.redis).length(),
"service_name": m_name
}
}
self.status_queue.publish(ServiceMessage(msg).as_primitives())
self.log.info(f"Sent service heartbeat: {msg['msg']}")
except Exception:
self.log.exception(
"An exception occurred while generating ServiceMessage")
else:
self.log.warning(
f"Skipping unknown counter: {m_name} [{m_type}] ==> {m_data}")