#!/usr/bin/env python
"""
HTTP server serving fuzzy JPEG image or Flash animation.

Project is currently specific to Firefox on Linux.
"""
from __future__ import with_statement

HOST = '127.0.0.1'
PORT = 8080
NB_FILES = 9
ROWS = 3
MAX_MEMORY = 300*1024*1024
HTML_REFRESH_TIMEOUT = 2.0

from fusil.application import Application
from optparse import OptionGroup
from fusil.file_tools import filenameExtension
from fusil.network.http_server import HttpServer
from fusil.process.attach import AttachProcess
from fusil.process.create import CreateProcess
from fusil.process.watch import WatchProcess
from fusil.process.tools import runCommand
from fusil.project_agent import ProjectAgent
from fusil.auto_mangle import AutoMangle
from fusil.dummy_mangle import DummyMangle
from datetime import datetime
from time import time
from fusil.x11 import sendKey, getDisplay, findWindowByNameRegex
from Xlib.keysymdef.miscellany import XK_F5
import re
from os.path import basename
from os import getenv

IMAGE_TEMPLATE = '<img src="%(url)s" alt="%(text)s">'
EMBEDED_TEMPLATE = '<embed src="%(url)s" name="%(text)s" type="%(mime)%s"></embed>'

FILE_EXTENSIONS = {
    # Imaged
    '.bmp': ('image/x-ms-bmp', IMAGE_TEMPLATE),
    '.gif': ('image/gif', IMAGE_TEMPLATE),
    '.ico': ('image/x-ico', IMAGE_TEMPLATE),
    '.jpg': ('image/jpeg', IMAGE_TEMPLATE),
    '.jpeg': ('image/jpeg', IMAGE_TEMPLATE),
    '.png': ('image/png', IMAGE_TEMPLATE),

    # Embeded
    '.svg': ('image/svg+xml', EMBEDED_TEMPLATE),
    '.swf': ('application/x-shockwave-flash', EMBEDED_TEMPLATE),
}

class Fuzzer(Application):
    NAME = "firefox"
    USAGE = "%prog [options] filename"
    NB_ARGUMENTS = 1

    def createFuzzerOptions(self, parser):
        options = OptionGroup(parser, "Firefox")
        options.add_option("--iceweasel", help="Firefox name is Iceweasel",
            action="store_true")
        options.add_option("--test", help="Test mode (no fuzzing, just make sure that the fuzzer works)",
            action="store_true")
        options.add_option("--port", help="TCP port using to listen on (default: %s)" % PORT,
            type="int", default=PORT)
        options.add_option("--nb-files", help="Number of generated files (default: %s)" % NB_FILES,
            type="int", default=NB_FILES)
        options.add_option("--rows", help="Number of HTML table rows (default: %s)" % ROWS,
            type="int", default=ROWS)
        options.add_option("--max-memory", help="Maximum memory in bytes (default: %s)" % MAX_MEMORY,
            type="int", default=MAX_MEMORY)
        options.add_option("--attach", help='Attach to an existing process using its name (eg. "firefox-bin")',
            type="str")
        return options

    def setupProject(self):
        project = self.project
        homepage_url = "http://%s:%s/" % (HOST, self.options.port)
        self.error("HTTP fuzzer homepage: %s" % homepage_url)

        if not self.options.attach:
            arguments = ["firefox", '-safe-mode', homepage_url]
            firefox = FirefoxProcess(project, arguments, timeout=None)
            # Firefox forks many times
            firefox.max_user_process = 50
            # On core dump, Firefox proposes to start a debugger
            # whereas the user is unable to answer to this question...
            # Possible workaround: close stdin?
            firefox.core_dump = False
            firefox.setupX11()
            firefox.max_memory = self.options.max_memory
            WatchProcess(firefox, default_score=1.0)
        else:
            fireboxbin = AttachProcess(project, self.options.attach)
            fireboxbin.max_memory = self.options.max_memory
            FirefoxOpenURL(project, homepage_url, self.options.iceweasel)

        orig_filename = self.arguments[0]
        filename_ext = filenameExtension(orig_filename)
        if self.options.test:
            DummyMangle(project, orig_filename, nb_file=self.options.nb_files)
        else:
            AutoMangle(project, orig_filename, nb_file=self.options.nb_files)
        FuzzyHttpServer(project, self.options, filename_ext, rows=self.options.rows)

        if self.options.iceweasel:
            regex = r"Iceweasel$"
        else:
            regex = r"Mozilla Firefox$"
        FirefoxWindow(project, regex)

class FirefoxProcess(CreateProcess):
    def on_project_start(self):
        CreateProcess.init(self)
        self.error("Start Firefox...")
        self.createProcess()
        self.error("Start Firefox... done")

    def on_project_done(self):
        CreateProcess.deinit(self)

    def init(self):
        if 1 < self.project().session_index:
            self.send('process_create', self)

    def deinit(self):
        pass

