# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: SageMath 9.5
#     language: sage
#     name: sagemath-9.5
# ---

# # Реализация модели защищенного канала передачи данных по спецификации TLS 1.3

IS_DEBUG = 1


def trace(*args, **kwargs):
    """
    Отладочная трассировка
    """
    
    global IS_DEBUG
    if IS_DEBUG:
        print('[TRACE]', end=' ')
        print(*args, **kwargs)


# ----

# +
# Криптография из различных ГОСТ'ов
from pygost import mgm
from pygost import gost3412
from pygost import gost3410
from pygost import gost34112012256

# Криптография из западных стандартов
from Crypto.Cipher import AES
from Crypto.Hash import SHA256

# HKDF
from Crypto.Protocol.KDF import HKDF

# КГПСЧ
from Crypto.Random import get_random_bytes, random
from Crypto.Util.number import getStrongPrime

# Различные утилиты
from collections import namedtuple
from Crypto.Util import number
from functools import partial
from binascii import hexlify
from enum import Enum
import uuid
import gmpy2
import pickle


# -

# ----
#
# ## Реализация оберток над криптографическими алгоритмами

# +
def multiply_by_scalar(curve, scalar, point):
    """
    Умножение точки эллиптической кривой на скаляр 
    """
    
    return curve.exp(scalar, *point)


def add_points(curve, lhs, rhs):
    """
    Сложение двух точек эллиптической кривой
    """
    
    return curve._add(*lhs, *rhs)


# -

# ---

# +
#
# Сертификат это простая структурка, поэтому обойдусь именованным кортежем
#

Certificate = namedtuple('Certificate', ['id', 'public_key', 'signature'])


# -

# ----

class KuznyechikMGM(object):
    """
    Класс-обертка над Кузнечиком в режиме MGM
    """
    
    KEY_SIZE = 32
    
    def __init__(self, key: bytes):
        """
        Инициализация - сохраняю ключ
        """

        self._key = key
    
    
    def encrypt(self, plaintext: bytes):
        """
        Зашифрование шифром Кузнечик в режиме MGM
        """
        
        encrypter = partial(KuznyechikMGM._encrypt_block, self._key)
        cipher = mgm.MGM(encrypter, KuznyechikMGM._block_size())
        
        nonce = KuznyechikMGM._generate_nonce()
        ct_with_tag = cipher.seal(nonce, plaintext, b'')
        
        return nonce, ct_with_tag[:-cipher.tag_size], ct_with_tag[-cipher.tag_size:]
    
    
    def decrypt(self, nonce: bytes, ciphertext: bytes, tag: bytes):
        """
        Расшифрование шифром Кузнечик в режиме MGM
        """
        
        encrypter = partial(KuznyechikMGM._encrypt_block, self._key)
        cipher = mgm.MGM(encrypter, KuznyechikMGM._block_size())
            
        ct_with_tag = ciphertext + tag
            
        return cipher.open(nonce, ct_with_tag, b'')
        
        
    @staticmethod
    def _block_size():
        return gost3412.GOST3412Kuznechik.blocksize
        
        
    @staticmethod
    def _encrypt_block(key: bytes, block: bytes):
        cipher = gost3412.GOST3412Kuznechik(key)
        return cipher.encrypt(block)
    
    
    @staticmethod
    def _generate_nonce():
        nonce = get_random_bytes(KuznyechikMGM._block_size())
        while bytearray(nonce)[0] & 0x80 > 0:
            nonce = get_random_bytes(KuznyechikMGM._block_size())
            
        return nonce


class AESGCM(object):
    """
    Класс-обертка над AES в режиме GCM
    """
    
    KEY_SIZE = 16
    
    def __init__(self, key: bytes):
        """
        Инициализация - сохраняю ключ
        """

        self._key = key
        
        
    def encrypt(self, plaintext: bytes):
        """
        Зашифрование шифром AES в режиме GCM
        """
        
        cipher = AES.new(self._key, AES.MODE_GCM)
        return cipher.nonce, *cipher.encrypt_and_digest(plaintext)
    
    
    def decrypt(self, nonce: bytes, ciphertext: bytes, tag: bytes):
        """
        Зашифрование шифром AES в режиме GCM
        """
        
        cipher = AES.new(self._key, AES.MODE_GCM, nonce=nonce)
        return cipher.decrypt_and_verify(ciphertext, tag)


# ----

class Streebog(gost34112012256.GOST34112012256):
    """
    pycryptodome-совместимая обертка над хэш-функцией Стрибог
    """
    
    digest_size = 32
    
    def __init__(self, data=None):
        """
        Инициализация и обновление состояния (если надо)
        """
        
        super(Streebog, self).__init__()
        if data is not None:
            self.update(data)
    
    @staticmethod
    def new(data=None):
        """
        Создание нового объекта
        """
        
        return Streebog(data)


def hash_digest(hasher_type, *args):
    """
    Вычисление хэша при помощи объекта-обертки
    """
    
    hasher = hasher_type.new()
    
    for arg in args:
        hasher.update(arg)
        
    return hasher.digest()


