#!/usr/bin/python3

import binascii
import select
import socket
import struct
import nacl.bindings
import nacl.utils

import util

from nacl.public import PrivateKey, PublicKey, Box

from exceptions import *

class SerializationError(ToxException):
    pass
class ParseError(ToxException):
    pass

def h(b):
    return binascii.hexlify(b)

class TcpPacket:
    """
    Holds generic packet to be sent via TCP
    """
    # Class variable holding callbacks: def callback(packet) will be called
    # if this packet is received
    listeners = {}

    def __init__(self):
        self.bytes = b""

    def send(self, socket=None):
        """
        Serialize and send packet via socket
        """
        serialized_b = self.serialize()
        print(binascii.hexlify(serialized_b))
        if socket is None:
            socket = self.relay.get_socket()
        socket.send(serialized_b)
        print ("sent %d bytes" % len(serialized_b))

    def encrypt(self, data):
        """
        Identity function - to be overridden
        """
        return data

    def serialize(self):
        """
        Return wire-ready serialized packet
        """
        serialized_b = self.encrypt(self.bytes)
        if len(serialized_b) >= 2**16:
            raise SerializationError('TCP packet too long: %d' % len(serialized_b))

        length_b = struct.pack(">H", len(serialized_b))
        return length_b + serialized_b

    def notify_listeners(self):
        if self.__class__ in self.listeners:
            for listener in self.listeners[self.__class__]:
                listener(self)

    @classmethod
    def register_listener(cls, listener):
        if cls not in TcpPacket.listeners:
            TcpPacket.listeners[cls] = []
        TcpPacket.listeners[cls].append(listener)
        print (cls, listener)


class EncryptedPacket(TcpPacket):
    def __init__(self, relay):
        """
        Relay needed for decrypting the incoming packet
        """
        self.relay = relay

    def encrypt(self, payload):
        """
        Encrypt with outgoing_public_key, sign with relay_private_key
        """
        relay = self.relay
        combined = nacl.bindings.crypto_box_beforenm(bytes(relay.get_outgoing_public_key()), bytes(relay.get_private_key()))
        nonce = relay.get_outgoing_nonce()
        print ("combined key", binascii.hexlify(combined))
        print ("outgoing nonce", h(nonce))
        encrypted = nacl.bindings.crypto_box_afternm(payload, nonce, combined)
        return encrypted

class TcpPingPacket(EncryptedPacket):
    """
    Ping packet (0x04)
    """
    packet_id = 0x04
    def __init__(self, relay=None):
        super().__init__(relay)
        self.bytes = b"\x04puretoxc"

    @classmethod
    def from_bytes(cls, packet_bytes):
        self = cls()
        assert (packet_bytes[0] == self.packet_id), "First byte not equal to packet id"
        self.bytes = packet_bytes
        self.ping_message = packet_bytes[1:]

        return self

class TcpPongPacket(EncryptedPacket):
    """
    Pong packet (0x05)
    """
    packet_id = 0x05
    def __init__(self, relay=None):
        super().__init__(relay)

    @classmethod
    def from_bytes(cls, packet_bytes):
        self = TcpPongPacket()
        assert (packet_bytes[0] == self.packet_id), "First byte does not equal to packet id"
        self.pong_message = packet_bytes[1:]

        self.bytes = packet_bytes

        return self

    def __repr__(self):
        return "<TcpPongPacket(%s)>" % self.pong_message

class TcpRoutingRequestPacket(EncryptedPacket):
    """
    Routing request (0x00)
    """
    packet_id = 0x00
    def __init__(self, relay=None):
        super().__init__(relay)
        

class TcpClientHelloPacket(TcpPacket):
    """
    The TCP packet sent just after opening the connection
    """
    def __init__(self, relay, session, identity):
        # The nonce used in hello packet only
        self.nonce = bytes(session.get_nonce())

        self.bytes = b""
        self.bytes += bytes(session.get_public_key())
        self.bytes += self.nonce
        self.bytes += self.get_payload(relay, session, identity)

        self.relay = relay

    def get_payload(self, relay, session, identity):
        payload = bytes(relay.get_incoming_public_key())
        payload += bytes(relay.get_outgoing_nonce_without_incrementing())

        combined = nacl.bindings.crypto_box_beforenm(bytes(relay.get_dht_public_key()), bytes(session.get_private_key()))
        print ("combined key", binascii.hexlify(combined))
        print ("relay incoming nonce", h(relay.get_outgoing_nonce_without_incrementing()))
        print ("self nonce", h(self.nonce))
        encrypted = nacl.bindings.crypto_box_afternm(payload, self.nonce, combined)
        print(encrypted)

        return bytes(encrypted)

    def serialize(self):
        """
        The hello packet doesn't add length header
        """
        return self.bytes

class TcpServerHelloPacket(TcpPacket):
    """
    TCP packet received in response to TcpClientHelloPacket
    """
    def parse(self, blob):
        if len(blob) != 24+72:
            raise ParseError("Packet length %d instead of %d" % (len(blob), 24+72))

        self.nonce = blob[:24]
        self.encrypted_payload = blob[24:]

    def decrypt_payload(self, relay, session):
        combined = nacl.bindings.crypto_box_beforenm(bytes(relay.get_dht_public_key()), bytes(session.get_private_key()))
        plaintext = nacl.bindings.crypto_box_open_afternm(self.encrypted_payload, self.nonce, combined)
        if len(plaintext) != 32+24:
            raise ParseError("Plaintext length %d instead of %d" % (len(plaintext), 32+24))
        self.relay_public_key = plaintext[:32]
        self.base_nonce = plaintext[32:32+24]
        


