54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import struct
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash
|
|
from ..registry import Header
|
|
from ..util import to_bytes, urlsafe_b64decode
|
|
|
|
|
|
__all__ = [
|
|
"derive_key_for_concat_kdf",
|
|
"u32be_len_input",
|
|
]
|
|
|
|
|
|
def derive_key_for_concat_kdf(
|
|
shared_key: bytes, header: Header, cek_size: int, key_size: int | None, tag: bytes | None = None
|
|
) -> bytes:
|
|
# PartyUInfo
|
|
apu_info = u32be_len_input(header.get("apu"), True)
|
|
# PartyVInfo
|
|
apv_info = u32be_len_input(header.get("apv"), True)
|
|
# SuppPubInfo
|
|
|
|
if key_size:
|
|
alg_id = u32be_len_input(header["alg"])
|
|
bit_size = key_size
|
|
else:
|
|
alg_id = u32be_len_input(header["enc"])
|
|
bit_size = cek_size
|
|
|
|
pub_info = struct.pack(">I", bit_size)
|
|
fixed_info = alg_id + apu_info + apv_info + pub_info
|
|
|
|
if tag:
|
|
cctag = u32be_len_input(tag)
|
|
fixed_info += cctag
|
|
|
|
ckdf = ConcatKDFHash(
|
|
algorithm=hashes.SHA256(),
|
|
length=bit_size // 8,
|
|
otherinfo=fixed_info,
|
|
)
|
|
return ckdf.derive(shared_key)
|
|
|
|
|
|
def u32be_len_input(s: bytes | str | None, use_base64: bool = False) -> bytes:
|
|
if not s:
|
|
return b"\x00\x00\x00\x00"
|
|
sb: bytes
|
|
if use_base64:
|
|
sb = urlsafe_b64decode(to_bytes(s))
|
|
else:
|
|
sb = to_bytes(s)
|
|
return struct.pack(">I", len(sb)) + sb
|