# ----

class DiffieHellmanProcessor(object):
    """
    Обработчик обмена ключами по протоколу Диффи-Хеллмана
    """
    
    def __init__(self, group=None, generator=None):
        """
        Инициализация - выбор группы, образующего элемента
        Возможно конструирование по готовой группе и образующему
        """

        self._group = Integers(getStrongPrime(512)) if group is None else group
        self._generator = random.randint(int(2), int(self._group.order() - 2)) if generator is None \
            else generator
        
        self._modulus = self._group.order()
        
        trace('[DiffieHellmanProcessor]', f'Using group of order {self._modulus}')
        
    
    def generate_private_key(self):
        """
        Генерирование закрытого ключа
        """
        
        return random.randint(int(2), int(self.order - 2))
    
    
    def get_public_key(self, private_key):
        """
        Получение открытого ключа по закрытому
        """
        
        return self.exponentiate(private_key, self._generator)
    
    
    def exponentiate(self, scalar, element):
        """
        Возведение в степень элемента группы
        """
        
        return self._group(int(gmpy2.powmod(int(element), int(scalar), self.order)))
    
    
    def element_to_bytes(self, element):
        """
        Конвертация элемента группы в байтовое представление
        """
        
        return number.long_to_bytes(int(element))
    
    
    @property
    def group(self):
        """
        Получение группы
        """
        
        return self._group
    
        
    @property
    def generator(self):
        """
        Получение образующего элемента
        """
        
        return self._generator
    
    
    @property
    def order(self):
        """
        Получение порядка группы
        """
        
        return self._modulus


class EllipticCurveDiffieHellmanProcessor(object):
    """
    Обработчик обмена ключами по протоколу Диффи-Хеллмана
    в группе точек эллиптической кривой
    """
    
    def __init__(self, curve=None, generator=None):
        """
        Инициализация - выбор кривой, образующего элемента
        Возможно конструирование по готовой кривой и образующему
        """
        
        #
        # Для простоты буду использовать кривую из ГОСТ 34.10 
        #
        
        self._curve = gost3410.CURVES['id-tc26-gost-3410-2012-512-paramSetA'] if curve is None else curve
        self._generator = (self._curve.x, self._curve.y) if generator is None else generator
        self._modulus = self._curve.p
        
        trace('[EllipticCurveDiffieHellmanProcessor]', 
              f'Using group of order {self._modulus} over curve {self._curve.name}')
        
    
    def generate_private_key(self):
        """
        Генерирование закрытого ключа
        """
        
        return random.randint(int(2), int(self.order - 2))
    
    
    def get_public_key(self, private_key):
        """
        Получение открытого ключа по закрытому
        """
        
        return self.exponentiate(private_key, self._generator)
    
    
    def exponentiate(self, scalar, element):
        """
        Умножение точки на скаляр
        """
        
        if not self._curve.contains(element):
            return None
        
        return multiply_by_scalar(self._curve, scalar, element)
    
    
    def element_to_bytes(self, element):
        """
        Конвертация элемента группы в байтовое представление
        """
        
        return gost3410.pub_marshal(element)
    
    
    @property
    def group(self):
        """
        Получение группы
        """
        
        return self._curve
    
    
    @property
    def generator(self):
        """
        Получение образующего элемента
        """
        
        return self._generator
    
    
    @property
    def order(self):
        """
        Получение порядка группы
        """
        
        return self._modulus

# ----

cipher_suites = ['ECDHE_KUZNYECHIK_MGM_STREEBOG', 'DHE_AES_GCM_SHA256']


def resolve_cipher_suite(cipher_suite):
    """
    Получение примитивов по строковому идентификатору
    """
        
    #
    # Пока захардкодил это все
    #
        
    if cipher_suite == 'ECDHE_KUZNYECHIK_MGM_STREEBOG':
        return EllipticCurveDiffieHellmanProcessor, KuznyechikMGM, Streebog
        
    elif cipher_suite == 'DHE_AES_GCM_SHA256':
        return DiffieHellmanProcessor, AESGCM, SHA256
        
    raise ValueError(f'Unsupported cipher suite {cipher_suite}')


# ----

def hmac(key, *args, hasher=Streebog):
    """
    Вычисление HMAC (согласно Р 50.1.113 - 2016)
    Информационная технология
    КРИПТОГРАФИЧЕСКАЯ ЗАЩИТА ИНФОРМАЦИИ
    Криптографические алгоритмы, сопутствующие применению алгоритмов 
    электронной цифровой подписи и функции хэширования
    """
    
    def xor(lhs, rhs):
        return bytes([l ^^ r for l, r in zip(lhs, rhs)])
    
    IPAD = b'\x36' * 64
    OPAD = b'\x5c' * 64
    
    K_star = key + (b'\x00' * (512 - len(key)))
    
    inner_hash = hash_digest(hasher, xor(K_star, IPAD), *args)
    return hash_digest(hasher, xor(K_star, OPAD), inner_hash)


