# encoding: utf-8

 """

 BLIP.py

 Created by Jens Alfke on 2008-06-03.

 Copyright notice and BSD license at end of file.

 """

 import asynchat

 import asyncore

 from cStringIO import StringIO

 import logging

 import socket

 import struct

 import sys

 import traceback

 import zlib

 # Connection status enumeration:

 kDisconnected = -1

 kClosed  = 0

 kOpening = 1

 kOpen    = 2

 kClosing = 3

 # 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

 ### LISTENER AND CONNECTION CLASSES:

 class Listener (asyncore.dispatcher):

     "BLIP listener/server class"

     def __init__(self, port, sslKeyFile=None, sslCertFile=None):

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

         self.sslKeyFile=sslKeyFile

         self.sslCertFile=sslCertFile

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

     def handle_accept( self ):

         socket,address = self.accept()

         if self.sslKeyFile:

             socket.ssl(socket,self.sslKeyFile,self.sslCertFile)

         conn = Connection(address, sock=socket, listener=self)

         conn.onRequest = self.onRequest

         if self.onConnected:

             self.onConnected(conn)

     def handle_error(self):

         (typ,val,trace) = sys.exc_info()

         log.error("Listener caught: %s %s\n%s", typ,val,traceback.format_exc())

         self.close()

 class Connection (asynchat.async_chat):

     def __init__( self, address, sock=None, listener=None, ssl=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."

         if sock:

             asynchat.async_chat.__init__(self,sock)

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

             self.status = kOpen

         else:

             asynchat.async_chat.__init__(self)

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

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

             self.status = kOpening

             if ssl:

                 ssl(self.socket)

             self.connect(address)

         self.address = address

         self.listener = listener

         self.onRequest = None

         self.pendingRequests = {}

         self.pendingResponses = {}

         self.outBox = []

         self.inMessage = None

         self.inNumRequests = self.outNumRequests = 0

         self.sending = False

         self._endOfFrame()

     def close(self):

         if self.status > kClosed:

             self.status = kClosing

             log.info("Connection closing...")

         asynchat.async_chat.close(self)

     def handle_connect(self):

         log.info("Connection open!")

         self.status = kOpen

     def handle_error(self):

         (typ,val,trace) = sys.exc_info()

         log.error("Connection caught: %s %s\n%s", typ,val,traceback.format_exc())

         self.discard_buffers()

         self.status = kDisconnected

         self.close()

     def handle_close(self):

         log.info("Connection closed!")

         self.pendingRequests = self.pendingResponses = None

         self.outBox = None

         if self.status == kClosing:

             self.status = kClosed

         else:

             self.status = kDisconnected

         asynchat.async_chat.handle_close(self)

     ### SENDING:

     @property

     def canSend(self):

         return self.status==kOpening or self.status==kOpen

     def _sendMessage(self, msg):

         if self.canSend:

             self._outQueueMessage(msg,True)

             if not self.sending:

                 log.debug("Waking up the output stream")

                 self.sending = True

                 self.push_with_producer(self)

             return True

         else:

             return False

     def _sendRequest(self, req):

         if self.canSend:

             requestNo = req.requestNo = self.outNumRequests = self.outNumRequests + 1

             response = req.response

             if response:

                 response.requestNo = requestNo

                 self.pendingResponses[requestNo] = response

                 log.debug("pendingResponses[%i] := %s",requestNo,response)

             return self._sendMessage(req)

         else:

             return False

     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.bytesSent==0:

                     break

                 index -= 1

             else:

                 index = 1

         self.outBox.insert(index,msg)

         if isNew:

             log.info("Queuing %s at index %i",msg,index)

         else:

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

     def more(self):

         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

             data = msg._sendNextFrame(frameSize)

             if msg._moreComing:

                 self._outQueueMessage(msg,isNew=False)

             else:

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

             return data

         else:

             log.debug("Nothing more to send")

             self.sending = False

             return None

     ### RECEIVING:

     def collect_incoming_data(self, data):

         if self.expectingHeader:

             if self.inHeader==None:

                 self.inHeader = data

             else:

                 self.inHeader += data

         elif self.inMessage:

             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: raise ConnectionException, "Incorrect frame magic number %x" %magic

             if frameLen < kFrameHeaderSize: raise ConnectionException,"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.pendingResponses[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

 ### MESSAGE CLASSES:

 class Message (object):

     "Abstract superclass of all request/response objects"

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

         self.connection = connection

         self.body = body

         self.properties = properties or {}

         self.requestNo = None

     @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[" %(type(self).__name__)

         if self.requestNo != None:

             s += "#%i" %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):

         """Received a frame header."""

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

     def _receivedData(self, data):

         """Received data from a frame."""

         self.frames.append(data)

     def _finished(self):

         """The entire message has been received; now decode it."""

         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, body=None, properties=None):

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

         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]

     @property

     def sent(self):

         return hasattr(self,'encoded')

     def _encode(self):

         "Generates the message's encoded form, prior to sending it."

         out = StringIO()

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

             def _writePropString(s):

                 out.write(str(s))    #FIX: Abbreviate

                 out.write('\000')

             _writePropString(key)

             _writePropString(value)

         propertiesSize = out.tell()

         assert propertiesSize<65536     #FIX: Return an error instead

         body = self.body

         if self.compressed:

             z = zlib.compressobj(6,zlib.DEFLATED,31)   # window size of 31 needed for gzip format

             out.write(z.compress(body))

             body = z.flush()

         out.write(body)

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

         out.close()

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

         self.bytesSent = 0

     def _sendNextFrame(self, maxLen):

         pos = self.bytesSent

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

         pos += len(payload)

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

         if not self._moreComing:

             self.encoded = None

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

         header = struct.pack(kFrameHeaderFormat, kFrameMagicNumber,

                                                    self.requestNo,

                                                    self.flags,

                                                    kFrameHeaderSize+len(payload))

         self.bytesSent = pos

         return header + payload

 class Request (object):

     @property

     def response(self):

         "The response object for this request."

         if self.noReply:

             return None

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

         if r==None:

             r = self._response = self._createResponse()

         return r

 class Response (Message):

     def _setRequest(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)

     def send(self):

         self._encode()

         return self.connection._sendRequest(self) and self.response

 class IncomingResponse (IncomingMessage, Response):

     def __init__(self, request):

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

         self._setRequest(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)

         self._setRequest(request)

     def send(self):

         self._encode()

         return self.connection._sendMessage(self)

 """

  Copyright (c) 2008, Jens Alfke <jens@mooseyard.com>. All rights reserved.

  Redistribution and use in source and binary forms, with or without modification, are permitted

  provided that the following conditions are met:

  * Redistributions of source code must retain the above copyright notice, this list of conditions

  and the following disclaimer.

  * Redistributions in binary form must reproduce the above copyright notice, this list of conditions

  and the following disclaimer in the documentation and/or other materials provided with the

  distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR

  IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND 

  FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRI-

  BUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES

  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 

   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 

  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF 

  THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 """