def Start(self): hmac = GetHMAC() for pathspec in self.args.pathspecs: policy = rdf_client.UploadPolicy( client_id=self.client_id, filename=pathspec.CollapsePath(), expires=rdfvalue.RDFDatetime.Now() + self.args.duration) serialized_policy = policy.SerializeToString() self.CallClient( standard_actions.UploadFile, rdf_client.UploadFileRequest( pathspec=pathspec, client_id=self.client_id, policy=serialized_policy, hmac=hmac.HMAC(serialized_policy)), next_state="Upload")
def testUpload(self): magic_string = "Hello world" test_file = os.path.join(self.temp_dir, "sample.txt") with open(test_file, "wb") as fd: fd.write(magic_string) args = rdf_client.UploadFileRequest() args.pathspec.path = test_file args.pathspec.pathtype = "OS" # Errors are logged on the server but not always provided to the client. We # check the server logs for the errors we inject. with test_lib.Instrument(logging, "error") as logger: # First do not provide a hmac at all. with self.assertRaises(IOError): self._UploadFile(args) self.assertRegexpMatches("HMAC not provided", str(logger.args)) logger.args[:] = [] # Now pass a rubbish HMAC but forget to give a policy. hmac = args.upload_token.GetHMAC() args.upload_token.hmac = hmac.HMAC("This is the wrong filename") with self.assertRaises(IOError): self._UploadFile(args) self.assertRegexpMatches("Policy not provided", str(logger.args)) logger.args[:] = [] # Ok - lets make an expired policy, Still wrong HMAC. policy = rdf_client.UploadPolicy(client_id=self.client_id, expires=1000) args.upload_token.SetPolicy(policy) with self.assertRaises(IOError): self._UploadFile(args) self.assertRegexpMatches("Signature did not match digest", str(logger.args)) logger.args[:] = [] # Ok lets hmac the policy now, but its still too old. args.upload_token.SetPolicy(policy) with self.assertRaises(IOError): self._UploadFile(args) # Make sure the file is not written yet. rootdir = config.CONFIG["FileUploadFileStore.root_dir"] target_filename = os.path.join( rootdir, self.client_id.Add(test_file).Path().lstrip(os.path.sep)) self.assertNotEqual(target_filename, test_file) with self.assertRaises(IOError): open(target_filename) self.assertRegexpMatches("Client upload policy is too old", str(logger.args)) logger.args[:] = [] # Lets expire the policy in the future. policy.expires = rdfvalue.RDFDatetime.Now() + 1000 args.upload_token.SetPolicy(policy) args.upload_token.GenerateHMAC() r = self._UploadFile(args) fs = file_store.FileUploadFileStore() # Make sure the file was uploaded correctly. fd = fs.OpenForReading(r.file_id) data = fd.read() self.assertEqual(data, magic_string)