def test_hmac():
    """
    Р 50.1.113 - 2016
    Информационная технология
    КРИПТОГРАФИЧЕСКАЯ ЗАЩИТА ИНФОРМАЦИИ
    Криптографические алгоритмы, сопутствующие применению алгоритмов 
    электронной цифровой подписи и функции хэширования
    
    4.1.1 HMAC_GOSTR3411_2012_256
    
    Приложение А. Контрольные примеры
    """
    
    K = b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f' + \
        b'\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
    
    T = b'\x01\x26\xbd\xb8\x78\x00\xaf\x21\x43\x41\x45\x65\x63\x78\x01\x00'
    
    expected = b'\xa1\xaa\x5f\x7d\xe4\x02\xd7\xb3\xd3\x23\xf2\x99\x1c\x8d\x45\x34' + \
        b'\x01\x31\x37\x01\x0a\x83\x75\x4f\xd0\xaf\x6d\x7c\xd4\x92\x2e\xd9'
    
    assert hmac(K, T, hasher=Streebog) == expected


# +
#
# Протестирую HMAC на контрольном примере из рекомендаций по стандартизации
#

test_hmac()


# -

# ----

class GOST3410Signer(object):
    """
    Объект-создатель подписей по ГОСТ 34.10 - 2012
    """
    
    def __init__(self):
        """
        Инициализация кривой и ключевой пары
        """
        
        self._curve = gost3410.CURVES['id-tc26-gost-3410-2012-512-paramSetA']
        self._private_key = gost3410.prv_unmarshal(get_random_bytes(64))
        self._public_key = gost3410.public_key(self._curve, self._private_key)
        
        
    def sign(self, *args):
        """
        Подпись произвольного набора байтовых объектов
        """
        
        digest = hash_digest(Streebog, *args)
        return gost3410.sign(self._curve, self._private_key, digest)
        
        
    @property
    def public_key(self):
        """
        Получение открытого ключа
        """
        
        return self._public_key
    
    
    @property
    def private_key(self):
        """
        Получение закрытого ключа
        """
        
        return self._private_key

    
    @property
    def curve(self):
        """
        Получение используемой кривой
        """
        
        return self._curve
    
    
    @property
    def generator(self):
        """
        Получение образующего группы точек кривой
        """
        
        return self._curve.x, self._curve.y


def verify_signature(curve, public_key, signature, *args):
    """
    Проверка подписи по ГОСТ 34.10 - 2012
    """
    
    digest = hash_digest(Streebog, *args)
    return gost3410.verify(curve, public_key, digest, signature)


# ----
#
# ## Реализация участников протокола

Offer = namedtuple('Offer', ['cipher_suite', 'group', 'generator'])


class GenericParticipant(object):
    """
    Класс-абстракция над участником протокола, который получает сертификат на свой открытый
    ключ и может осуществлять подпись сообщений
    """
    
    def __init__(self, ca, *, fake_data=None, demo_cert_error=False):
        """
        Инициализация - получение идентификатора, генерирование ключевой пары, 
        получение сертификата
        """
        
        self._id = uuid.uuid4()
        self._signer = GOST3410Signer()
        
        if fake_data is not None:
            #
            # Демонстрация неудачного выполнения протокола
            #
            
            if demo_cert_error:
                #
                # Тут демонстрация неудачи получения серификата
                #
                
                self._signer._public_key = fake_data
                self._certificate = ca.issue_certificate(self)
            else:
                #
                # Тут как бы сертификат и открытый ключ правильные, а закрытый ключ противник не угадал
                #
                
                self._certificate = ca.issue_certificate(self)
                self._signer._public_key = fake_data
        else:
            #
            # Стандартное выполнение
            #
            
            self._certificate = ca.issue_certificate(self)
        
        self._n = None
        
        trace('[GenericParticipant]', f'Certificate = {self._certificate}')
        
        
    def sign(self, *args):
        """
        Подписывание произвольных данных
        """
        
        return self._signer.sign(*args)
        
        
    @property
    def public_key(self):
        """
        Получение открытого ключа подписи
        """
        
        return self._signer.public_key
    
    
    @property
    def identifier(self):
        """
        Получение идентификатора
        """
        
        return self._id
    
    
    @property
    def certificate(self):
        """
        Получение сертификата
        """
        
        return self._certificate
    
    
    @property
    def curve(self):
        """
        Получение кривой, используемой для подписи
        """
        
        return self._signer.curve
    
    
    def _get_point_for_schnorr_proof(self):
        #
        # Костыль, но все же работает
        #
        
        self._n = gost3410.prv_unmarshal(get_random_bytes(64))
        return multiply_by_scalar(self._signer.curve, 
                                  self._n, self._signer.generator)
    
    
    def _get_response(self, chellenge):
        return self._n + chellenge * self._signer.private_key


