class Config(object):
  def GetExtra(cls, name: str='') -> Dict[str, Any]:
  def LoadExtra(cls, filename: str) -> bool:
  def LoadExtraData(cls, json_string: bytes) -> None:
  def ClearExtra(cls) -> None: class Recipe(object):
  def __init__(self,
               description: str,
               contents: Dict[str, Any],
               args: List[Tuple[str, str, Any]]) -> None:
  def GetHelpString(self) -> str: T = TypeVar("T", bound="interface.AttributeContainer")  # pylint: disable=invalid-name,line-too-long
NEW_ISSUE_URL = 'https://github.com/log2timeline/dftimewolf/issues/new'
class DFTimewolfState(object):
  def __init__(self, config: Type[Config]) -> None:
  def _InvokeModulesInThreads(self, callback: Callable[[Any], Any]) -> None:
  def ImportRecipeModules(self, module_locations: Dict[str, str]) -> None:
  def LoadRecipe(self,
                 recipe: Dict[str, Any],
                 module_locations: Dict[str, str]) -> None:
  def FormatExecutionPlan(self) -> str:
  def LogExecutionPlan(self) -> None:
  def AddToCache(self, name: str, value: Any) -> None:
  def GetFromCache(self, name: str, default_value: Any=None) -> Any:
  def StoreContainer(self, container: "interface.AttributeContainer") -> None:
  def GetContainers(self,
                    container_class: Type[T],
                    pop: bool=False) -> Sequence[T]:
  def DedupeContainers(self, container_class: Type[T]) -> None:
  def _SetupModuleThread(self, module_definition: Dict[str, str]) -> None:
  def SetupModules(self) -> None:
  def _RunModuleThread(self, module_definition: Dict[str, str]) -> None:
  def RunPreflights(self) -> None:
  def CleanUpPreflights(self) -> None:
  def InstantiateModule(self, module_name: str) -> "dftw_module.BaseModule":
  def RunModules(self) -> None:
  def RegisterStreamingCallback(
      self,
      target: Callable[["interface.AttributeContainer"], Any],
      container_type: Type["interface.AttributeContainer"]) -> None:
  def StreamContainer(self, container: "interface.AttributeContainer") -> None:
  def AddError(self, error: DFTimewolfError) -> None:
  def CleanUp(self) -> None:
  def CheckErrors(self, is_global: bool=False) -> None: class FSPath(interface.AttributeContainer):
