def test_simpleencryptio():
key = b"123"
buf = os.urandom(1024 * 1024 * 50)
bio = io.BytesIO(buf)
c = SimpleEncryptIO(bio, key, len(buf))
assert total_len(c) == len(buf) + PADDED_ENCRYPT_HEAD_WITH_SALT_LEN
enc = c.read()
d = to_decryptio(io.BytesIO(enc), key)
assert total_len(d) == len(buf)
dec = d.read()
assert buf == dec
def test_simpleencryptio():
key = "123"
nonce_or_iv = os.urandom(16)
buf = os.urandom(1024 * 1024 * 50)
bio = io.BytesIO(buf)
c = SimpleEncryptIO(bio, key, nonce_or_iv, len(buf))
assert total_len(c) == len(buf) + ENCRYPT_HEAD_LEN
enc = c.read()
d = to_decryptio(io.BytesIO(enc), key)
assert total_len(d) == len(buf)
dec = d.read()
assert buf == dec
def test_linked_crypted_io():
key = os.urandom(32)
buf = os.urandom(1024 * 50 + 14)
raw_len = len(buf)
raw_io = io.BytesIO(buf)
eio = SimpleEncryptIO(raw_io, key, raw_len)
dio = to_decryptio(eio, key)
eio = ChaCha20EncryptIO(dio, key, raw_len)
dio = to_decryptio(eio, key)
eio = AES256CBCEncryptIO(dio, key, raw_len)
dio = to_decryptio(eio, key)
length = 0
dbuf = b""
while True:
d = dio.read(1)
if not d:
break
dbuf += d
assert len(d) == 1
length += 1
assert length == raw_len
assert dbuf == buf
def encrypt_io(self, io: IO, encrypt_key: Any, nonce_or_iv: Any = None):
io_len = total_len(io)
if self == EncryptType.No:
return io
elif self == EncryptType.Simple:
return SimpleEncryptIO(io, encrypt_key, nonce_or_iv
or os.urandom(16), io_len)
elif self == EncryptType.ChaCha20:
return ChaCha20EncryptIO(io, encrypt_key, nonce_or_iv
or os.urandom(16), io_len)
elif self == EncryptType.AES265CBC:
return AES256CBCEncryptIO(io, encrypt_key, nonce_or_iv
or os.urandom(16), io_len)
else:
raise ValueError(f"Unknown EncryptType: {self}")