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)))