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, )