import socketserver
import pickle
import simplecrypto
from uuid import uuid4

BUFSIZE = 8192
BLOCKSIZE = 4096

PACKET_TYPE_RECV_RESULT = 'recv_result'
PACKET_TYPE_DATA_FRAGMENT = 'data_fragment'
PACKET_TYPE_SVC_MESSAGE = 'svc_message'
PACKET_TYPE_NEW_CONNECTION = 'new_connection'

DATA_TYPE_FILE_CHUNK = 'file_chunk'
DATA_TYPE_CMD = 'cmd'
DATA_TYPE_SEARCH_QUERY = 'search_query'
DATA_TYPE_SEARCH_RESULT = 'search_result'
DATA_TYPE_PEER_PUBKEY = 'peer_pubkey'

SVC_MESSAGE_BAD_PACKET = 'bad_packet'
SVC_MESSAGE_DECRYPT_ERROR = 'decrypt_error'
SVC_MESSAGE_YOU_ARE_STRANGER = 'you_are_stranger'

RECV_OK_CODE = 0
RECV_ERROR_CODE = 1

HEADER_SVC_YOU_ARE_STRANGER = {'type': PACKET_TYPE_SVC_MESSAGE, 'msg': SVC_MESSAGE_YOU_ARE_STRANGER}

HEADER_RECV_OK = {'type': PACKET_TYPE_RECV_RESULT, 'code': RECV_OK_CODE}
HEADER_RECV_ERROR = {'type': PACKET_TYPE_RECV_RESULT, 'code': RECV_ERROR_CODE}

key = None

peers = {}

class Peer:
    def __init__(self, peer_addr, pubkey):
        self.addr = peer_addr
        self.pubkey = pubkey

def get_key():
    global key
    if not key:
        key = simplecrypto.RsaKeypair()
    return key

def pickle_data(data):
    return pickle.dumps(data, protocol=4)

def write_svc_msg(fobj, svc_msg, extra_data={}):
    MSG = {'type': PACKET_TYPE_SVC_MESSAGE, 'msg': svc_msg}
    if extra_data:
        MSG.update(extra_data)
    fobj.write(pickle_data(MSG))

class UDPRequestHandler(socketserver.DatagramRequestHandler):
    def handle(self):
        datagram = self.rfile.read(BUFSIZE)
        if self.client_address not in peers:
            write_svc_msg(self.wfile, SVC_MESSAGE_YOU_ARE_STRANGER)
            return
        try:
            unpickled_datagram = pickle.dumps(datagram)
        except pickle.UnpicklingError:
            write_svc_msg(self.wfile, SVC_MESSAGE_BAD_PACKET)
            return
        pk = get_key()
        try:
            data = pickle.loads(pk.decrypt_raw(unpickled_datagram['packet']))
        except:
            write_svc_msg(self.wfile, SVC_MESSAGE_DECRYPT_ERROR, {'packet_id': unpickled_datagram['packet_id']})
            return
        block_id = data['block_id']
        fragment_num = data['num']
        cmd = data['']