241 lines
7.9 KiB
Python
241 lines
7.9 KiB
Python
import os
|
|
import pickle
|
|
import hmac
|
|
import copy
|
|
import json
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
|
|
from cryptography.exceptions import InvalidTag
|
|
from cryptography.hazmat.primitives import hashes, serialization
|
|
from cryptography.hazmat.primitives.asymmetric import ec
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
|
|
|
CURVE = ec.SECP256R1()
|
|
|
|
@dataclass
|
|
class MessageHeader:
|
|
dh: object
|
|
pn: int
|
|
n: int
|
|
|
|
def GENERATE_DH():
|
|
return ec.generate_private_key(CURVE)
|
|
|
|
def DH(dh_pair, dh_pub):
|
|
return dh_pair.exchange(ec.ECDH(), dh_pub)
|
|
|
|
def KDF_RK(rk, dh_out):
|
|
output = HKDF(
|
|
algorithm=hashes.SHA256(),
|
|
length=64,
|
|
salt=rk,
|
|
info=b"",
|
|
).derive(dh_out)
|
|
return output[:32], output[32:]
|
|
|
|
def KDF_CK(ck):
|
|
return (
|
|
hmac.new(ck, b"\x01", "sha256").digest(),
|
|
hmac.new(ck, b"\x02", "sha256").digest(),
|
|
)
|
|
|
|
def ENCRYPT(mk, plaintext, associated_data):
|
|
nonce = os.urandom(12)
|
|
return nonce + AESGCM(mk).encrypt(nonce, plaintext, associated_data)
|
|
|
|
def DECRYPT(mk, ciphertext, associated_data):
|
|
if mk is None or not isinstance(ciphertext, bytes) or len(ciphertext) < 12:
|
|
return None
|
|
return AESGCM(mk).decrypt(ciphertext[:12], ciphertext[12:], associated_data)
|
|
|
|
def HEADER(dh_pair, pn, n):
|
|
public_key = dh_pair.public_key()
|
|
return MessageHeader(public_key, int(pn), int(n))
|
|
|
|
def CONCAT(ad, header):
|
|
header_tup = pickle.dumps(
|
|
(header.dh.public_bytes(serialization.Encoding.X962,serialization.PublicFormat.UncompressedPoint), header.pn, header.n),
|
|
)
|
|
return len(ad).to_bytes(8) + ad + header_tup
|
|
|
|
class MessengerServer:
|
|
def __init__(self, server_signing_key, server_decryption_key):
|
|
self.server_signing_key = server_signing_key
|
|
self.server_decryption_key = server_decryption_key
|
|
|
|
def decryptReport(self, ct):
|
|
data = pickle.loads(ct)
|
|
key = ec.EllipticCurvePublicKey.from_encoded_point(
|
|
CURVE, data["bytes"]
|
|
)
|
|
shared = DH(self.server_decryption_key, key)
|
|
derived_key = HKDF(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=None,
|
|
info=b""
|
|
).derive(shared)
|
|
aesgcm = AESGCM(derived_key)
|
|
plaintext = aesgcm.decrypt(data["nonce"], data["ciphertext"], data["bytes"])
|
|
return plaintext.decode("latin1")
|
|
|
|
def signCert(self, cert):
|
|
return self.server_signing_key.sign(cert, ec.ECDSA(hashes.SHA256()))
|
|
|
|
class MessengerClient:
|
|
|
|
def __init__(self, name, server_signing_pk, server_encryption_pk):
|
|
self.name = name
|
|
self.server_signing_pk = server_signing_pk
|
|
self.server_encryption_pk = server_encryption_pk
|
|
self.conns = {}
|
|
self.certs = {}
|
|
self.dh = GENERATE_DH()
|
|
self.certificate = None
|
|
|
|
@staticmethod
|
|
def RatchetInitAlice(state, SK, bob_dh_public_key, alice_dh_key_pair=None):
|
|
state.DHs = alice_dh_key_pair or GENERATE_DH()
|
|
state.DHr = bob_dh_public_key
|
|
state.RK, state.CKs = KDF_RK(SK, DH(state.DHs, state.DHr))
|
|
state.CKr = None
|
|
state.Ns = 0
|
|
state.Nr = 0
|
|
state.PN = 0
|
|
state.MKSKIPPED = {}
|
|
|
|
@staticmethod
|
|
def RatchetInitBob(state, SK, bob_dh_key_pair):
|
|
state.DHs = bob_dh_key_pair
|
|
state.DHr = None
|
|
state.RK = SK
|
|
state.CKs = None
|
|
state.CKr = None
|
|
state.Ns = 0
|
|
state.Nr = 0
|
|
state.PN = 0
|
|
state.MKSKIPPED = {}
|
|
|
|
@staticmethod
|
|
def RatchetSendKey(state):
|
|
state.CKs, mk = KDF_CK(state.CKs)
|
|
Ns = state.Ns
|
|
state.Ns += 1
|
|
return Ns, mk
|
|
|
|
@staticmethod
|
|
def RatchetEncrypt(state, plaintext, AD):
|
|
Ns, mk = MessengerClient.RatchetSendKey(state)
|
|
header = HEADER(state.DHs, state.PN, Ns)
|
|
return header, ENCRYPT(mk, plaintext, CONCAT(AD, header))
|
|
|
|
@staticmethod
|
|
def RatchetReceiveKey(state, header):
|
|
if not isinstance(header, MessageHeader) or header.dh is None:
|
|
return None
|
|
if header.pn < 0 or header.n < 0:
|
|
return None
|
|
if state.DHr is None or header.dh.public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint) != state.DHr.public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint):
|
|
MessengerClient.DHRatchet(state, header)
|
|
if header.n != state.Nr:
|
|
return None
|
|
state.CKr, mk = KDF_CK(state.CKr)
|
|
state.Nr += 1
|
|
return mk
|
|
|
|
@staticmethod
|
|
def RatchetDecrypt(state, header, ciphertext, AD):
|
|
mk = MessengerClient.RatchetReceiveKey(state, header)
|
|
return DECRYPT(mk, ciphertext, CONCAT(AD, header))
|
|
|
|
@staticmethod
|
|
def DHRatchet(state, header):
|
|
state.PN = state.Ns
|
|
state.Ns = 0
|
|
state.Nr = 0
|
|
state.DHr = header.dh
|
|
state.RK, state.CKr = KDF_RK(state.RK, DH(state.DHs, state.DHr))
|
|
state.DHs = GENERATE_DH()
|
|
state.RK, state.CKs = KDF_RK(state.RK, DH(state.DHs, state.DHr))
|
|
|
|
def generateCertificate(self):
|
|
self.certificate = pickle.dumps((self.name, self.dh.public_key().public_bytes(serialization.Encoding.X962,
|
|
serialization.PublicFormat.UncompressedPoint)))
|
|
return self.certificate
|
|
|
|
def receiveCertificate(self, certificate, signature):
|
|
self.server_signing_pk.verify(signature, certificate, ec.ECDSA(hashes.SHA256()))
|
|
name, public_bytes = pickle.loads(certificate)
|
|
public_key = ec.EllipticCurvePublicKey.from_encoded_point(CURVE, public_bytes)
|
|
self.certs[name] = public_key
|
|
|
|
def sendMessage(self, name, message):
|
|
if name not in self.certs:
|
|
return None
|
|
state = self.conns.get(name)
|
|
if state is None:
|
|
state = SimpleNamespace()
|
|
initial_secret = DH(self.dh, self.certs[name])
|
|
self.RatchetInitAlice(
|
|
state, initial_secret, self.certs[name], self.dh
|
|
)
|
|
self.conns[name] = state
|
|
plaintext = message.encode("utf-8")
|
|
return self.RatchetEncrypt(state, plaintext, b"")
|
|
|
|
def receiveMessage(self, name, header, ciphertext):
|
|
if name not in self.certs:
|
|
return None
|
|
state = self.conns.get(name)
|
|
if state is None:
|
|
state = SimpleNamespace()
|
|
initial_secret = DH(self.dh, self.certs[name])
|
|
self.RatchetInitBob(state, initial_secret, self.dh)
|
|
temp_state = copy.copy(state) # fixes the issue with bad ciphertexts corrupting the state
|
|
try:
|
|
plaintext = self.RatchetDecrypt(temp_state, header, ciphertext, b"")
|
|
message = plaintext.decode("utf-8")
|
|
except (
|
|
InvalidTag,
|
|
ValueError,
|
|
TypeError,
|
|
AttributeError,
|
|
IndexError,
|
|
UnicodeDecodeError,
|
|
pickle.UnpicklingError,
|
|
EOFError,
|
|
KeyError,
|
|
):
|
|
return None
|
|
self.conns[name] = temp_state
|
|
return message
|
|
|
|
|
|
def report(self, name, message):
|
|
report = (name, message)
|
|
plaintext = pickle.dumps(report)
|
|
temp_key = GENERATE_DH()
|
|
temp_bytes = temp_key.public_key().public_bytes(
|
|
serialization.Encoding.X962,
|
|
serialization.PublicFormat.UncompressedPoint,
|
|
)
|
|
shared_secret = DH(temp_key, self.server_encryption_pk)
|
|
hkdf = HKDF(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=None,
|
|
info=b""
|
|
)
|
|
key = hkdf.derive(shared_secret)
|
|
nonce = os.urandom(12)
|
|
aesgcm = AESGCM(key)
|
|
ciphertext = aesgcm.encrypt(nonce, plaintext, temp_bytes) # not sure if passing in temp_bytes is correct, but it seems like it should be
|
|
data = {
|
|
"bytes": temp_bytes,
|
|
"nonce": nonce,
|
|
"ciphertext": ciphertext
|
|
}
|
|
return plaintext.decode("latin1"), pickle.dumps(data)
|