import hashlib
import lib.serpent
from lib.serpent import Serpent
from lib.serpent import serpent_cbc_encrypt 
from lib.serpent import serpent_cbc_decrypt 
from commands import BROADCAST
from commands import DIRECT 
from commands import IGNORE 
from lib.message import Message
import base64
import binascii
import time
import struct
import sys
import hmac
import random
import os 
import pprint
import logging
pp = pprint.PrettyPrinter(indent=4)

PACKET_SIZE = 496 
MAX_SPEAKER_SIZE = 32
TS_ACCEPTABLE_SKEW = 60 * 15
BLACK_PACKET_FORMAT = "<448s48s"
RED_PACKET_FORMAT = "<16sBBxB428s"
RED_PACKET_LENGTH_WITH_PADDING = 448
MESSAGE_PACKET_FORMAT = "<q32s32s32s324s"
MAX_MESSAGE_LENGTH = 428
MAX_BOUNCES = 3
STALE_PACKET = 0
DUPLICATE_PACKET = 1
MALFORMED_PACKET = 2
INVALID_SIGNATURE = 3
IGNORED = 4

class Infosec(object): 
    def __init__(self, state=None):
        self.state = state

    def message(self, message):
        # if we are not rebroadcasting we need to set the timestamp

        if message.timestamp == None:
            message.original = True
            message.timestamp = int(time.time())
        else:
            message.original = False

        target_peer = (self.state.get_peer_by_handle(message.handle)
                        if message.command == DIRECT
                        else None)

        if target_peer and not target_peer.get_key():
            logging.debug("No key for peer associated with %s" % message.handle)
            return

        if message.command == DIRECT and target_peer == None:
            logging.debug("Aborting message: unknown handle: %s" % message.handle)
            return

        message_bytes = self.get_message_bytes(message, target_peer)
        if message.command != IGNORE:
            message_hash = binascii.hexlify(hashlib.sha256(message_bytes).digest())
            logging.debug("generated message_hash: %s" % message_hash)
            self.state.add_to_dedup_queue(message_hash)
            self.state.log(message.speaker, message_bytes, target_peer)

        if message.command == DIRECT:
            signed_packet_bytes = self.pack(target_peer, message, message_bytes)
            target_peer.send(signed_packet_bytes)
        elif message.command == BROADCAST or message.command == IGNORE:
            for peer in self.state.get_keyed_peers():
                
                # we don't want to send a broadcast back to the originator

                if message.peer and (peer.peer_id == message.peer.peer_id):
                    next

                signed_packet_bytes = self.pack(peer, message, message_bytes)
                peer.send(signed_packet_bytes)
        else:
            pass

    def get_message_bytes(self, message, peer=None):
        timestamp = message.timestamp
        command = message.command
        speaker = self._pad(message.speaker, MAX_SPEAKER_SIZE)

        # let's generate the self_chain value from the last message or set it to zero if 
        # there this is the first message
        
        if message.original:
            if command == DIRECT:
                self_chain = self.state.get_last_message_hash(message.speaker, peer.peer_id)
            elif command == BROADCAST:
                self_chain = self.state.get_last_message_hash(message.speaker)
            elif command == IGNORE:
                self_chain = "\x00" * 32
            net_chain = "\x00" * 32
        else:
            self_chain = message.self_chain 
            net_chain = message.net_chain

        # pack message bytes

        if message.command != IGNORE:
            logging.debug("packing message bytes: %s" % message.body)
        else:
            logging.debug("packing rubbish message bytes: %s" % binascii.hexlify(message.body))

        message_bytes = struct.pack(MESSAGE_PACKET_FORMAT, message.timestamp, self_chain, net_chain, speaker, message.body)
        return message_bytes

    def pack(self, peer, message, message_bytes):
        key_bytes = base64.b64decode(peer.get_key())
        signing_key = key_bytes[:32]
        cipher_key = key_bytes[32:]

        # pack packet bytes

        nonce = self._generate_nonce(16)
        bounces = message.bounces 
        version = 0xfe
        red_packet_bytes = struct.pack(RED_PACKET_FORMAT, nonce, bounces, version, message.command, self._pad(message_bytes, MAX_MESSAGE_LENGTH))

        # encrypt packet

        serpent = Serpent(cipher_key)
        black_packet_bytes = serpent_cbc_encrypt(cipher_key, red_packet_bytes)

        # sign packet

        signature_bytes = hmac.new(signing_key, black_packet_bytes, hashlib.sha384).digest()

        # pack the signed black packet

        signed_packet_bytes = struct.pack(BLACK_PACKET_FORMAT, black_packet_bytes, signature_bytes)

        return signed_packet_bytes
    
    def unpack(self, peer, black_packet):
        # unpack the black packet

        key_bytes = base64.b64decode(peer.get_key())
        signing_key = key_bytes[:32]
        cipher_key = key_bytes[32:]

        try:
            black_packet_bytes, signature_bytes = struct.unpack(BLACK_PACKET_FORMAT, black_packet)
        except:
            logging.error("Discarding malformed black packet from %s" % peer.get_key())
            return Message({ "error_code": MALFORMED_PACKET })

        # check signature

        signature_check_bytes = hmac.new(signing_key, black_packet_bytes, hashlib.sha384).digest()

        if(signature_check_bytes != signature_bytes):
            return Message({ "error_code": INVALID_SIGNATURE })

        # try to decrypt black packet
        
        serpent = Serpent(cipher_key)
        red_packet_bytes = serpent_cbc_decrypt(cipher_key, black_packet_bytes)

        # unpack red packet

        nonce, bounces, version, command, message_bytes = struct.unpack(RED_PACKET_FORMAT, red_packet_bytes)

        # compute message_hash

        message_hash = binascii.hexlify(hashlib.sha256(message_bytes).digest())

        # unpack message

        int_ts, self_chain, net_chain, speaker, body = struct.unpack(MESSAGE_PACKET_FORMAT, message_bytes)
        
        # remove padding from speaker

        for index, byte in enumerate(speaker):
            if byte == '\x00':
                speaker = speaker[0:index]
                break 

        # remove padding from body

        for index, byte in enumerate(body):
            if byte == '\x00':
                body = body[0:index]
                break 

        # nothing to be done for an IGNORE command

        if command == IGNORE:
            return Message({"speaker": speaker, "error_code": IGNORED})

        # check timestamp

        if(int_ts not in self._ts_range()):
            return Message({ "error_code": STALE_PACKET })

        # check self_chain

        if command == DIRECT:
            self_chain_check = self.state.get_last_message_hash(speaker, peer.peer_id)
        elif command == BROADCAST:
            self_chain_check = self.state.get_last_message_hash(speaker)

        self_chain_valid = (self_chain_check == self_chain)

        # log this message for use in the self_chain check

        self.state.log(speaker, message_bytes, peer if (command == DIRECT) else None)

        # build message object

        message = Message({ 
                        "peer": peer,
                        "body": body.rstrip(),
                        "timestamp": int_ts,
                        "command": command,
                        "speaker": speaker,
                        "bounces": bounces,
                        "self_chain": self_chain,
                        "net_chain": net_chain,
                        "self_chain_valid": self_chain_valid,
                        "message_hash": message_hash
                        })

        # check for duplicates

        if(self.state.is_duplicate_message(message_hash)):
            message.error_code = DUPLICATE_PACKET
            return message

        return message

    def _pad(self, text, size):
        return text.ljust(size, "\x00")

    def _ts_range(self):
        current_ts = int(time.time())
        return range(current_ts - TS_ACCEPTABLE_SKEW, current_ts + TS_ACCEPTABLE_SKEW)

    def _generate_nonce(self, length=8):
        """Generate pseudorandom number."""
        return ''.join([str(random.randint(0, 9)) for i in range(length)])

    def gen_rubbish_body(self):
        return os.urandom(MAX_MESSAGE_LENGTH)