class TLS13Participant(GenericParticipant):
    """
    Обобщенный участник общения по TLS 1.3
    """
    
    def __init__(self, ca, *, fake_data=None, demo_cert_error=False):
        """
        Инициализация - по факту просто сохраняю ссылку на СА
        """
        
        super(TLS13Participant, self).__init__(ca, 
                                               fake_data=fake_data, 
                                               demo_cert_error=demo_cert_error)
        self._ca = ca
    
    
    def check_certificate(self, identifier: uuid.UUID, public_key, cert: Certificate):
        """
        Проверка полученного сертификата
        """
        
        if self._ca.is_revoked(cert):
            trace('[TLS13Participant]', f'Revoked certificate')
            return False
            
        id_as_bytes = str(identifier).encode()
        pk_as_bytes = gost3410.pub_marshal(public_key)
        
        if not verify_signature(self._ca.curve, 
                                self._ca.public_key, 
                                cert.signature, 
                                id_as_bytes, 
                                pk_as_bytes):
            trace('[TLS13Participant]', f'Invalid certificate')
            return False
        
        return True
        
        
    @staticmethod
    def check_mac(key, mac, *args):
        """
        Проверка имитовставки
        """
        
        return mac == hmac(key, *args)
    
    
    @staticmethod
    def check_signature(curve, public_key, signature, *args):
        """
        Проверка подписи
        """
        
        return verify_signature(curve, public_key, signature, *args)


