twistedその４

HTMLのフィルタリング機能を実装する。
参考コードはこちら
http://blog.somethingaboutcode.com/?p=155

+-------+
|reactor|
+-------+---------+  +-----+  +-----------------+
+10080 HTTPFactory|->|Proxy|->|UpperProxyRequest|
+-----------------+  +-----+  +-------------------------------------------+
A                                   |process-> FilteringProxyClientFactory|
|                                   +------------------V------------------+----------+
+-------+                           |FilteringProxyClient                            |
|browser|                           |   Proxyから取ってきてHTMLにフィルタをかける    |
+-------+                           +------------------------------------------------+

htmlfilter

# coding: utf-8

import re

RE=re.compile('(<body[^>]*>)', re.IGNORECASE)

def filter(url, html):
    return RE.sub(r'\g<1><h1>Filtered!</h1>', html)

bodyタグの直後にFiltered!と追加する。
後でなんとかするためのプレースホルダ。

from twisted.web import http, proxy
from twisted.internet import reactor
from twisted.python import log
import urlparse
import htmlfilter
from cStringIO import StringIO

###############################################################################
# upper proxy
###############################################################################
PROXY_HOST='localhost'
PROXY_PORT=8080

class UpperProxyRequest(proxy.ProxyRequest):

    def process(self):
        parsed = urlparse.urlparse(self.uri)
        protocol = parsed[0]
        host = parsed[1]
        print self.method

        port = self.ports[protocol]
        if ':' in host:
            host, port = host.split(':')
            port = int(port)
        rest = urlparse.urlunparse(('', '') + parsed[2:])
        if not rest:
            rest = rest + '/'
        class_ = self.protocols[protocol]
        headers = self.getAllHeaders().copy()
        if 'host' not in headers:
            headers['host'] = host
        self.content.seek(0, 0)
        s = self.content.read()

        clientFactory = class_(
                self.method, self.uri, self.clientproto, headers, s, self)
        assert(PROXY_HOST)
        assert(PROXY_PORT)
        self.reactor.connectTCP(PROXY_HOST, PROXY_PORT, clientFactory)


proxy.Proxy.requestFactory=UpperProxyRequest

###############################################################################
# filtering
###############################################################################
class FilteringProxyClient(proxy.ProxyClient):
    def __init__(self, *args, **kwargs):
        proxy.ProxyClient.__init__(self, *args, **kwargs)
        self.__buffer=None
 
    def handleHeader(self, key, value):
        if key.lower()=="content-type" and value.startswith("text/html"):
            self.__buffer=StringIO()
        else:
            proxy.ProxyClient.handleHeader(self, key, value)
 
    def handleEndHeaders(self):
        if self.__buffer:
            pass #Need to calculate and send Content-Length first
        else:
            proxy.ProxyClient.handleEndHeaders(self)
 
    def handleResponsePart(self, buffer):
        if self.__buffer:
            self.__buffer.write(buffer)
        else:
            proxy.ProxyClient.handleResponsePart(self, buffer)
 
    def handleResponseEnd(self):
        if self.__buffer:
            b=self.__buffer.getvalue()
            self.__buffer=None
            buffer=htmlfilter.filter(self.father.uri, b)
            self.father.responseHeaders.setRawHeaders(
                    "Content-Length", ["%d" % len(buffer)])
            proxy.ProxyClient.handleEndHeaders(self)
            proxy.ProxyClient.handleResponsePart(self, buffer)
        proxy.ProxyClient.handleResponseEnd(self)

class FilteringProxyClientFactory(proxy.ProxyClientFactory):
    protocol=FilteringProxyClient

proxy.Proxy.requestFactory.protocols={'http': FilteringProxyClientFactory}


if __name__=="__main__":
    import sys
    log.startLogging(sys.stdout)

    f=http.HTTPFactory()
    f.protocol=proxy.Proxy

    reactor.listenTCP(10080, f)
    reactor.run()

ProxyRequest.protocolsのもつProxyClientFactoryに細工して、Filtering機能を持つProxyClientを生成するように改造。FilteringProxyClientは、content-typeがtext/htmlの時だけ本体を確保しておいて全部取得してからHTMLフィルターをかけてcontent-lengthを再計算するようにしている。