class HtmlPage:
    def __init__(self):
        self.title = None
        self.headers = []
        self.body = ''

    def __str__(self):
        html = ['<html>\n']
        if self.title or self.headers:
            html.append('  <head>\n')
            if self.title:
                html.append('    <title>%s</title>\n' % self.title)
            html.append('  </head>\n')
        html.append('<body>\n')
        html.append(self.body+"\n")
        html.append('</body>\n')
        html.append('</html>\n')
        return ''.join(html)

class FuzzyHttpServer(HttpServer):
    def __init__(self, project, options, filename_ext, rows=5):
        HttpServer.__init__(self, project, options.port)
        self.options = options

        if filename_ext:
            ext = filename_ext.lower()
        else:
            ext = None
        try:
            self.content_type, self.file_html_format = FILE_EXTENSIONS[ext]
        except KeyError:
            raise ValueError("Unknown file extension: %r" % filename_ext)
        self.file_url_format = "file-%u-%u" + ext
        self.file_url_match = re.compile("file-[0-9]+-([0-9]+)" + re.escape(ext)).match
        self.timeout = HTML_REFRESH_TIMEOUT
        self.rows = rows

    def init(self):
        HttpServer.init(self)
        self.pages = set()
#        self.filenames = tuple()
        self.filenames = None
        self.is_done = False
        self.done_at = None

    def on_mangle_filenames(self, filenames):
        self.filenames = filenames
        self.send('http_server_ready')

    def serveRequest(self, client, request):
        url = request.uri[1:]
        if not url:
            url = "index.html"

        page = url
        self.error("Serve URL=%r" % url)
        match = self.file_url_match(url)
        if match:
            file_index = int(match.group(1))
            filename = self.filenames[file_index]
            data = open(filename, 'rb').read()
            self.serveData(client, 200, "OK", data=data, content_type=self.content_type)
        elif url == "index.html":
            self.createHomepage(client)
        else:
            page = None
            self.error("Error 404! %r" % url)
            self.error404(client, url)
        if page:
            self.pages.add(page)

        if (1 + self.options.nb_files) <= len(self.pages) and not self.is_done:
            self.is_done = True
            self.done_at = time() + self.timeout

    def createHomepage(self, client):
        self.error("Serve homepage")
        page = HtmlPage()
        page.title = 'Fusil HTTP server'
        page.body = '<h1>Fuzzing</h1>\n'
        page.body += '<table border="1">\n'
        tr_open = False
        count = len(self.filenames)
        session_index = self.project().session_index
        for index in xrange(count):
            url = self.file_url_format % (session_index, index)

            if (index % self.rows) == 0:
                if tr_open:
                    page.body += '  </tr>\n'
                page.body += '  <tr>\n'
                tr_open = True
            span = ''
            if index == (count-1):
                colspan = (index+1) % self.rows
                if colspan:
                    span += ' colspan="%s"' % (self.rows - colspan + 1)
            content = self.file_html_format % {
                'url': url,
                'text': "[%s]" % url,
                'mime': self.content_type}
            page.body += '    <td%s>%s</td>\n' % (span, content)
        page.body += '  </tr>\n'
        page.body += '</table>\n'
        page.body += '<p>Created: %s</p>\n' % datetime.now()
        page.body += '<p>Session: %s</p>\n' % self.project().session_index

        # Write the HTML in a file
        page = str(page)
        filename = self.session().createFilename("index.html")
        with open(filename, 'wb') as fp:
            fp.write(page)

        # Send bytes to the client
        self.serveData(client, 200, "OK", data=page, content_type="text/html")

    def live(self):
        HttpServer.live(self)
        if not self.is_done:
            return
        if time() < self.done_at:
            return
        self.error("DONE")
        self.is_done = False
        self.send('session_stop')

class FirefoxOpenURL(ProjectAgent):
    def __init__(self, project, homepage_url, iceweasel):
        ProjectAgent.__init__(self, project, "open_url")
        self.first = True
        self.homepage_url = homepage_url
        if iceweasel:
            self.program = "iceweasel"
        else:
            self.program = "firefox"

    def on_http_server_ready(self):
        if not self.first:
            return
        self.first = False
        env = {
            'HOME': getenv('HOME'),
            'DISPLAY': getenv('DISPLAY'),
        }
        runCommand(self, (self.program, self.homepage_url), options={'env': env})

class FirefoxWindow(ProjectAgent):
    def __init__(self, project, regex):
        ProjectAgent.__init__(self, project, "firefox_window")
        self.display = getDisplay()
        self.root_window = self.display.screen().root
        self.F5_keycode = self.display.keysym_to_keycode(XK_F5)
        self.window = None
        self.regex = regex

    def findWindow(self):
        if self.window:
            return
        self.window = findWindowByNameRegex(self.root_window, self.regex)

    def on_http_server_ready(self):
        if self.project().session_index == 1:
            return
        self.error("HTTP SERVER READY")
        self.findWindow()
        self.error("Send key F5 (%s) to Firefox window!" % self.F5_keycode)
        sendKey(self.window, self.F5_keycode, released=False) # 71=keycode of "F5" key (reload page)

if __name__ == "__main__":
    Fuzzer().main()

