#!/usr/bin/env python

# encoding: utf-8

"""

BLIP.py

Created by Jens Alfke on 2008-06-03.

Copyright (c) 2008 Jens Alfke. All rights reserved.

"""

import asynchat

import asyncore

from cStringIO import StringIO

import logging

import socket

import struct

import sys

import traceback

import unittest

import zlib

# INTERNAL CONSTANTS -- NO TOUCHIES!

kFrameMagicNumber   = 0x9B34F205

kFrameHeaderFormat  = '!LLHH'

kFrameHeaderSize    = 12

kMsgFlag_TypeMask   = 0x000F

kMsgFlag_Compressed = 0x0010

kMsgFlag_Urgent     = 0x0020

kMsgFlag_NoReply    = 0x0040

kMsgFlag_MoreComing = 0x0080

kMsgType_Request    = 0

kMsgType_Response   = 1

kMsgType_Error      = 2

log = logging.getLogger('BLIP')

log.propagate = True

class MessageException(Exception):

    pass

class ConnectionException(Exception):

    pass

class Listener (asyncore.dispatcher):

    "BLIP listener/server class"

    def __init__(self, port):

        "Create a listener on a port"

        asyncore.dispatcher.__init__(self)

        self.onConnected = self.onRequest = None

        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)

        self.bind( ('',port) )

        self.listen(5)

        log.info("Listening on port %u", port)

    def handle_accept( self ):

        client,address = self.accept()

        conn = Connection(address,client)

        conn.onRequest = self.onRequest

        if self.onConnected:

            self.onConnected(conn)

class Connection (asynchat.async_chat):

    def __init__( self, address, conn=None ):

        "Opens a connection with the given address. If a connection/socket object is provided it'll use that,"

        "otherwise it'll open a new outgoing socket."

        asynchat.async_chat.__init__(self,conn)

        self.address = address

        if conn:

            log.info("Accepted connection from %s",address)

        else:

            log.info("Opening connection to %s",address)

            self.create_socket(socket.AF_INET, socket.SOCK_STREAM)

            self.connect(address)

        self.onRequest = None

        self.pendingRequests = {}

        self.pendingResponses = {}

        self.outBox = []

        self.inMessage = None

        self.inNumRequests = 0

        self._endOfFrame()

    #def handle_error(self,x):

    #    log.error("Uncaught exception: %s",x)

    #    self.close()

    def _fatal(self, error):

        log.error("Fatal BLIP connection error: %s",error)

        self.close()

    ### SENDING:

    def _outQueueMessage(self, msg,isNew=True):

        n = len(self.outBox)

        index = n

        if msg.urgent and n>1:

            while index > 0:

                otherMsg = self.outBox[index-1]

                if otherMsg.urgent:

                    if index<n:

                        index += 1

                    break

                elif isNew and otherMsg._bytesWritten==0:

                    break

                index -= 1

            else:

                index = 1

        self.outBox.insert(index,msg)

        if isNew:

            log.info("Queuing outgoing message at index %i",index)

            if n==0:

                self._sendNextFrame()

        else:

            log.debug("Re-queueing outgoing message at index %i of %i",index,len(self.outBox))

    def _sendNextFrame(self):

        while self.outBox:              #FIX: Don't send everything at once; only as space becomes available!

            n = len(self.outBox)

            if n > 0:

                msg = self.outBox.pop(0)

                frameSize = 4096

                if msg.urgent or n==1 or not self.outBox[0].urgent:

                    frameSize *= 4

                if msg._sendNextFrame(self,frameSize):

                    self._outQueueMessage(msg,isNew=False)

                else:

                    log.info("Finished sending %s",msg)

    ### RECEIVING:

    def collect_incoming_data(self, data):

        if self.expectingHeader:

            if self.inHeader==None:

                self.inHeader = data

            else:

                self.inHeader += data

        else:

            self.inMessage._receivedData(data)

    def found_terminator(self):

        if self.expectingHeader:

            # Got a header:

            (magic, requestNo, flags, frameLen) = struct.unpack(kFrameHeaderFormat,self.inHeader)

            self.inHeader = None

            if magic!=kFrameMagicNumber: self._fatal("Incorrect frame magic number %x" %magic)

            if frameLen < kFrameHeaderSize: self._fatal("Invalid frame length %u" %frameLen)

            frameLen -= kFrameHeaderSize

            log.debug("Incoming frame: type=%i, number=%i, flags=%x, length=%i",

                        (flags&kMsgFlag_TypeMask),requestNo,flags,frameLen)

            self.inMessage = self._inMessageForFrame(requestNo,flags)

            if frameLen > 0:

                self.expectingHeader = False

                self.set_terminator(frameLen)

            else:

                self._endOfFrame()

        else:

            # Got the frame's payload:

            self._endOfFrame()

    def _inMessageForFrame(self, requestNo,flags):

        message = None

        msgType = flags & kMsgFlag_TypeMask

        if msgType==kMsgType_Request:

            message = self.pendingRequests.get(requestNo)

            if message==None and requestNo == self.inNumRequests+1:

                message = IncomingRequest(self,requestNo,flags)

                assert message!=None

                self.pendingRequests[requestNo] = message

                self.inNumRequests += 1

        elif msgType==kMsgType_Response or msgType==kMsgType_Error:

            message = self.pendingResponses.get(requestNo)

        if message != None:

            message._beginFrame(flags)

        else:

            log.warning("Ignoring unexpected frame with type %u, request #%u", msgType,requestNo)

        return message

    def _endOfFrame(self):

        msg = self.inMessage

        self.inMessage = None

        self.expectingHeader = True

        self.inHeader = None

        self.set_terminator(kFrameHeaderSize) # wait for binary header

        if msg:

            log.debug("End of frame of %s",msg)

            if not msg.moreComing:

                self._receivedMessage(msg)

    def _receivedMessage(self, msg):

        log.info("Received: %s",msg)

        # Remove from pending:

        if msg.isResponse:

            del self.pendingReplies[msg.requestNo]

        else:

            del self.pendingRequests[msg.requestNo]

        # Decode:

        try:

            msg._finished()

            if not msg.isResponse:

                self.onRequest(msg)

        except Exception, x:

            log.error("Exception handling incoming message: %s", traceback.format_exc())

            #FIX: Send an error reply

