def setUp(self):
super().setUp()
task_id = str(uuid.uuid4())
dir_manager = DirManager(self.path)
res_dir = dir_manager.get_task_temporary_dir(task_id)
out_dir = os.path.join(res_dir, 'out_dir')
out_dir_file = os.path.join(out_dir, 'dir_file')
out_file = os.path.join(res_dir, 'out_file')
os.makedirs(out_dir, exist_ok=True)
with open(out_file, 'w') as f:
f.write("File contents")
with open(out_dir_file, 'w') as f:
f.write("Dir file contents")
self.dir_manager = dir_manager
self.task_id = task_id
self.secret = FileEncryptor.gen_secret(10, 20)
self.disk_files = [out_file, out_dir_file]
disk_file_names = [os.path.basename(f) for f in self.disk_files]
self.all_files = disk_file_names
self.res_dir = res_dir
self.out_dir = out_dir
self.out_path = os.path.join(self.out_dir, str(uuid.uuid4()))
def setUp(self):
TestDirFixture.setUp(self)
self.task_id = str(uuid.uuid4())
self.dir_manager = DirManager(self.path)
self.secret = FileEncryptor.gen_secret(10, 20)
MockDirContents.populate(self, self.dir_manager, self.task_id)
def setUp(self):
super().setUp()
task_id = str(uuid.uuid4())
dir_manager = DirManager(self.path)
res_dir = dir_manager.get_task_temporary_dir(task_id)
out_dir = os.path.join(res_dir, 'out_dir')
os.makedirs(out_dir, exist_ok=True)
self.dir_manager = dir_manager
self.task_id = task_id
self.secret = FileEncryptor.gen_secret(10, 20)
# Create directory structure:
# |-- directory
# |-- directory2
# | |-- directory3
# | | `-- file.txt
# | `-- file.txt
# `-- file.txt
file_path = os.path.join(res_dir, "file.txt")
directory_path = os.path.join(res_dir, "directory")
directory2_path = os.path.join(res_dir, "directory2/")
directory2_file_path = os.path.join(directory2_path, "file.txt")
directory3_path = os.path.join(directory2_path, "directory3/")
directory3_file_path = os.path.join(directory3_path, "file.txt")
os.makedirs(directory_path)
os.makedirs(directory2_path)
os.makedirs(directory3_path)
with open(file_path, 'w') as out:
out.write("content")
with open(directory2_file_path, 'w') as out:
out.write("content")
with open(directory3_file_path, 'w') as out:
out.write("content")
self.disk_files = [
file_path,
directory_path,
directory2_path,
]
self.expected_results = [
os.path.basename(file_path),
os.path.basename(directory_path),
os.path.relpath(directory2_path, res_dir),
os.path.relpath(directory3_path, res_dir),
os.path.relpath(directory2_file_path, res_dir),
os.path.relpath(directory3_file_path, res_dir)
]
self.res_dir = res_dir
self.out_dir = out_dir
self.out_path = os.path.join(self.out_dir, str(uuid.uuid4()))
def test_decrypt(self):
""" Test decryption procedure """
secret = FileEncryptor.gen_secret(10, 20)
decrypted_path = self.test_file_path + ".dec"
if os.path.exists(self.enc_file_path):
os.remove(self.enc_file_path)
AESFileEncryptor.encrypt(self.test_file_path,
self.enc_file_path,
secret)
AESFileEncryptor.decrypt(self.enc_file_path,
decrypted_path,
secret)
self.assertEqual(os.path.getsize(self.test_file_path),
os.path.getsize(decrypted_path))
with open(self.test_file_path, 'rb') as f1, \
open(decrypted_path, 'rb') as f2:
while True:
chunk1 = f1.read(32)
chunk2 = f2.read(32)
if chunk1 != chunk2:
raise ValueError("Invalid decrypted file chunk")
elif not chunk1 and not chunk2:
break
AESFileEncryptor.decrypt(self.enc_file_path,
decrypted_path,
secret + b"0")
decrypted = True
if os.path.getsize(self.test_file_path) != os.path.getsize(decrypted_path):
decrypted = False
else:
with open(self.test_file_path, 'rb') as f1, \
open(decrypted_path, 'rb') as f2:
while True:
chunk1 = f1.read(32)
chunk2 = f2.read(32)
if chunk1 != chunk2:
decrypted = False
break
elif not chunk1 and not chunk2:
break
self.assertFalse(decrypted)
def test_get_key_and_iv(self):
""" Test helper methods: gen_salt and get_key_and_iv """
salt = AESFileEncryptor.gen_salt(AESFileEncryptor.block_size)
self.assertEqual(len(salt), AESFileEncryptor.block_size - AESFileEncryptor.salt_prefix_len)
secret = FileEncryptor.gen_secret(10, 20)
self.assertGreaterEqual(len(secret), 10)
self.assertLessEqual(len(secret), 20)
key_len = 32
iv_len = AESFileEncryptor.block_size
key, iv = AESFileEncryptor.get_key_and_iv(secret, salt, key_len, iv_len)
self.assertEqual(len(key), key_len)
self.assertEqual(len(iv), iv_len)
def test_encrypt(self):
""" Test encryption procedure """
secret = FileEncryptor.gen_secret(10, 20)
if os.path.exists(self.enc_file_path):
os.remove(self.enc_file_path)
AESFileEncryptor.encrypt(self.test_file_path,
self.enc_file_path,
secret)
self.assertTrue(os.path.exists(self.enc_file_path))
with open(self.enc_file_path, 'rb') as f:
encrypted = f.read()
self.assertEqual(
len(encrypted) % AESFileEncryptor.block_size, 0,
"Incorrect ciphertext size: {}. Should be multiple of {}".format(len(encrypted),
AESFileEncryptor.block_size))
def gen_secret(self):
return FileEncryptor.gen_secret(
self.min_secret_len,
self.max_secret_len,
)