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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258
|
# type: ignore
import base64
import datetime
import decimal
import hashlib
import uuid
from typing import Any, Optional
import ormar
import pytest
from ormar import ModelDefinitionError, NoMatch
from ormar.fields.sqlalchemy_encrypted import EncryptedString
from tests.lifespan import init_tests
from tests.settings import create_config
base_ormar_config = create_config()
default_fernet = dict(
encrypt_secret="asd123", encrypt_backend=ormar.EncryptBackends.FERNET
)
class DummyBackend(ormar.fields.EncryptBackend):
def _initialize_backend(self, secret_key: bytes) -> None:
pass
def encrypt(self, value: Any) -> str:
return value
def decrypt(self, value: Any) -> str:
return value
class Author(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="authors")
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100, **default_fernet)
uuid_test = ormar.UUID(default=uuid.uuid4, uuid_format="string")
uuid_test2 = ormar.UUID(nullable=True, uuid_format="string")
password: str = ormar.String(
max_length=128,
encrypt_secret="udxc32",
encrypt_backend=ormar.EncryptBackends.HASH,
)
birth_year: int = ormar.Integer(
nullable=True,
encrypt_secret="secure89key%^&psdijfipew",
encrypt_backend=ormar.EncryptBackends.FERNET,
)
test_text: str = ormar.Text(default="", **default_fernet)
test_bool: bool = ormar.Boolean(nullable=False, **default_fernet)
test_float: float = ormar.Float(**default_fernet)
test_float2: Optional[float] = ormar.Float(nullable=True, **default_fernet)
test_datetime = ormar.DateTime(default=datetime.datetime.now, **default_fernet)
test_date = ormar.Date(default=datetime.date.today, **default_fernet)
test_time = ormar.Time(default=datetime.time, **default_fernet)
test_json = ormar.JSON(default={}, **default_fernet)
test_bigint: int = ormar.BigInteger(default=0, **default_fernet)
test_smallint: int = ormar.SmallInteger(default=0, **default_fernet)
test_decimal = ormar.Decimal(scale=2, precision=10, **default_fernet)
test_decimal2 = ormar.Decimal(max_digits=10, decimal_places=2, **default_fernet)
test_bytes = ormar.LargeBinary(max_length=100, **default_fernet)
test_b64bytes = ormar.LargeBinary(
max_length=100, represent_as_base64_str=True, **default_fernet
)
custom_backend: str = ormar.String(
max_length=200,
encrypt_secret="asda8",
encrypt_backend=ormar.EncryptBackends.CUSTOM,
encrypt_custom_backend=DummyBackend,
)
class Hash(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="hashes")
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(
max_length=128,
encrypt_secret="udxc32",
encrypt_backend=ormar.EncryptBackends.HASH,
)
class Filter(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="filters")
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100, **default_fernet)
hash = ormar.ForeignKey(Hash)
class Report(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="reports")
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100)
filters = ormar.ManyToMany(Filter)
create_test_database = init_tests(base_ormar_config)
def test_error_on_encrypted_pk():
with pytest.raises(ModelDefinitionError):
class Wrong(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="wrongs")
id: int = ormar.Integer(
primary_key=True,
encrypt_secret="asd123",
encrypt_backend=ormar.EncryptBackends.FERNET,
)
def test_error_on_encrypted_relation():
with pytest.raises(ModelDefinitionError):
class Wrong2(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="wrongs2")
id: int = ormar.Integer(primary_key=True)
author = ormar.ForeignKey(
Author,
encrypt_secret="asd123",
encrypt_backend=ormar.EncryptBackends.FERNET,
)
def test_error_on_encrypted_m2m_relation():
with pytest.raises(ModelDefinitionError):
class Wrong3(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="wrongs3")
id: int = ormar.Integer(primary_key=True)
author = ormar.ManyToMany(
Author,
encrypt_secret="asd123",
encrypt_backend=ormar.EncryptBackends.FERNET,
)
def test_wrong_backend():
with pytest.raises(ModelDefinitionError):
class Wrong3(ormar.Model):
ormar_config = base_ormar_config.copy(tablename="wrongs3")
id: int = ormar.Integer(primary_key=True)
author = ormar.Integer(
encrypt_secret="asd123",
encrypt_backend=ormar.EncryptBackends.CUSTOM,
encrypt_custom_backend="aa",
)
def test_db_structure():
assert Author.ormar_config.table.c.get("name").type.__class__ == EncryptedString
@pytest.mark.asyncio
async def test_save_and_retrieve():
async with base_ormar_config.database:
test_uuid = uuid.uuid4()
await Author(
name="Test",
birth_year=1988,
password="test123",
uuid_test=test_uuid,
test_float=1.2,
test_bool=True,
test_decimal=3.57,
test_decimal2=decimal.Decimal(5.5),
test_json=dict(aa=12),
custom_backend="test12",
test_bytes=b"test",
test_b64bytes=b"test2",
).save()
author = await Author.objects.get()
assert author.name == "Test"
assert author.birth_year == 1988
password = (
"03e4a4d513e99cb3fe4ee3db282c053daa3f3572b849c3868939a306944ad5c08"
"22b50d4886e10f4cd418c3f2df3ceb02e2e7ac6e920ae0c90f2dedfc8fa16e2"
)
assert author.password == password
assert author.uuid_test == test_uuid
assert author.uuid_test2 is None
assert author.test_datetime.date() == datetime.date.today()
assert author.test_date == datetime.date.today()
assert author.test_text == ""
assert author.test_float == 1.2
assert author.test_float2 is None
assert author.test_bigint == 0
assert author.test_json == {"aa": 12}
assert float(author.test_decimal) == 3.57
assert author.test_decimal2 == 5.5
assert author.custom_backend == "test12"
assert author.test_bytes == "test".encode("utf-8")
assert author.test_b64bytes == "dGVzdDI="
assert base64.b64decode(author.test_b64bytes) == b"test2"
@pytest.mark.asyncio
async def test_fernet_filters_nomatch():
async with base_ormar_config.database:
await Filter(name="test1").save()
await Filter(name="test1").save()
filters = await Filter.objects.all()
assert filters[0].name == filters[1].name == "test1"
with pytest.raises(NoMatch):
await Filter.objects.get(name="test1")
assert await Filter.objects.get_or_none(name="test1") is None
@pytest.mark.asyncio
async def test_hash_filters_works():
async with base_ormar_config.database:
await Hash(name="test1").save()
await Hash(name="test2").save()
secret = hashlib.sha256("udxc32".encode()).digest()
secret = base64.urlsafe_b64encode(secret)
hashed_test1 = hashlib.sha512(secret + "test1".encode()).hexdigest()
hash1 = await Hash.objects.get(name="test1")
assert hash1.name == hashed_test1
with pytest.raises(NoMatch):
await Filter.objects.get(name__icontains="test")
@pytest.mark.asyncio
async def test_related_model_fields_properly_decrypted():
async with base_ormar_config.database:
hash1 = await Hash(name="test1").save()
report = await Report.objects.create(name="Report1")
await report.filters.create(name="test1", hash=hash1)
await report.filters.create(name="test2")
report2 = await Report.objects.select_related("filters").get()
assert report2.filters[0].name == "test1"
assert report2.filters[1].name == "test2"
secret = hashlib.sha256("udxc32".encode()).digest()
secret = base64.urlsafe_b64encode(secret)
hashed_test1 = hashlib.sha512(secret + "test1".encode()).hexdigest()
report2 = await Report.objects.select_related("filters__hash").get()
assert report2.filters[0].name == "test1"
assert report2.filters[0].hash.name == hashed_test1
|