from typing import Literal, Union, overload
from .constants import AEADID, HPKEMode
from .context import ContextRecipient, ContextSender
from .primitives.aead import AEADBase
from .primitives.kdf import KDFBase
from .primitives.kem import KEMBase
from .types import KEMPrivateKey, KEMPublicKey
from .utils import I2OSP, concat
[docs]
class HPKESetup:
"""
HPKE Setup functions (RFC 9180 ยง5.1).
Handles the setup phase of HPKE operations, creating encryption contexts
for different HPKE modes.
Parameters
----------
kem : KEMBase
Key Encapsulation Mechanism instance.
kdf : KDFBase
Key Derivation Function instance.
aead : AEADBase
Authenticated Encryption with Associated Data instance.
Attributes
----------
kem : KEMBase
Key Encapsulation Mechanism instance.
kdf : KDFBase
Key Derivation Function instance.
aead : AEADBase
Authenticated Encryption with Associated Data instance.
suite_id : bytes
HPKE suite identifier.
"""
[docs]
def __init__(self, kem: KEMBase, kdf: KDFBase, aead: AEADBase) -> None:
self.kem = kem
self.kdf = kdf
self.aead = aead
self.suite_id = concat(
b"HPKE",
I2OSP(kem.kem_id, 2),
I2OSP(kdf.kdf_id, 2),
I2OSP(aead.aead_id, 2),
)
@overload
def key_schedule(
self,
role: Literal["S"],
mode: HPKEMode,
shared_secret: bytes,
info: bytes,
psk: bytes,
psk_id: bytes,
) -> ContextSender: ...
@overload
def key_schedule(
self,
role: Literal["R"],
mode: HPKEMode,
shared_secret: bytes,
info: bytes,
psk: bytes,
psk_id: bytes,
) -> ContextRecipient: ...
[docs]
def key_schedule(
self,
role: Literal["S", "R"],
mode: HPKEMode,
shared_secret: bytes,
info: bytes,
psk: bytes,
psk_id: bytes,
) -> Union[ContextSender, ContextRecipient]:
"""
Perform key schedule to derive encryption context.
Parameters
----------
role : str
Context role ('S' for sender, 'R' for recipient).
mode : HPKEMode
HPKE mode.
shared_secret : bytes
Shared secret from KEM.
info : bytes
Application-supplied information.
psk : bytes
Pre-shared key.
psk_id : bytes
Pre-shared key identifier.
Returns
-------
ContextSender or ContextRecipient
Encryption context for the specified role.
Raises
------
ValueError
If PSK inputs are inconsistent with the mode.
"""
self.verify_psk_inputs(mode, psk, psk_id)
psk_id_hash = self.kdf.labeled_extract(
salt=b"",
label="psk_id_hash",
ikm=psk_id,
suite_id=self.suite_id,
)
info_hash = self.kdf.labeled_extract(
salt=b"",
label="info_hash",
ikm=info,
suite_id=self.suite_id,
)
key_schedule_context = concat(
I2OSP(mode, 1),
psk_id_hash,
info_hash,
)
secret = self.kdf.labeled_extract(
salt=shared_secret,
label="secret",
ikm=psk,
suite_id=self.suite_id,
)
key = b""
base_nonce = b""
if self.aead.aead_id != AEADID.EXPORT_ONLY:
key = self.kdf.labeled_expand(
prk=secret,
label="key",
info=key_schedule_context,
L=self.aead.Nk,
suite_id=self.suite_id,
)
base_nonce = self.kdf.labeled_expand(
prk=secret,
label="base_nonce",
info=key_schedule_context,
L=self.aead.Nn,
suite_id=self.suite_id,
)
exporter_secret = self.kdf.labeled_expand(
prk=secret,
label="exp",
info=key_schedule_context,
L=self.kdf.Nh,
suite_id=self.suite_id,
)
if role == "S":
return ContextSender(
aead=self.aead,
kdf=self.kdf,
key=key,
base_nonce=base_nonce,
exporter_secret=exporter_secret,
suite_id=self.suite_id,
)
return ContextRecipient(
aead=self.aead,
kdf=self.kdf,
key=key,
base_nonce=base_nonce,
exporter_secret=exporter_secret,
suite_id=self.suite_id,
)
[docs]
def setup_base_sender(self, pkR: KEMPublicKey, info: bytes) -> tuple[bytes, ContextSender]:
"""
Setup Base mode for sender.
Parameters
----------
pkR : Key Object
Recipient's public key.
info : bytes
Application-supplied information.
Returns
-------
tuple
Tuple of (encapsulated key, sender context).
"""
shared_secret, enc = self.kem.encap(pkR)
ctx = self.key_schedule("S", HPKEMode.MODE_BASE, shared_secret, info, psk=b"", psk_id=b"")
return enc, ctx
[docs]
def setup_base_recipient(self, enc: bytes, skR: KEMPrivateKey, info: bytes) -> ContextRecipient:
"""
Setup Base mode for recipient.
Parameters
----------
enc : bytes
Encapsulated public key.
skR : Key Object
Recipient's private key.
info : bytes
Application-supplied information.
Returns
-------
ContextRecipient
Recipient context.
"""
shared_secret = self.kem.decap(enc, skR)
return self.key_schedule("R", HPKEMode.MODE_BASE, shared_secret, info, psk=b"", psk_id=b"")
[docs]
def setup_psk_sender(
self, pkR: KEMPublicKey, info: bytes, psk: bytes, psk_id: bytes
) -> tuple[bytes, ContextSender]:
"""
Setup PSK mode for sender.
Parameters
----------
pkR : Key Object
Recipient's public key.
info : bytes
Application-supplied information.
psk : bytes
Pre-shared key (must have at least 32 bytes of entropy).
psk_id : bytes
Pre-shared key identifier.
Returns
-------
tuple
Tuple of (encapsulated key, sender context).
Raises
------
ValueError
If PSK has insufficient entropy.
"""
if len(psk) < 32:
raise ValueError("PSK must have at least 32 bytes of entropy")
shared_secret, enc = self.kem.encap(pkR)
ctx = self.key_schedule("S", HPKEMode.MODE_PSK, shared_secret, info, psk=psk, psk_id=psk_id)
return enc, ctx
[docs]
def setup_psk_recipient(
self, enc: bytes, skR: KEMPrivateKey, info: bytes, psk: bytes, psk_id: bytes
) -> ContextRecipient:
"""
Setup PSK mode for recipient.
Parameters
----------
enc : bytes
Encapsulated public key.
skR : Key Object
Recipient's private key.
info : bytes
Application-supplied information.
psk : bytes
Pre-shared key (must have at least 32 bytes of entropy).
psk_id : bytes
Pre-shared key identifier.
Returns
-------
ContextRecipient
Recipient context.
Raises
------
ValueError
If PSK has insufficient entropy.
"""
if len(psk) < 32:
raise ValueError("PSK must have at least 32 bytes of entropy")
shared_secret = self.kem.decap(enc, skR)
return self.key_schedule(
"R", HPKEMode.MODE_PSK, shared_secret, info, psk=psk, psk_id=psk_id
)
[docs]
def setup_auth_sender(
self, pkR: KEMPublicKey, info: bytes, skS: KEMPrivateKey
) -> tuple[bytes, ContextSender]:
"""
Setup Auth mode for sender.
Parameters
----------
pkR : Key Object
Recipient's public key.
info : bytes
Application-supplied information.
skS : Key Object
Sender's private key.
Returns
-------
tuple
Tuple of (encapsulated key, sender context).
"""
shared_secret, enc = self.kem.auth_encap(pkR, skS)
ctx = self.key_schedule("S", HPKEMode.MODE_AUTH, shared_secret, info, psk=b"", psk_id=b"")
return enc, ctx
[docs]
def setup_auth_recipient(
self, enc: bytes, skR: KEMPrivateKey, info: bytes, pkS: KEMPublicKey
) -> ContextRecipient:
"""
Setup Auth mode for recipient.
Parameters
----------
enc : bytes
Encapsulated public key.
skR : Key Object
Recipient's private key.
info : bytes
Application-supplied information.
pkS : Key Object
Sender's public key.
Returns
-------
ContextRecipient
Recipient context.
"""
shared_secret = self.kem.auth_decap(enc, skR, pkS)
return self.key_schedule("R", HPKEMode.MODE_AUTH, shared_secret, info, psk=b"", psk_id=b"")
[docs]
def setup_auth_psk_sender(
self, pkR: KEMPublicKey, info: bytes, psk: bytes, psk_id: bytes, skS: KEMPrivateKey
) -> tuple[bytes, ContextSender]:
"""
Setup AuthPSK mode for sender.
Parameters
----------
pkR : Key Object
Recipient's public key.
info : bytes
Application-supplied information.
psk : bytes
Pre-shared key (must have at least 32 bytes of entropy).
psk_id : bytes
Pre-shared key identifier.
skS : Key Object
Sender's private key.
Returns
-------
tuple
Tuple of (encapsulated key, sender context).
Raises
------
ValueError
If PSK has insufficient entropy.
"""
if len(psk) < 32:
raise ValueError("PSK must have at least 32 bytes of entropy")
shared_secret, enc = self.kem.auth_encap(pkR, skS)
ctx = self.key_schedule(
"S", HPKEMode.MODE_AUTH_PSK, shared_secret, info, psk=psk, psk_id=psk_id
)
return enc, ctx
[docs]
def setup_auth_psk_recipient(
self,
enc: bytes,
skR: KEMPrivateKey,
info: bytes,
psk: bytes,
psk_id: bytes,
pkS: KEMPublicKey,
) -> ContextRecipient:
"""
Setup AuthPSK mode for recipient.
Parameters
----------
enc : bytes
Encapsulated public key.
skR : Key Object
Recipient's private key.
info : bytes
Application-supplied information.
psk : bytes
Pre-shared key (must have at least 32 bytes of entropy).
psk_id : bytes
Pre-shared key identifier.
pkS : Key Object
Sender's public key.
Returns
-------
ContextRecipient
Recipient context.
Raises
------
ValueError
If PSK has insufficient entropy.
"""
if len(psk) < 32:
raise ValueError("PSK must have at least 32 bytes of entropy")
shared_secret = self.kem.auth_decap(enc, skR, pkS)
return self.key_schedule(
"R", HPKEMode.MODE_AUTH_PSK, shared_secret, info, psk=psk, psk_id=psk_id
)