Source code for l2tscaffolder.scaffolders.turbinia

# -*- coding: utf-8 -*-
"""Turbinia component scaffolder."""
import datetime
import os
import logging

from typing import Dict
from typing import Iterator
from typing import Tuple

from l2tscaffolder.lib import definitions
from l2tscaffolder.lib import mapping_helper
from l2tscaffolder.scaffolders import interface
from l2tscaffolder.scaffolders import manager

[docs]class TurbiniaJobTaskScaffolder(interface.Scaffolder): """The Turbinia base scaffolder interface. Attributes: class_name (str): class name of the Turbinia job and task to be generated. """ # The name of the scaffolder plugin. NAME = 'turbinia_job_and_task' # One liner describing what the scaffolder provides. DESCRIPTION = ( 'Provides a scaffolder to generate a Turbinia job and task plugins.') # Define which project this particular scaffolder belongs to. PROJECT = definitions.DEFINITION_TURBINIA # Filenames of templates. TEMPLATE_JOB_FILE = 'turbinia_job.jinja2' TEMPLATE_TASK_FILE = 'turbinia_task.jinja2' def __init__(self): """Initializes the Turbinia scaffolder.""" super(TurbiniaJobTaskScaffolder, self).__init__() self._job_path = os.path.join('turbinia', 'jobs') self._task_path = os.path.join('turbinia', 'workers') self._mapping_helper = mapping_helper.MappingHelper() self.class_name = '' def _GenerateJobFile(self) -> str: """Generates the job job file.""" return self._mapping_helper.RenderTemplate( self.TEMPLATE_JOB_FILE, self.GetJinjaContext()) def _GenerateTaskFile(self) -> str: """Generates the task file.""" return self._mapping_helper.RenderTemplate( self.TEMPLATE_TASK_FILE, self.GetJinjaContext())
[docs] def GetInitFileChanges(self) -> Iterator[Tuple[str, str]]: """Generate a list of init files that need changing and the changes to them. Yields: Tuple[str, str]: path to the init file and the entry to add to it. """ python_init_path = self._job_path.replace(os.sep, '.') job_string = 'from {0:s} import {1:s}\n'.format( python_init_path, self._output_name) job_init_path = os.path.join(self._job_path, '') yield job_init_path, job_string
[docs] def GetJinjaContext(self) -> Dict[str, object]: """Returns a dict that can be used as a context for Jinja2 templates. Returns: dict: containing: str: name of Jinja argument. object: Jinja argument value. """ context = super(TurbiniaJobTaskScaffolder, self).GetJinjaContext() context['class_name'] = self.class_name context['plugin_name'] = self._output_name time_now = datetime.datetime.utcnow() context['year'] = time_now.year return context
[docs] def GenerateFiles(self) -> Iterator[Tuple[str, str]]: """Generates all the files required for a Turbinia component. Yields: list[tuple]: containing: str: file name. str: file content. """ plugin_name = '{0:s}.py'.format(self._output_name) self.class_name = self._mapping_helper.GenerateClassName( self._output_name) try: job_path = os.path.join(self._job_path, plugin_name) job_content = self._GenerateJobFile() yield job_path, job_content except SyntaxError as exception: logging.error(( 'Syntax error while attempting to generate component, error ' 'message: {0!s}').format(exception)) try: task_path = os.path.join(self._task_path, plugin_name) task_content = self._GenerateTaskFile() yield task_path, task_content except SyntaxError as exception: logging.error(( 'Syntax error while attempting to generate component, error ' 'message: {0!s}').format(exception))