def __init__(self, **kargs):
self.path_dirs = PathDirs(**kargs)
self.manifest = join(self.path_dirs.meta_dir, 'plugin_manifest.cfg')
self.p_helper = PluginHelper(**kargs)
self.d_client = docker.from_env()
self.logger = Logger(__name__)
self.plugin_config_file = self.path_dirs.plugin_config_file
def while_waiting(self):
""" Update fields periodically if nothing is happening """
# give a little extra time for file descriptors to close
time.sleep(0.1)
self.addfield.value = Timestamp()
self.addfield.display()
self.addfield2.value = Uptime()
self.addfield2.display()
self.addfield3.value = str(len(Containers())) + ' running'
if len(Containers()) > 0:
self.addfield3.labelColor = 'GOOD'
else:
self.addfield3.labelColor = 'DEFAULT'
self.addfield3.display()
# if file drop location changes deal with it
logger = Logger(__name__)
if self.file_drop.value != DropLocation()[1]:
logger.info('Starting: file drop restart')
try:
self.file_drop.value = DropLocation()[1]
logger.info('Path given: ' + str(self.file_drop.value))
except Exception as e: # pragma no cover
logger.error('file drop restart failed with error: ' + str(e))
logger.info('Finished: file drop restart')
self.file_drop.display()
return
def __init__(self, action_dict=None, action_name=None, *args, **keywords):
api_action = Tools()
action = {'api_action': api_action}
if action_dict:
action.update(action_dict)
logger = Logger(action_name)
InventoryForm.__init__(self, action, logger, *args, **keywords)
def __init__(self, *args, **kwargs):
self.d_client = docker.from_env()
self.path_dirs = PathDirs(**kwargs)
self.manifest = join(self.path_dirs.meta_dir,
'plugin_manifest.cfg')
self.vent_config = self.path_dirs.cfg_file
self.startup_file = self.path_dirs.startup_file
self.logger = Logger(__name__)
self._auto_install()
def __init__(self,
version='HEAD',
branch='master',
user=None,
pw=None,
*args,
**kwargs):
self.version = version
self.branch = branch
self.user = user
self.pw = pw
self.d_client = docker.from_env()
self.path_dirs = PathDirs(**kwargs)
self.path_dirs.host_config()
self.manifest = join(self.path_dirs.meta_dir, 'plugin_manifest.cfg')
self.logger = Logger(__name__)
def __init__(self, *args, **keywords):
""" Initialize tool form objects """
self.logger = Logger(__name__)
self.logger.info(str(keywords['names']))
self.api_action = Action()
self.m_helper = MenuHelper()
action = {'api_action': self.api_action}
self.tools_tc = {}
self.repo_widgets = {}
if keywords['action_dict']:
action.update(keywords['action_dict'])
if keywords['names']:
i = 1
for name in keywords['names']:
action['action_object' + str(i)] = getattr(
self.api_action, name)
i += 1
self.action = action
# get list of all possible group views to display
self.views = deque()
possible_groups = set()
manifest = Template(self.api_action.plugin.manifest)
if self.action['cores']:
tools = self.api_action.inventory(choices=['core'])[1]['core']
else:
tools = self.api_action.inventory(choices=['tools'])[1]['tools']
for tool in tools:
groups = manifest.option(tool, 'groups')[1].split(',')
for group in groups:
# don't do core because that's the purpose of all in views
if group != '' and group != 'core':
possible_groups.add(group)
self.manifest = manifest
self.views += possible_groups
self.views.append('all groups')
self.no_instance = ['build', 'remove']
super(ToolForm, self).__init__(*args, **keywords)
def __init__(self, **kargs):
self.api_action = Action(**kargs)
self.plugin = self.api_action.plugin
self.p_helper = self.api_action.p_helper
self.logger = Logger(__name__)
def __init__(self, manifest):
self.manifest = manifest
self.d_client = docker.from_env()
self.logger = Logger(__name__)
from os.path import expanduser
from os.path import join
from subprocess import check_output
from subprocess import PIPE
from subprocess import Popen
from subprocess import STDOUT
import docker
import pkg_resources
import requests
from vent.helpers.logs import Logger
from vent.helpers.paths import PathDirs
from vent.helpers.templates import Template
logger = Logger(__name__)
def Logs(c_type=None, grep_list=None):
""" Generically filter logs stored in log containers """
def get_logs(logs, log_entries):
try:
for log in logs:
if str(container.name) in log_entries:
log_entries[str(container.name)].append(log)
else:
log_entries[str(container.name)] = [log]
except Exception as e: # pragma: no cover
logger.error('Unable to get logs for ' +
str(container.name) +
' because: ' + str(e))
class UpdateToolsForm(npyscreen.ActionForm):
""" For picking which tools to update """
tools_tc = {}
triggered = 0
logger = Logger(__name__)
def create(self):
self.add_handlers({"^T": self.change_forms, "^Q": self.quit})
self.add(
npyscreen.TitleText,
name='Select which tools to update (only plugin tools are shown):',
editable=False)
def while_waiting(self):
""" Update with current tools that are not cores """
if not self.triggered:
i = 4
api_action = Action()
response = api_action.inventory(choices=['repos', 'tools', 'core'])
if response[0]:
inventory = response[1]
for repo in inventory['repos']:
if repo != 'https://github.com/cyberreboot/vent':
repo_name = repo.rsplit("/", 2)[1:]
self.tools_tc[repo] = {}
title_text = self.add(npyscreen.TitleText,
name='Plugin: ' + repo,
editable=False,
rely=i,
relx=5)
title_text.display()
i += 1
for tool in inventory['tools']:
r_name = tool[0].split(":")
if repo_name[0] == r_name[0] and repo_name[
1] == r_name[1]:
core = False
for item in inventory['core']:
if tool[0] == item[0]:
core = True
t = tool[1]
if t == "":
t = "/"
if not core:
t += ":" + ":".join(
tool[0].split(":")[-2:])
self.tools_tc[repo][t] = self.add(
npyscreen.CheckBox,
name=t,
value=True,
relx=10)
self.tools_tc[repo][t].display()
i += 1
i += 2
self.triggered = 1
return
def quit(self, *args, **kwargs):
self.parentApp.switchForm("MAIN")
def on_ok(self):
"""
Take the tool selections and update them
"""
def diff(first, second):
"""
Get the elements that exist in the first list and not in the second
"""
second = set(second)
return [item for item in first if item not in second]
def popup(original, orig_type, thr, title):
"""
Start the thread and display a popup of info
until the thread is finished
"""
thr.start()
info_str = ""
while thr.is_alive():
if orig_type == 'containers':
info = diff(Containers(), original)
elif orig_type == 'images':
info = diff(Images(), original)
if info:
info_str = ""
for entry in info:
# TODO limit length of info_str to fit box
info_str += entry[0] + ": " + entry[1] + "\n"
npyscreen.notify_wait(info_str, title=title)
time.sleep(1)
return
original_containers = Containers()
api_action = Action()
for repo in self.tools_tc:
for tool in self.tools_tc[repo]:
self.logger.info(tool)
if self.tools_tc[repo][tool].value:
t = tool
if t.startswith('/:'):
t = " " + t[1:]
t = t.split(":")
thr = threading.Thread(target=api_action.update,
args=(),
kwargs={
'name': t[0],
'branch': t[1],
'version': t[2]
})
popup(original_containers, "containers", thr,
'Please wait, updating containers...')
npyscreen.notify_confirm("Done updating containers.",
title='Updated containers')
self.quit()
def on_cancel(self):
self.quit()
def change_forms(self, *args, **keywords):
""" Toggles to main """
change_to = "MAIN"
# Tell the VentApp object to change forms.
self.parentApp.change_form(change_to)
def __init__(self, es_host="elasticsearch", rmq_host="rabbitmq"):
""" initialize host information """
self.es_host = es_host
self.rmq_host = rmq_host
self.logger = Logger(__name__)
def while_waiting(self):
""" Update fields periodically if nothing is happening """
# give a little extra time for file descriptors to close
time.sleep(0.1)
self.addfield.value = Timestamp()
self.addfield.display()
self.addfield2.value = Uptime()
self.addfield2.display()
self.addfield3.value = str(len(Containers())) + ' running'
if len(Containers()) > 0:
self.addfield3.labelColor = 'GOOD'
else:
self.addfield3.labelColor = 'DEFAULT'
self.addfield3.display()
# update core tool status
self.addfield5.value, values = MainForm.t_status(True)
if values[0] + values[1] == 0:
color = 'DANGER'
self.addfield4.labelColor = 'CAUTION'
self.addfield4.value = 'Idle'
elif values[0] >= int(values[2]):
color = 'GOOD'
self.addfield4.labelColor = color
self.addfield4.value = 'Ready to start jobs'
else:
color = 'CAUTION'
self.addfield4.labelColor = color
self.addfield4.value = 'Ready to start jobs'
self.addfield5.labelColor = color
# update plugin tool status
plugin_str, values = MainForm.t_status(False)
plugin_str += ', ' + str(values[3]) + ' plugin(s) installed'
self.addfield6.value = plugin_str
# get jobs
jobs = Jobs()
# number of jobs, number of tool containers
self.addfield7.value = str(jobs[0]) + ' jobs running (' + str(jobs[1])
self.addfield7.value += ' tool containers), ' + str(jobs[2])
self.addfield7.value += ' completed jobs'
if jobs[0] > 0:
self.addfield4.labelColor = 'GOOD'
self.addfield4.value = 'Processing jobs'
self.addfield7.labelColor = 'GOOD'
else:
self.addfield7.labelColor = 'DEFAULT'
self.addfield4.display()
self.addfield5.display()
self.addfield6.display()
self.addfield7.display()
# if file drop location changes deal with it
logger = Logger(__name__)
status = (False, None)
if self.file_drop.value != DropLocation()[1]:
logger.info('Starting: file drop restart')
try:
self.file_drop.value = DropLocation()[1]
logger.info('Path given: ' + str(self.file_drop.value))
# restart if the path is valid
if DropLocation()[0]:
status = self.api_action.clean(name='file_drop')
status = self.api_action.prep_start(name='file_drop')
else:
logger.error('file drop path name invalid' +
DropLocation()[1])
if status[0]:
tool_d = status[1]
status = self.api_action.start(tool_d)
logger.info('Status of file drop restart: ' +
str(status[0]))
except Exception as e: # pragma no cover
logger.error('file drop restart failed with error: ' + str(e))
logger.info('Finished: file drop restart')
self.file_drop.display()
return
def __init__(self, **kargs):
self.path_dirs = PathDirs(**kargs)
self.manifest = os.path.join(self.path_dirs.meta_dir,
"plugin_manifest.cfg")
self.d_client = docker.from_env()
self.logger = Logger(__name__)
def __init__(self, **kargs):
self.plugin = Plugin(**kargs)
self.d_client = self.plugin.d_client
self.vent_config = os.path.join(self.plugin.path_dirs.meta_dir,
"vent.cfg")
self.logger = Logger(__name__)
def file_queue(path, template_path="/vent/", r_host="redis"):
"""
Processes files that have been added from the rq-worker, starts plugins
that match the mime type for the new file.
"""
import ConfigParser
import ast
import docker
import json
import requests
import os
import sys
from redis import Redis
from rq import Queue
from subprocess import check_output, Popen, PIPE
from string import punctuation
from vent.helpers.logs import Logger
status = (True, None)
images = []
configs = {}
logger = Logger(__name__)
if os.path.isdir("/root/.vent"):
template_path = "/root/.vent/"
try:
d_client = docker.from_env()
# get the correct path for binding
vent_config = ConfigParser.RawConfigParser()
vent_config.optionxform = str
vent_config.read(template_path + 'vent.cfg')
if (vent_config.has_section('main')
and vent_config.has_option('main', 'files')):
files = vent_config.get('main', 'files')
else:
files = '/'
# deal with ~
files = os.path.expanduser(files)
chars = set(punctuation)
chars.discard('/')
chars.discard('_')
chars.discard('-')
file_name = ''
# escape any funky symbols to allow users FREEDOM of directory name
for char in files:
if char in chars:
if char == '\\':
file_name += '\\' + char
else:
file_name += '\\\\' + char
else:
file_name += char
files = file_name
_, path = path.split('_', 1)
directory = path.rsplit('/', 1)[0]
path = path.replace('/files', files, 1)
path_copy = path
# read in configuration of plugins to get the ones that should run
# against the path.
# keep track of images that failed getting configurations for
failed_images = set()
config = ConfigParser.RawConfigParser()
config.optionxform = str
config.read(template_path + 'plugin_manifest.cfg')
sections = config.sections()
name_maps = {}
orig_path_d = {}
path_cmd = {}
labels_d = {}
# get all name maps
for section in sections:
link_name = config.get(section, 'link_name')
image_name = config.get(section, 'image_name')
name_maps[link_name] = image_name.replace(':',
'-').replace('/', '-')
for section in sections:
path = path_copy
orig_path = ''
repo = config.get(section, 'repo')
t_type = config.get(section, 'type')
labels = {
'vent-plugin': '',
'file': path,
'vent.section': section,
'vent.repo': repo,
'vent.type': t_type
}
image_name = config.get(section, 'image_name')
link_name = config.get(section, 'link_name')
# doesn't matter if it's a repository or registry because both in manifest
if config.has_option(section, 'groups'):
if 'replay' in config.get(section, 'groups'):
try:
# read the vent.cfg file to grab the network-mapping
# specified. For replay_pcap
n_name = 'network-mapping'
n_map = []
if vent_config.has_section(n_name):
# make sure that the options aren't empty
if vent_config.options(n_name):
options = vent_config.options(n_name)
for option in options:
if vent_config.get(n_name, option):
n_map.append(
vent_config.get(n_name, option))
orig_path = path
path = str(n_map[0]) + " " + path
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if config.has_option(section, 'service'):
try:
options_dict = json.loads(config.get(section, 'service'))
for option in options_dict:
value = options_dict[option]
labels[option] = value
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if config.has_option(section, 'settings'):
try:
options_dict = json.loads(config.get(section, 'settings'))
in_base = directory == '/files'
# process base by default
process_file = in_base
# check if this tool shouldn't process the base by default
if 'process_base' in options_dict:
if options_dict['process_base'] == 'no':
process_file = False
# check if this tool should look at subdirs created by
# other tools' output
if 'process_from_tool' in options_dict and not in_base:
for tool in options_dict['process_from_tool'].split(
','):
if tool.replace(' ', '-') in directory:
process_file = True
if 'ext_types' in options_dict and process_file:
ext_types = options_dict['ext_types'].split(',')
for ext_type in ext_types:
if path.endswith(ext_type):
images.append(image_name)
configs[image_name] = {}
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if image_name in configs:
if config.has_option(section, 'docker'):
try:
options_dict = ast.literal_eval(
config.get(section, 'docker'))
for option in options_dict:
try:
configs[image_name][option] = ast.literal_eval(
options_dict[option])
except Exception as e: # pragma: no cover
configs[image_name][option] = options_dict[
option]
if 'links' in configs[image_name]:
for link in configs[image_name]['links']:
if link in name_maps:
configs[image_name]['links'][
name_maps[link]] = configs[image_name][
'links'].pop(link)
# TODO network_mode
# TODO volumes_from
# TODO external services
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if config.has_option(section, 'gpu') and image_name in configs:
try:
options_dict = json.loads(config.get(section, 'gpu'))
if 'enabled' in options_dict:
enabled = options_dict['enabled']
if enabled == 'yes':
configs[image_name]['gpu_options'] = options_dict
labels['vent.gpu'] = 'yes'
if 'dedicated' in options_dict:
labels['vent.gpu.dedicated'] = options_dict[
'dedicated']
if 'device' in options_dict:
labels['vent.gpu.device'] = options_dict[
'device']
if 'mem_mb' in options_dict:
labels['vent.gpu.mem_mb'] = options_dict[
'mem_mb']
port = ''
host = ''
if (vent_config.has_section('nvidia-docker-plugin')
and vent_config.has_option(
'nvidia-docker-plugin', 'port')):
port = vent_config.get('nvidia-docker-plugin',
'port')
else:
port = '3476'
if (vent_config.has_section('nvidia-docker-plugin')
and vent_config.has_option(
'nvidia-docker-plugin', 'host')):
host = vent_config.get('nvidia-docker-plugin',
'host')
else:
# grab the default gateway
try:
route = Popen(('/sbin/ip', 'route'),
stdout=PIPE)
h = check_output(
('awk', '/default/ {print$3}'),
stdin=route.stdout)
route.wait()
host = h.strip()
except Exception as e: # pragma no cover
logger.error("Default gateway "
"went wrong " + str(e))
nd_url = 'http://' + host + ':' + port + '/v1.0/docker/cli'
params = {'vol': 'nvidia_driver'}
try:
r = requests.get(nd_url, params=params)
if r.status_code == 200:
options = r.text.split()
for option in options:
if option.startswith(
'--volume-driver='):
configs[image_name][
'volume_driver'] = option.split(
"=", 1)[1]
elif option.startswith('--volume='):
vol = option.split("=",
1)[1].split(":")
if 'volumes' in configs[
image_name]:
# !! TODO handle if volumes is a list
configs[image_name]['volumes'][
vol[0]] = {
'bind': vol[1],
'mode': vol[2]
}
else:
configs[image_name][
'volumes'] = {
vol[0]: {
'bind': vol[1],
'mode': vol[2]
}
}
elif option.startswith('--device='):
dev = option.split("=", 1)[1]
if 'devices' in configs[
image_name]:
configs[image_name][
'devices'].append(dev +
":" +
dev +
":rwm")
else:
configs[image_name][
'devices'] = [
dev + ":" + dev +
":rwm"
]
else:
# unable to parse option provided by
# nvidia-docker-plugin
pass
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
print("Failure with nvidia-docker-plugin: " +
str(e))
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
print("Unable to process gpu options: " + str(e))
path_cmd[image_name] = path
orig_path_d[image_name] = orig_path
labels_d[image_name] = labels
# TODO get syslog address rather than hardcode
# TODO add group label
# TODO get group and name for syslog tag
log_config = {
'type': 'syslog',
'config': {
'syslog-address': 'tcp://0.0.0.0:514',
'syslog-facility': 'daemon',
'tag': path.rsplit('.', 1)[-1]
}
}
# setup gpu queue
can_queue_gpu = True
try:
q = Queue(connection=Redis(host=r_host), default_timeout=86400)
except Exception as e: # pragma: no cover
can_queue_gpu = False
print("Unable to connect to redis: " + str(e))
# start containers
for image in images:
if image not in failed_images:
orig_path = orig_path_d[image]
labels = labels_d[image]
if orig_path:
# replay_pcap is special so we can't bind it like normal
# since the plugin takes in an additional argument
dir_path = orig_path.rsplit('/', 1)[0]
else:
dir_path = path.rsplit('/', 1)[0]
volumes = {dir_path: {'bind': dir_path, 'mode': 'rw'}}
if 'volumes' in configs[image]:
for volume in volumes:
configs[image]['volumes'][volume] = volumes[volume]
else:
configs[image]['volumes'] = volumes
if 'vent.gpu' in labels and labels['vent.gpu'] == 'yes':
if can_queue_gpu:
# queue up containers requiring a gpu
q_str = json.dumps({
'image': image,
'command': path_cmd[image],
'labels': labels,
'detach': True,
'log_config': log_config,
'configs': configs[image]
})
q.enqueue('watch.gpu_queue', q_str, ttl=2592000)
else:
failed_images.add(image)
else:
if 'gpu_options' in configs[image]:
del configs[image]['gpu_options']
d_client.containers.run(image=image,
command=path_cmd[image],
labels=labels,
detach=True,
log_config=log_config,
**configs[image])
if failed_images:
status = (False, failed_images)
else:
status = (True, images)
except Exception as e: # pragma: no cover
status = (False, str(e))
print('Error on line {}'.format(sys.exc_info()[-1].tb_lineno))
print("Failed to process job: " + str(e))
print(str(status))
return status
def file_queue(path, template_path='/vent/', r_host='redis'):
"""
Processes files that have been added from the rq-worker, starts plugins
that match the mime type for the new file.
"""
import configparser
import ast
import docker
import json
import requests
import os
import sys
import time
import uuid
from redis import Redis
from redis import StrictRedis
from rq import Queue
from subprocess import check_output, Popen, PIPE
from string import punctuation
from vent.helpers.logs import Logger
status = (True, None)
images = []
configs = {}
logger = Logger(__name__)
if (os.path.isfile('/root/.vent/vent.cfg')
and os.path.isfile('/root/.vent/plugin_manifest.cfg')):
template_path = '/root/.vent/'
try:
d_client = docker.from_env()
# get the correct path for binding
vent_config = configparser.ConfigParser(interpolation=None)
vent_config.optionxform = str
vent_config.read(template_path + 'vent.cfg')
if (vent_config.has_section('main')
and vent_config.has_option('main', 'files')):
files = vent_config.get('main', 'files')
else:
files = '/'
# deal with ~
files = os.path.expanduser(files)
chars = set(punctuation)
chars.discard('/')
chars.discard('_')
chars.discard('-')
file_name = ''
# escape any funky symbols to allow users FREEDOM of directory name
for char in files:
if char in chars:
if char == '\\':
file_name += '\\' + char
else:
file_name += '\\\\' + char
else:
file_name += char
files = file_name
_, path = path.split('_', 1)
directory = path.rsplit('/', 1)[0]
path = path.replace('/files', files, 1)
path_copy = path
# read in configuration of plugins to get the ones that should run
# against the path.
# keep track of images that failed getting configurations for
failed_images = set()
config = configparser.ConfigParser(interpolation=None)
config.optionxform = str
logger.debug('Path to manifest: ' + template_path +
'plugin_manifest.cfg')
config.read(template_path + 'plugin_manifest.cfg')
sections = config.sections()
name_maps = {}
orig_path_d = {}
path_cmd = {}
labels_d = {}
# get all name maps
for section in sections:
link_name = config.get(section, 'link_name')
image_name = config.get(section, 'image_name')
name_maps[link_name] = image_name.replace(':',
'-').replace('/', '-')
for section in sections:
path = path_copy
orig_path = ''
repo = config.get(section, 'repo')
t_type = config.get(section, 'type')
labels = {
'vent-plugin': '',
'file': path,
'vent.section': section,
'vent.repo': repo,
'vent.type': t_type
}
image_name = config.get(section, 'image_name')
link_name = config.get(section, 'link_name')
if config.has_option(section, 'service'):
try:
options_dict = json.loads(config.get(section, 'service'))
for option in options_dict:
value = options_dict[option]
labels[option] = value
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if config.has_option(section, 'settings'):
try:
options_dict = json.loads(config.get(section, 'settings'))
in_base = directory == '/files'
# process base by default
process_file = in_base
# check if this tool shouldn't process the base by default
if 'process_base' in options_dict:
if options_dict['process_base'] == 'no':
process_file = False
# check if this tool should look at subdirs created by
# other tools' output
if 'process_from_tool' in options_dict and not in_base:
for tool in options_dict['process_from_tool'].split(
','):
dir_pieces = directory.split('/')
dir_check = dir_pieces
for dir_piece in dir_pieces:
if 'UTC' in dir_piece:
dir_check = dir_piece
if tool.replace(' ', '-') in dir_check:
process_file = True
if 'ext_types' in options_dict and process_file:
ext_types = options_dict['ext_types'].split(',')
for ext_type in ext_types:
if path.endswith(ext_type):
images.append(image_name)
configs[image_name] = {}
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if image_name in configs:
if config.has_option(section, 'docker'):
try:
options_dict = ast.literal_eval(
config.get(section, 'docker'))
for option in options_dict:
try:
configs[image_name][option] = ast.literal_eval(
options_dict[option])
except Exception as e: # pragma: no cover
configs[image_name][option] = options_dict[
option]
if 'links' in configs[image_name]:
for link in configs[image_name]['links']:
if link in name_maps:
configs[image_name]['links'][
name_maps[link]] = configs[image_name][
'links'].pop(link)
# TODO volumes_from
# TODO external services
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
if config.has_option(section, 'gpu') and image_name in configs:
try:
options_dict = json.loads(config.get(section, 'gpu'))
if 'enabled' in options_dict:
enabled = options_dict['enabled']
if enabled == 'yes':
configs[image_name]['gpu_options'] = options_dict
labels['vent.gpu'] = 'yes'
if 'dedicated' in options_dict:
labels['vent.gpu.dedicated'] = options_dict[
'dedicated']
if 'device' in options_dict:
labels['vent.gpu.device'] = options_dict[
'device']
if 'mem_mb' in options_dict:
labels['vent.gpu.mem_mb'] = options_dict[
'mem_mb']
port = ''
host = ''
if (vent_config.has_section('nvidia-docker-plugin')
and vent_config.has_option(
'nvidia-docker-plugin', 'port')):
port = vent_config.get('nvidia-docker-plugin',
'port')
else:
port = '3476'
if (vent_config.has_section('nvidia-docker-plugin')
and vent_config.has_option(
'nvidia-docker-plugin', 'host')):
host = vent_config.get('nvidia-docker-plugin',
'host')
else:
# grab the default gateway
try:
route = Popen(('/sbin/ip', 'route'),
stdout=PIPE)
host = check_output(
('awk', '/default/ {print$3}'),
stdin=route.stdout).strip().decode(
'utf-8')
route.wait()
except Exception as e: # pragma no cover
logger.error('Default gateway '
'went wrong ' + str(e))
nd_url = 'http://' + host + ':' + port + '/v1.0/docker/cli'
params = {'vol': 'nvidia_driver'}
try:
r = requests.get(nd_url, params=params)
if r.status_code == 200:
options = r.text.split()
for option in options:
if option.startswith(
'--volume-driver='):
configs[image_name][
'volume_driver'] = option.split(
'=', 1)[1]
elif option.startswith('--volume='):
vol = option.split('=',
1)[1].split(':')
if 'volumes' in configs[
image_name]:
# !! TODO handle if volumes is a list
configs[image_name]['volumes'][
vol[0]] = {
'bind': vol[1],
'mode': vol[2]
}
else:
configs[image_name][
'volumes'] = {
vol[0]: {
'bind': vol[1],
'mode': vol[2]
}
}
elif option.startswith('--device='):
dev = option.split('=', 1)[1]
if 'devices' in configs[
image_name]:
configs[image_name][
'devices'].append(dev +
':' +
dev +
':rwm')
else:
configs[image_name][
'devices'] = [
dev + ':' + dev +
':rwm'
]
else:
# unable to parse option provided by
# nvidia-docker-plugin
pass
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
logger.error(
'Failure with nvidia-docker-plugin: ' +
str(e))
except Exception as e: # pragma: no cover
failed_images.add(image_name)
status = (False, str(e))
logger.error('Unable to process gpu options: ' + str(e))
path_cmd[image_name] = path
orig_path_d[image_name] = orig_path
labels_d[image_name] = labels
# TODO get syslog address rather than hardcode
# TODO add group label
# TODO get group and name for syslog tag
log_config = {
'type': 'syslog',
'config': {
'syslog-address': 'tcp://127.0.0.1:514',
'syslog-facility': 'daemon',
'tag': '{{.Name}}'
}
}
# setup gpu queue
can_queue_gpu = True
try:
q = Queue(connection=Redis(host=r_host), default_timeout=86400)
r = StrictRedis(host=r_host, port=6379, db=0)
except Exception as e: # pragma: no cover
can_queue_gpu = False
logger.error('Unable to connect to redis: ' + str(e))
# start containers
for image in images:
if image not in failed_images:
orig_path = orig_path_d[image]
labels = labels_d[image]
configs[image]['auto_remove'] = True
name = image.replace('/', '-').replace(':', '-') + '_' + \
str(int(time.time()))+'_'+str(uuid.uuid4())[:4]
if orig_path:
# replay_pcap is special so we can't bind it like normal
# since the plugin takes in an additional argument
dir_path = orig_path.rsplit('/', 1)[0]
else:
dir_path = path.rsplit('/', 1)[0]
volumes = {dir_path: {'bind': dir_path, 'mode': 'rw'}}
if 'volumes' in configs[image]:
for volume in volumes:
configs[image]['volumes'][volume] = volumes[volume]
else:
configs[image]['volumes'] = volumes
command = path_cmd[image]
if 'command' in configs[image]:
command = configs[image]['command'] + ' ' + command
del configs[image]['command']
if 'vent.gpu' in labels and labels['vent.gpu'] == 'yes':
if can_queue_gpu:
# queue up containers requiring a gpu
q_str = json.dumps({
'image': image,
'command': command,
'labels': labels,
'detach': True,
'name': name,
'network': 'vent',
'log_config': log_config,
'configs': configs[image]
})
q.enqueue('watch.gpu_queue', q_str, ttl=2592000)
else:
failed_images.add(image)
else:
if 'gpu_options' in configs[image]:
del configs[image]['gpu_options']
# TODO check for links
change_networking = False
links = []
network_name = ''
configs[image]['network'] = 'vent'
if 'links' in configs[image]:
for link in configs[image]['links']:
links.append((link, configs[image]['links'][link]))
if 'network' in configs[image]:
network_name = configs[image]['network']
del configs[image]['network']
del configs[image]['links']
change_networking = True
cont = d_client.containers.create(image=image,
command=command,
labels=labels,
detach=True,
name=name,
log_config=log_config,
**configs[image])
cont_id = cont.id
if change_networking:
network_to_attach = d_client.networks.list(
names=[network_name])
if len(network_to_attach) > 0:
logger.info(
'Attaching to network: "{0}" with the following links: {1}'
.format(network_name, links))
network_to_attach[0].connect(cont_id, links=links)
logger.info('Detaching from network: bridge')
network_to_detach = d_client.networks.list(
names=['bridge'])
network_to_detach[0].disconnect(cont_id)
cont.start()
try:
r.hincrby('vent_plugin_counts', image)
except Exception as e: # pragma: no cover
logger.error(
'Failed to update count of plugins because: {0}'.
format(str(e)))
if failed_images:
status = (False, failed_images)
else:
status = (True, images)
except Exception as e: # pragma: no cover
status = (False, str(e))
logger.error('Error on line {}'.format(sys.exc_info()[-1].tb_lineno))
logger.error('Failed to process job: ' + str(e))
logger.info(str(status))
return status
def __init__(self, manifest, *args, **kwargs):
self.path_dirs = PathDirs(**kwargs)
self.manifest = manifest
self.d_client = docker.from_env()
self.logger = Logger(__name__)