ldapscan/lib/Protocol/Http/Message.py

"""
@author Georg Hopp

"""
from ..Message import Message as BaseMessage

class Message(BaseMessage):
    START_READY   = 0x01
    HEADERS_READY = 0x02
    BODY_READY    = 0x04

    METHODS = ('OPTIONS','GET','HEAD','POST','PUT','DELETE','TRACE','CONNECT')
    METHOD_OPTIONS = METHODS.index('OPTIONS')
    METHOD_GET     = METHODS.index('GET')
    METHOD_HEAD    = METHODS.index('HEAD')
    METHOD_POST    = METHODS.index('POST')
    METHOD_PUT     = METHODS.index('PUT')
    METHOD_DELETE  = METHODS.index('DELETE')
    METHOD_TRACE   = METHODS.index('TRACE')
    METHOD_CONNECT = METHODS.index('CONNECT')

    def __init__(self, remote):
        super(Message, self).__init__(remote)
        self.state = 0

        self._chunk_size = 0
        self._chunked    = False

        self._headers = {}
        self._body    = ''

        self._http    = None
        self._method  = None
        self._uri     = None
        self._code    = None
        self._message = None

    """
    cleaner
    =====================================================================
    """
    def resetStartLine(self):
        self._http    = None
        self._uri     = None
        self._code    = None
        self._message = None
        self.state &= ~Message.START_READY

    def resetHeaders(self):
        self._headers = {}
        self.state &= ~Message.HEADERS_READY

    def resetBody(self):
        self._body = ''
        self.state &= ~Message.BODY_READY
        self._chunked = False
        self._chunk_size = 0

    def reset(self):
        self.resetStartLine()
        self.resetHeaders()
        self.resetBody()

    def removeHeadersByKey(self, key):
        """
        Remove HTTP headers to a given key. This will remove all headers right
        now associated to that key. Keys are alwasys stored lower case and
        cenverted to title case during composition.

        returns None

        @key: str The header key to remove.
        """
        if key.lower() in self._headers:
            del(self._headers[key.lower()])

    def removeHeader(self, header):
        """
        Remove a header.

        returns None

        @header: tuple Holds key and value of the header to remove.
        """
        key = header[0].lower()
        if key in self._headers:
            if header[1] in self._headers[key]:
                self._headers[key].remove(header[1])

    """
    setter
    =====================================================================
    """
    def setRequestLine(self, method, uri, http):
        if self.isResponse():
            raise Exception('try to make a request from a response')
        self._method = method
        self._uri    = uri
        self._http   = http

    def setStateLine(self, http, code, message):
        if self.isRequest():
            raise Exception('try to make a response from a request')
        self._http    = http
        self._code    = code
        self._message = message

    def setHeader(self, key, value):
        """
        Add a header to the message.
        Under some circumstances HTTP allows to have multiple headers with
        the same key. Thats the reason why the values are handled in a list
        here.

        Returns None

        key:   The header key (The part before the colon :).
        value: The header value (The part behind the colon :).
               Value might also be a list a values for this key.
        """
        key = key.lower()
        if key in self._headers:
            self._headers[key] += [v.strip() for v in value.split(',')
                                   if v.strip() not in self._headers[key]]
        else:
            self._headers[key.lower()] = [v.strip() for v in value.split(',')]

    def replaceHeader(self, key, value):
        self._headers[key.lower()] = [v.strip() for v in value.split(',')]

    def setHeaders(self, headers):
        """
        This sets a bunch of headers at once. It will add the headers and not
        override anything. It is neccessary to clear the headers before calling
        this if only the headers given here should be in the message.

        Returns None

        headers: Either a list of tuples [(key,value),...] or
                 a dictionary {key:value,...}.
                 In both cases the values should be a list again.
        """
        if type(headers) == dict:
            headers = headers.items()

        for h in headers:
            self.setHeader(h[0], h[1])

    def setBody(self, data):
        """
        Set the body of a message. Currently we do not support sending 
        chunked message so this is simple...

        Returns None

        data: The data to set in the message body.
        """
        self.replaceHeader('Content-Length', '%d'%len(data))
        self._body = data

    """
    getter
    =====================================================================
    """
    def getHttpVersion(self):
        return self._http

    def getMethod(self):
        return self._method

    def getUri(self):
        return self._uri

    def getResponseCode(self):
        return self._code

    def getResponseMessage(self):
        return self._message

    def getStartLine(self):
        line = ''
        if self.isRequest():
            method = Message.METHODS[self._method]
            line = ' '.join((method, self._uri, self._http))
        elif self.isResponse():
            line = ' '.join((self._http, str(self._code), self._message))
        return line

    def getHeaders(self):
        return [(k, self.getHeader(k)) for k in self._headers]

    def getHeader(self, key):
        """
        Get all values currently associated to this header key.

        returns list All values to the given key.

        @key: str The key to get values for.
        """
        key = key.lower()
        if key not in self._headers: return ''
        return ', '.join(self._headers[key])

    def getBody(self):
        return self._body


    """
    checker
    =====================================================================
    """
    def headerKeyExists(self, key):
        return key.lower() in self._headers

    def startlineReady(self):
        return Message.START_READY == self.state & Message.START_READY

    def headersReady(self):
        return Message.HEADERS_READY == self.state & Message.HEADERS_READY

    def bodyReady(self):
        return Message.BODY_READY == self.state & Message.BODY_READY

    def ready(self):
        return self.headersReady() and self.bodyReady()

    def isRequest(self):
        return self._method is not None

    def isResponse(self):
        return self._code is not None

    def isCloseMessage(self):
        if self.isRequest():
            # HTTP always expects a response to be send, so a request is
            # never the close message.
            return False
        else:
            con_header = self.getHeader('Connection').lower()
            if self._http == 'HTTP/1.0':
                return 'keep-alive' not in con_header
            else:
                return 'close' in con_header

    def isUpgradeMessage(self):
        con_header = self.getHeader('Connection').lower()
        return 'upgrade' in con_header

    def isOptions(self):
        return Message.METHOD_OPTIONS == self.getMethod()

    def isGet(self):
        return Message.METHOD_GET == self.getMethod()

    def isHead(self):
        return Message.METHOD_HEAD == self.getMethod()

    def isPost(self):
        return Message.METHOD_POST == self.getMethod()

    def isPut(self):
        return Message.METHOD_PUT == self.getMethod()

    def isDelete(self):
        return Message.METHOD_DELETE == self.getMethod()

    def isTrace(self):
        return Message.METHOD_TRACE == self.getMethod()

    def isConnect(self):
        return Message.METHOD_CONNECT == self.getMethod()

# vim: set ft=python et ts=8 sw=4 sts=4: