intercepting_proxy.py

# The MIT License (MIT)
# 
# Copyright (c) 2015 Context Information Security Ltd.
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

try:
    from cStringIO import StringIO
except:
    from io import BytesIO as StringIO
    
try:
    from urlparse import urlparse, urlunparse
except:
    from urllib.parse import urlparse, urlunparse
    
from twisted.internet import reactor
from twisted.python.log import startLogging
from twisted.web import server, resource, proxy, http
from twisted.python.filepath import FilePath

# add timeout features to the ProxyClient requests
from twisted.protocols.policies import TimeoutMixin


#ProxyClient(twisted.web.http.HTTPClient)
class InterceptingProxyClient(proxy.ProxyClient, TimeoutMixin):
    def __init__(self, command, rest, version, headers, data, father):
        proxy.ProxyClient.__init__(self, command, rest, version, headers, data, father)
        
        # Define variables to track a terminated request and set a timeout
        # set a callback handler for when a "finish" is reached
        self.finished = False
        self.setTimeout(20)
        d = father.notifyFinish()
        d.addBoth(self.onRequestFinished)
        
    def onRequestFinished(self, message=None):
        self.finish()
        
    def finish(self):
        # terminate the connection the proper way
        # if the parent had not been terminated, terminate it
        if not self.father.finished and not self.father._disconnected:
            self.father.finish()
            
        #if you have not been terminated, terminate it
        if not self.finished:
            self.transport.loseConnection()
            self.setTimeout(None)
            self.finished = True
                
    def connectionMade(self):
        # Edge case where the connection terminated prior to the event notification being setup 
        if self.father._disconnected:
            self.finish()
            return

        # Send request to web server
        proxy.ProxyClient.connectionMade(self)
        
    def timeoutConnection(self):
        # Recognize a timeout 
        TimeoutMixin.timeoutConnection(self)

    def handleResponsePart(self, buffer):
        # Buffer the output if we intend to modify it
        if self.father.has_response_modifiers():
            self.father.response_buffer.write(buffer)
        else:
            proxy.ProxyClient.handleResponsePart(self, buffer)
 
    def handleResponseEnd(self):
        # Process the buffered output if we are modifying it
        if self.father.has_response_modifiers():
            if not self._finished:
                # Replace the StringIO with a string for the modifiers
                data = self.father.response_buffer.getvalue()
                self.father.response_buffer.close()
                self.father.response_buffer = data

                # Do editing of response headers / content here
                self.father.run_response_modifiers()
                self.father.responseHeaders.setRawHeaders('content-length', [len(self.father.response_buffer)])
                self.father.write(self.father.response_buffer)
        proxy.ProxyClient.handleResponseEnd(self)


class InterceptingProxyClientFactory(proxy.ProxyClientFactory):
    noisy = False
    protocol = InterceptingProxyClient

#ProxyRequest(twisted.web.http.Request)
class InterceptingProxyRequest(proxy.ProxyRequest):
    def __init__(self, *args, **kwargs):
        proxy.ProxyRequest.__init__(self, *args, **kwargs)
        self.response_buffer = StringIO()
        self.request_buffer = StringIO()
        self.modifiers = self.channel.factory.modifiers
        
    def run_request_modifiers(self):
        if not self.has_request_modifiers():
            return
            
        if self.requestHeaders.hasHeader('content-length'):
            self.request_buffer = self.content.read()
            
        for m in self.modifiers:
            m.modify_request(self)
            
        if self.requestHeaders.hasHeader('content-length'):
            self.content.seek(0,0)
            self.content.write(self.request_buffer)
            self.content.truncate()
            self.requestHeaders.setRawHeaders('content-length', [len(self.request_buffer)])

    def has_request_modifiers(self):
        ret = False
        for m in self.modifiers:
            if m.will_modify_request(self):
                ret = True
        return ret
        
    def has_response_modifiers(self):
        ret = False
        for m in self.modifiers:
            if m.will_modify_response(self):
                ret = True
        return ret

    def run_response_modifiers(self):
        for m in self.modifiers:
            m.modify_response(self)

    def has_response_server(self):
        for m in self.modifiers:
            if m.will_serve_response(self):
                return True
        return False
        
    def serve_resource(self):
        body = None
        for m in self.modifiers:
            if m.will_serve_response(self):
                body = m.get_response(self)
                break
        if not body:
            raise Exception('Nothing served a resource')
        self.setHeader('content-length', str(len(body)))
        if self.method == 'HEAD':
            self.write('')
        else:
            self.write(body)
        self.finish()
        
    def process(self):
        host = None
        port = None

        if not self.uri.startswith("http://") and not self.uri.startswith("https://"):
            self.uri = "http://" + self.getHeader("Host") + self.uri

        parsed_uri = urlparse(self.uri)
        self.uri = urlunparse(('', '', parsed_uri.path, parsed_uri.params, parsed_uri.query, parsed_uri.fragment)) or "/"

        if self.has_request_modifiers():
            self.run_request_modifiers()
            
        if self.has_response_server():
            self.serve_resource()
            return 
            
        protocol = parsed_uri.scheme or 'http'
        host = host or parsed_uri.netloc
        port = port or parsed_uri.port or self.ports[protocol]

        headers = self.getAllHeaders().copy()
        if 'host' not in headers:
            headers['host'] = host
             
        if ':' in host:
            host,_ = host.split(':')

        self.content.seek(0, 0)
        content = self.content.read()
    
        clientFactory = InterceptingProxyClientFactory(self.method, self.uri, self.clientproto, headers, content, self)
        self.reactor.connectTCP(host, port, clientFactory)

#Proxy(twisted.web.http.HTTPChannel)
class InterceptingProxy(proxy.Proxy):
    requestFactory = InterceptingProxyRequest

class InterceptingProxyFactory(http.HTTPFactory):
    protocol = InterceptingProxy
    
    def add_modifier(self, m):
        self.modifiers.append(m)
        
    def __init__(self, modifier, *args, **kwargs):
        http.HTTPFactory.__init__(self, *args, **kwargs)
        self.modifiers = []
        self.add_modifier(modifier)