223 lines
6.9 KiB
Python
223 lines
6.9 KiB
Python
"""ChaCha20-Blake3 AEAD boxes, counter-nonce channels, and a fast CSPRNG.
|
|
Everything supports in-place operation into caller buffers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hmac
|
|
import secrets
|
|
import threading
|
|
from collections.abc import Buffer
|
|
from typing import overload
|
|
|
|
import blake3
|
|
from cryptography.hazmat.primitives.ciphers import Cipher
|
|
from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20
|
|
|
|
KEY_LEN = 32
|
|
NONCE_LEN = 12
|
|
_MAC_KEY_BLOCK = 64
|
|
|
|
|
|
def nonce_increment(nonce: Buffer) -> None:
|
|
"""Increment the nonce in place, little-endian."""
|
|
n = memoryview(nonce).cast("B")
|
|
for i in range(len(n)):
|
|
if n[i] < 255:
|
|
n[i] += 1
|
|
return
|
|
n[i] = 0
|
|
|
|
|
|
def _session(key: Buffer, nonce: Buffer, ad: Buffer | None, maclen: int):
|
|
"""ChaCha20 stream + keyed BLAKE3 with ad absorbed (None when maclen=0)."""
|
|
key = memoryview(key).cast("B")
|
|
nonce = memoryview(nonce).cast("B")
|
|
if len(key) != KEY_LEN:
|
|
raise ValueError("key must be 32 bytes")
|
|
if len(nonce) != NONCE_LEN:
|
|
raise ValueError("nonce must be 12 bytes")
|
|
if not 0 <= maclen <= 32:
|
|
raise ValueError("maclen must be in 0..32")
|
|
# IETF layout: 32-bit little-endian block counter || 96-bit nonce
|
|
stream = Cipher(ChaCha20(key, bytes(4) + nonce), mode=None).encryptor()
|
|
if maclen == 0:
|
|
return stream, None
|
|
tag_key = stream.update(bytes(_MAC_KEY_BLOCK))[:KEY_LEN]
|
|
mac = blake3.blake3(key=tag_key)
|
|
ad = memoryview(ad).cast("B") if ad else memoryview(b"")
|
|
mac.update(len(ad).to_bytes(8, "little"))
|
|
mac.update(ad)
|
|
return stream, mac
|
|
|
|
|
|
def _tag(mac, ct_len: int, maclen: int) -> bytes:
|
|
m = mac.copy()
|
|
m.update(ct_len.to_bytes(8, "little"))
|
|
return m.digest()[:maclen]
|
|
|
|
|
|
def encrypt_into(
|
|
out: Buffer,
|
|
msg: Buffer,
|
|
key: Buffer,
|
|
nonce: Buffer | None = None,
|
|
*,
|
|
ad: Buffer | None = None,
|
|
maclen: int = 16,
|
|
) -> memoryview:
|
|
"""Encrypt message under key, with optional nonce, ad and mac, into out buffer."""
|
|
msg = memoryview(msg).cast("B")
|
|
hdr = NONCE_LEN if nonce is None else 0
|
|
used = hdr + len(msg) + maclen
|
|
mv = memoryview(out).cast("B")
|
|
if len(mv) < used:
|
|
raise TypeError("output buffer too small")
|
|
out = mv[:used]
|
|
if nonce is None:
|
|
nonce = out[:NONCE_LEN]
|
|
random(nonce)
|
|
stream, mac = _session(key, nonce, ad, maclen)
|
|
ct = out[hdr : hdr + len(msg)]
|
|
stream.update_into(msg, ct)
|
|
if mac is not None:
|
|
mac.update(ct)
|
|
out[hdr + len(msg) :] = _tag(mac, len(msg), maclen)
|
|
return out
|
|
|
|
|
|
def decrypt_into(
|
|
out: Buffer,
|
|
ct: Buffer,
|
|
key: Buffer,
|
|
nonce: Buffer | None = None,
|
|
*,
|
|
ad: Buffer | None = None,
|
|
maclen: int = 16,
|
|
) -> memoryview:
|
|
"""Decrypt ciphertext into out buffer. Other arguments must match those at encryption point."""
|
|
ct = memoryview(ct).cast("B")
|
|
if nonce is None:
|
|
if len(ct) < NONCE_LEN + maclen:
|
|
raise ValueError("ciphertext buffer too short")
|
|
nonce, ct = ct[:NONCE_LEN], ct[NONCE_LEN:]
|
|
elif len(ct) < maclen:
|
|
raise ValueError("ciphertext buffer too short")
|
|
used = len(ct) - maclen
|
|
mv = memoryview(out).cast("B")
|
|
if len(mv) < used:
|
|
raise TypeError("output buffer too small")
|
|
out = mv[:used]
|
|
stream, mac = _session(key, nonce, ad, maclen)
|
|
if mac is not None:
|
|
mac.update(ct[:used])
|
|
if not hmac.compare_digest(bytes(ct[used:]), _tag(mac, used, maclen)):
|
|
raise ValueError("invalid authentication tag")
|
|
stream.update_into(ct[:used], out)
|
|
return out
|
|
|
|
|
|
def encrypt(
|
|
msg: Buffer,
|
|
key: Buffer,
|
|
nonce: Buffer | None = None,
|
|
*,
|
|
ad: Buffer | None = None,
|
|
maclen: int = 16,
|
|
) -> bytearray:
|
|
"""Encrypt message under key, with optional nonce, ad and mac, returning ciphertext."""
|
|
msg = memoryview(msg).cast("B")
|
|
out = bytearray(len(msg) + maclen + (NONCE_LEN if nonce is None else 0))
|
|
encrypt_into(out, msg, key, nonce, ad=ad, maclen=maclen)
|
|
return out
|
|
|
|
|
|
def decrypt(
|
|
ct: Buffer,
|
|
key: Buffer,
|
|
nonce: Buffer | None = None,
|
|
*,
|
|
ad: Buffer | None = None,
|
|
maclen: int = 16,
|
|
) -> bytearray:
|
|
"""Decrypt and return message. Other arguments must match those at encryption point."""
|
|
ct = memoryview(ct).cast("B")
|
|
hdr = NONCE_LEN if nonce is None else 0
|
|
if len(ct) < hdr + maclen:
|
|
raise ValueError("ciphertext buffer too short")
|
|
out = bytearray(len(ct) - hdr - maclen)
|
|
decrypt_into(out, ct, key, nonce, ad=ad, maclen=maclen)
|
|
return out
|
|
|
|
|
|
class _Sequence:
|
|
def __init__(self, key: Buffer, *, maclen: int = 16):
|
|
key = memoryview(key).cast("B")
|
|
if len(key) != KEY_LEN:
|
|
raise ValueError("key must be 32 bytes")
|
|
self._key = key
|
|
self._maclen = maclen
|
|
self._nonce = bytearray(NONCE_LEN)
|
|
|
|
|
|
class Sender(_Sequence):
|
|
"""Encrypting end of a message channel: messages are sent in order,
|
|
each individually authenticated. Pair with Receiver on the same key."""
|
|
|
|
def encrypt_into(
|
|
self, out: Buffer, msg: Buffer, *, ad: Buffer | None = None
|
|
) -> memoryview:
|
|
ret = encrypt_into(out, msg, self._key, self._nonce, ad=ad, maclen=self._maclen)
|
|
nonce_increment(self._nonce)
|
|
return ret
|
|
|
|
def encrypt(self, msg: Buffer, *, ad: Buffer | None = None) -> bytearray:
|
|
msg = memoryview(msg).cast("B")
|
|
out = bytearray(len(msg) + self._maclen)
|
|
self.encrypt_into(out, msg, ad=ad)
|
|
return out
|
|
|
|
|
|
class Receiver(_Sequence):
|
|
"""Decrypting end of a message channel: messages are verified and
|
|
returned in order. Pair with Sender on the same key."""
|
|
|
|
def decrypt_into(
|
|
self, out: Buffer, ct: Buffer, *, ad: Buffer | None = None
|
|
) -> memoryview:
|
|
ret = decrypt_into(out, ct, self._key, self._nonce, ad=ad, maclen=self._maclen)
|
|
nonce_increment(self._nonce)
|
|
return ret
|
|
|
|
def decrypt(self, ct: Buffer, *, ad: Buffer | None = None) -> bytearray:
|
|
ct = memoryview(ct).cast("B")
|
|
if len(ct) < self._maclen:
|
|
raise ValueError("ciphertext buffer too short")
|
|
out = bytearray(len(ct) - self._maclen)
|
|
self.decrypt_into(out, ct, ad=ad)
|
|
return out
|
|
|
|
|
|
_random = threading.local()
|
|
|
|
|
|
@overload
|
|
def random(nbytes: int, /) -> bytearray: ...
|
|
@overload
|
|
def random(buf: Buffer, /) -> None: ...
|
|
def random(arg: int | Buffer, /):
|
|
"""random(nbytes) -> nbytes as bytearray; random(buf) fills in place.
|
|
Seeded from OS random, thread safe."""
|
|
if isinstance(arg, int):
|
|
mv = memoryview(ret := bytearray(arg))
|
|
else:
|
|
mv, ret = memoryview(arg).cast("B"), None
|
|
zeroes = bytes(len(mv))
|
|
try:
|
|
_random.enc.update_into(zeroes, mv)
|
|
except (AttributeError, ValueError):
|
|
_random.enc = Cipher(
|
|
ChaCha20(secrets.token_bytes(KEY_LEN), bytes(16)), mode=None
|
|
).encryptor()
|
|
_random.enc.update_into(zeroes, mv)
|
|
return ret
|