class TLS13Client(TLS13Participant):
    """
    Клиент для TLS 1.3
    """
        
    S2C = 0
    C2S = 1
    
    class AuthMode(Enum):
        """
        Режим аутентицикации
        """
        
        OWA = 1  # Односторонняя
        TWA = 2  # Двусторонняя
    
    
    def __init__(self, ca, cipher_suite, *, fake_data=None, demo_cert_error=False):
        """
        Инициализация - разбор набора примитивов
        """
        
        super(TLS13Client, self).__init__(ca, fake_data=fake_data, 
                                          demo_cert_error=demo_cert_error)
        
        self._key_exchange_type, self._cipher_type, self._hasher = resolve_cipher_suite(cipher_suite)
        self._cipher_suite = cipher_suite
        self._session_keys = {}
        
        trace('[TLS13Client]', f'Using cipher suite {cipher_suite}')
        
        
    def establish_connection(self, server, auth_mode: AuthMode):
        """
        Установление соединения с сервером
        """
        
        key_exchange = self._key_exchange_type()
        
        ephemeral_dh_sk = key_exchange.generate_private_key()
        ephemeral_dh_pk = key_exchange.get_public_key(ephemeral_dh_sk)
        client_nonce    = get_random_bytes(64)
        
        trace('[TLS13Client]', f'Ephemeral DH public key: {ephemeral_dh_pk}')
        trace('[TLS13Client]', f'Nonce: {hexlify(client_nonce)}')
        
        offer = Offer(self._cipher_suite, 
                      key_exchange.group, 
                      key_exchange.generator)
        
        #
        # А теперь запрос на сервер
        #
        
        is_one_way = auth_mode == TLS13Client.AuthMode.OWA
        server_hello = server.get_hello(self.identifier, ephemeral_dh_pk, client_nonce, 
                                        offer, is_one_way)
        
        if server_hello is None:
            trace('[TLS13Client]', f'Server refused the connection')
            return False
        
        #
        # Распакую ответ
        #
        
        ephemeral_dh_server_pk, server_nonce, cipher_suite, c1, c2, c3, c4 = server_hello
        
        if cipher_suite != self._cipher_suite:
            trace('[TLS13Client]', f'Unsupported cipher suite {cipher_suite}')
            return False
        
        #
        # Общий секрет DH
        #
        
        ephemeral_dh_common = key_exchange.exponentiate(ephemeral_dh_sk, ephemeral_dh_server_pk)
        if ephemeral_dh_common is None:
            trace(f'[TLS13Client]', 'Server ECDH public key is not on curve')
        
        #
        # Вырабатываю ключи для симметричного шифрования и имитовставки
        #
        
        client_packet = b''.join([key_exchange.element_to_bytes(ephemeral_dh_pk),
                                  client_nonce, str(offer).encode()])
        
        server_packet = b''.join([key_exchange.element_to_bytes(ephemeral_dh_server_pk),
                                  server_nonce, cipher_suite.encode()])
        
        key_derivation_data = b''.join([key_exchange.element_to_bytes(ephemeral_dh_common),
                                        client_packet, server_packet])
        
        encryption_key, mac_key = HKDF(key_derivation_data, self._cipher_type.KEY_SIZE, 
                                       b'\x00' * 16, self._hasher, 2)
        
        #
        # Расшифровываю полученные пакеты
        #
        
        cipher = self._cipher_type(encryption_key)
        
        try:
            cert_requested     = cipher.decrypt(*c1)
            server_certificate = cipher.decrypt(*c2)
            server_signature   = cipher.decrypt(*c3)
            server_mac         = cipher.decrypt(*c4)
        except ValueError:
            trace('[TLS13Client]', f'Malformed server response')
            return False
        
        #
        # Проверяю полученнные данные
        #
        
        if is_one_way == cert_requested:
            #
            # Запросили сертификат в ходе односторонней аутентификации или наоборот
            #
            
            trace('[TLS13Client]', f'Invalid client certificate request')
            return False
        
        if not self.check_certificate(server.identifier, server.public_key, 
                                      pickle.loads(server_certificate)):
            trace('[TLS13Client]', f'Invalid server certificate')
            return False
        
        if not TLS13Participant.check_signature(server.curve, server.public_key, server_signature, 
                                                client_packet, server_packet, c1[1], c2[1]):
            trace('[TLS13Client]', f'Malformed server response')
            return False
        
        if not TLS13Participant.check_mac(mac_key, server_mac, 
                                          client_packet, server_packet, c1[1], c2[1], c3[1]):
            trace('[TLS13Client]', f'Malformed server response')
            return False
            
        #
        # Если требуется двусторонняя аутентификация, то зашлю на сервер сертификат
        #
        
        if not is_one_way:
            c5 = cipher.encrypt(pickle.dumps(self.certificate))
            
            signature = self.sign(client_packet, server_packet, c1[1], c2[1], c3[1], c4[1], c5[1])
            c6 = cipher.encrypt(signature)
            
            mac = hmac(mac_key, client_packet, server_packet, c1[1], c2[1], c3[1], c4[1], c5[1], c6[1])
            c7 = cipher.encrypt(mac)
            
            if not server.process_client_certificate(self, c5, c6, c7):
                trace('[TLS13Client]', f'Server refused the connection')
                return False
        
        #
        # Теперь я генерирую ключи шифрования
        #
        
        key_derivation_data = b''.join([key_exchange.element_to_bytes(ephemeral_dh_common),
                                        client_packet, server_packet, c1[1], c2[1], c3[1], c4[1]])
            
        self._session_keys[server.identifier] = HKDF(key_derivation_data, self._cipher_type.KEY_SIZE,
                                                     b'\x00' * 16, self._hasher, 2)
        
        trace('[TLS13Client]', f'Server accepted the connection')
        
        trace('[TLS13Client]', 
              f's -> c key: {hexlify(self._session_keys[server.identifier][TLS13Client.S2C])}')
        trace('[TLS13Client]', 
              f'c -> s key: {hexlify(self._session_keys[server.identifier][TLS13Client.C2S])}')
        
        return True
    
    
    def send_wait_receive_message(self, server, message: str):
        """
        Отправка сообщения на сервер и получение ответа
        """
        
        if server.identifier not in self._session_keys:
            trace('[TLS13Client]', f'Connection with server {server.identifier} is not established yet')
            return
        
        cipher = self._cipher_type(self._session_keys[server.identifier][TLS13Client.C2S])
        
        response = server._receive_message_and_respond(self, cipher.encrypt(message.encode()))
        if response is None:
            trace('[TLS13Client]', f'Malformed server response')
        
        try:
            cipher = self._cipher_type(self._session_keys[server.identifier][TLS13Client.S2C])
            return cipher.decrypt(*response).decode()
        except ValueError:
            trace('[TLS13Client]', f'Malformed server response')
            
            
    def change_keys(self, server):
        """
        Смена ключей по инициативе клиента
        """
        
        if server.identifier not in self._session_keys:
            trace('[TLS13Client]', f'Connection with server {server.identifier} is not established yet')
            return
        
        #
        # Меняю у себя
        #
        
        self._change_keys_internal(server.identifier)
        
        #
        # Меняю на сервере
        #
        
        server._change_keys_internal(self.identifier)
        
        
    def _change_keys_internal(self, server_id):
        s2c_key = HKDF(self._session_keys[server_id][TLS13Client.S2C], self._cipher_type.KEY_SIZE, 
                       b'\x00' * 16, self._hasher, 1)
        
        c2s_key = HKDF(self._session_keys[server_id][TLS13Client.C2S], self._cipher_type.KEY_SIZE, 
                       b'\x00' * 16, self._hasher, 1)
        
        trace('[TLS13Client]', f'New s -> c key: {hexlify(s2c_key)}')
        trace('[TLS13Client]', f'New c -> s key: {hexlify(c2s_key)}')
        
        self._session_keys[server_id] = (s2c_key, c2s_key)


