def explain_state(self, task):
log.info('Got explain_state message.')
nq = NamedQueue(task.watch_queue)
submission = self.entries.get(task.sid, None)
if submission:
has_timeout = False
for v in self.ack_timeout.itervalues():
for t in v:
if t.sid == task.sid:
has_timeout = True
for v in self.service_timeout.itervalues():
for t in v:
if t.sid == task.sid:
has_timeout = True
if not has_timeout:
nq.push({
'srl': 0,
'message': 'No timeouts for this submission!',
'depth': 0,
})
for entry in submission.itervalues():
if not entry.task.psrl:
explain(entry, nq)
nq.push(False)
def outstanding_services(self, task):
nq = NamedQueue(task.watch_queue)
outstanding = {}
submission = self.entries.get(task.sid, None)
if submission:
for entry in submission.itervalues():
get_outstanding_services(entry, outstanding)
nq.push(outstanding)
def wait_for_networking(timeout):
uid = uuid.uuid4().get_hex()
for _each_second in xrange(timeout):
try:
q = NamedQueue('hostagent-redischeck-%s' % uid)
q.push('can i reach you')
q.pop(timeout=1, blocking=False)
return True
except Exception as e:
print('waiting for redis reachability. %s ' % str(e))
return False
def sysprep(self):
"""Basic prep and return."""
self._init_registration()
self._init_queues()
self.log.info('performing sysprep')
self._clear_tempdir()
self._wait_for_networking(20)
self._check_time_drift()
if not self.registration:
raise ProvisioningError('Host registration not found.')
if not self.registration.get('enabled', None):
raise ProvisioningError('Host explicitly disabled.')
profile_name = self.registration.get('profile', None)
if not profile_name:
raise ProvisioningError('Host has no assigned profile.')
if 'profile_definition' not in self.registration:
self.host_profile = self.store.get_profile(profile_name)
if not self.host_profile:
raise ProvisioningError(
'Host profile does not appear to exist in datastore: %s.',
profile_name)
self.log.info('Our profile: %s', pprint.pformat(self.host_profile))
else:
self.host_profile = self.registration.get('profile_definition', {})
self.log.info('Our profile: %s', pprint.pformat(self.host_profile))
vm_config = self.host_profile.get('virtual_machines', {})
if vm_config and not profile_name.startswith('flex'):
from assemblyline.al.common.vm import VmManager
self.vm_manager = VmManager(vm_config)
self.vm_manager.sysprep()
# if we are are running within a VM. patch hosts files.
if self.is_a_vm():
nq = NamedQueue('vm-%s' % self.mac, db=DATABASE_NUM)
nq.push(self.registration)
def watch(self, task):
queue = task.watch_queue
sid = task.sid
# Make sure this submission exists.
watchers = self.watchers.get(sid, {})
# Bail if we have a watcher with the same name for this sid.
if queue in watchers:
return
ttl = 0
try:
ttl = config.core.dispatcher.timeouts.watch_queue
except: # pylint: disable=W0702
pass
w = NamedQueue(queue, ttl=ttl)
errors = self.errors.get(sid, None)
results = self.results.get(sid, None)
if results is None and errors is None:
# TODO: Should we send UNKNOWN.
w.push({'status': 'STOP'})
return
watchers[queue] = w
self.watchers[sid] = watchers
# Send all cache keys to the newly created queue.
# Afterward they will be sent as they are received.
w.push({'status': 'START'})
if results:
w.push(*[{'status': 'OK', 'cache_key': c} for c in results])
if errors:
w.push(*[{'status': 'FAIL', 'cache_key': c} for c in errors])
def writer(self):
queue = {}
store = forge.get_datastore()
while self.running:
try:
msg = self.storage_queue.pop(timeout=1)
if not msg:
if self.drain:
break
continue
response = None
if msg['type'] == 'complete':
key = msg['filescore_key']
if key:
store.save_filescore(
key, msg['expiry'], {
'psid': msg['psid'],
'sid': msg['sid'],
'score': msg['score'],
'time': msg['now'],
})
elif msg['type'] == 'error':
name, response = msg['name'], msg['response']
response.cache_key = \
store.save_error(name, None, None, response)
q.send_raw(response.as_dispatcher_response())
elif msg['type'] == 'finalize':
store.finalize_submission(msg['sid'],
msg['classification'],
msg['errors'], msg['results'],
msg['score'])
completed_queue = msg['completed_queue']
if completed_queue:
cq = queue.get(completed_queue, None)
if not cq:
cq = NamedQueue(completed_queue)
queue[completed_queue] = cq
cq.push(msg['raw'])
# Send complete message to any watchers.
for w in msg['watchers'].itervalues():
w.push({'status': 'STOP'})
else:
log.warning("Unhandled message type: %s",
msg.get('type', '<unknown>'))
except riak.RiakError:
msg['retries'] = retries = msg.get('retries', 0) + 1
if retries > 5:
log.exception("Max retries exceeded")
continue
self.storage_queue.push(msg)
log.exception("Problem doing %s", msg.get('type', 'unknown'))
except Exception: # pylint:disable=W0702
log.exception('Problem in writer')
# TODO: Should we sleep for a bit here to avoid flailing?
store.close()
def outstanding_submissions(self, task):
nq = NamedQueue(task.watch_queue)
nq.push({'sids': self.entries.keys()})
def get_system_time(self, task):
nq = NamedQueue(task.watch_queue)
nq.push({'time': time.time()})
def _init_registration(self):
if self.is_a_vm():
nq = NamedQueue('vm-%s' % self.mac, db=DATABASE_NUM)
reg = nq.pop()
nq.push(reg)
self.log.info('Updating our registration.')
reg['hostname'] = net.get_hostname()
reg['ip'] = self.ip
reg['machine_info'] = sysinfo.get_machine_info()
reg['last_checkin'] = isotime.now_as_iso()
reg['platform'] = sysinfo.get_platform()
reg['updated'] = time.asctime()
reg['system_name'] = config.system.name
if 'roles' not in reg:
reg['roles'] = []
if "hostagent" not in reg["roles"]:
reg['roles'].append("hostagent")
else:
reg = self.store.get_node(self.mac)
if not reg:
self.log.info(
'This appears to be our first run on this host. Registering ourselves.'
)
reg = DEFAULT_REGISTRATION.copy()
reg['hostname'] = net.get_hostname()
reg['ip'] = self.ip
reg['mac_address'] = self.mac
reg['machine_info'] = sysinfo.get_machine_info()
reg['last_checkin'] = isotime.now_as_iso()
reg['platform'] = sysinfo.get_platform()
reg['profile'] = 'idle'
reg['created'] = time.asctime()
if 'roles' not in reg:
reg['roles'] = []
if "controller" not in reg["roles"]:
reg['roles'].append("controller")
if "hostagent" not in reg["roles"]:
reg['roles'].append("hostagent")
self.store.save_node(self.mac, reg)
else:
# Just do an update of the extra info in registration.
self.log.info('Updating our registration.')
reg['hostname'] = net.get_hostname()
reg['ip'] = self.ip
if not reg.get('profile', None):
reg['profile'] = config.workers.default_profile
reg['machine_info'] = sysinfo.get_machine_info()
reg['last_checkin'] = isotime.now_as_iso()
reg['platform'] = sysinfo.get_platform()
reg['updated'] = time.asctime()
reg['system_name'] = config.system.name
if 'roles' not in reg:
reg['roles'] = []
if "controller" not in reg["roles"] and not reg.get(
'is_vm', False):
reg['roles'].append("controller")
if "hostagent" not in reg["roles"]:
reg['roles'].append("hostagent")
self.store.save_node(self.mac, reg)
self.registration = reg
msgs = forge.apply_overrides(reg.get('config_overrides', None))
if msgs:
self.log.info("Using %s.", " and ".join(msgs))
self.log.info('Our registration: %s',
pprint.pformat(self.registration))
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)
dn = auth.get('dn', None)
otp = auth.get('otp', 0)
u2f_response = auth.get('u2f_response', None)
u2f_challenge = ses.pop('_u2f_challenge_', None)
password = auth.get('password', None)
uname = auth.get('username', None) or dn
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 = apikey_handler(uname, apikey, storage)
if validated_user:
return validated_user, priv
validated_user, priv = userpass_handler(uname, password, storage)
if validated_user:
validate_2fa(validated_user, otp, u2f_challenge, u2f_response,
storage)
return validated_user, priv
validated_user, priv = dn_handler(dn, storage)
if validated_user:
validate_2fa(validated_user, otp, u2f_challenge, u2f_response,
storage)
return validated_user, priv
except AuthenticationException as ae:
# 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")