class RemoteFSPath(FSPath):
class Report(interface.AttributeContainer):
class GCPLogs(interface.AttributeContainer):
class AWSLogs(interface.AttributeContainer):
class AzureLogs(interface.AttributeContainer):
class ThreatIntelligence(interface.AttributeContainer):
class YaraRule(interface.AttributeContainer):
class TicketAttribute(interface.AttributeContainer):
class File(interface.AttributeContainer):
class Directory(interface.AttributeContainer):
class ForensicsVM(interface.AttributeContainer):
class URL(interface.AttributeContainer):
class GCEDisk(interface.AttributeContainer):
class GCEDiskEvidence(interface.AttributeContainer):
class GCEImage(interface.AttributeContainer):
class DataFrame(interface.AttributeContainer):
class Host(interface.AttributeContainer):
class GrrFlow(interface.AttributeContainer):
class WorkspaceLogs(interface.AttributeContainer):
class GCSObject(interface.AttributeContainer):
class AWSS3Object(interface.AttributeContainer):
class AWSVolume(interface.AttributeContainer):
class AWSSnapshot(interface.AttributeContainer):
class OsqueryQuery(interface.AttributeContainer):
  CONTAINER_TYPE = 'fspath'
  CONTAINER_TYPE = 'remotefspath'
  CONTAINER_TYPE = 'report'
  CONTAINER_TYPE = 'gcp_logs'
  CONTAINER_TYPE = 'aws_logs'
  CONTAINER_TYPE = 'azure_logs'
  CONTAINER_TYPE = 'threat_intelligence'
  CONTAINER_TYPE = 'ticketattribute'
  CONTAINER_TYPE = 'file'
  CONTAINER_TYPE = 'directory'
  CONTAINER_TYPE = 'forensics_vm'
  CONTAINER_TYPE = 'url'
  CONTAINER_TYPE = 'gcedisk'
  CONTAINER_TYPE = 'gcediskevidence'
  CONTAINER_TYPE = 'gceimage'
  CONTAINER_TYPE = 'data_frame'
  CONTAINER_TYPE = 'host'
  CONTAINER_TYPE = 'grr_flow'
  CONTAINER_TYPE = 'workspace_logs'
  CONTAINER_TYPE = 'gcs_object'
  CONTAINER_TYPE = 'aws_s3_object'
  CONTAINER_TYPE = 'aws_volume'
  CONTAINER_TYPE = 'aws_snapshot'
  CONTAINER_TYPE = 'osquery_query'
  def __init__(self, path: str) -> None:
  def __init__(self, path: str, hostname: str) -> None:
  def __init__(
      self,
      module_name: str,
      text: str,
      text_format: str = 'plaintext',
      attributes: Optional[List[Dict[str, Any]]] = None) -> None:
  def __init__(self, path: str, filter_expression: str, project_name: str):
  def __init__(
      self, path: str, profile_name: Optional[str], query_filter: Optional[str],
      start_time: Optional[datetime], end_time: Optional[datetime]) -> None:
  def __init__(
      self, path: str, filter_expression: Optional[str],
      subscription_id: Optional[str]) -> None:
  def __init__(self, name: str, indicator: Optional[str], path: str) -> None:
  def __init__(self, name: str, rule_text: str) -> None:
  def __init__(self, type_: str, name: str, value: str) -> None:
  def __init__(
      self, name: str, path: str, description: Optional[str] = None) -> None:
  def __init__(
      self, name: str, path: str, description: Optional[str] = None) -> None:
  def __init__(
      self, name: str, evidence_disk: Union["GoogleComputeDisk", "AWSVol",
                                            "AZComputeDisk"],
      platform: str) -> None:
  def __init__(self, path: str) -> None:
  def __init__(self, name: str) -> None:
  def __eq__(self, other: GCEDisk) -> bool:
  def __init__(self, name: str, project: str) -> None:
  def __init__(self, name: str) -> None:
  def __init__(
      self,
      data_frame: "pandas.DataFrame",
      description: str,
      name: str,
      source: Optional[str] = None) -> None:
  def __init__(self, hostname: str, platform: str = 'unknown') -> None:
  def __init__(self, hostname: str, flow: str) -> None:
  def __init__(
      self,
      application_name: str,
      path: str,
      filter_expression: str,
      user_key: Optional[str ]= '',
      start_time: Optional[str]='',
      end_time: Optional[str]='') -> None:
  def __init__(self, path: str):
  def __init__(self, path: str):
  def __init__(self, vol_id: str) -> None:
  def __init__(self, snap_id: str) -> None:
  def __init__(self, query: str) -> None: class ModulesManager(object):
  ALLOW_MODULE_OVERRIDE = False
  def DeregisterModule(cls, module_class: Type['module.BaseModule']) -> None:
  def GetModuleByName(cls, name: str) -> Type['module.BaseModule']:
  def RegisterModule(cls, module_class: Type['module.BaseModule']) -> None:
  def RegisterModules(
      cls, module_classes: List[Type['module.BaseModule']]) -> None: class RecipesManager(object):
  ALLOW_RECIPE_OVERRIDE = False
  def _ReadRecipeFromFileObject(
      self, file_object: Union[StringIO, TextIOWrapper]) -> Recipe:
  def DeregisterRecipe(self, recipe: Recipe) -> None:
  def GetRecipes(self) -> List[Recipe]:
  def ReadRecipeFromFile(self, path: str) -> None:
  def ReadRecipesFromDirectory(self, path: str) -> None:
  def RegisterRecipe(self, recipe: Recipe) -> None:
  def RegisterRecipes(self, recipes: List[Recipe]) -> None:
  def Recipes(self) -> Dict[str, resources.Recipe]: class DFTimewolfError(Exception):
class BadConfigurationError(DFTimewolfError):
class RecipeParseError(DFTimewolfError):
class CommandLineParseError(DFTimewolfError):
class CriticalError(DFTimewolfError):
  def __init__(self,
               message: Optional[str]=None,
               name: Optional[str]=None,
               stacktrace: Optional[str]=None,
               critical: bool=False,
               unexpected: bool=False) -> None: class DummyModule1(module.BaseModule):