class TLS13Server(TLS13Participant):
    """
    Сервер для TLS 1.3
    """
    
    CIPHER_TYPE            = 0
    HASHER                 = 1
    HANDSHAKE_ENC_KEY      = 2
    HANDSHAKE_MAC_KEY      = 3
    HANDSHAKE_VALUES       = 4
    HANDSHAKE_KEY_EXCHANGE = 5
    
    S2C = 0
    C2S = 1
    
    
    def __init__(self, ca, cipher_suites):
        """
        Инициализация - сохранение настроек
        """
        
        super(TLS13Server, self).__init__(ca)
        
        self._cipher_suites = cipher_suites
        self._session_keys  = {}
        self._mapped_states = {}
        
        
    def get_hello(self, client_id, ephemeral_dh_client_pk, client_nonce, offer, is_one_way):
        """
        Получение ServerHello
        """
        
        #
        # Провверяю, что поддерживается вообще набор примитивов
        #
        
        cipher_suite, group, generator = offer
        
        if cipher_suite not in self._cipher_suites:
            trace('[TLS13Server]', f'Unsupported cipher suite {cipher_suite}')
            return None
        
        trace('[TLS13Server]', f'Using cipher suite {cipher_suite} with client {client_id}')
        
        #
        # Получаю примитивы
        #
        
        key_exchange_type, cipher_type, hasher = resolve_cipher_suite(cipher_suite)
        key_exchange = key_exchange_type(group, generator)
        
        #
        # Вырабатываю ключи DH и nonce, после чего получаю общий секрет DH
        #
        
        ephemeral_dh_sk = key_exchange.generate_private_key()
        ephemeral_dh_pk = key_exchange.get_public_key(ephemeral_dh_sk)
        server_nonce    = get_random_bytes(64)
        
        trace('[TLS13Server]', f'Ephemeral DH public key: {ephemeral_dh_pk}')
        trace('[TLS13Server]', f'Nonce: {hexlify(client_nonce)}')
        
        ephemeral_dh_common = key_exchange.exponentiate(ephemeral_dh_sk, ephemeral_dh_client_pk)
        if ephemeral_dh_common is None:
            trace(f'[TLS13Client]', 'Client ECDH public key is not on curve')
        
        #
        # Вырабатываю ключи для симметричного шифрования и имитовставки
        #
        
        client_packet = b''.join([key_exchange.element_to_bytes(ephemeral_dh_client_pk),
                                  client_nonce, str(offer).encode()])
        
        server_packet = b''.join([key_exchange.element_to_bytes(ephemeral_dh_pk),
                                  server_nonce, cipher_suite.encode()])
        
        key_derivation_data = b''.join([key_exchange.element_to_bytes(ephemeral_dh_common),
                                        client_packet, server_packet])
        
        encryption_key, mac_key = HKDF(key_derivation_data, cipher_type.KEY_SIZE, 
                                       b'\x00' * 16, hasher, 2)
        
        #
        # Рассчитываю нужные значения
        # 
        
        cipher = cipher_type(encryption_key)
        
        cert_request = bytes([not is_one_way])
        c1 = cipher.encrypt(cert_request)
        
        cert = pickle.dumps(self.certificate)
        c2 = cipher.encrypt(cert)
        
        signature = self.sign(client_packet, server_packet, c1[1], c2[1])
        c3 = cipher.encrypt(signature)
        
        mac = hmac(mac_key, client_packet, server_packet, c1[1], c2[1], c3[1])
        c4 = cipher.encrypt(mac)
        
        #
        # Тут будут сохранены нужные значения
        #
        
        self._mapped_states[client_id] = [None, None, None, None, None, None]
        
        #
        # Теперь выработаю ключи (если это односторонняя аутентификация)
        #
        
        if is_one_way:
            self._session_keys[client_id] = TLS13Server._generate_keys(key_exchange, cipher_type.KEY_SIZE, 
                                                                       hasher, ephemeral_dh_common, 
                                                                       client_packet, server_packet, 
                                                                       c1, c2, c3, c4)
        
        #
        # Сохраняю остальное и возвращаю клиенту
        #
        
        self._mapped_states[client_id][TLS13Server.CIPHER_TYPE]       = cipher_type
        self._mapped_states[client_id][TLS13Server.HASHER]            = hasher
        self._mapped_states[client_id][TLS13Server.HANDSHAKE_ENC_KEY] = encryption_key
        self._mapped_states[client_id][TLS13Server.HANDSHAKE_MAC_KEY] = mac_key
        self._mapped_states[client_id][TLS13Server.HANDSHAKE_VALUES]  = (ephemeral_dh_common, client_packet, 
                                                                         server_packet, c1, c2, c3, c4)
        self._mapped_states[client_id][TLS13Server.HANDSHAKE_KEY_EXCHANGE] = key_exchange
        
        return ephemeral_dh_pk, server_nonce, cipher_suite, c1, c2, c3, c4
    
    
    def process_client_certificate(self, client, c5, c6, c7):
        """
        Проверка сертификата клиента
        """
        
        client_id = client.identifier
        
        #
        # Достану ранее сохраненные параметры
        #
        
        cipher_type    = self._mapped_states[client_id][TLS13Server.CIPHER_TYPE]
        hasher         = self._mapped_states[client_id][TLS13Server.HASHER]
        encryption_key = self._mapped_states[client_id][TLS13Server.HANDSHAKE_ENC_KEY]
        mac_key        = self._mapped_states[client_id][TLS13Server.HANDSHAKE_MAC_KEY]
        
        ephemeral_dh_common, client_packet, server_packet, c1, c2, c3, c4 = \
            self._mapped_states[client_id][TLS13Server.HANDSHAKE_VALUES]
        
        key_exchange   = self._mapped_states[client_id][TLS13Server.HANDSHAKE_KEY_EXCHANGE]
        
        #
        # Расшифрую полученный пакет
        #
        
        cipher = cipher_type(encryption_key)
        
        try:
            client_certificate = cipher.decrypt(*c5)
            client_signature   = cipher.decrypt(*c6)
            client_mac         = cipher.decrypt(*c7)
        except ValueError:
            trace('[TLS13Server]', f'Malformed client response')
            return False
        
        #
        # проверю все полученные данные
        #
        
        if not self.check_certificate(client.identifier, client.public_key, 
                                      pickle.loads(client_certificate)):
            trace('[TLS13Server]', f'Invalid client certificate')
            return False
        
        if not TLS13Participant.check_signature(client.curve, client.public_key, client_signature, 
                                                client_packet, server_packet, 
                                                c1[1], c2[1], c3[1], c4[1], c5[1]):
            trace('[TLS13Server]', f'Malformed client response')
            return False
        
        if not TLS13Participant.check_mac(mac_key, client_mac, client_packet, server_packet, 
                                          c1[1], c2[1], c3[1], c4[1], c5[1], c6[1]):
            trace('[TLS13Server]', f'Malformed client response')
            return False
        
        #
        # Все чудесно - генерирую ключи
        #
        
        key_derivation_data = b''.join([key_exchange.element_to_bytes(ephemeral_dh_common),
                                        client_packet, server_packet, c1[1], c2[1], c3[1], c4[1]])
            
        self._session_keys[client_id] = TLS13Server._generate_keys(key_exchange, cipher_type.KEY_SIZE, 
                                                                   hasher, ephemeral_dh_common, client_packet, 
                                                                   server_packet, c1, c2, c3, c4)
        
        return True
            
            
    def change_keys(self, client):
        """
        Смена ключей по инициативе сервера
        """
        
        if client.identifier not in self._session_keys:
            trace('[TLS13Client]', f'Connection with client {client.identifier} is not established yet')
            return
        
        #
        # Меняю у себя
        #
        
        self._change_keys_internal(client.identifier)
        
        #
        # Меняю на сервере
        #
        
        client._change_keys_internal(self.identifier)
        
        
    def _change_keys_internal(self, client_id):
        s2c_key = HKDF(self._session_keys[client_id][TLS13Client.S2C], 
                       self._mapped_states[client_id][TLS13Server.CIPHER_TYPE].KEY_SIZE, 
                       b'\x00' * 16, self._mapped_states[client_id][TLS13Server.HASHER], 1)
        
        c2s_key = HKDF(self._session_keys[client_id][TLS13Client.C2S], 
                       self._mapped_states[client_id][TLS13Server.CIPHER_TYPE].KEY_SIZE, 
                       b'\x00' * 16, self._mapped_states[client_id][TLS13Server.HASHER], 1)
        
        trace('[TLS13Client]', f'New s -> c key: {hexlify(s2c_key)}')
        trace('[TLS13Client]', f'New c -> s key: {hexlify(c2s_key)}')
        
        self._session_keys[client_id] = (s2c_key, c2s_key)
    
    
    def _receive_message_and_respond(self, client, encrypted_message):
        client_id = client.identifier
        
        try:
            cipher = self._mapped_states[client_id][TLS13Server.CIPHER_TYPE](
                self._session_keys[client_id][TLS13Server.C2S])
            
            message = cipher.decrypt(*encrypted_message).decode()
            print(f'Received message "{message}"')
            
            cipher = self._mapped_states[client_id][TLS13Server.CIPHER_TYPE](
                self._session_keys[client_id][TLS13Server.S2C])
            
            return cipher.encrypt(f'Hello, {message}!'.encode())
        except ValueError:
            trace('[TLS13Server]', f'Malformed client message')
    
    
    @staticmethod
    def _generate_keys(key_exchange, key_size, hasher, ephemeral_dh_common, 
                       client_packet, server_packet, c1, c2, c3, c4):
        key_derivation_data = b''.join([key_exchange.element_to_bytes(ephemeral_dh_common),
                                        client_packet, server_packet, c1[1], c2[1], c3[1], c4[1]])
            
        s2c_key, c2s_key = HKDF(key_derivation_data, key_size, 
                                b'\x00' * 16, hasher, 2)
        
        trace('[TLS13Server]', f's -> c key: {hexlify(s2c_key)}')
        trace('[TLS13Server]', f'c -> s key: {hexlify(c2s_key)}')
        
        return s2c_key, c2s_key


