From c78533816f8f732e7ee9e6f91caaa786a4c249ee Mon Sep 17 00:00:00 2001 From: Tyson Smith Date: Fri, 13 Dec 2024 13:56:57 -0800 Subject: [PATCH] feat: add support for marionette --- src/ffpuppet/bootstrapper.py | 56 ++++++++++++++++++++++++++----- src/ffpuppet/core.py | 25 ++++++++++++-- src/ffpuppet/main.py | 15 +++++++++ src/ffpuppet/resources/testff.py | 1 + src/ffpuppet/test_bootstrapper.py | 43 ++++++++++++++++++++---- src/ffpuppet/test_ffpuppet.py | 22 ++++++++++++ src/ffpuppet/test_main.py | 6 ++++ 7 files changed, 151 insertions(+), 17 deletions(-) diff --git a/src/ffpuppet/bootstrapper.py b/src/ffpuppet/bootstrapper.py index c4eddca..639ec34 100644 --- a/src/ffpuppet/bootstrapper.py +++ b/src/ffpuppet/bootstrapper.py @@ -8,14 +8,16 @@ from select import select from socket import SO_REUSEADDR, SOL_SOCKET, socket from time import perf_counter, sleep -from typing import Callable - -from .exceptions import BrowserTerminatedError, BrowserTimeoutError, LaunchError +from typing import TYPE_CHECKING, Callable # as of python 3.10 socket.timeout was made an alias of TimeoutError # pylint: disable=ungrouped-imports,wrong-import-order from socket import timeout as socket_timeout # isort: skip +from .exceptions import BrowserTerminatedError, BrowserTimeoutError, LaunchError + +if TYPE_CHECKING: + from collections.abc import Iterable LOG = getLogger(__name__) @@ -63,6 +65,18 @@ def __enter__(self) -> Bootstrapper: def __exit__(self, *exc: object) -> None: self.close() + @classmethod + def check_port(cls, value: int) -> bool: + """Verify port value is in valid range. + + Args: + None + + Returns: + bool + """ + return value == 0 or 1024 <= value <= 65535 + def close(self) -> None: """Close listening socket. @@ -85,9 +99,35 @@ def create(cls, attempts: int = 50, port: int = 0) -> Bootstrapper: Returns: Bootstrapper. """ + sock = cls.create_socket(attempts=attempts, port=port) + if sock is None: + raise LaunchError("Could not find available port") + return cls(sock) + + @classmethod + def create_socket( + cls, + attempts: int = 50, + blocked: Iterable[int] | None = BLOCKED_PORTS, + port: int = 0, + ) -> socket | None: + """Create a listening socket. + + Args: + attempts: Number of times to attempt to bind. + blocked: Ports that cannot be used. + port: Port to use. Use 0 for system select. + + Returns: + A listening socket. + """ assert attempts > 0 - assert port == 0 or 1024 <= port <= 65535 - assert port not in cls.BLOCKED_PORTS + if not cls.check_port(port): + LOG.debug("requested invalid port: %d", port) + return None + if blocked and port in blocked: + LOG.debug("requested blocked port: %d", port) + return None for _ in range(attempts): sock = socket() sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1) @@ -100,14 +140,14 @@ def create(cls, attempts: int = 50, port: int = 0) -> Bootstrapper: sleep(0.1) continue # avoid blocked ports - if sock.getsockname()[1] in cls.BLOCKED_PORTS: + if blocked and sock.getsockname()[1] in blocked: LOG.debug("bound to blocked port, retrying...") sock.close() continue break else: - raise LaunchError("Could not find available port") - return cls(sock) + return None + return sock @property def location(self) -> str: diff --git a/src/ffpuppet/core.py b/src/ffpuppet/core.py index 3f2f998..fc19ca8 100644 --- a/src/ffpuppet/core.py +++ b/src/ffpuppet/core.py @@ -109,8 +109,9 @@ class FFPuppet: "_logs", "_proc_tree", "_profile_template", - "_xvfb", "_working_path", + "_xvfb", + "marionette", "profile", "reason", ) @@ -136,6 +137,7 @@ def __init__( self._profile_template = use_profile self._xvfb: Xvfb | None = None self._working_path = working_path + self.marionette: int | None = None self.profile: Profile | None = None self.reason: Reason | None = Reason.CLOSED @@ -583,6 +585,7 @@ def close(self, force_close: bool = False) -> None: # reset remaining to closed state try: + self.marionette = None self._proc_tree = None self._logs.close() self._checks = [] @@ -677,6 +680,7 @@ def launch( launch_timeout: int = 300, location: str | None = None, log_limit: int = 0, + marionette: int | None = None, memory_limit: int = 0, prefs_js: Path | None = None, extension: list[Path] | None = None, @@ -764,7 +768,6 @@ def launch( "network.proxy.failover_direct": "false", "privacy.partition.network_state": "false", } - self.profile.add_prefs(prefs) launch_args = [bootstrapper.location] is_windows = system() == "Windows" @@ -772,6 +775,23 @@ def launch( # disable launcher process launch_args.append("-no-deelevate") launch_args.append("-wait-for-browser") + + if marionette is not None: + # find/validate port to use + free_sock = Bootstrapper.create_socket(port=marionette) + if free_sock is None: + if marionette == 0: + LOG.error("Cannot find available port for marionette") + else: + LOG.error("Marionette cannot use port: %d", marionette) + raise LaunchError("Debugging server port unavailable") + self.marionette = free_sock.getsockname()[1] + free_sock.close() + launch_args.append("-marionette") + prefs["marionette.port"] = str(self.marionette) + + self.profile.add_prefs(prefs) + cmd = self.build_launch_cmd(str(bin_path), additional_args=launch_args) # open logs @@ -830,6 +850,7 @@ def launch( if self._proc_tree is None: # only clean up here if a launch was not attempted or Popen failed LOG.debug("process not launched") + self.marionette = None self.profile.remove() self.profile = None self.reason = Reason.CLOSED diff --git a/src/ffpuppet/main.py b/src/ffpuppet/main.py index d740739..2886fc0 100644 --- a/src/ffpuppet/main.py +++ b/src/ffpuppet/main.py @@ -13,6 +13,7 @@ from tempfile import mkdtemp from time import sleep, strftime +from .bootstrapper import Bootstrapper from .core import Debugger, FFPuppet, Reason from .exceptions import BrowserExecutionError from .helpers import certutil_available, certutil_find @@ -134,6 +135,15 @@ def parse_args(argv: list[str] | None = None) -> Namespace: nargs="?", help="Headless mode. 'default' uses browser's built-in headless mode.", ) + cfg_group.add_argument( + "--marionette", + const=0, + default=None, + nargs="?", + type=int, + help="Enable marionette. If a port is provided it is used otherwise " + "a random port is selected. (default: disabled)", + ) cfg_group.add_argument( "-p", "--prefs", @@ -267,6 +277,8 @@ def parse_args(argv: list[str] | None = None) -> Namespace: value = int(Path(settings).read_bytes()) if value > 1: parser.error(f"rr needs {settings} <= 1, but it is {value}") + if args.marionette is not None and not Bootstrapper.check_port(args.marionette): + parser.error("--marionette must be 0 or > 1024 and < 65536") if not args.logs.is_dir(): parser.error(f"Log output directory is invalid '{args.logs}'") args.log_level = log_level_map[args.log_level] @@ -317,6 +329,7 @@ def main(argv: list[str] | None = None) -> None: location=args.url, launch_timeout=args.launch_timeout, log_limit=args.log_limit, + marionette=args.marionette, memory_limit=args.memory, prefs_js=args.prefs, extension=args.extension, @@ -326,6 +339,8 @@ def main(argv: list[str] | None = None) -> None: assert ffp.profile is not None assert ffp.profile.path is not None Profile.check_prefs(ffp.profile.path / "prefs.js", args.prefs) + if ffp.marionette is not None: + LOG.info("Marionette listening on port: %d", ffp.marionette) LOG.info("Running Firefox (pid: %d)...", ffp.get_pid()) while ffp.is_healthy(): sleep(args.poll_interval) diff --git a/src/ffpuppet/resources/testff.py b/src/ffpuppet/resources/testff.py index 7890611..4e60da0 100755 --- a/src/ffpuppet/resources/testff.py +++ b/src/ffpuppet/resources/testff.py @@ -29,6 +29,7 @@ def main() -> int: parser = ArgumentParser(prog="testff", description="Fake Firefox for testing") parser.add_argument("url") parser.add_argument("-headless", action="store_true", help="ignored") + parser.add_argument("-marionette", nargs="?", type=int, help="ignored") parser.add_argument("-new-instance", action="store_true", help="ignored") parser.add_argument("-no-deelevate", action="store_true", help="ignored") parser.add_argument("-wait-for-browser", action="store_true", help="ignored") diff --git a/src/ffpuppet/test_bootstrapper.py b/src/ffpuppet/test_bootstrapper.py index 5e95f6d..28ac02c 100644 --- a/src/ffpuppet/test_bootstrapper.py +++ b/src/ffpuppet/test_bootstrapper.py @@ -184,23 +184,22 @@ def _fake_browser(port, payload_size=5120): @mark.parametrize( - "bind, attempts, raised", + "bind, attempts", [ # failed to bind (OSError) - ((OSError(0, "foo1"),), 1, LaunchError), + ((OSError(0, "foo1"),), 1), # failed to bind (PermissionError) - multiple attempts - (repeat(PermissionError(10013, "foo2"), 4), 4, LaunchError), + (repeat(PermissionError(10013, "foo2"), 4), 4), ], ) -def test_bootstrapper_08(mocker, bind, attempts, raised): - """test Bootstrapper.create() - failures""" +def test_bootstrapper_08(mocker, bind, attempts): + """test Bootstrapper.create_socket() - failures""" mocker.patch("ffpuppet.bootstrapper.sleep", autospec=True) fake_sock = mocker.MagicMock(spec_set=socket) fake_sock.bind.side_effect = bind mocker.patch("ffpuppet.bootstrapper.select", return_value=([fake_sock], None, None)) mocker.patch("ffpuppet.bootstrapper.socket", return_value=fake_sock) - with raises(raised), Bootstrapper.create(attempts=attempts): - pass + assert Bootstrapper.create_socket(attempts=attempts) is None assert fake_sock.bind.call_count == attempts assert fake_sock.close.call_count == attempts @@ -216,3 +215,33 @@ def test_bootstrapper_09(mocker): with Bootstrapper.create(attempts=2): pass assert fake_sock.close.call_count == 2 + + +def test_bootstrapper_10(mocker): + """test Bootstrapper.create() - failure""" + mocker.patch("ffpuppet.bootstrapper.Bootstrapper.create_socket", return_value=None) + with raises(LaunchError), Bootstrapper.create(): + pass + + +@mark.parametrize("value", [123, 5555]) +def test_bootstrapper_11(value): + """test Bootstrapper.create_socket() - unusable ports""" + assert Bootstrapper.create_socket(blocked=[5555], port=value) is None + + +@mark.parametrize( + "value, result", + [ + (0, True), + (1337, True), + (32768, True), + (-1, False), + (1, False), + (1023, False), + (65536, False), + ], +) +def test_bootstrapper_12(value, result): + """test Bootstrapper.check_port()""" + assert Bootstrapper.check_port(value) == result diff --git a/src/ffpuppet/test_ffpuppet.py b/src/ffpuppet/test_ffpuppet.py index aebbde3..444c9d4 100644 --- a/src/ffpuppet/test_ffpuppet.py +++ b/src/ffpuppet/test_ffpuppet.py @@ -963,3 +963,25 @@ def test_ffpuppet_32(mocker): with raises(NotImplementedError): ffp.dump_coverage() ffp._proc_tree = None + + +def test_ffpuppet_33(): + """test FFPuppet.launch() with marionette""" + with FFPuppet() as ffp, HTTPTestServer() as srv: + assert ffp.marionette is None + ffp.launch(TESTFF_BIN, marionette=0, location=srv.get_addr()) + assert ffp.marionette is not None + ffp.close() + assert ffp.marionette is None + + +@mark.parametrize("port", [0, 123]) +def test_ffpuppet_34(mocker, port): + """test FFPuppet.launch() with marionette failure""" + fake_bts = mocker.patch("ffpuppet.core.Bootstrapper", autospec=True) + fake_bts.create.return_value.location = "" + fake_bts.create_socket.return_value = None + with FFPuppet() as ffp, HTTPTestServer() as srv: + with raises(LaunchError): + ffp.launch(TESTFF_BIN, marionette=port, location=srv.get_addr()) + assert ffp.marionette is None diff --git a/src/ffpuppet/test_main.py b/src/ffpuppet/test_main.py index 50db908..617e88c 100644 --- a/src/ffpuppet/test_main.py +++ b/src/ffpuppet/test_main.py @@ -72,6 +72,12 @@ def test_parse_args_01(capsys, mocker, tmp_path): with raises(SystemExit): parse_args([str(fake_bin), "--log-limit", "-1"]) assert "error: --log-limit must be >= 0" in capsys.readouterr()[-1] + # invalid marionette port + with raises(SystemExit): + parse_args([str(fake_bin), "--marionette", "123"]) + assert ( + "error: --marionette must be 0 or > 1024 and < 65536" in capsys.readouterr()[-1] + ) # invalid memory limit with raises(SystemExit): parse_args([str(fake_bin), "--memory", "-1"])