mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 22:16:35 +00:00
Oleksandr Bezdieniezhnykh
95 lines
2.4 KiB
Python
95 lines
2.4 KiB
Python
import os
|
|
|
|
import pytest
|
|
|
|
from security import Security
|
|
|
|
|
|
def test_bt_enc_01_roundtrip_1024_random_bytes():
|
|
key = "test-key"
|
|
data = os.urandom(1024)
|
|
enc = Security.encrypt_to(data, key)
|
|
assert Security.decrypt_to(enc, key) == data
|
|
|
|
|
|
def test_bt_enc_02_roundtrip_onnx_model(fixture_onnx_model):
|
|
key = Security.get_model_encryption_key()
|
|
data = fixture_onnx_model
|
|
enc = Security.encrypt_to(data, key)
|
|
assert Security.decrypt_to(enc, key) == data
|
|
|
|
|
|
def test_bt_enc_03_roundtrip_empty_input():
|
|
key = "k"
|
|
data = b""
|
|
enc = Security.encrypt_to(data, key)
|
|
assert Security.decrypt_to(enc, key) == b""
|
|
|
|
|
|
def test_bt_enc_04_roundtrip_single_zero_byte():
|
|
key = "k"
|
|
data = b"\x00"
|
|
enc = Security.encrypt_to(data, key)
|
|
assert Security.decrypt_to(enc, key) == b"\x00"
|
|
|
|
|
|
def test_bt_enc_05_same_data_different_keys_different_ciphertext():
|
|
data = b"payload"
|
|
a = Security.encrypt_to(data, "key-a")
|
|
b = Security.encrypt_to(data, "key-b")
|
|
assert a != b
|
|
|
|
|
|
def test_bt_enc_06_decrypt_wrong_key_not_equal_original():
|
|
original = b"secret"
|
|
enc = Security.encrypt_to(original, "key-a")
|
|
out = Security.decrypt_to(enc, "key-b")
|
|
assert out != original
|
|
|
|
|
|
@pytest.mark.resilience
|
|
def test_rt_enc_01_corrupted_ciphertext():
|
|
key = "k"
|
|
original = b"hello world"
|
|
enc = Security.encrypt_to(original, key)
|
|
corrupted = bytearray(enc)
|
|
corrupted[len(corrupted) // 2] ^= 0xFF
|
|
try:
|
|
result = Security.decrypt_to(bytes(corrupted), key)
|
|
except Exception:
|
|
return
|
|
assert result != original
|
|
|
|
|
|
@pytest.mark.security
|
|
def test_st_enc_01_same_data_same_key_two_encryptions_differ():
|
|
key = "k"
|
|
data = b"x" * 64
|
|
a = Security.encrypt_to(data, key)
|
|
b = Security.encrypt_to(data, key)
|
|
assert a != b
|
|
|
|
|
|
@pytest.mark.security
|
|
def test_st_enc_02_wrong_key_cannot_recover_plaintext():
|
|
original = b"data"
|
|
enc = Security.encrypt_to(original, "key-one")
|
|
out = Security.decrypt_to(enc, "key-two")
|
|
assert out != original
|
|
|
|
|
|
@pytest.mark.security
|
|
def test_st_enc_03_model_encryption_key_deterministic():
|
|
a = Security.get_model_encryption_key()
|
|
b = Security.get_model_encryption_key()
|
|
assert a == b
|
|
|
|
|
|
@pytest.mark.resource_limit
|
|
def test_rl_enc_01_encrypted_size_at_most_plaintext_plus_32():
|
|
key = "k"
|
|
for n in (0, 1, 15, 16, 17, 1024, 4096):
|
|
data = os.urandom(n)
|
|
enc = Security.encrypt_to(data, key)
|
|
assert len(enc) <= n + 32
|