class TLS13CertificationAuthority(object):
    """
    Удостоверющий центр для TLS 1.3
    """
    
    def __init__(self):
        """
        Инициализация - генерирование ключевой пары
        """
        
        self._signer = GOST3410Signer()
        
        self._certificates = set()
        self._revoked_certificates = set()
        
        trace('[TLS13CertificationAuthority]', 
              f'Public key = {self._signer.public_key}')
        
        
    @property
    def public_key(self):
        """
        Получение открытого ключа
        """
        
        return self._signer.public_key
    
    
    @property
    def curve(self):
        """
        Получение кривой, используемой для подписи
        """
        
        return self._signer.curve
        
        
    def issue_certificate(self, participant: GenericParticipant):
        """
        Выдача сертификата на открытый ключ
        """
        
        #
        # Проверю, что клиент действительно владеет закрытым ключом
        #
        
        if not self._check_participant(participant):
            trace('[TLS13CertificationAuthority]', 
                  f'Participant cannot prove a valid private key possession')
            
            return None
        
        #
        # Все чудесно, подписываю ключ и отдаю сертификат
        #
        
        id_as_bytes = str(participant.identifier).encode()
        pk_as_bytes = gost3410.pub_marshal(participant.public_key)
        
        cert = Certificate(participant.identifier, participant.public_key, 
                           self._signer.sign(id_as_bytes, pk_as_bytes))
        
        self._certificates.add(cert)
        
        trace('[TLS13CertificationAuthority]', 
              f'Issued certificate for {participant.identifier}. Total certificates: {len(self._certificates)}')
        
        return cert
    
    
    def revoke_certificate(self, participant: GenericParticipant):
        """
        Отзыв сертификата
        """
        
        #
        # А можно ли вообще отзывать сертификат?
        #
        
        if not self._check_participant(participant):
            trace('[TLS13CertificationAuthority]', 
                  f'Participant cannot prove a valid private key possession')
            
            return 
        
        if participant.certificate not in self._certificates:
            #
            # А отзывать то и нечего
            #
            
            return
            
        self._revoked_certificates.add(participant.certificate)
        self._certificates.discard(participant.certificate)
        
        trace('[TLS13CertificationAuthority]', 
              f'List of revoked certificates now has {len(self._revoked_certificates)} elements')
        
        
    def is_revoked(self, cert: Certificate):
        """
        Проверка на факт отзыва сертификата
        """
        
        return cert in self._revoked_certificates
    
    
    def _check_participant(self, participant: GenericParticipant) -> bool:
        Q = participant._get_point_for_schnorr_proof()
        c = gost3410.prv_unmarshal(get_random_bytes(64))
        t = participant._get_response(c)
        
        tP  = multiply_by_scalar(self._signer.curve, t, self._signer.generator)
        cPK = multiply_by_scalar(self._signer.curve, c, participant.public_key)
        
        return tP == add_points(self._signer.curve, Q, cPK)

