Source code
Revision control
Copy as Markdown
Other Tools
# 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
import json
import os
import signal
import socket
import sys
import time
import mozinfo
import six
from mozprocess import ProcessHandler
from mozproxy.backends.base import Playback
from mozproxy.recordings import RecordingFile
from mozproxy.utils import (
LOG,
download_file_from_url,
get_available_port,
tooltool_download,
transform_platform,
)
here = os.path.dirname(__file__)
mitm_folder = os.path.dirname(os.path.realpath(__file__))
# maximal allowed runtime of a mitmproxy command
MITMDUMP_COMMAND_TIMEOUT = 30
class Mitmproxy(Playback):
def __init__(self, config):
self.config = config
self.host = (
"127.0.0.1" if "localhost" in self.config["host"] else self.config["host"]
)
self.port = None
self.mitmproxy_proc = None
self.mitmdump_path = None
self.record_mode = config.get("record", False)
self.recording = None
self.playback_files = []
self.browser_path = ""
if config.get("binary", None):
self.browser_path = os.path.normpath(config.get("binary"))
self.policies_dir = None
self.ignore_mitmdump_exit_failure = config.get(
"ignore_mitmdump_exit_failure", False
)
if self.record_mode:
if "recording_file" not in self.config:
LOG.error(
"recording_file value was not provided. Proxy service wont' start "
)
raise Exception("Please provide a playback_files list.")
if not isinstance(self.config.get("recording_file"), six.string_types):
LOG.error("recording_file argument type is not str!")
raise Exception("recording_file argument type invalid!")
if not os.path.splitext(self.config.get("recording_file"))[1] == ".zip":
LOG.error(
"Recording file type (%s) should be a zip. "
"Please provide a valid file type!"
% self.config.get("recording_file")
)
raise Exception("Recording file type should be a zip")
if os.path.exists(self.config.get("recording_file")):
LOG.error(
"Recording file (%s) already exists."
"Please provide a valid file path!"
% self.config.get("recording_file")
)
raise Exception("Recording file already exists.")
if self.config.get("playback_files", False):
LOG.error("Record mode is True and playback_files where provided!")
raise Exception("playback_files specified during record!")
if self.config.get("playback_version") is None:
LOG.error(
"mitmproxy was not provided with a 'playback_version' "
"Please provide a valid playback version"
)
raise Exception("playback_version not specified!")
# mozproxy_dir is where we will download all mitmproxy required files
# when running locally it comes from obj_path via mozharness/mach
if self.config.get("obj_path") is not None:
self.mozproxy_dir = self.config.get("obj_path")
else:
# in production it is ../tasks/task_N/build/, in production that dir
# is not available as an envvar, however MOZ_UPLOAD_DIR is set as
# ../tasks/task_N/build/blobber_upload_dir so take that and go up 1 level
self.mozproxy_dir = os.path.dirname(
os.path.dirname(os.environ["MOZ_UPLOAD_DIR"])
)
self.mozproxy_dir = os.path.join(self.mozproxy_dir, "testing", "mozproxy")
self.upload_dir = os.environ.get("MOZ_UPLOAD_DIR", self.mozproxy_dir)
LOG.info(
"mozproxy_dir used for mitmproxy downloads and exe files: %s"
% self.mozproxy_dir
)
# setting up the MOZPROXY_DIR env variable so custom scripts know
# where to get the data
os.environ["MOZPROXY_DIR"] = self.mozproxy_dir
LOG.info("Playback tool: %s" % self.config["playback_tool"])
LOG.info("Playback tool version: %s" % self.config["playback_version"])
def download_mitm_bin(self):
# Download and setup mitm binaries
manifest = os.path.join(
here,
"manifests",
"mitmproxy-rel-bin-%s-{platform}.manifest"
% self.config["playback_version"],
)
transformed_manifest = transform_platform(manifest, self.config["platform"])
# generate the mitmdump_path
self.mitmdump_path = os.path.normpath(
os.path.join(
self.mozproxy_dir,
"mitmdump-%s" % self.config["playback_version"],
"mitmdump",
)
)
# Check if mitmproxy bin exists
if os.path.exists(self.mitmdump_path):
LOG.info("mitmproxy binary already exists. Skipping download")
else:
# Download and unpack mitmproxy binary
download_path = os.path.dirname(self.mitmdump_path)
LOG.info("create mitmproxy %s dir" % self.config["playback_version"])
if not os.path.exists(download_path):
os.makedirs(download_path)
LOG.info("downloading mitmproxy binary")
tooltool_download(
transformed_manifest, self.config["run_local"], download_path
)
def download_manifest_file(self, manifest_path):
# Manifest File
# we use one pageset for all platforms
LOG.info("downloading mitmproxy pageset")
tooltool_download(manifest_path, self.config["run_local"], self.mozproxy_dir)
with open(manifest_path) as manifest_file:
manifest = json.load(manifest_file)
for file in manifest:
zip_path = os.path.join(self.mozproxy_dir, file["filename"])
LOG.info("Adding %s to recording list" % zip_path)
self.playback_files.append(RecordingFile(zip_path))
def download_playback_files(self):
# Detect type of file from playback_files and download accordingly
if "playback_files" not in self.config:
LOG.error(
"playback_files value was not provided. Proxy service wont' start "
)
raise Exception("Please provide a playback_files list.")
if not isinstance(self.config["playback_files"], list):
LOG.error("playback_files should be a list")
raise Exception("playback_files should be a list")
for playback_file in self.config["playback_files"]:
if playback_file.startswith("https://") and "mozilla.com" in playback_file:
# URL provided
dest = os.path.join(self.mozproxy_dir, os.path.basename(playback_file))
download_file_from_url(playback_file, self.mozproxy_dir, extract=False)
# Add Downloaded file to playback_files list
LOG.info("Adding %s to recording list" % dest)
self.playback_files.append(RecordingFile(dest))
continue
if not os.path.exists(playback_file):
LOG.error(
"Zip or manifest file path (%s) does not exist. Please provide a valid path!"
% playback_file
)
raise Exception("Zip or manifest file path does not exist")
if os.path.splitext(playback_file)[1] == ".zip":
# zip file path provided
LOG.info("Adding %s to recording list" % playback_file)
self.playback_files.append(RecordingFile(playback_file))
elif os.path.splitext(playback_file)[1] == ".manifest":
# manifest file path provided
self.download_manifest_file(playback_file)
def download(self):
"""Download and unpack mitmproxy binary and pageset using tooltool"""
if not os.path.exists(self.mozproxy_dir):
os.makedirs(self.mozproxy_dir)
self.download_mitm_bin()
if self.record_mode:
self.recording = RecordingFile(self.config["recording_file"])
else:
self.download_playback_files()
def stop(self):
LOG.info("Mitmproxy stop!!")
self.stop_mitmproxy_playback()
if self.record_mode:
LOG.info("Record mode ON. Generating zip file ")
self.recording.generate_zip_file()
def wait(self, timeout=1):
"""Wait until the mitmproxy process has terminated."""
# We wait using this method to allow Windows to respond to the Ctrl+Break
# signal so that we can exit cleanly from the command-line driver.
while True:
returncode = self.mitmproxy_proc.wait(timeout)
if returncode is not None:
return returncode
def start(self):
# go ahead and download and setup mitmproxy
self.download()
# mitmproxy must be started before setup, so that the CA cert is available
self.start_mitmproxy(self.mitmdump_path, self.browser_path)
# In case the setup fails, we want to stop the process before raising.
try:
self.setup()
except Exception:
try:
self.stop()
except Exception:
LOG.error("MitmProxy failed to STOP.", exc_info=True)
LOG.error("Setup of MitmProxy failed.", exc_info=True)
raise
def start_mitmproxy(self, mitmdump_path, browser_path):
"""Startup mitmproxy and replay the specified flow file"""
if self.mitmproxy_proc is not None:
raise Exception("Proxy already started.")
self.port = get_available_port()
LOG.info("mitmdump path: %s" % mitmdump_path)
LOG.info("browser path: %s" % browser_path)
# mitmproxy needs some DLL's that are a part of Firefox itself, so add to path
env = os.environ.copy()
env["PATH"] = os.path.dirname(browser_path) + os.pathsep + env["PATH"]
command = [mitmdump_path]
if self.config.get("verbose", False):
# Generate mitmproxy verbose logs
command.extend(["-v"])
# add proxy host and port options
command.extend(["--listen-host", self.host, "--listen-port", str(self.port)])
# record mode
if self.record_mode:
# generate recording script paths
command.extend(
[
"--save-stream-file",
os.path.normpath(self.recording.recording_path),
"--set",
"websocket=false",
]
)
if "inject_deterministic" in self.config.keys():
command.extend(
[
"--scripts",
os.path.join(mitm_folder, "scripts", "inject-deterministic.py"),
]
)
self.recording.set_metadata(
"proxy_version", self.config["playback_version"]
)
else:
# playback mode
if len(self.playback_files) > 0:
if self.config["playback_version"] == "8.1.1":
command.extend(
[
"--set",
"websocket=false",
"--set",
"connection_strategy=lazy",
"--set",
"alt_server_replay_nopop=true",
"--set",
"alt_server_replay_kill_extra=true",
"--set",
"alt_server_replay_order_reversed=true",
"--set",
"alt_server_replay={}".format(
",".join(
[
os.path.normpath(playback_file.recording_path)
for playback_file in self.playback_files
]
)
),
"--scripts",
os.path.normpath(
os.path.join(
mitm_folder, "scripts", "alt-serverplayback.py"
)
),
]
)
elif self.config["playback_version"] in [
"4.0.4",
"5.1.1",
"6.0.2",
]:
command.extend(
[
"--set",
"upstream_cert=false",
"--set",
"upload_dir=" + os.path.normpath(self.upload_dir),
"--set",
"websocket=false",
"--set",
"server_replay_files={}".format(
",".join(
[
os.path.normpath(playback_file.recording_path)
for playback_file in self.playback_files
]
)
),
"--scripts",
os.path.normpath(
os.path.join(
mitm_folder, "scripts", "alternate-server-replay.py"
)
),
]
)
else:
raise Exception("Mitmproxy version is unknown!")
else:
raise Exception(
"Mitmproxy can't start playback! Playback settings missing."
)
# mitmproxy needs some DLL's that are a part of Firefox itself, so add to path
env = os.environ.copy()
if not os.path.dirname(self.browser_path) in env["PATH"]:
env["PATH"] = os.path.dirname(self.browser_path) + os.pathsep + env["PATH"]
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)
self.mitmproxy_proc = ProcessHandler(
command,
logfile=os.path.join(self.upload_dir, "mitmproxy.log"),
env=env,
storeOutput=False,
)
self.mitmproxy_proc.run()
end_time = time.time() + MITMDUMP_COMMAND_TIMEOUT
ready = False
while time.time() < end_time:
ready = self.check_proxy(host=self.host, port=self.port)
if ready:
LOG.info(
"Mitmproxy playback successfully started on %s:%d as pid %d"
% (self.host, self.port, self.mitmproxy_proc.pid)
)
return
time.sleep(0.25)
# cannot continue as we won't be able to playback the pages
LOG.error("Aborting: Mitmproxy process did not startup")
self.stop_mitmproxy_playback()
sys.exit(1) # XXX why do we need to do that? a raise is not enough?
def stop_mitmproxy_playback(self):
"""Stop the mitproxy server playback"""
if self.mitmproxy_proc is None or self.mitmproxy_proc.poll() is not None:
return
LOG.info(
"Stopping mitmproxy playback, killing process %d" % self.mitmproxy_proc.pid
)
# On Windows, mozprocess brutally kills mitmproxy with TerminateJobObject
# The process has no chance to gracefully shutdown.
# Here, we send the process a break event to give it a chance to wrapup.
# See the signal handler in the alternate-server-replay-4.0.4.py script
if mozinfo.os == "win":
LOG.info("Sending CTRL_BREAK_EVENT to mitmproxy")
os.kill(self.mitmproxy_proc.pid, signal.CTRL_BREAK_EVENT)
time.sleep(2)
exit_code = self.mitmproxy_proc.kill()
self.mitmproxy_proc = None
if exit_code != 0:
if exit_code is None:
LOG.error("Failed to kill the mitmproxy playback process")
return
if mozinfo.os == "win":
from mozprocess.winprocess import ( # noqa
ERROR_CONTROL_C_EXIT,
ERROR_CONTROL_C_EXIT_DECIMAL,
)
if exit_code in [ERROR_CONTROL_C_EXIT, ERROR_CONTROL_C_EXIT_DECIMAL]:
LOG.info(
"Successfully killed the mitmproxy playback process"
" with exit code %d" % exit_code
)
return
log_func = LOG.error
if self.ignore_mitmdump_exit_failure:
log_func = LOG.info
log_func("Mitmproxy exited with error code %d" % exit_code)
else:
LOG.info("Successfully killed the mitmproxy playback process")
def check_proxy(self, host, port):
"""Check that mitmproxy process is working by doing a socket call using the proxy settings
:param host: Host of the proxy server
:param port: Port of the proxy server
:return: True if the proxy service is working
"""
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
s.connect((host, port))
s.shutdown(socket.SHUT_RDWR)
s.close()
return True
except socket.error:
return False