#!/usr/bin/env python2
"""
HTTP_Transaction is a container of useful HTTP Transaction information to
simplify code both in the framework and the plugins.
"""

import cgi
import logging
import StringIO
import gzip
import zlib
from httplib import responses as response_messages
import json

from framework.lib.general import *
from cookies import Cookie


class HTTP_Transaction(object):
    def __init__(self, Timer):
        self.Timer = Timer
        self.New = False

    def ScopeToStr(self):
        return str(self.IsInScope)[0]

    def InScope(self):
        return(self.IsInScope)

    def Start(self, url, data, method, is_in_scope):
        self.IsInScope = is_in_scope
        self.StartRequest()
        self.URL = url
        self.InitData(data)
        self.Method = DeriveHTTPMethod(method, data)
        self.Found = None
        self.RawRequest = ''
        self.ResponseHeaders = []
        self.ResponseSize = ''
        self.Status = ''
        self.ID = ''
        self.HTMLLinkToID = ''
        self.New = True  # Flag new transaction.

    def InitData(self, data):
        self.Data = data
        if self.Data is None:
            # This simplifies other code later, no need to cast to str if None, etc.
            self.Data = ''

    def StartRequest(self):
        self.Timer.start_timer('Request')
        self.Time = self.TimeHuman = ''

    def EndRequest(self):
        self.Time = self.Timer.get_elapsed_time_as_str('Request')
        self.TimeHuman = self.Time
        self.LocalTimestamp = self.Timer.get_current_date_time()

    def SetTransaction(self, found, request, response):
        # Response can be "Response" for 200 OK or "Error" for everything else, we don't care here.
        if self.URL != response.url:
            if response.code not in [302, 301]:  # No way, error in hook.
                # Mark as a redirect, dirty but more accurate than 200 :P
                self.Status = "%s Found" % str(302)
                self.Status += " --Redirect--> %s " % str(response.code)
                self.Status += response.msg
            # Redirect differs in schema (i.e. https instead of http).
            if self.URL.split(':')[0] != response.url.split(':')[0]:
                pass
            self.URL = response.url
        else:
            self.Status = "%s %s" % (str(response.code), response.msg)
        self.RawRequest = request
        self.Found = found
        self.ResponseHeaders = response.headers
        self.ResponseContents = response.read()
        # a new self.Decodedcontent is added if the received response is in compressed format
        self.checkIfCompressed(response, self.ResponseContents)
        self.EndRequest()

    def SetTransactionFromDB(self, id, url, method, status, time, time_human, local_timestamp, request_data,
                             raw_request, response_headers, response_size, response_body):
        self.ID = id
        self.New = False  # Flag NOT new transaction.
        self.URL = url
        self.Method = method
        self.Status = status
        self.Found = (self.Status == "200 OK")
        self.Time = time
        self.TimeHuman = time_human
        self.LocalTimestamp = local_timestamp
        self.Data = request_data
        self.RawRequest = raw_request
        self.ResponseHeaders = response_headers
        self.ResponseSize = response_size
        self.ResponseContents = response_body

    def GetSessionTokens(self):
        cookies = []
        try:  # parsing may sometimes fail
            for cookie in self.Cookies_list:
                cookies.append(Cookie.from_string(cookie).to_dict())
        except:
            pass
        return json.dumps(cookies)

    def SetError(self, error_message):
        # Only called for unknown errors, 404 and other HTTP stuff handled on self.SetResponse.
        self.ResponseContents = error_message
        self.EndRequest()

    def GetID(self):
        return (self.ID)

    def SetID(self, id, html_link_to_id):
        self.ID = id
        self.HTMLLinkToID = html_link_to_id
        # Only for new transactions, not when retrieved from DB, etc.
        if self.New:
            log = logging.getLogger('general')
            log.info("New owtf HTTP Transaction: %s",
                     " - ".join([self.ID, self.TimeHuman, self.Status, self.Method, self.URL]))

    def GetHTMLLink(self, link_name=''):
        if '' == link_name:
            link_name = "Transaction %s" % self.ID
        return self.HTMLLinkToID.replace('@@@PLACE_HOLDER@@@', link_name)

    def GetHTMLLinkWithTime(self, link_name=''):
        return "%s (%s)" % (self.GetHTMLLink(link_name), self.TimeHuman)

    def GetRawEscaped(self):
        return "<pre>%s</pre>" % cgi.escape(self.GetRaw())

    def GetRaw(self):
        return "%s\n\n%s" % (self.GetRawRequest(), self.GetRawResponse())

    def GetRawRequest(self):
        return self.RawRequest

    def GetStatus(self):
        return self.Status

    def GetResponseHeaders(self):
        return self.ResponseHeaders

    def GetRawResponse(self, with_status=True):
        try:
            return "%s\r\n%s\n\n%s" % (self.GetStatus(), str(self.ResponseHeaders), self.ResponseContents)
        except UnicodeDecodeError:
            return "%s\r\n%s\n\n[Binary Content]" % (self.GetStatus(), str(self.ResponseHeaders))

    def GetRawResponseHeaders(self, with_status=True):
        return "%s\r\n%s" % (self.GetStatus(), str(self.ResponseHeaders))

    def GetRawResponseBody(self):
        return self.ResponseContents

    def ImportProxyRequestResponse(self, request, response):
        self.IsInScope = request.in_scope
        self.URL = request.url
        self.InitData(request.body)
        self.Method = request.method
        try:
            self.Status = "%s %s" % (str(response.code), response_messages[int(response.code)])
        except KeyError:
            self.Status = "%s Unknown Error" % str(response.code)
        self.RawRequest = request.raw_request
        self.ResponseHeaders = response.header_string
        self.ResponseContents = response.body
        self.ResponseSize = len(self.ResponseContents)
        self.Time = str(response.request_time)
        self.TimeHuman = self.Timer.get_time_human(self.Time)
        self.LocalTimestamp = request.local_timestamp
        self.Found = (self.Status == "200 OK")
        self.Cookies_list = response.cookies
        self.New = True
        self.ID = ''
        self.HTMLLinkToID = ''

    def getDecodedResponse(self):
        return self.DecodedResponse

    def checkIfCompressed(self, response, content):
        if response.info().get('Content-Encoding') == 'gzip':  # check for gzip compression
            compressedFile = StringIO.StringIO()
            compressedFile.write(content)
            compressedFile.seek(0)
            f = gzip.GzipFile(fileobj=compressedFile, mode='rb')
            self.DecodedContent = f.read()
        elif response.info().get('Content-Encoding') == 'deflate':  # check for deflate compression
            self.DecodedContent = zlib.decompress(content)
        else:
            self.DecodedContent = content  # else the no compression
