habraproxy.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Требования: нужна библиотека beautifulsoup4 версии не ниже 4.4.1:
# "pip install beautifulsoup4"
#
# Запуск возможен с тремя параметрами:
# -host - хост сервера, по умолчанию localhost
# -port - порт сервера, по умолчанию 8000
# -site - адрес сайта, без "http://". По умолчанию: habrahabr.ru
#
# После запуска открывается браузер с обработанной главной страницей сайта

import BaseHTTPServer
import requests
import re
import argparse
import webbrowser
from bs4 import BeautifulSoup
from bs4.element import Comment

parser = argparse.ArgumentParser()
parser.add_argument('-host', default='localhost')
parser.add_argument('-port', default=8000)
parser.add_argument('-site', default='habrahabr.ru')
args = parser.parse_args()

HOST = args.host
PORT = int(args.port)
SITE = args.site


def add_tm(match):
    """
    Добавляет HTML-сущность "TM" к слову
    """
    return u'{}{}'.format(match.group(0), u'™')


def replace_content(html):
    """
    Тут происходит замена слов из шести букв на слова с "TM"
    """
    soup = BeautifulSoup(html, 'html.parser')
    body = soup.find('body')
    for text in body.find_all(string=True):
        if text.parent.name != 'script' and not isinstance(text, Comment):
            fixed_text = re.sub(r'\b\w{6}\b', add_tm, text, 0, re.UNICODE)
            text.replace_with(fixed_text)
    return soup.decode('utf-8', formatter=None)


class HabraProxyServer(BaseHTTPServer.HTTPServer):
    def server_activate(self):
        BaseHTTPServer.HTTPServer.server_activate(self)
        self.RequestHandlerClass.after_activate(*self.server_address)


class HabraProxyRequestHandler(BaseHTTPServer.BaseHTTPRequestHandler):
    def do_GET(self):
        """
        Обрабатывает запрос
        """
        url = 'http://{host}{path}'.format(**{'host': SITE, 'path': self.path})
        response = requests.get(url)

        self.send_response(response.status_code)
        self.send_header('Content-type', response.headers['content-type'])
        self.end_headers()

        if 'text/html' in response.headers['content-type']:
            content = replace_content(response.text)
            self.wfile.write(content.encode(response.encoding))
        else:
            self.wfile.write(response.content)

    @classmethod
    def after_activate(cls, host, port):
        """
        Открытие браузера с главной страницей сайта после запуска сервера
        """
        url = u'http://{host}{port}/'.format(
            **{'host': host, 'port': '' if port == 80 else ':' + str(port)})
        webbrowser.open_new_tab(url)


def main():
    httpd = HabraProxyServer((HOST, PORT), HabraProxyRequestHandler)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        httpd.server_close()


if __name__ == '__main__':
    main()