def testConfigFileInclusionWindowsPaths(self): one = r""" Config.includes: - 2.yaml Section1.int: 1 """ two = r""" Section1.int: 2 SecondaryFileIncluded: true """ config_path = "C:\\Windows\\System32\\GRR" def MockedWindowsOpen(filename, _=None): basename = ntpath.basename(filename) dirname = ntpath.dirname(filename) # Make sure we only try to open files from this directory. if dirname != config_path: raise IOError("Tried to open wrong file %s" % filename) if basename == "1.yaml": return StringIO.StringIO(one) if basename == "2.yaml": return StringIO.StringIO(two) raise IOError("File not found %s" % filename) # We need to also use the nt path manipulation modules. with utils.MultiStubber((__builtin__, "open", MockedWindowsOpen), (os, "path", ntpath)): conf = self._GetNewConf() conf.Initialize(filename=ntpath.join(config_path, "1.yaml")) self.assertEqual(conf["Section1.int"], 2) self.assertEqual(conf["SecondaryFileIncluded"], True)
def setUp(self): # pylint: disable=invalid-name """Set up test method.""" super(ApiRegressionTest, self).setUp() if not self.__class__.api_method: raise ValueError("%s.api_method has to be set." % self.__class__.__name__) if not self.__class__.handler: raise ValueError("%s.handler has to be set." % self.__class__.__name__) self.checks = [] self.syscalls_stubber = utils.MultiStubber( (socket, "gethostname", lambda: "test.host"), (os, "getpid", lambda: 42)) self.syscalls_stubber.Start() self.token.username = "******" testing_startup.TestInit() # Force creation of new APIAuthorizationManager. api_auth_manager.APIACLInit.InitApiAuthManager()
def Start(self): """Install the stubs.""" modules = { "_winreg": mock.MagicMock(), "ctypes": mock.MagicMock(), "ctypes.wintypes": mock.MagicMock(), # Requires mocking because exceptions.WindowsError does not exist "exceptions": mock.MagicMock(), } self.module_patcher = mock.patch.dict("sys.modules", modules) self.module_patcher.start() # pylint: disable= g-import-not-at-top from grr_response_client.vfs_handlers import registry import exceptions import _winreg # pylint: enable=g-import-not-at-top fixture = RegistryFake() self.stubber = utils.MultiStubber( (registry, "KeyHandle", RegistryFake.FakeKeyHandle), (registry, "OpenKey", fixture.OpenKey), (registry, "QueryValueEx", fixture.QueryValueEx), (registry, "QueryInfoKey", fixture.QueryInfoKey), (registry, "EnumValue", fixture.EnumValue), (registry, "EnumKey", fixture.EnumKey)) self.stubber.Start() # Add the Registry handler to the vfs. vfs.VFSInit().Run() _winreg.HKEY_USERS = "HKEY_USERS" _winreg.HKEY_LOCAL_MACHINE = "HKEY_LOCAL_MACHINE" exceptions.WindowsError = IOError
def testIndexIsUsedWhenRandomAccessIsUsed(self): with utils.MultiStubber( (aff4.PackedVersionedCollection, "COMPACTION_BATCH_SIZE", 100), (aff4.PackedVersionedCollection, "INDEX_INTERVAL", 1)): with aff4.FACTORY.Create(self.collection_urn, "PackedVersionedCollection", mode="w", token=self.token): pass for i in range(20): with aff4.FACTORY.Open( self.collection_urn, "PackedVersionedCollection", mode="w", token=self.token) as fd: fd.Add(rdf_flows.GrrMessage(request_id=i)) with aff4.FACTORY.OpenWithLock( self.collection_urn, "PackedVersionedCollection", token=self.token) as fd: fd.Compact() collection = aff4.FACTORY.Open(self.collection_urn, token=self.token) item_size = collection.fd.size / len(collection) # There's no seek expected for the first element for i in range(1, 20): seek_ops = [] old_seek = collection.fd.Seek def SeekStub(offset): seek_ops.append(offset) # pylint: disable=cell-var-from-loop old_seek(offset) # pylint: disable=cell-var-from-loop # Check that the stream is seeked to a correct byte offset on every # GenerateItems() call with an offset specified. with utils.Stubber(collection.fd, "Seek", SeekStub): _ = list(collection.GenerateItems(offset=i)) self.assertListEqual([item_size * i], seek_ops)
def _ConfigStub(self, sections=None): mock = GetConfigMockClass(sections) return utils.MultiStubber( (config.CONFIG, "GetRaw", mock["GetRaw"]), (config.CONFIG, "Get", mock["Get"]), (config.CONFIG, "type_infos", mock["type_infos"]))
def MakeRequest(self, instrumentor, manager, path, verify_cb=lambda x: True): with utils.MultiStubber((requests, "request", instrumentor.request), (time, "sleep", instrumentor.sleep)): return manager.OpenServerEndpoint(path, verify_cb=verify_cb)
def MakeRequest(self, instrumentor, manager, path, verify_cb=lambda x: True): with utils.MultiStubber( (urllib2, "install_opener", instrumentor.install_opener), (urllib2, "urlopen", instrumentor.urlopen), (time, "sleep", instrumentor.sleep)): return manager.OpenServerEndpoint(path, verify_cb=verify_cb)
def setUp(self): super(GRRBaseTest, self).setUp() tmpdir = os.environ.get("TEST_TMPDIR") or config.CONFIG["Test.tmpdir"] if platform.system() != "Windows": # Make a temporary directory for test files. self.temp_dir = tempfile.mkdtemp(dir=tmpdir) else: self.temp_dir = tempfile.mkdtemp() config.CONFIG.SetWriteBack( os.path.join(self.temp_dir, "writeback.yaml")) logging.info("Starting test: %s.%s", self.__class__.__name__, self._testMethodName) self.last_start_time = time.time() try: # Clear() is much faster than init but only supported for FakeDataStore. data_store.DB.Clear() except AttributeError: self.InitDatastore() aff4.FACTORY.Flush() # Create a Foreman and Filestores, they are used in many tests. aff4_grr.GRRAFF4Init().Run() filestore.FileStoreInit().Run() hunts_results.ResultQueueInitHook().Run() email_alerts.EmailAlerterInit().RunOnce() # Stub out the email function self.emails_sent = [] def SendEmailStub(to_user, from_user, subject, message, **unused_kwargs): self.emails_sent.append((to_user, from_user, subject, message)) self.mail_stubber = utils.MultiStubber( (email_alerts.EMAIL_ALERTER, "SendEmail", SendEmailStub), (email.utils, "make_msgid", lambda: "<message id stub>")) self.mail_stubber.Start() self.nanny_stubber = utils.Stubber( client_utils_linux.NannyController, "StartNanny", lambda unresponsive_kill_period=None, nanny_logfile=None: True) self.nanny_stubber.Start() # We don't want to send actual email in our tests self.smtp_patcher = mock.patch("smtplib.SMTP") self.mock_smtp = self.smtp_patcher.start() def DisabledSet(*unused_args, **unused_kw): raise NotImplementedError( "Usage of Set() is disabled, please use a configoverrider in tests." ) self.config_set_disable = utils.Stubber(config.CONFIG, "Set", DisabledSet) self.config_set_disable.Start()