def UploadFile(self, args):
"""Just copy the file into the filestore."""
file_fd = vfs.VFSOpen(args.pathspec)
fs = file_store.FileUploadFileStore()
fd = fs.CreateFileStoreFile()
while True:
data = file_fd.read(self.BUFFER_SIZE)
if not data:
break
fd.write(data)
file_id = fd.Finalize()
return [rdf_client.UploadedFile(stat_entry=file_fd.Stat(), file_id=file_id)]
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)