Ejemplo n.º 1
0
def test_recording_generation(*args):
    test_file = os.path.join(here, "files", "new_file.zip")
    file = RecordingFile(test_file)
    with open(file.recording_path, "w") as recording:
        recording.write("This is a recording")

    file.set_metadata("test_file", True)
    file.generate_zip_file()

    assert os.path.exists(test_file)
    os.remove(test_file)
Ejemplo n.º 2
0
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
            inject_deterministic = os.path.join(
                mitm_folder,
                "scripts",
                "inject-deterministic.py",
            )
            http_protocol_extractor = os.path.join(
                mitm_folder,
                "scripts",
                "http_protocol_extractor.py",
            )

            args = [
                "--save-stream-file",
                normalize_path(self.recording.recording_path),
                "--set",
                "websocket=false",
                "--scripts",
                inject_deterministic,
                "--scripts",
                http_protocol_extractor,
            ]
            command.extend(args)
            self.recording.set_metadata("proxy_version",
                                        self.config["playback_version"])
        else:
            # playback mode
            if len(self.playback_files) > 0:
                script = os.path.join(
                    mitm_folder,
                    "scripts",
                    "alternate-server-replay.py",
                )

                if self.config["playback_version"] in [
                        "4.0.4", "5.1.1", "6.0.2"
                ]:
                    args = [
                        "--set",
                        "upstream_cert=false",
                        "--set",
                        "upload_dir=" + normalize_path(self.upload_dir),
                        "--set",
                        "websocket=false",
                        "--set",
                        "server_replay_files={}".format(",".join([
                            normalize_path(playback_file.recording_path)
                            for playback_file in self.playback_files
                        ])),
                        "--scripts",
                        normalize_path(script),
                    ]
                    command.extend(args)
                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,
            processStderrLine=LOG.error,
            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 ERROR_CONTROL_C_EXIT  # noqa

                if exit_code == ERROR_CONTROL_C_EXIT:
                    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