### MESSAGES:

class Message (object):

    "Abstract superclass of all request/response objects"

    def __init__(self, connection, properties=None, body=None):

        self.connection = connection

        self.properties = properties or {}

        self.body = body

    @property

    def flags(self):

        if self.isResponse:

            flags = kMsgType_Response

        else:

            flags = kMsgType_Request

        if self.urgent:     flags |= kMsgFlag_Urgent

        if self.compressed: flags |= kMsgFlag_Compressed

        if self.noReply:    flags |= kMsgFlag_NoReply

        if self.moreComing: flags |= kMsgFlag_MoreComing

        return flags

    def __str__(self):

        s = "%s[#%i" %(type(self).__name__,self.requestNo)

        if self.urgent:     s += " URG"

        if self.compressed: s += " CMP"

        if self.noReply:    s += " NOR"

        if self.moreComing: s += " MOR"

        if self.body:       s += " %i bytes" %len(self.body)

        return s+"]"

    def __repr__(self):

        s = str(self)

        if len(self.properties): s += repr(self.properties)

        return s

    @property 

    def isResponse(self):

        "Is this message a response?"

        return False

    @property 

    def contentType(self):

        return self.properties.get('Content-Type')

    def __getitem__(self, key):     return self.properties.get(key)

    def __contains__(self, key):    return key in self.properties

    def __len__(self):              return len(self.properties)

    def __nonzero__(self):          return True

    def __iter__(self):             return self.properties.__iter__()

class IncomingMessage (Message):

    "Abstract superclass of incoming messages."

    def __init__(self, connection, requestNo, flags):

        super(IncomingMessage,self).__init__(connection)

        self.requestNo  = requestNo

        self.urgent     = (flags & kMsgFlag_Urgent) != 0

        self.compressed = (flags & kMsgFlag_Compressed) != 0

        self.noReply    = (flags & kMsgFlag_NoReply) != 0

        self.moreComing = (flags & kMsgFlag_MoreComing) != 0

        self.frames     = []

    def _beginFrame(self, flags):

        if (flags & kMsgFlag_MoreComing)==0:

            self.moreComing = False

    def _receivedData(self, data):

        self.frames.append(data)

    def _finished(self):

        encoded = "".join(self.frames)

        self.frames = None

        # Decode the properties:

        if len(encoded) < 2: raise MessageException, "missing properties length"

        propSize = 2 + struct.unpack('!H',encoded[0:2])[0]

        if propSize>len(encoded): raise MessageException, "properties too long to fit"

        if propSize>2 and encoded[propSize-1] != '\000': raise MessageException, "properties are not nul-terminated"

        proplist = encoded[2:propSize-1].split('\000')

        encoded = encoded[propSize:]

        if len(proplist) & 1: raise MessageException, "odd number of property strings"

        for i in xrange(0,len(proplist),2):

            def expand(str):

                if len(str)==1:

                    str = IncomingMessage.__expandDict.get(str,str)

                return str

            self.properties[ expand(proplist[i])] = expand(proplist[i+1])

        # Decode the body:

        if self.compressed and len(encoded)>0:

            try:

                encoded = zlib.decompress(encoded,31)   # window size of 31 needed for gzip format

            except zlib.error:

                raise MessageException, sys.exc_info()[1]

        self.body = encoded

    __expandDict= {'\x01' : "Content-Type",

                   '\x02' : "Profile",

                   '\x03' : "application/octet-stream",

                   '\x04' : "text/plain; charset=UTF-8",

                   '\x05' : "text/xml",

                   '\x06' : "text/yaml",

                   '\x07' : "Channel",

                   '\x08' : "Error-Code",

                   '\x09' : "Error-Domain"}

