testing/talos/talos/mitmproxy/mitmproxy.py
author Rob Wood <rwood@mozilla.com>
Tue, 15 Aug 2017 15:30:23 -0400
changeset 375466 2619edc827ada43379ff483c26271284cfdd199a
parent 363025 69a0d95d19bbee728d6e255bb63706da1fd36ce8
child 376336 ebbe081e95547d66c02333090689727dafe77cc8
permissions -rw-r--r--
Bug 1387265 - Expand talos tp6 to macosx; r=jmaher MozReview-Commit-ID: 78LGHEJNGz4

'''This helps loading mitmproxy's cert and change proxy settings for Firefox.'''
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import os
import psutil
import sys
import subprocess
import time

import mozinfo

from mozlog import get_proxy_logger

here = os.path.dirname(os.path.realpath(__file__))
LOG = get_proxy_logger()

# path for mitmproxy certificate, generated auto after mitmdump is started
# on local machine it is 'HOME', however it is different on production machines
try:
    DEFAULT_CERT_PATH = os.path.join(os.getenv('HOME'),
                                     '.mitmproxy', 'mitmproxy-ca-cert.cer')
except:
    DEFAULT_CERT_PATH = os.path.join(os.getenv('HOMEDRIVE'), os.getenv('HOMEPATH'),
                                     '.mitmproxy', 'mitmproxy-ca-cert.cer')

MITMPROXY_SETTINGS = '''// Start with a comment
// Load up mitmproxy cert
var Cc = Components.classes;
var Ci = Components.interfaces;
var certdb = Cc["@mozilla.org/security/x509certdb;1"].getService(Ci.nsIX509CertDB);
var certdb2 = certdb;

try {
certdb2 = Cc["@mozilla.org/security/x509certdb;1"].getService(Ci.nsIX509CertDB2);
} catch (e) {}

cert = "%(cert)s";
certdb2.addCertFromBase64(cert, "C,C,C", "");

// Use mitmdump as the proxy
// Manual proxy configuration
pref("network.proxy.type", 1);
pref("network.proxy.http", "127.0.0.1");
pref("network.proxy.http_port", 8080);
pref("network.proxy.ssl", "127.0.0.1");
pref("network.proxy.ssl_port", 8080);
'''


def configure_mitmproxy(fx_install_dir,
                        scripts_path,
                        certificate_path=DEFAULT_CERT_PATH):
    # scripts_path is path to mozharness on test machine; needed so can import
    if scripts_path is not False:
        sys.path.insert(1, scripts_path)
        sys.path.insert(1, os.path.join(scripts_path, 'mozharness'))
    from mozharness.mozilla.firefox.autoconfig import write_autoconfig_files
    certificate = _read_certificate(certificate_path)
    write_autoconfig_files(fx_install_dir=fx_install_dir,
                           cfg_contents=MITMPROXY_SETTINGS % {
                              'cert': certificate})


def _read_certificate(certificate_path):
    ''' Return the certificate's hash from the certificate file.'''
    # NOTE: mitmproxy's certificates do not exist until one of its binaries
    #       has been executed once on the host
    with open(certificate_path, 'r') as fd:
        contents = fd.read()
    return ''.join(contents.splitlines()[1:-1])


def is_mitmproxy_cert_installed():
    """Verify mitmxproy CA cert was added to Firefox"""
    # TODO: Bug 1366071
    return True


def install_mitmproxy_cert(mitmproxy_proc, browser_path, scripts_path):
    """Install the CA certificate generated by mitmproxy, into Firefox"""
    LOG.info("Installing mitmxproxy CA certficate into Firefox")
    # browser_path is exe, we want install dir
    browser_install = os.path.dirname(browser_path)
    # on macosx we need to remove the last folders 'Content/MacOS'
    if mozinfo.os == 'mac':
        browser_install = browser_install[:-14]

    LOG.info('Calling configure_mitmproxy with browser folder: %s' % browser_install)
    configure_mitmproxy(browser_install, scripts_path)
    # cannot continue if failed to add CA cert to Firefox, need to check
    if not is_mitmproxy_cert_installed():
        LOG.error('Aborting: failed to install mitmproxy CA cert into Firefox')
        stop_mitmproxy_playback(mitmproxy_proc)
        sys.exit()


def start_mitmproxy_playback(mitmdump_path,
                             mitmproxy_recording_path,
                             mitmproxy_recordings_list,
                             browser_path):
    """Startup mitmproxy and replay the specified flow file"""
    mitmproxy_recordings = []
    # recording names can be provided in comma-separated list; build py list including path
    for recording in mitmproxy_recordings_list:
        mitmproxy_recordings.append(os.path.join(mitmproxy_recording_path, recording))

    # cmd line to start mitmproxy playback using custom playback script is as follows:
    # <path>/mitmdump -s "<path>mitmdump-alternate-server-replay/alternate-server-replay.py
    #  <path>recording-1.mp <path>recording-2.mp..."
    param = os.path.join(here, 'alternate-server-replay.py')
    env = os.environ.copy()

    # this part is platform-specific
    if mozinfo.os == 'win':
        param2 = '""' + param.replace('\\', '\\\\\\') + ' ' + \
                 ' '.join(mitmproxy_recordings).replace('\\', '\\\\\\') + '""'
        sys.path.insert(1, mitmdump_path)
        # mitmproxy needs some DLL's that are a part of Firefox itself, so add to path
        env["PATH"] = os.path.dirname(browser_path) + ";" + env["PATH"]
    elif mozinfo.os == 'mac':
        param2 = param + ' ' + ' '.join(mitmproxy_recordings)
        env["PATH"] = os.path.dirname(browser_path)
    else:
        # TODO: support other platforms, Bug 1366355
        LOG.error('Aborting: talos mitmproxy is currently only supported on Windows and Mac')
        sys.exit()

    command = [mitmdump_path, '-k', '-s', param2]

    LOG.info("Starting mitmproxy playback using env path: %s" % env["PATH"])
    LOG.info("Starting mitmproxy playback using command: %s" % ' '.join(command))
    # to turn off mitmproxy log output, use these params for Popen:
    # Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
    mitmproxy_proc = subprocess.Popen(command, env=env)
    time.sleep(10)
    data = mitmproxy_proc.poll()
    if data is None:
        LOG.info("Mitmproxy playback successfully started as pid %d" % mitmproxy_proc.pid)
        return mitmproxy_proc
    # cannot continue as we won't be able to playback the pages
    LOG.error('Aborting: mitmproxy playback process failed to start, poll returned: %s' % data)
    sys.exit()


def stop_mitmproxy_playback(mitmproxy_proc):
    """Stop the mitproxy server playback"""
    LOG.info("Stopping mitmproxy playback, klling process %d" % mitmproxy_proc.pid)
    if mozinfo.os == 'mac':
        mitmproxy_proc.terminate()
    else:
        mitmproxy_proc.kill()
    time.sleep(10)
    if mitmproxy_proc.pid in psutil.pids():
        # I *think* we can still continue, as process will be automatically
        # killed anyway when mozharness is done (?) if not, we won't be able
        # to startup mitmxproy next time if it is already running
        LOG.error("Failed to kill the mitmproxy playback process")
    else:
        LOG.info("Successfully killed the mitmproxy playback process")