def __init__(self, jobs_denylist=None, jobs_allowlist=None):
"""Initialization for celery worker.
Args:
jobs_denylist (Optional[list[str]]): Jobs we will exclude from running
jobs_allowlist (Optional[list[str]]): The only Jobs we will include to run
"""
super(TurbiniaCeleryWorker, self).__init__()
# Deregister jobs from denylist/allowlist.
job_manager.JobsManager.DeregisterJobs(jobs_denylist, jobs_allowlist)
disabled_jobs = list(config.DISABLED_JOBS) if config.DISABLED_JOBS else []
disabled_jobs = [j.lower() for j in disabled_jobs]
# Only actually disable jobs that have not been allowlisted.
if jobs_allowlist:
disabled_jobs = list(set(disabled_jobs) - set(jobs_allowlist))
if disabled_jobs:
log.info(
'Disabling non-allowlisted jobs configured to be disabled in the '
'config file: {0:s}'.format(', '.join(disabled_jobs)))
job_manager.JobsManager.DeregisterJobs(jobs_denylist=disabled_jobs)
# Check for valid dependencies/directories.
dependencies = config.ParseDependencies()
if config.DOCKER_ENABLED:
check_docker_dependencies(dependencies)
check_system_dependencies(dependencies)
check_directory(config.MOUNT_DIR_PREFIX)
check_directory(config.OUTPUT_DIR)
check_directory(config.TMP_DIR)
jobs = job_manager.JobsManager.GetJobNames()
log.info(
'Dependency check complete. The following jobs will be enabled '
'for this worker: {0:s}'.format(','.join(jobs)))
self.worker = self.task_manager.celery.app
def testParseDependencies(self):
"""Tests a valid config for the ParseDependencies() method."""
smpl_depends = 'DEPENDENCIES = [{"job": "PlasoJob","programs": ["test"], "docker_image": "test"}]'
self.WriteConfig(smpl_depends)
config.LoadConfig()
smpl_out = {'plasojob': {'programs': ['test'], 'docker_image': 'test'}}
smpl_test = config.ParseDependencies()
self.assertEqual(smpl_out, smpl_test)
def __init__(self, jobs_denylist=None, jobs_allowlist=None):
"""Initialization for PSQ Worker.
Args:
jobs_denylist (Optional[list[str]]): Jobs we will exclude from running
jobs_allowlist (Optional[list[str]]): The only Jobs we will include to run
"""
setup()
psq_publisher = pubsub.PublisherClient()
psq_subscriber = pubsub.SubscriberClient()
datastore_client = datastore.Client(project=config.TURBINIA_PROJECT)
try:
self.psq = psq.Queue(
psq_publisher,
psq_subscriber,
config.TURBINIA_PROJECT,
name=config.PSQ_TOPIC,
storage=psq.DatastoreStorage(datastore_client))
except exceptions.GoogleCloudError as e:
msg = 'Error creating PSQ Queue: {0:s}'.format(str(e))
log.error(msg)
raise TurbiniaException(msg)
# Deregister jobs from denylist/allowlist.
job_manager.JobsManager.DeregisterJobs(jobs_denylist, jobs_allowlist)
disabled_jobs = list(
config.DISABLED_JOBS) if config.DISABLED_JOBS else []
disabled_jobs = [j.lower() for j in disabled_jobs]
# Only actually disable jobs that have not been allowlisted.
if jobs_allowlist:
disabled_jobs = list(set(disabled_jobs) - set(jobs_allowlist))
if disabled_jobs:
log.info(
'Disabling non-allowlisted jobs configured to be disabled in the '
'config file: {0:s}'.format(', '.join(disabled_jobs)))
job_manager.JobsManager.DeregisterJobs(jobs_denylist=disabled_jobs)
# Check for valid dependencies/directories.
dependencies = config.ParseDependencies()
if config.DOCKER_ENABLED:
try:
check_docker_dependencies(dependencies)
except TurbiniaException as e:
log.warning(
"DOCKER_ENABLED=True is set in the config, but there is an error checking for the docker daemon: {0:s}"
).format(str(e))
check_system_dependencies(dependencies)
check_directory(config.MOUNT_DIR_PREFIX)
check_directory(config.OUTPUT_DIR)
check_directory(config.TMP_DIR)
register_job_timeouts(dependencies)
jobs = job_manager.JobsManager.GetJobNames()
log.info('Dependency check complete. The following jobs are enabled '
'for this worker: {0:s}'.format(','.join(jobs)))
log.info('Starting PSQ listener on queue {0:s}'.format(self.psq.name))
self.worker = psq.Worker(queue=self.psq)
def __init__(self, jobs_blacklist=None, jobs_whitelist=None):
"""Initialization for PSQ Worker.
Args:
jobs_blacklist (Optional[list[str]]): Jobs we will exclude from running
jobs_whitelist (Optional[list[str]]): The only Jobs we will include to run
"""
config.LoadConfig()
psq_publisher = pubsub.PublisherClient()
psq_subscriber = pubsub.SubscriberClient()
datastore_client = datastore.Client(project=config.TURBINIA_PROJECT)
try:
self.psq = psq.Queue(
psq_publisher,
psq_subscriber,
config.TURBINIA_PROJECT,
name=config.PSQ_TOPIC,
storage=psq.DatastoreStorage(datastore_client))
except exceptions.GoogleCloudError as e:
msg = 'Error creating PSQ Queue: {0:s}'.format(str(e))
log.error(msg)
raise TurbiniaException(msg)
# Deregister jobs from blacklist/whitelist.
disabled_jobs = list(
config.DISABLED_JOBS) if config.DISABLED_JOBS else []
job_manager.JobsManager.DeregisterJobs(jobs_blacklist, jobs_whitelist)
if disabled_jobs:
log.info(
'Disabling jobs that were configured to be disabled in the '
'config file: {0:s}'.format(', '.join(disabled_jobs)))
job_manager.JobsManager.DeregisterJobs(
jobs_blacklist=disabled_jobs)
# Check for valid dependencies/directories.
dependencies = config.ParseDependencies()
if config.DOCKER_ENABLED:
check_docker_dependencies(dependencies)
check_system_dependencies(dependencies)
check_directory(config.MOUNT_DIR_PREFIX)
check_directory(config.OUTPUT_DIR)
check_directory(config.TMP_DIR)
jobs = job_manager.JobsManager.GetJobNames()
log.info(
'Dependency check complete. The following jobs will be enabled '
'for this worker: {0:s}'.format(','.join(jobs)))
log.info('Starting PSQ listener on queue {0:s}'.format(self.psq.name))
self.worker = psq.Worker(queue=self.psq)
def setup(self, jobs_denylist=None, jobs_allowlist=None, *args, **kwargs):
"""Does setup of Task manager and its dependencies.
Args:
jobs_denylist (list): Jobs that will be excluded from running
jobs_allowlist (list): The only Jobs will be included to run
"""
self._backend_setup(*args, **kwargs)
job_names = jobs_manager.JobsManager.GetJobNames()
if jobs_denylist or jobs_allowlist:
selected_jobs = jobs_denylist or jobs_allowlist
for job in selected_jobs:
if job.lower() not in job_names:
msg = (
'Error creating server. Job {0!s} is not found in registered '
'jobs {1!s}.'.format(job, job_names))
log.error(msg)
raise TurbiniaException(msg)
log.info('Filtering Jobs with allowlist {0!s} and denylist {1!s}'.
format(jobs_allowlist, jobs_denylist))
job_names = jobs_manager.JobsManager.FilterJobNames(
job_names, jobs_denylist, jobs_allowlist)
# Disable any jobs from the config that were not previously allowlisted.
disabled_jobs = list(
config.DISABLED_JOBS) if config.DISABLED_JOBS else []
disabled_jobs = [j.lower() for j in disabled_jobs]
if jobs_allowlist:
disabled_jobs = list(set(disabled_jobs) - set(jobs_allowlist))
if disabled_jobs:
log.info(
'Disabling non-allowlisted jobs configured to be disabled in the '
'config file: {0:s}'.format(', '.join(disabled_jobs)))
job_names = jobs_manager.JobsManager.FilterJobNames(
job_names, disabled_jobs, [])
self.jobs = [
job for _, job in jobs_manager.JobsManager.GetJobs(job_names)
]
dependencies = config.ParseDependencies()
job_utils.register_job_timeouts(dependencies)
log.debug('Registered job list: {0:s}'.format(str(job_names)))
def __init__(self, jobs_denylist=None, jobs_allowlist=None):
"""Initialization for Turbinia Worker.
Args:
jobs_denylist (Optional[list[str]]): Jobs we will exclude from running
jobs_allowlist (Optional[list[str]]): The only Jobs we will include to run
"""
setup()
# Deregister jobs from denylist/allowlist.
job_manager.JobsManager.DeregisterJobs(jobs_denylist, jobs_allowlist)
disabled_jobs = list(
config.DISABLED_JOBS) if config.DISABLED_JOBS else []
disabled_jobs = [j.lower() for j in disabled_jobs]
# Only actually disable jobs that have not been allowlisted.
if jobs_allowlist:
disabled_jobs = list(set(disabled_jobs) - set(jobs_allowlist))
if disabled_jobs:
log.info(
'Disabling non-allowlisted jobs configured to be disabled in the '
'config file: {0:s}'.format(', '.join(disabled_jobs)))
job_manager.JobsManager.DeregisterJobs(jobs_denylist=disabled_jobs)
# Check for valid dependencies/directories.
dependencies = config.ParseDependencies()
if config.DOCKER_ENABLED:
try:
check_docker_dependencies(dependencies)
except TurbiniaException as e:
log.warning(
"DOCKER_ENABLED=True is set in the config, but there is an error checking for the docker daemon: {0:s}"
).format(str(e))
check_system_dependencies(dependencies)
check_directory(config.MOUNT_DIR_PREFIX)
check_directory(config.OUTPUT_DIR)
check_directory(config.TMP_DIR)
job_utils.register_job_timeouts(dependencies)
jobs = job_manager.JobsManager.GetJobNames()
log.info('Dependency check complete. The following jobs are enabled '
'for this worker: {0:s}'.format(','.join(jobs)))
