def CheckFreeGRRTempSpace(self, _): """Mock out the driver loading code to pass the memory image.""" path = tempfiles.GetDefaultGRRTempDirectory() reply = rdf_client.DiskUsage(path=path, total=10 * 1024 * 1024 * 1024, used=5 * 1024 * 1024 * 1024, free=5 * 1024 * 1024 * 1024) return [reply]
def _SaveSignatureShard(self, scan_request): """Writes a YaraSignatureShard received from the server to disk. Args: scan_request: The YaraProcessScanRequest sent by the server. Returns: The full Yara signature, if all shards have been received. Otherwise, None is returned. """ def GetShardName(shard_index, num_shards): return "shard_%02d_of_%02d" % (shard_index, num_shards) signature_dir = os.path.join(tempfiles.GetDefaultGRRTempDirectory(), "Sig_%s" % self.session_id.Basename()) # Create the temporary directory and set permissions, if it does not exist. tempfiles.EnsureTempDirIsSane(signature_dir) shard_path = os.path.join( signature_dir, GetShardName(scan_request.signature_shard.index, scan_request.num_signature_shards)) with io.open(shard_path, "wb") as f: f.write(scan_request.signature_shard.payload) dir_contents = set(os.listdir(signature_dir)) all_shards = [ GetShardName(i, scan_request.num_signature_shards) for i in range(scan_request.num_signature_shards) ] if dir_contents.issuperset(all_shards): # All shards have been received; delete the temporary directory and # return the full signature. full_signature = io.BytesIO() for shard in all_shards: with io.open(os.path.join(signature_dir, shard), "rb") as f: full_signature.write(f.read()) shutil.rmtree(signature_dir, ignore_errors=True) return full_signature.getvalue().decode("utf-8") else: return None
def cleanup(): shutil.rmtree(tempfiles.GetDefaultGRRTempDirectory())
def tearDown(self): super(GRRTempFileTestFilename, self).tearDown() # The actual GRR temp dir. os.rmdir(tempfiles.GetDefaultGRRTempDirectory()) self.tempdir_overrider.Stop()