class TcpRelay:
    """
    Represents a connection to TCP relay
    """
    def __init__(self, hostname, port, public_key):
        """
        public_key is a hex-encoded string containing relay's public DHT key
        """
        self.hostname = hostname
        self.port = int(port)
        self.dht_public_key = PublicKey(binascii.unhexlify(public_key))
        self.tcp_incoming_nonce = None
        # temporary key that will be used for encryption during the connection and will be discarded after
        self.private_key = PrivateKey.generate()

        self.tcp_outgoing_nonce = nacl.utils.random(Box.NONCE_SIZE)
        # a temporary public key tied to this connection
        self.outgoing_public_key = None

        # The socked used to connect to this relay
        self._socket = None

    def connect(self):
        """
        Establish a TCP connection
        """
        self._socket = socket.socket()
        self._socket.connect((self.hostname, self.port))
        return self._socket

    def get_socket(self):
        """
        Return a socket to TCP relay, connect if necessary
        """
        if self._socket is not None:
            return self._socket
        return self.connect()

    def get_client_hello_packet(self, session, identity):
        return TcpClientHelloPacket(self, session, identity)

    def get_incoming_nonce(self):
        """
        Return nonce for decrypting incoming packets (TCP relay -> me), incrementing the nonce
        """
        nonce = self.tcp_incoming_nonce
        self.increment_incoming_nonce()
        return nonce

    def increment_incoming_nonce(self):
        self.tcp_incoming_nonce = util.increment_nonce(self.tcp_incoming_nonce)

    def get_incoming_public_key(self):
        return self.private_key.public_key

    def get_private_key(self):
        return self.private_key

    def get_dht_public_key(self):
        """
        Return relay's DHT public key
        """
        return self.dht_public_key

    def set_incoming_nonce(self, nonce):
        self.tcp_incoming_nonce = nonce

    def set_outgoing_public_key(self, key):
        """
        Set the ephemeral relay's public key with which this connection is encrypted
        """
        self.tcp_outgoing_public_key = PublicKey(key)

    def get_outgoing_nonce(self):
        """
        Return the nonce used for encrypting outgoing packets (me -> TCP relay)
        Increments the nonce.
        """
        rv = self.tcp_outgoing_nonce
        self.increment_outgoing_nonce()
        return rv

    def get_outgoing_nonce_without_incrementing(self):
        return self.tcp_outgoing_nonce

    def increment_outgoing_nonce(self):
        self.tcp_outgoing_nonce = util.increment_nonce(self.tcp_outgoing_nonce)

    def get_outgoing_public_key(self):
        return self.tcp_outgoing_public_key

    def negotiate_forward_secret_keys(self, session, identity):
        """
        Given our session (DHT key) and identity (long term key), send
        hello packets to the TCP relay and agree on nonce and ephemeral
        keys.
        This is a high-level function.
        """
        hello_packet = self.get_client_hello_packet(session, identity)
        hello_packet.send()

        sock = self.get_socket()
        received = sock.recv(4096)

        packet = TcpServerHelloPacket()
        packet.parse(received)
        packet.decrypt_payload(self, session)
        print ("relay key", h(packet.relay_public_key))
        print ("base nonce", h(packet.base_nonce))
        self.set_incoming_nonce(packet.base_nonce)
        self.set_outgoing_public_key(packet.relay_public_key)
        self._socket.setblocking(0)

class IncomingStreamHandler:
    """
    Split incoming TCP stream into packets
    """
    def __init__(self, relay):
        """
        sock - tcp socket to read from
        relay - TcpRelay instance needed for decryption
        """
        self.relay = relay
        self.socket = relay.get_socket()
        self.buf = b""
        self.packet_types = {}
        self.build_packet_types_table()

    def tick(self):
        """
        To be called in a loop. Reads bytes from the stream, parses them
        """
        ready_to_read, _, _ = select.select([self.socket], [], [], 0)
        if self.socket not in ready_to_read:
            return

        self.buf += self.socket.recv(4096)

        packets = []
        while True:
            packet_ciphertext = self.try_build_packet()
            if not packet_ciphertext:
                break

            packet_cleartext = self.decrypt(packet_ciphertext)
            packet = self.parse_packet(packet_cleartext)
            if packet:
                packets.append(packet)

        return packets
            
    def decrypt(self, ciphertext):
        combined = nacl.bindings.crypto_box_beforenm(bytes(self.relay.get_outgoing_public_key()), bytes(self.relay.get_private_key()))
        nonce = self.relay.get_incoming_nonce()
        print ("IncomingStreamHandler.decrypt:")
        print ("combined key", binascii.hexlify(combined))
        print ("incoming nonce", h(nonce))
        plaintext = nacl.bindings.crypto_box_open_afternm(ciphertext, nonce, combined)
        
        return plaintext

    def try_build_packet(self):
        length = int.from_bytes(self.buf[:2], byteorder='big')
        try:
            packet = self.buf[2:2+length]
            self.buf = self.buf[2+length:]
            return packet
        except IndexError:
            return None

    def build_packet_types_table(self):
        for cls in EncryptedPacket.__subclasses__():
            self.packet_types[cls.packet_id] = cls

        print (self.packet_types)

    def parse_packet(self, packet_bytes):
        print (packet_bytes)
        packet_id = packet_bytes[0]
        if packet_id in self.packet_types:
            packet = self.packet_types[packet_id].from_bytes(packet_bytes)
            packet.notify_listeners()

            return packet
        else:
            print("Unknown packet type %d" % packet_id)