class DummyModule2(module.BaseModule):
class DummyPreflightModule(module.PreflightModule):
  def __init__(self, state, name=None):
  def SetUp(self, runtime_value=None):  # pylint: disable=arguments-differ
  def Callback(self, container):
  def Process(self):
  def __init__(self, state, name=None):
  def SetUp(self, runtime_value=None):  # pylint: disable=arguments-differ
  def Process(self):
  def __init__(self, state, name=None):
  def SetUp(self, args):  # pylint: disable=arguments-differ
  def Process(self):
  def CleanUp(self): class TestContainer(interface.AttributeContainer):
class TestContainerTwo(interface.AttributeContainer):
class TestContainerThree(interface.AttributeContainer):
class ContainerGeneratorModule(module.BaseModule):
class ThreadAwareConsumerModule(module.ThreadAwareModule):
class Issue503Module(module.ThreadAwareModule):
  CONTAINER_TYPE = 'test_container'
  CONTAINER_TYPE = 'test_container_two'
  CONTAINER_TYPE = 'test_container_three'
  def __init__(self, value: str) -> None:
  def __eq__(self, other: object) -> bool:
  def __init__(self, value: str) -> None:
  def __init__(self, value: str) -> None:
  def __init__(self, state, name=None):
  def SetUp(self, runtime_value=None): # pylint: disable=arguments-differ
  def Process(self):
  def __init__(self, state, name=None):
  def SetUp(self): # pylint: disable=arguments-differ
  def Process(self, container) -> None:
  def GetThreadOnContainerType():
  def GetThreadPoolSize(self):
  def PreProcess(self) -> None:
  def PostProcess(self) -> None:
  def __init__(self, state, name=None):
  def SetUp(self): # pylint: disable=arguments-differ
  def Process(self, container) -> None:
  def GetThreadOnContainerType():
  def GetThreadPoolSize(self):
  def PreProcess(self) -> None:
  def PostProcess(self) -> None:
  def KeepThreadedContainersInState(self) -> bool:
import unittest
import mock
from dftimewolf import config
from dftimewolf.lib import resources
from dftimewolf.lib import state
from dftimewolf.lib.containers import containers
from dftimewolf.lib.modules import manager as modules_manager
from dftimewolf.lib.recipes import manager as recipes_manager
from dftimewolf.lib import errors
from tests.test_modules import modules
from tests.test_modules import thread_aware_modules
from tests.test_modules import test_recipe
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Tests State."""





TEST_MODULES = {
  'DummyModule1': 'tests.test_modules.modules',
  'DummyModule2': 'tests.test_modules.modules',
  'DummyPreflightModule': 'tests.test_modules.modules',
  'ContainerGeneratorModule': 'tests.test_modules.thread_aware_modules',
  'ThreadAwareConsumerModule': 'tests.test_modules.thread_aware_modules',
  'Issue503Module': 'tests.test_modules.thread_aware_modules'
}

class StateTest(unittest.TestCase):
  """Tests for the DFTimewolfState class."""

  def setUp(self):
    """Registers the dummy modules and recipe to be used in tests."""
    modules_manager.ModulesManager.RegisterModules([
        modules.DummyModule1,
        modules.DummyModule2,
        modules.DummyPreflightModule,
        thread_aware_modules.ContainerGeneratorModule,
        thread_aware_modules.ThreadAwareConsumerModule,
        thread_aware_modules.Issue503Module])

    self._recipe = resources.Recipe(
        test_recipe.__doc__, test_recipe.contents, test_recipe.args)
    self._threaded_recipe = resources.Recipe(
        test_recipe.__doc__,
        test_recipe.threaded_no_preflights,
        test_recipe.args)
    self._recipes_manager = recipes_manager.RecipesManager()
    self._recipes_manager.RegisterRecipe(self._recipe)
    self._recipes_manager.RegisterRecipe(self._threaded_recipe)

  def tearDown(self):
    """Deregister the recipe used in tests."""
    self._recipes_manager.DeregisterRecipe(self._recipe)
    self._recipes_manager.DeregisterRecipe(self._threaded_recipe)

    modules_manager.ModulesManager.DeregisterModule(modules.DummyModule1)
    modules_manager.ModulesManager.DeregisterModule(modules.DummyModule2)
    modules_manager.ModulesManager.DeregisterModule(

        modules.DummyPreflightModule)