make messageheader a dataclass

This commit is contained in:
hinuiiik
2026-10-22 22:48:21 -04:00
parent 305e2a2b28
commit cd3dcaf4a3
+3 -2
View File
@@ -2,6 +2,8 @@ import os
import pickle import pickle
import string import string
import hmac import hmac
from dataclasses import dataclass
from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM from cryptography.hazmat.primitives.ciphers.aead import AESGCM
@@ -9,6 +11,7 @@ from cryptography.hazmat.primitives.kdf.hkdf import HKDF
CURVE = ec.SECP256R1() CURVE = ec.SECP256R1()
@dataclass
class MessageHeader: class MessageHeader:
dh: object dh: object
pn: int pn: int
@@ -29,7 +32,6 @@ def KDF_RK(rk, dh_out):
).derive(dh_out) ).derive(dh_out)
return output[:32], output[32:] return output[:32], output[32:]
def KDF_CK(ck): def KDF_CK(ck):
return ( return (
hmac.new(ck, b"\x01", "sha256").digest(), hmac.new(ck, b"\x01", "sha256").digest(),
@@ -40,7 +42,6 @@ def ENCRYPT(mk, plaintext, associated_data):
nonce = os.urandom(12) nonce = os.urandom(12)
return nonce + AESGCM(mk).encrypt(nonce, plaintext, associated_data) return nonce + AESGCM(mk).encrypt(nonce, plaintext, associated_data)
def DECRYPT(mk, ciphertext, associated_data): def DECRYPT(mk, ciphertext, associated_data):
return AESGCM(mk).decrypt(ciphertext[:12], ciphertext[12:], associated_data) return AESGCM(mk).decrypt(ciphertext[:12], ciphertext[12:], associated_data)