def _testReadWhole(self, data: bytes): # pylint: disable=invalid-name
key = os.urandom(32)
encrypted = aead.Encrypt(io.BytesIO(data), key)
decrypted = aead.Decrypt(encrypted, key)
self.assertEqual(decrypted.read(), data)
def testRandomFile(self):
content = os.urandom(1024)
response = responses.Response(responses.POST, "https://foo.bar/quux")
response.status = 201
response.headers = {
"Location": "https://foo.bar/norf",
}
responses.add(response)
handler = gcs_test_lib.FakeUploadHandler()
responses.add_callback(responses.PUT, "https://foo.bar/norf", handler)
with temp.AutoTempFilePath() as temp_path:
with open(temp_path, mode="wb") as temp_file:
temp_file.write(content)
flow_id = self._Collect(path=temp_path,
signed_url="https://foo.bar/quux")
state = flow_test_lib.GetFlowState(self.client_id, flow_id)
self.assertNotEmpty(state.encryption_key)
encrypted_buf = io.BytesIO(handler.content)
decrypted_buf = aead.Decrypt(encrypted_buf, state.encryption_key)
self.assertEqual(decrypted_buf.read(), content)
def _testFile(self, content: bytes): # pylint: disable=invalid-name
key = os.urandom(32)
response = responses.Response(responses.POST, "https://foo.bar/quux")
response.status = 201
response.headers = {
"Location": "https://foo.bar/norf",
}
responses.add(response)
handler = gcs_test_lib.FakeUploadHandler()
responses.add_callback(responses.PUT, "https://foo.bar/norf", handler)
with temp.AutoTempDirPath(remove_non_empty=True) as tempdir:
with open(os.path.join(tempdir, "file"), mode="wb") as file:
file.write(content)
args = rdf_large_file.CollectLargeFileArgs()
args.signed_url = "https://foo.bar/quux"
args.encryption_key = key
args.path_spec.pathtype = rdf_paths.PathSpec.PathType.OS
args.path_spec.path = os.path.join(tempdir, "file")
results = list(large_file.CollectLargeFile(args))
self.assertLen(results, 1)
self.assertEqual(results[0].session_uri, "https://foo.bar/norf")
encrypted_buf = io.BytesIO(handler.content)
decrypted_buf = aead.Decrypt(encrypted_buf, key)
self.assertEqual(decrypted_buf.read(), content)
def testIncorrectNonceLength(self):
key = os.urandom(32)
buf = io.BytesIO()
nonce = os.urandom(aead._AEAD_NONCE_SIZE - 1)
buf.write(nonce)
buf.seek(0, io.SEEK_SET)
with self.assertRaisesRegex(EOFError, "nonce length"):
aead.Decrypt(buf, key).read()
def testIncorrectTag(self):
key = os.urandom(32)
aesgcm = crypto_aead.AESGCM(key)
buf = io.BytesIO()
nonce = os.urandom(aead._AEAD_NONCE_SIZE)
buf.write(nonce)
buf.write(aesgcm.encrypt(nonce, b"foo", b"QUUX"))
buf.seek(0, io.SEEK_SET)
with self.assertRaises(crypto_exceptions.InvalidTag):
aead.Decrypt(buf, key).read()
def testIncorrectData(self):
key = os.urandom(32)
aesgcm = crypto_aead.AESGCM(key)
buf = io.BytesIO()
nonce = os.urandom(aead._AEAD_NONCE_SIZE)
adata = aead._AEAD_ADATA_FORMAT.pack(0, True)
buf.write(nonce)
buf.write(aesgcm.encrypt(nonce, b"foo", adata))
buf.getbuffer()[-1] ^= 0b10101010 # Corrupt last byte.
buf.seek(0, io.SEEK_SET)
with self.assertRaises(crypto_exceptions.InvalidTag):
aead.Decrypt(buf, key).read()
def testReadExact(self):
count = 2048
key = os.urandom(32)
buf = io.BytesIO(b"foobarbazquxnorf" * count)
encrypted = aead.Encrypt(buf, key)
decrypted = aead.Decrypt(encrypted, key)
for _ in range(count):
self.assertEqual(decrypted.read(3), b"foo")
self.assertEqual(decrypted.read(3), b"bar")
self.assertEqual(decrypted.read(3), b"baz")
self.assertEqual(decrypted.read(3), b"qux")
self.assertEqual(decrypted.read(4), b"norf")
self.assertEqual(decrypted.read(), b"")
def testReadExact(self):
key = os.urandom(32)
data = os.urandom(1024)
encrypted = aead.Encrypt(io.BytesIO(data), key)
chunk_1 = encrypted.read(1)
self.assertLen(chunk_1, 1)
chunk_23 = encrypted.read(23)
self.assertLen(chunk_23, 23)
chunk_71 = encrypted.read(71)
self.assertLen(chunk_71, 71)
chunk_107 = encrypted.read(107)
self.assertLen(chunk_107, 107)
chunk_rest = encrypted.read()
buf = io.BytesIO(chunk_1 + chunk_23 + chunk_71 + chunk_107 + chunk_rest)
decrypted = aead.Decrypt(buf, key)
self.assertEqual(decrypted.read(), data)
def Stream(url: str, encryption_key: bytes) -> Iterator[IO[bytes]]:
"""Streams a decrypted large file from the given URL."""
with requests.get(url, stream=True) as streamer:
stream = ioutil.Unchunk(streamer.iter_content(io.DEFAULT_BUFFER_SIZE))
yield aead.Decrypt(stream, encryption_key)