from donna25519 import *
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
import struct
import socket
import hashlib
import hmac
from math import ceil

O_NUM = 10 # number of initial One-Time Pre Keys

hash_len = 32
def hmac_sha256(key, data):
    return hmac.new(key, data, hashlib.sha256).digest()

def hkdf(length, ikm, salt = b""):
    """Parameters:
    length: output length in bytes
    ikm: input key material
    """
    prk = hmac_sha256(salt, ikm)
    t = b""
    okm = b""
    for i in range(int(ceil(1.0 * length / hash_len))):
        t = hmac_sha256(prk, t + bytes([1+i]))
        okm += t
    return okm[:length]

class ClientSession:
    def __init__(self, root_key, chain_key_s, chain_key_r, eph_own, eph_peer):
        self.root_key = root_key
        self.chain_key_s = chain_key_s
        self.chain_key_r = chain_key_r
        self.eph_own = eph_own
        self.eph_peer = eph_peer

    def update_keys(self, eph_peer):
        """Performs vertical ratcheting
        Updates root & chain keys to match peer's (for decrypting its messages)
        Then updates root & chain keys again with fresh ephemereal key (for encrypting
        my messages)
        """
        self.eph_peer = eph_peer
        # Update RootKey & receiving ChainKey
        # TODO 2 Obtain new root_key and chain_key_r from hkdf over
        # DH(eph_own, eph_peer) and root_key (can be used as salt parameter in hkdf)
        self.root_key = ""
        self.chain_key_r = ""

        # Update RootKey & sending ChainKey
        self.eph_own = None # TODO 2 Generate new ephemereal key pair
        # TODO 2 Obtain new root_key and chain_key_s from hkdf over
        # DH(eph_own, eph_peer) and root_key (can be used as salt parameter in hkdf)
        self.root_key = ""
        self.chain_key_s = ""

