import abc import nacl.bindings from Cryptodome.Cipher import AES from .key_schedule import tls_sha256, tls_sha384 class TLS_AEAD_Cipher(abc.ABC): NONCE_LEN = 12 @property @abc.abstractmethod def KEY_LEN(self): "" @property @abc.abstractmethod def MAC_LEN(self): "" @property @abc.abstractmethod def tls_hash(self): "" @abc.abstractmethod def cipher(self): "" def __init__(self, secret): self.reset(secret) def reset(self, secret): self.secret = secret self.key = self.tls_hash.derive_key(self.secret, self.KEY_LEN) self.iv = int.from_bytes( self.tls_hash.derive_iv(self.secret, self.NONCE_LEN), "big" ) self.sequence_number = 0 def next_application_traffic_secret(self): return self.tls_hash.hkdf_expand_label( self.secret, b"traffic upd", b"", self.tls_hash.hash_len ) def update_traffic_secret(self): self.reset(self.next_application_traffic_secret()) def verify_data(self, msg): return self.tls_hash.verify_data(self.secret, msg) def get_nonce(self): nonce = self.sequence_number ^ self.iv nonce = nonce.to_bytes(self.NONCE_LEN, "big") self.sequence_number += 1 return nonce def decrypt(self, ciphertext, associated_data): cipher = self.cipher() cipher.update(associated_data) return cipher.decrypt_and_verify( ciphertext[: -self.MAC_LEN], ciphertext[-self.MAC_LEN :] ) def encrypt(self, plaintext, associated_data): cipher = self.cipher() cipher.update(associated_data) ciphertext, tag = cipher.encrypt_and_digest(plaintext) return ciphertext + tag def tls_ciphertext(self, plaintext): head = b"\x17\x03\x03" + (len(plaintext) + self.MAC_LEN).to_bytes(2, "big") return head + self.encrypt(plaintext, head) class TLS_CHACHA20_POLY1305_SHA256(TLS_AEAD_Cipher): KEY_LEN = 32 MAC_LEN = 16 tls_hash = tls_sha256 def cipher(self): "" def decrypt(self, ciphertext, associated_data): nonce = self.get_nonce() return nacl.bindings.crypto_aead_chacha20poly1305_ietf_decrypt( bytes(ciphertext), associated_data, nonce, self.key ) def encrypt(self, plaintext, associated_data): nonce = self.get_nonce() return nacl.bindings.crypto_aead_chacha20poly1305_ietf_encrypt( bytes(plaintext), associated_data, nonce, self.key ) class TLS_AES_128_GCM_SHA256(TLS_AEAD_Cipher): KEY_LEN = 16 MAC_LEN = 16 tls_hash = tls_sha256 def cipher(self): return AES.new( self.key, AES.MODE_GCM, nonce=self.get_nonce(), mac_len=self.MAC_LEN ) class TLS_AES_256_GCM_SHA384(TLS_AEAD_Cipher): KEY_LEN = 32 MAC_LEN = 16 tls_hash = tls_sha384 def cipher(self): return AES.new( self.key, AES.MODE_GCM, nonce=self.get_nonce(), mac_len=self.MAC_LEN ) class TLS_AES_128_CCM_SHA256(TLS_AEAD_Cipher): KEY_LEN = 16 MAC_LEN = 16 tls_hash = tls_sha256 def cipher(self): return AES.new( self.key, AES.MODE_CCM, nonce=self.get_nonce(), mac_len=self.MAC_LEN ) class TLS_AES_128_CCM_8_SHA256(TLS_AEAD_Cipher): tls_hash = tls_sha256 KEY_LEN = 16 MAC_LEN = 8 def cipher(self): return AES.new( self.key, AES.MODE_CCM, nonce=self.get_nonce(), mac_len=self.MAC_LEN )