From d08e78b89d293c825ee524da268fd997fe93c524 Mon Sep 17 00:00:00 2001 From: Tyson Smith Date: Mon, 13 Nov 2023 09:20:09 -0800 Subject: [PATCH] Use Path objects in AssetManager --- grizzly/common/storage.py | 6 +- grizzly/common/test_storage.py | 13 ++-- grizzly/reduce/test_reduce.py | 2 +- grizzly/replay/replay.py | 7 +- grizzly/replay/test_replay.py | 2 +- grizzly/target/assets.py | 93 ++++++++------------------ grizzly/target/puppet_target.py | 8 +-- grizzly/target/test_assets.py | 97 ++++++++++------------------ grizzly/target/test_puppet_target.py | 34 +++++----- 9 files changed, 96 insertions(+), 166 deletions(-) diff --git a/grizzly/common/storage.py b/grizzly/common/storage.py index 18744bee..5be8daaf 100644 --- a/grizzly/common/storage.py +++ b/grizzly/common/storage.py @@ -175,7 +175,7 @@ def clone(self): result.assets = dict(self.assets) if result.assets: assert self.assets_path - org_path = Path(self.assets_path) + org_path = self.assets_path try: # copy asset data from test case result.assets_path = result.root / org_path.relative_to(self.root) @@ -273,7 +273,7 @@ def dump(self, dst_path, include_details=False): # save target assets and update meta data if self.assets: assert isinstance(self.assets, dict) - assert isinstance(self.assets_path, (Path, str)) + assert isinstance(self.assets_path, Path) info["assets"] = self.assets info["assets_path"] = "_assets_" copytree(self.assets_path, dst_path / info["assets_path"]) @@ -439,7 +439,7 @@ def load_single(cls, path, adjacent=False, copy=True): ) # load all adjacent data from directory if adjacent: - for entry in Path(entry_point.parent).rglob("*"): + for entry in entry_point.parent.rglob("*"): if not entry.is_file(): continue # ignore asset path diff --git a/grizzly/common/test_storage.py b/grizzly/common/test_storage.py index 415d421b..602b5bab 100644 --- a/grizzly/common/test_storage.py +++ b/grizzly/common/test_storage.py @@ -5,7 +5,6 @@ from itertools import chain from json import dumps, loads -from pathlib import Path from zipfile import ZIP_DEFLATED, ZipFile from pytest import mark, raises @@ -219,8 +218,8 @@ def test_testcase_08(tmp_path): (tmp_path / "src" / "nested" / "empty").mkdir() dst_dir = tmp_path / "dst" # build test case - with AssetManager(base_path=str(tmp_path)) as assets: - assets.add("example", str(asset_file)) + with AssetManager(base_path=tmp_path) as assets: + assets.add("example", asset_file) with TestCase("target.bin", "test-adapter") as src: src.env_vars["TEST_ENV_VAR"] = "100" src.add_from_file(entry_point) @@ -238,7 +237,7 @@ def test_testcase_08(tmp_path): with TestCase.load_single(dst_dir) as dst: assert "example" in dst.assets assert dst.assets_path is not None - assert (Path(dst.assets_path) / "asset.bin").is_file() + assert (dst.assets_path / "asset.bin").is_file() assert "_assets_/asset.bin" not in (x.file_name for x in dst._files.optional) assert dst.entry_point == "target.bin" assert "target.bin" in (x.file_name for x in dst._files.required) @@ -293,7 +292,7 @@ def test_testcase_10(tmp_path): org.time_limit = 10 org.add_from_bytes(b"a", "a.html") org.assets = {"sample": asset.name} - org.assets_path = str(asset_path) + org.assets_path = asset_path org.dump(working, include_details=True) assert (working / "_assets_" / asset.name).is_file() with TestCase.load_single(working, adjacent=False) as loaded: @@ -335,8 +334,8 @@ def test_testcase_14(tmp_path): nested.mkdir() asset_file = tmp_path / "example_asset" asset_file.touch() - with AssetManager(base_path=str(tmp_path)) as assets: - assets.add("example", str(asset_file)) + with AssetManager(base_path=tmp_path) as assets: + assets.add("example", asset_file) with TestCase("target.bin", "test-adapter") as src: src.assets = assets.assets src.assets_path = assets.path diff --git a/grizzly/reduce/test_reduce.py b/grizzly/reduce/test_reduce.py index 62193120..4830f7d8 100644 --- a/grizzly/reduce/test_reduce.py +++ b/grizzly/reduce/test_reduce.py @@ -857,7 +857,7 @@ def submit(test_cases, report, force=False): target.filtered_environ.return_value = {"test": "abc"} with AssetManager(base_path=str(tmp_path)) as assets: (tmp_path / "example_asset").touch() - assets.add("example", str(tmp_path / "example_asset"), copy=False) + assets.add("example", tmp_path / "example_asset", copy=False) target.assets = assets try: mgr = ReduceManager( diff --git a/grizzly/replay/replay.py b/grizzly/replay/replay.py index f11295c8..2585cd22 100644 --- a/grizzly/replay/replay.py +++ b/grizzly/replay/replay.py @@ -179,9 +179,7 @@ def load_testcases(cls, path, subset=None): env_vars = None for test in testcases: if assets is None and test.assets and test.assets_path: - assets = AssetManager.load( - test.assets, str(test.root / test.assets_path) - ) + assets = AssetManager.load(test.assets, test.root / test.assets_path) if not env_vars and test.env_vars: env_vars = dict(test.env_vars) test.env_vars.clear() @@ -658,8 +656,7 @@ def main(cls, args): LOG.debug("adding environment loaded from test case") target.merge_environment(env_vars) - # TODO: support overriding existing assets - # prioritize specified assets over included + # TODO: prioritize specified assets over included target.assets.add_batch(args.asset) target.process_assets() diff --git a/grizzly/replay/test_replay.py b/grizzly/replay/test_replay.py index 7a38ad1a..e44bd496 100644 --- a/grizzly/replay/test_replay.py +++ b/grizzly/replay/test_replay.py @@ -637,7 +637,7 @@ def test_replay_21(tmp_path): # build test case with AssetManager() as assets: (tmp_path / "prefs.js").touch() - assets.add("prefs", str(tmp_path / "prefs.js"), copy=False) + assets.add("prefs", tmp_path / "prefs.js", copy=False) with TestCase(data.name, "foo") as src: src.env_vars = {"foo": "bar"} src.assets = dict(assets.assets) diff --git a/grizzly/target/assets.py b/grizzly/target/assets.py index 02e953c8..341a74b0 100644 --- a/grizzly/target/assets.py +++ b/grizzly/target/assets.py @@ -2,9 +2,7 @@ # 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/. from logging import getLogger -from os import makedirs, unlink -from os.path import abspath, basename, exists, isdir, isfile -from os.path import join as pathjoin +from pathlib import Path from shutil import copyfile, copytree, move, rmtree from tempfile import mkdtemp @@ -19,13 +17,12 @@ class AssetError(Exception): """Raised by AssetManager""" -# TODO: use pathlib class AssetManager: __slots__ = ("assets", "path") def __init__(self, base_path=None): self.assets = {} - self.path = mkdtemp(prefix="assets_", dir=base_path) + self.path = Path(mkdtemp(prefix="assets_", dir=base_path)) def __enter__(self): return self @@ -38,41 +35,36 @@ def add(self, asset, path, copy=True): Args: asset (str): Name of asset. - path (str): Location on disk. + path (Path): Location on disk. copy (bool): Copy or move the content. Returns: str: Path to the asset on the filesystem. """ assert isinstance(asset, str) - assert isinstance(path, str) + assert isinstance(path, Path) assert self.path, "cleanup() was called" - if not exists(path): - raise OSError(f"{path!r} does not exist") - path = abspath(path) - # only copy files from outside working path - if path.startswith(self.path): - raise AssetError(f"Cannot add existing asset content {path!r}") - dst_path = pathjoin(self.path, basename(path)) - # avoid overwriting data that is part of an existing asset - if exists(dst_path): - raise AssetError(f"{basename(path)!r} is an existing asset") + if not path.exists(): + raise OSError(f"'{path}' does not exist") + dst = self.path / path.name # remove existing asset with the same name if asset in self.assets: LOG.debug("asset %r exists, removing existing", asset) self.remove(asset) + # avoid overwriting data that is part of an existing asset + if dst.exists(): + raise AssetError(f"{asset}: '{path.name}' already exists") if copy: - if isfile(path): - copyfile(path, dst_path) + if path.is_file(): + copyfile(path, dst) else: - copytree(path, dst_path) + copytree(path, dst) else: - move(path, self.path) - self.assets[asset] = basename(path) - LOG.debug( - "added asset %r %s to %r", asset, "copied" if copy else "moved", dst_path - ) - return dst_path + # TODO: move() only accepts str in Python 3.8 + move(str(path), str(self.path)) + self.assets[asset] = path.name + LOG.debug("%s asset %r to '%s'", "copied" if copy else "moved", asset, dst) + return dst def add_batch(self, assets): """Add collection of assets to the AssetManager. @@ -84,7 +76,7 @@ def add_batch(self, assets): None """ for asset, path in assets: - self.add(asset, path) + self.add(asset, Path(path)) def cleanup(self): """Remove asset files from filesystem. @@ -100,33 +92,6 @@ def cleanup(self): self.assets.clear() self.path = None - def dump(self, dst_path, subdir="_assets_"): - """Copy assets content to a given path. - - Args: - dst_path (str): Path to copy assets content to. - subdir (str): Create and use as destination if given. - - Returns: - dict: Collection asset paths keyed by asset name. - """ - dumped = {} - if self.assets: - if subdir: - dst_path = pathjoin(dst_path, subdir) - makedirs(dst_path, exist_ok=not subdir) - for name, path in self.assets.items(): - dumped[name] = path - src = pathjoin(self.path, path) - if isfile(src): - copyfile(src, pathjoin(dst_path, path)) - elif isdir(src): - copytree(src, pathjoin(dst_path, path)) - else: - dumped.pop(name) - LOG.warning("Failed to dump asset %r from %r", name, src) - return dumped - def get(self, asset): """Get path to content on filesystem for given asset. @@ -134,12 +99,10 @@ def get(self, asset): asset (str): Asset to lookup. Returns: - str: Path to asset content or None if asset does not exist. + Path: Path to asset content or None if asset does not exist. """ item = self.assets.get(asset, None) - if item is not None: - return pathjoin(self.path, item) - return None + return self.path / item if item else None def is_empty(self): """Check if AssetManager contains entries. @@ -159,7 +122,7 @@ def load(cls, assets, src_path, base_path=None): Args: asset (dict): Asset paths on filesystem relative to src_path, keyed on asset name. - src_path (str): Path to scan for assets. + src_path (Path): Path to scan for assets. base_path (str): Base path to use to create local storage. Returns: @@ -167,7 +130,7 @@ def load(cls, assets, src_path, base_path=None): """ obj = cls(base_path=base_path) for asset, src_name in assets.items(): - obj.add(asset, pathjoin(src_path, src_name)) + obj.add(asset, src_path / src_name) return obj def remove(self, asset): @@ -179,10 +142,10 @@ def remove(self, asset): Returns: None """ - path = self.assets.pop(asset, None) - if path: - path = pathjoin(self.path, path) - if isfile(path): - unlink(path) + local_path = self.assets.pop(asset, None) + if local_path: + path = self.path / local_path + if path.is_file(): + path.unlink() else: rmtree(path, ignore_errors=True) diff --git a/grizzly/target/puppet_target.py b/grizzly/target/puppet_target.py index 365fd713..14426fd5 100644 --- a/grizzly/target/puppet_target.py +++ b/grizzly/target/puppet_target.py @@ -383,11 +383,11 @@ def process_assets(self): prefs = Path(tmp_path) / "prefs.js" template = PrefPicker.lookup_template("browser-fuzzing.yml") PrefPicker.load_template(template).create_prefsjs(prefs) - self._prefs = self.assets.add("prefs", str(prefs), copy=False) + self._prefs = self.assets.add("prefs", prefs, copy=False) abort_tokens = self.assets.get("abort-tokens") if abort_tokens: LOG.debug("loading 'abort tokens' from %r", abort_tokens) - with (Path(self.assets.path) / abort_tokens).open() as in_fp: + with (self.assets.path / abort_tokens).open() as in_fp: for line in in_fp: line = line.strip() if line: @@ -402,7 +402,7 @@ def process_assets(self): opts.load_options(self.environ.get(var_name, "")) if self.assets.get(asset): # use suppression file if provided as asset - opts.add("suppressions", repr(self.assets.get(asset)), overwrite=True) + opts.add("suppressions", f"'{self.assets.get(asset)}'", overwrite=True) elif opts.get("suppressions"): supp_file = opts.pop("suppressions") if SanitizerOptions.is_quoted(supp_file): @@ -412,7 +412,7 @@ def process_assets(self): LOG.debug("using %r from environment", asset) opts.add( "suppressions", - repr(self.assets.add(asset, supp_file)), + f"'{self.assets.add(asset, Path(supp_file))}'", overwrite=True, ) else: diff --git a/grizzly/target/test_assets.py b/grizzly/target/test_assets.py index 265ebd9a..dd1f457a 100644 --- a/grizzly/target/test_assets.py +++ b/grizzly/target/test_assets.py @@ -11,30 +11,25 @@ def test_asset_manager_01(tmp_path): """test AssetManager()""" - with AssetManager(base_path=str(tmp_path)) as assets: + with AssetManager(base_path=tmp_path) as assets: assert not assets.assets assert assets.path assert assets.is_empty() # add file (move) example = tmp_path / "example.txt" example.write_text("foo") - asset_location = assets.add("example_file", str(example), copy=False) + asset_location = assets.add("example_file", example, copy=False) assert assets.assets["example_file"] == "example.txt" - assert str(Path(assets.path) / "example.txt") == asset_location + assert assets.path / "example.txt" == asset_location assert assets.get("example_file") == asset_location assert len(assets.assets) == 1 assert not assets.is_empty() assert not example.is_file() - assert (Path(assets.path) / "example.txt").is_file() - # add file - file name collision - example.write_text("foo") - with raises(AssetError, match="'example.txt' is an existing asset"): - assets.add("example_file", str(example)) - assert len(assets.assets) == 1 + assert (assets.path / "example.txt").is_file() # add existing asset - update asset (copy) example = tmp_path / "example_2.txt" example.write_text("foo") - assets.add("example_file", str(example)) + assets.add("example_file", example) assert example.is_file() assert len(assets.assets) == 1 # add directory @@ -44,23 +39,23 @@ def test_asset_manager_01(tmp_path): (example / "a" / "1.txt").write_text("1") (example / "b").mkdir() (example / "b" / "2.txt").write_text("2") - assets.add("example_path", str(example)) + assets.add("example_path", example) rmtree(example) - assert (Path(assets.path) / "example_path/a/1.txt").is_file() - assert (Path(assets.path) / "example_path/b/2.txt").is_file() + assert (assets.path / "example_path/a/1.txt").is_file() + assert (assets.path / "example_path/b/2.txt").is_file() assert len(assets.assets) == 2 # get - assert Path(assets.path) / "example_2.txt" == Path(assets.get("example_file")) - assert Path(assets.path) / "example_path" == Path(assets.get("example_path")) + assert (assets.path / "example_2.txt").samefile(assets.get("example_file")) + assert (assets.path / "example_path").samefile(assets.get("example_path")) # remove directory assets.remove("example_path") assert len(assets.assets) == 1 - assert Path(assets.path).is_dir() - assert not (Path(assets.path) / "example_path").is_dir() + assert assets.path.is_dir() + assert not (assets.path / "example_path").is_dir() # remove file assets.remove("example_file") assert len(assets.assets) == 0 - assert not any(Path(assets.path).iterdir()) + assert not any(assets.path.iterdir()) # cleanup assets.cleanup() assert not assets.assets @@ -74,59 +69,33 @@ def test_asset_manager_02(tmp_path): assert assets.get("missing") is None # add missing file with raises(OSError, match="'missing' does not exist"): - assets.add("a", "missing") + assets.add("a", Path("missing")) assert not assets.assets # remove invalid asset assets.remove("missing") # add file example = tmp_path / "example.txt" - example.write_text("example") - assets.add("example_file", str(example), copy=True) - # add existing file under as different asset (collision) - with raises(AssetError, match="'example.txt' is an existing asset"): - assets.add("collide", str(example)) + example.touch() + assets.add("example_file", example, copy=True) + # add with asset with name and file collision + assets.add("example_file", example, copy=True) + # add with file name collision as different asset + with raises(AssetError, match="collide: 'example.txt' already exists"): + assets.add("collide", example) assert "collide" not in assets.assets - example.unlink() - # add file from existing asset - example = Path(assets.path) / "example.txt" - with raises(AssetError, match="Cannot add existing asset content"): - assets.add("existing", str(example)) - assert "existing" not in assets.assets - # add file from existing asset with asset name collision - example = Path(assets.path) / "example.txt" - with raises(AssetError, match="Cannot add existing asset content"): - assets.add("example_file", str(example)) - assert "example_file" in assets.assets def test_asset_manager_03(tmp_path): - """test AssetManager() dump/load""" - # test dump() - with AssetManager(base_path=str(tmp_path)) as assets: - # add file - example = tmp_path / "example.txt" - example.write_text("example") - assets.add("example_file", str(example), copy=False) - # add directory - example = tmp_path / "example_path" - example.mkdir() - (example / "a").mkdir() - (example / "a" / "1.txt").write_text("1") - assets.add("example_path", str(example)) - # invalid entry - assets.assets["invalid"] = "bad/path" - # dump - dump_path = tmp_path / "dump" - dumped = assets.dump(str(dump_path), subdir="sub") - assert len(dumped) == 2 - assert (dump_path / "sub" / "example.txt").is_file() - assert (dump_path / "sub" / "example_path" / "a" / "1.txt").is_file() - # test load() - with AssetManager.load( - dumped, str(dump_path / "sub"), base_path=str(tmp_path) - ) as assets: - assert len(assets.assets) == 2 - assert (Path(assets.path) / assets.assets["example_file"]).is_file() + """test AssetManager() load""" + src = tmp_path / "src" + src.mkdir() + (src / "b.txt").touch() + + with AssetManager.load({"a": "b.txt"}, src, base_path=str(tmp_path)) as assets: + assert len(assets.assets) == 1 + assert "a" in assets.assets + assert assets.assets["a"] == "b.txt" + assert (assets.path / assets.assets["a"]).is_file() def test_asset_manager_04(tmp_path): @@ -136,12 +105,12 @@ def test_asset_manager_04(tmp_path): # add file example = tmp_path / "example.txt" example.write_text("example") - batch.append(["example_file", str(example)]) + batch.append(["example_file", example]) # add directory example = tmp_path / "example_path" example.mkdir() (example / "a").mkdir() (example / "a" / "1.txt").write_text("1") - batch.append(["example_path", str(example)]) + batch.append(["example_path", example]) assets.add_batch(batch) assert len(assets.assets) == 2 diff --git a/grizzly/target/test_puppet_target.py b/grizzly/target/test_puppet_target.py index fc725792..298d4c4b 100644 --- a/grizzly/target/test_puppet_target.py +++ b/grizzly/target/test_puppet_target.py @@ -3,7 +3,6 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. # pylint: disable=protected-access from itertools import count -from pathlib import Path from platform import system from ffpuppet import BrowserTerminatedError, BrowserTimeoutError, Debugger, Reason @@ -304,28 +303,26 @@ def test_puppet_target_08(mocker, tmp_path): with PuppetTarget(fake_file, 300, 25, 5000) as target: assert target.assets.get("prefs") is None target.process_assets() - assert (Path(target.assets.path) / target.assets.get("prefs")).is_file() - assert target.assets.get("prefs").endswith("prefs.js") + assert target.assets.get("prefs").is_file() + assert target.assets.get("prefs").name == "prefs.js" # prefs file provided with AssetManager(base_path=str(tmp_path)) as assets: - assets.add("prefs", str(fake_file)) + assets.add("prefs", fake_file) with PuppetTarget(fake_file, 300, 25, 5000, assets=assets) as target: target.process_assets() - assert (Path(target.assets.path) / target.assets.get("prefs")).is_file() - assert target.assets.get("prefs").endswith("fake") + assert target.assets.get("prefs").is_file() + assert target.assets.get("prefs").name == "fake" # abort tokens file provided with AssetManager(base_path=str(tmp_path)) as assets: - assets.add("abort-tokens", str(fake_file)) + assets.add("abort-tokens", fake_file) with PuppetTarget(fake_file, 300, 25, 5000, assets=assets) as target: # ignore E1101: (pylint 2.9.3 bug?) # Method 'add_abort_token' has no 'call_count' member (no-member) # pylint: disable=no-member assert target._puppet.add_abort_token.call_count == 0 target.process_assets() - assert ( - Path(target.assets.path) / target.assets.get("abort-tokens") - ).is_file() - assert target.assets.get("abort-tokens").endswith("fake") + assert target.assets.get("abort-tokens").is_file() + assert target.assets.get("abort-tokens").name == "fake" assert target._puppet.add_abort_token.call_count == 2 @@ -383,10 +380,10 @@ def test_puppet_target_10(tmp_path, asset, env): supp_asset = tmp_path / "supp_asset" supp_env = tmp_path / "supp_env" with AssetManager(base_path=str(tmp_path)) as assets: - assets.add("prefs", str(fake_file)) + assets.add("prefs", fake_file) if asset: supp_asset.touch() - assets.add("lsan-suppressions", str(supp_asset)) + assets.add("lsan-suppressions", supp_asset) with PuppetTarget(fake_file, 300, 25, 5000, assets=assets) as target: target.environ["TSAN_OPTIONS"] = "a=1" if env: @@ -396,10 +393,15 @@ def test_puppet_target_10(tmp_path, asset, env): target.environ["LSAN_OPTIONS"] = "suppressions='missing'" target.process_assets() if asset: - assert supp_asset.name in target.environ["LSAN_OPTIONS"] + assert ( + f"suppressions='{target.assets.path / supp_asset.name}'" + in target.environ["LSAN_OPTIONS"] + ) elif env: - assert supp_env.name in target.environ["LSAN_OPTIONS"] - assert assets.get("lsan-suppressions") + assert ( + f"suppressions='{target.assets.path / supp_env.name}'" + in target.environ["LSAN_OPTIONS"] + ) else: assert not assets.get("lsan-suppressions")