async def ping_loop(ctx, ping_interval, cycle_time, initial_ping_timeout,
ping_retries, backoff, loop, inventory_router_url):
"""
:param ctx:
:param ping_interval:
:param cycle_time:
:param initial_ping_timeout:
:param ping_retries:
:param backoff:
:param loop:
:param inventory_router_url:
:return:
"""
# load the queue
inventory_client = InventoryClient(inventory_router_url)
while True:
if stop_ping_loop:
log.info('Stopping ping loop')
break
log.debug('Looking for work')
now = time.time()
for mercury_id, data in list(active_state.items(
)): # copy to list because the list length could change
# out from under us
if now - data['last_ping'] > ping_interval and not data['pinging']:
log.debug('Scheduling ping for {}'.format(mercury_id))
active_state[mercury_id]['pinging'] = True
asyncio.ensure_future(ping(data, ctx, initial_ping_timeout,
ping_retries, backoff,
inventory_client),
loop=loop)
await asyncio.sleep(cycle_time)
def reacquire(inventory_url, backend_name):
"""
:param inventory_url:
:param backend_name:
:return:
"""
# Onetime use synchronous client
log.info('Attempting to reacquire active agents')
log.debug('Inventory Router: {}'.format(inventory_url))
inventory_client = InventoryClient(inventory_url,
# TODO: Add these to configuration
response_timeout=60,
rcv_retry=10)
existing_documents = inventory_client.query({'active': {'$ne': None},
'origin.name': backend_name},
projection={'mercury_id': 1,
'active': 1})
if existing_documents.get('error'): # Transport Error
log.error('[BACKEND CRITICAL] '
'Error communicating with inventory service, could not '
'reacquire: <{}>'.format(existing_documents.get('message')))
# Return without reacquiring any nodes. Once communication is
# reestablished, agents will begin to re-register
return
for doc in existing_documents['message']['items']:
if not BackendController.validate_agent_info(doc['active']):
log.error('Found junk in document {} expunging'.format(
doc['mercury_id']))
inventory_client.update_one(doc['mercury_id'], {'active': None})
log.info('Attempting to reacquire %s : %s' % (
doc['mercury_id'], doc['active']['rpc_address']))
add_active_record(doc)
log.info('Reacquire operation complete')
inventory_client.close()
def reacquire(self):
# Onetime use synchronous client
inventory_client = InventoryClient(self.inventory_router_url)
existing_documents = inventory_client.query({'active': {
'$ne': None
}},
projection={
'mercury_id': 1,
'active': 1
})
for doc in existing_documents['items']:
if not self.controller.validate_agent_info(doc['active']):
log.error('Found junk in document {} expunging'.format(
doc['mercury_id']))
inventory_client.update_one(doc['mercury_id'],
{'active': None})
log.info('Attempting to reacquire %s : %s' %
(doc['mercury_id'], doc['active']['rpc_address']))
add_active_record(doc)
inventory_client.close()
def __init__(self, *args, **kwargs):
super(DiscoverView, self).__init__(*args, **kwargs)
inventory_url = configuration.inventory.inventory_router
self.inventory_client = InventoryClient(inventory_url)
class DiscoverView(MethodView):
"""
Discover method view
"""
def __init__(self, *args, **kwargs):
super(DiscoverView, self).__init__(*args, **kwargs)
inventory_url = configuration.inventory.inventory_router
self.inventory_client = InventoryClient(inventory_url)
@staticmethod
def render_agent_script():
with open('scripts/agent.ipxe') as fp:
template = fp.read()
return pystache.render(template, **configuration)
@staticmethod
def plain(message):
return Response(message, content_type='text/plain')
def get(self, mac_address):
""" Attempt to relate a device using the provided mac address """
result = self.inventory_client.query(
{'interfaces.address': mac_address},
projection={
'boot': 1,
'mercury_id': 1,
'dmi': 1
})
if result.get('error'):
abort(500, result)
message = result['message']
if not message['total']:
log.info('New device: {}'.format(mac_address))
return self.plain(self.render_agent_script())
if message['total'] > 1:
log.error('DUPLICATE MAC ADDRESS: {}'.format(mac_address))
abort(500, 'Duplicate mac addresses in inventory')
inventory_data = message['items'][0]
boot_info = inventory_data.get('boot', {})
if boot_info.get('script'):
# Dangerous, can potentially leak entire configuration into image
# if an attacker knows the key names
return pystache.render(boot_info['script'],
dict(**inventory_data, **configuration))
boot_state = boot_info.get('state', 'agent')
if boot_state == 'local':
log.info('Booting {} from hard drive'.format(
inventory_data['mercury_id']))
return self.plain('#!ipxe\nexit\n')
elif boot_info == 'rescue':
log.info('Booting {} to rescue mode'.format(
inventory_data['mercury_id']))
return self.plain('Boot rescue iPXE script here')
log.info('Booting {} to agent'.format(inventory_data['mercury_id']))
print(request.base_url)
return self.plain(self.render_agent_script())
import bson
import logging
from bottle import route, run, request, HTTPResponse
from mercury.common.clients.inventory import InventoryClient
from mercury.common.clients.rpc.frontend import RPCFrontEndClient
from mercury_api.configuration import api_configuration
log = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)
inventory_router_url = api_configuration['inventory']['inventory_router']
rpc_router_url = api_configuration['rpc']['rpc_router']
inventory_client = InventoryClient(inventory_router_url)
rpc_client = RPCFrontEndClient(rpc_router_url)
def http_error(message, code=500):
return HTTPResponse({'error': True, 'message': message}, status=code)
def validate_json(f):
def wrapper(*args, **kwargs):
try:
if not request.json:
return http_error('JSON request is missing', code=400)
except ValueError:
log.debug('JSON request is malformed: {}'.format(
request.body.read()))