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 
import base64
import binascii
import time
import struct
import sys
import hmac
import random
import pprint
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

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

    def pack(self, peer, message):
        try:
            timestamp = message["timestamp"]
        except:
            timestamp = None
        command = message["command"]
        speaker = self._pad(message["speaker"], MAX_SPEAKER_SIZE)

        # if we are rebroadcasting we need to use the original timestamp

        if(timestamp == None):
            int_ts = int(time.time())
        else:
            int_ts = timestamp

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

        # 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.server.state.get_last_message_hash(message["speaker"], peer.peer_id)
            elif command == BROADCAST:
                self_chain = self.server.state.get_last_message_hash(message["speaker"])
            net_chain = "\x00" * 32
        else:
            self_chain = message["self_chain"] 
            net_chain = message["net_chain"]

        # pack message bytes

        message_bytes = struct.pack(MESSAGE_PACKET_FORMAT, int_ts, self_chain, net_chain, speaker, message["body"].encode("ascii"))

        # log messages 

        if message["original"]:
            if command == DIRECT:
                self.server.state.log(message["speaker"], message_bytes, peer.peer_id)
            elif command == BROADCAST:
                self.server.state.log(message["speaker"], message_bytes)

        # pack packet bytes

        nonce = self._generate_nonce(16)
        bounces = message["bounces"] 
        version = 0xfe
        red_packet_bytes = struct.pack(RED_PACKET_FORMAT, nonce, bounces, version, 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)

        # we want to ignore this ts if it is sent back to us

        self.server.recent.insert(int_ts)
        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:
            self.server.print_error("Discarding malformed black packet from %s" % peer.get_key())
            return { "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 { "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)

        # unpack message

        int_ts, self_chain, net_chain, speaker, message = struct.unpack(MESSAGE_PACKET_FORMAT, message_bytes)
        speaker = speaker.strip()
        
        # check timestamp

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

        if(self.server.recent.has(int_ts)):
            return { "error_code": DUPLICATE_PACKET } 
        else:
            self.server.recent.insert(int_ts)

        # check self_chain

        if command == DIRECT:
            self_chain_check = self.server.state.get_last_message_hash(speaker, peer.peer_id)
        elif command == BROADCAST:
            self_chain_check = self.server.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.server.state.log(speaker, message_bytes, peer.peer_id if (command == DIRECT) else None)

        # remove padding from message bytes

        for index, byte in enumerate(message):
            if binascii.hexlify(byte) == "00":
                unpadded_message = message[0:index]
                break 

        # return the message, timestamp, and command (command replaces propagate)

        return { "body": unpadded_message.rstrip(),
                 "timestamp": int_ts,
                 "command": command,
                 "speaker": speaker,
                 "bounces": bounces,
                 "self_chain": self_chain,
                 "net_chain": net_chain,
                 "self_chain_valid": self_chain_valid,
                 "error_code": None}

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

    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)])