class OutgoingMessage (Message):

    "Abstract superclass of outgoing requests/responses."

    def __init__(self, connection, properties=None, body=None):

        Message.__init__(self,connection,properties,body)

        self.urgent = self.compressed = self.noReply = False

        self.moreComing = True

    def __setitem__(self, key,val):

        self.properties[key] = val

    def __delitem__(self, key):

        del self.properties[key]

    def send(self):

        "Sends this message."

        log.info("Sending %s",self)

        out = StringIO()

        for (key,value) in self.properties.iteritems():

            def _writePropString(str):

                out.write(str)    #FIX: Abbreviate

                out.write('\000')

            _writePropString(key)

            _writePropString(value)

        self.encoded = struct.pack('!H',out.tell()) + out.getvalue()

        out.close()

        body = self.body

        if self.compressed:

            body = zlib.compress(body,5)

        self.encoded += body

        log.debug("Encoded %s into %u bytes", self,len(self.encoded))

        self.bytesSent = 0

        self.connection._outQueueMessage(self)

    def _sendNextFrame(self, conn,maxLen):

        pos = self.bytesSent

        payload = self.encoded[pos:pos+maxLen]

        pos += len(payload)

        self.moreComing = (pos < len(self.encoded))

        log.debug("Sending frame of %s; bytes %i--%i", self,pos-len(payload),pos)

        conn.push( struct.pack(kFrameHeaderFormat, kFrameMagicNumber,

                                                   self.requestNo,

                                                   self.flags,

                                                   kFrameHeaderSize+len(payload)) )

        conn.push( payload )

        self.bytesSent = pos

        return self.moreComing

class Request (object):

    @property

    def response(self):

        "The response object for this request."

        r = self.__dict__.get('_response')

        if r==None:

            r = self._response = self._createResponse()

        return r

class Response (Message):

    def __init__(self, request):

        assert not request.noReply

        self.request = request

        self.requestNo = request.requestNo

        self.urgent = request.urgent

    @property

    def isResponse(self):

        return True

class IncomingRequest (IncomingMessage, Request):

    def _createResponse(self):

        return OutgoingResponse(self)

class OutgoingRequest (OutgoingMessage, Request):

    def _createResponse(self):

        return IncomingResponse(self)

class IncomingResponse (IncomingMessage, Response):

    def __init__(self, request):

        IncomingMessage.__init__(self,request.connection,request.requestNo,0)

        Response.__init__(self,request)

        self.onComplete = None

    def _finished(self):

        super(IncomingResponse,self)._finished()

        if self.onComplete:

            try:

                self.onComplete(self)

            except Exception, x:

                log.error("Exception dispatching response: %s", traceback.format_exc())

class OutgoingResponse (OutgoingMessage, Response):

    def __init__(self, request):

        OutgoingMessage.__init__(self,request.connection)

        Response.__init__(self,request)

### UNIT TESTS:

class BLIPTests(unittest.TestCase):

    def setUp(self):

        def handleRequest(request):

            logging.info("Got request!: %r",request)

            body = request.body

            assert len(body)<32768

            assert request.contentType == 'application/octet-stream'

            assert int(request['Size']) == len(body)

            assert request['User-Agent'] == 'BLIPConnectionTester'

            for i in xrange(0,len(request.body)):

                assert ord(body[i]) == i%256

            response = request.response

            response.body = request.body

            response['Content-Type'] = request.contentType

            response.send()

        listener = Listener(46353)

        listener.onRequest = handleRequest

    def testListener(self):

        logging.info("Waiting...")

        asyncore.loop()

if __name__ == '__main__':

    logging.basicConfig(level=logging.INFO)

    unittest.main()