# ----
# ## Демонстрация удачного выполнения протоколов

# Создание удостоверяющего центра
ca = TLS13CertificationAuthority()

# Создание сервера
srv = TLS13Server(ca, cipher_suites)

# Первый валидный пользователь (использует ECDH, Кузнечик в режиме MGM и Стрибог)
alice = TLS13Client(ca, cipher_suite=cipher_suites[0])

# Второй валидный пользователь (использует DH, AES в режиме GCM и SHA256)
bob = TLS13Client(ca, cipher_suite=cipher_suites[1])

# Устанавливаю соединение с односторонней аутентификацией
alice.establish_connection(srv, TLS13Client.AuthMode.OWA)

# Устанавливаю соединение с двусторонней аутентификацией
bob.establish_connection(srv, TLS13Client.AuthMode.TWA)

# Отправлю сообщение
alice.send_wait_receive_message(srv, 'world')

# Сменю ключи по инициативе клиента
alice.change_keys(srv)

# И еще раз отправлю сообщение, чтобы было видно, что ключи успешно поменялись
alice.send_wait_receive_message(srv, 'world')

# Отправлю сообщение другим клиентом
bob.send_wait_receive_message(srv, 'crypto')

# Смена ключей по инициативе сервера
srv.change_keys(bob)

# Все по-прежнему успешно отправляется
bob.send_wait_receive_message(srv, 'crypto')

# ----
# ## Отзыв сертификата

# Третий валидный пользователь
carol = TLS13Client(ca, cipher_suite=cipher_suites[0])

# Отзову сертификат Кэрол
ca.revoke_certificate(carol)

# Попробую подключиться, используя двустороннюю аутентификацию
carol.establish_connection(srv, TLS13Client.AuthMode.TWA)

# ----
# ## Демонстрация неудачного выполнения протоколов

# Тут Мэллори пытается установить подключение с двусторонней аутентификацией,
# используя сертификат не на свой ключ (я просто подменю тут открытый ключ для простоты)
mallory = TLS13Client(ca, cipher_suite=cipher_suites[0], 
                      fake_data=bob.public_key)

mallory.establish_connection(srv, TLS13Client.AuthMode.TWA)

# Мэллори теперь пытается получить сертификат на ключ Боба, не зная его закрытого ключа
mallory = TLS13Client(ca, cipher_suite=cipher_suites[0], 
                      fake_data=bob.public_key, 
                      demo_cert_error=True)

StarcoderdataJupyterScriptsDedupFiltered