import secrets
import sys
from struct import *
from tools import *
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256
from Crypto.Random import get_random_bytes
from Crypto.Cipher import PKCS1_v1_5
from Crypto.Hash import HMAC


class Info:

    def build_client_hello(self):
        body = bytes()

        size = 32 + 2 + 3
        head = pack('BH', 0x80, size)

        random = get_random_bytes(32)
        a = 0b00000001
        cipherSuite = a.to_bytes(2, 'big')

        body = head + random + cipherSuite
        return body

    def build_server_hello(self):
        body = bytes()

        size = 32 + 2 + 3
        head = pack('BH', 0x81, size)

        random = get_random_bytes(32)
        a = 0b00000001
        cipherSuite = a.to_bytes(2, 'big')

        body=head+random+cipherSuite
        return body

    def build_server_certificate(self, public_key):
        size = len(public_key)+3
        head = pack('BH', 0x82, size)

        body = head+public_key

        return bytes(body)

    def build_client_certificate(self, public_key):
        size = len(public_key)+3
        head = pack('BH', 0x83, size)

        body = head+public_key

        return bytes(body)

    def build_certificate_verify(self, key, client_hello, server_hello, server_certificate):

        temp = bitwise_or_bytes(client_hello, server_hello)
        message = bitwise_or_bytes(server_certificate, temp)
        h = SHA256.new(message)
        signature = pkcs1_15.new(key).sign(h)

        size = len(signature)+3
        head = pack('BH', 0x84, size)

        body = head+signature
        return body

    def build_client_key_exchange(self, key):
        session_key = get_random_bytes(48)
        cipher_rsa = PKCS1_v1_5.new(key)
        enc_session_key = cipher_rsa.encrypt(session_key)

        size = len(enc_session_key)+3
        head = pack('BH', 0x85, size)

        body = head+enc_session_key
        return session_key,body

    def build_server_finished(self, messages, master_secret):
        handshake_messages = bitwise_or_bytes(messages[0], messages[1])
        temp = bitwise_or_bytes(SHA256.new(messages[2]).digest(), SHA256.new(messages[3]).digest())
        temp1 = bitwise_or_bytes(messages[4], messages[5])
        temp2 = bitwise_or_bytes(temp, temp1)
        handshake_messages = bitwise_or_bytes(handshake_messages, temp2)

        finish_label = b'SERVER'
        msg = bitwise_or_bytes(finish_label, SHA256.new(handshake_messages).digest())
        hmac = HMAC.new(master_secret, msg, digestmod=SHA256).digest()

        size = 32 + 3
        head = pack('BH', 0x86, size)

        body = head+hmac

        return body

    def build_client_finished(self, messages, master_secret):
        handshake_messages = bitwise_or_bytes(messages[0], messages[1])
        temp = bitwise_or_bytes(SHA256.new(messages[2]).digest(), SHA256.new(messages[3]).digest())
        temp1 = bitwise_or_bytes(messages[4], messages[5])
        temp2 = bitwise_or_bytes(temp, temp1)
        handshake_messages = bitwise_or_bytes(handshake_messages, temp2)

        finish_label = b'CLIENT'
        msg = bitwise_or_bytes(finish_label, SHA256.new(handshake_messages).digest())
        hmac = HMAC.new(master_secret, msg, digestmod=SHA256).digest()

        size = 32 + 3
        head = pack('BH', 0x87, size)

        body = head+hmac

        return body

    def build_error_message(self, error):
        size = 2 + 3
        error=int.to_bytes(error,1,'big')

        head = pack('BH', 0x88, size)
        body = head + error

        return body