class Client:
    def __init__(self, server_ip, server_port):
        self.I = PrivateKey() # Identity Key Pair
        self.S = PrivateKey() # Signed Pre Key
        self.O_queue = [PrivateKey() for i in range(O_NUM)] # One-Time Pre Keys

        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.s.connect((server_ip, server_port))
        self.user_id = self.register()

        # self.existing_sessions[user_id] = session with that user (ClientSession)
        self.existing_sessions = {} # initially no existing session

    def register(self):
        """Lets the server know about its presence"""
        # TODO 1.1 send public keys of I, S and the list of O to the server
        # e.g. self.s.send(self.I.get_public().public)
        # You may also print a log message, e.g. print "Sent I = " + self.I.get_public().public.hex()

        # You can send the list of O by first sending its length and then the keys one by one
        # use struct.pack('!i', my_int) to send 4 byte signed integers

        # receive your user id from the server (might need later)
        # use struct.unpack('!i', received_int)[0] to get an integer
        # e.g. struct.unpack('!i', self.s.recv(4))[0]

        user_id = 0 # TODO 1.1 change with received user id
        print("Got User ID = %d" % user_id)
        return user_id

    def get_send_message_key(self, peer_user_id):
        """Gets sending MessageKey using ChainKey from peer_user_id's ClientSession"""
        session = self.existing_sessions[peer_user_id]

        # TODO 2 Generate message_key from chain_key_s (session.chain_key_s)
        # message_key = HMAC_SHA256(ChainKey, 0x01)
        message_key = None

        # TODO 2 Update chain_key_s
        # ChainKey = HMAC_SHA256(ChainKey, 0x02)
        session.chain_key_s = None

        return message_key

    def get_recv_message_key(self, peer_user_id):
        """Gets MessageKey using ChainKey from peer_user_id's ClientSession"""
        session = self.existing_sessions[peer_user_id]

        # TODO 2 Generate message_key from chain_key_r (session.chain_key_r)
        # message_key = HMAC_SHA256(ChainKey, 0x01)
        message_key = None

        # TODO 2 Update chain_key_r
        # ChainKey = HMAC_SHA256(ChainKey, 0x02)
        session.chain_key = None

        return message_key

    def pad_message(self, message):
        pad_number = (-len(message)) % 16
        if pad_number == 0:
            pad_number += 16

        message += ''.join([chr(pad_number) for i in xrange(pad_number)])
        return message

    def unpad_message(self, message):
        pad_number = ord(message[-1])
        message = message[:-pad_number]
        return message

    def send_message(self, to_user_id, message):
        """Sends a message to user with id <to_user_id>
        If there is no existing ClientSession between them, they create one
        """
        if to_user_id not in self.existing_sessions:
            self.setup_session(to_user_id)

        self.s.send(b"SEND")
        self.s.send(struct.pack('!i', to_user_id))

        # Get MessageKey for next message
        message_key = self.get_send_message_key(to_user_id)

        # Encrypt message using MessageKey
        iv = os.urandom(16)
        cipher = Cipher(algorithms.AES(message_key), modes.CBC(iv))
        encryptor = cipher.encryptor()
        message = self.pad_message(message)
        enc_message = encryptor.update(message) + encryptor.finalize()

        # TODO 2 Send own ephemereal public key (get it from the ClientSession with to_user_id)
        session = self.existing_sessions[to_user_id]
        raw_eph_own = session.eph_own.get_public().public

        # TODO 2 Send encrypted message length to server

        # TODO 2 Send encrypted message to server

    def recv_message(self):
        """Receives a message
        It can be either from a new session exchange (tag=NEWS), or an actual message(tag=MESG)"""
        tag = self.s.recv(4)
        if tag == "NEWS": # NEW Session
            ini_user_id = struct.unpack('!i', self.s.recv(4))[0]
            self.setup_session_rec(ini_user_id)
            tag = self.s.recv(4)

        if tag == "MESG": # message incoming
            # Get sender user id
            sender_user_id = struct.unpack('!i', self.s.recv(4))[0]

            # Get ephemereal key
            raw_eph_peer = self.s.recv(32)


            session = self.existing_sessions[sender_user_id]
            # TOD2 3 If received eph_peer is unknown (different than the one stored in the
            # ClientSession), update the keys
            # The first time current eph_peer will be None, so make sure you update in that
            # case also
            if False: # change with actual condition
                eph_peer = PublicKey(raw_eph_peer)
                session.update_keys(eph_peer)

            # Get message length
            mlen = struct.unpack('!i', self.s.recv(4))[0]

            # Get message
            enc_message = self.s.recv(mlen)

            # Decrypt message using MessageKey
            message_key = self.get_recv_message_key(sender_user_id)
            iv = enc_message[:16]
            enc_message = enc_message[16:]
            cipher = Cipher(algorithms.AES(message_key), modes.CBC(iv))
            decryptor = cipher.decryptor()
            msg = decryptor.update(enc_message) + decryptor.finalize()
            msg = self.unpad_message(msg)

            return msg

    def record_new_client_session(self, master_secret, client_id, eph_peer = None):
        """Forges RootKey and ChainKey from master_secret and stores them in a ClientSession
        which is then put in the existing_sessions dictionary"""

        root_key = master_secret # master_secret is considered as the base root key
        if eph_peer: # initiator considers Srec as initial peer ephemereal key
            eph_own = PrivateKey() # Generate new send ephemereal key

            # TODO 2 Get root_key and chain_key_s with hkdf over DH(eph_own, eph_peer) and
            # master secret
            root_key = ""
            chain_key_s = ""
            chain_key_r = None # Will initialize with first received message
        else: # the non-initiator (first receiver) will enter this branch
            eph_own = self.S # initiator considers receiver's S as first "ephemereal" key
            chain_key_s = None
            chain_key_r = None

        new_client_session = ClientSession(root_key, chain_key_s, chain_key_r, eph_own, eph_peer)
        self.existing_sessions[client_id] = new_client_session

    def setup_session(self, to_user_id):
        """Initiates a new session with to_user_id"""
        self.s.send(b"GENS") # GENerate Session command
        self.s.send(struct.pack('!i', to_user_id))
        resp = self.s.recv(4)
        if resp == "FAIL":
            print("User does not exist")
            return False

        # TODO 1.1 Get Irec (done), Srec, Orec from server
        Irec = PublicKey(self.s.recv(32))
        # Srec = ...
        # Prec = ...

        # TODO 1.1 Generate Eini (ephemeral private key of initiator)
        # Eini = ... (see above for function to retrieve ephemeral private key)

        # TODO 1.1 Compute master_secret from DH(Iini, Srec), DH(Eini, Irec), DH(Eini, Srec), DH(Eini, Orec)
        # See method do_exchange in donna library: https://github.com/Muterra/donna25519
        # Use '+' to concatenate output of various DH exchanges
        master_secret = ""
        print(master_secret.hex())

        # TODO 1.2 Send Eini to the other client (through server)

        # Generate new ClientSession using master_secret
        self.record_new_client_session(master_secret, to_user_id, Srec)

    def setup_session_rec(self, from_user_id):
        """Receives a new session with from_user_id"""

        # TODO 1.2 Receive initiator's public keys Eini, Iini and Orec_pub
        Eini = PublicKey(self.s.recv(32))
        # Iini = ...
        # Orec_pub = ...

        # TODO 1.2 Get private Orec from my O_queue based on received public Orec and then remove from O_queue

        # TODO 1.2 Compute master_secret from DH(Iini, Srec), DH(Eini, Irec), DH(Eini, Srec), DH(Eini, Orec)
        # See above, in method setup_session
        master_secret = ""
        print(master_secret.hex())

        # Generate new ClientSession using master_secret
        self.record_new_client_session(master_secret, from_user_id)