blob: 83fbe41b04fc235706480d4e63fcc1d028b28e8c (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
|
# test_crypto.py
import unittest
from io import BytesIO
from cryptography.fernet import Fernet
from crypto import generate_key, load_key, encrypt, decrypt
class TestCryptoFunctions(unittest.TestCase):
def setUp(self):
self.key = generate_key()
with open('vault.sqlite', 'wb') as vault:
vault.write(b'Hello, world!')
def tearDown(self):
try:
os.remove('vault.sqlite')
except FileNotFoundError:
pass
def test_generate_key(self):
new_key = generate_key()
self.assertNotEqual(new_key, self.key)
def test_load_key(self):
key_file = 'key.bin'
with open(key_file, 'wb') as key:
key.write(self.key)
loaded_key = load_key(key_file)
self.assertEqual(loaded_key, self.key)
def test_encrypt_decrypt(self):
encrypt(self.key, filename='vault.sqlite')
with open('vault.sqlite', 'rb') as vault:
encrypted_data = vault.read()
decrypt(self.key)
with open('vault.sqlite', 'rb') as vault:
decrypted_data = vault.read()
f = Fernet(self.key)
self.assertEqual(f.decrypt(encrypted_data), b'Hello, world!')
def test_load_invalid_key(self):
key_file = 'key.bin'
try:
load_key(key_file)
self.fail('Expected FileNotFoundError')
except FileNotFoundError:
pass
if __name__ == '__main__':
unittest.main()
|
