diff --git a/src/fuzzfetch/extract.py b/src/fuzzfetch/extract.py index a456d141..b86e89a1 100644 --- a/src/fuzzfetch/extract.py +++ b/src/fuzzfetch/extract.py @@ -35,8 +35,8 @@ def extract_zip(zip_fn: PathArg, path: PathArg = ".") -> None: """ dest_path = Path(path) - def _extract_file(zip_fp: zipfile.ZipFile, info: zipfile.ZipInfo) -> None: - """Extract files while explicitly setting the proper permissions""" + def _extract_entry(zip_fp: zipfile.ZipFile, info: zipfile.ZipInfo) -> None: + """Extract entries while explicitly setting the proper permissions""" rel_path = Path(info.filename) # strip leading "firefox" from path @@ -47,9 +47,12 @@ def _extract_file(zip_fp: zipfile.ZipFile, info: zipfile.ZipInfo) -> None: out_path = dest_path / rel_path - out_path.parent.mkdir(parents=True, exist_ok=True) - with zip_fp.open(info) as zip_member_fp, out_path.open("wb") as out_fp: - shutil.copyfileobj(zip_member_fp, out_fp) + if info.is_dir(): + out_path.mkdir(parents=True, exist_ok=True) + else: + out_path.parent.mkdir(parents=True, exist_ok=True) + with zip_fp.open(info) as zip_member_fp, out_path.open("wb") as out_fp: + shutil.copyfileobj(zip_member_fp, out_fp) perm = info.external_attr >> 16 perm |= stat.S_IREAD # make sure we're not accidentally setting this to 0 @@ -57,7 +60,7 @@ def _extract_file(zip_fp: zipfile.ZipFile, info: zipfile.ZipInfo) -> None: with zipfile.ZipFile(zip_fn) as zip_fp: for info in zip_fp.infolist(): - _extract_file(zip_fp, info) + _extract_entry(zip_fp, info) def _is_within_directory(directory: PathArg, target: PathArg) -> bool: diff --git a/tests/test_extract.py b/tests/test_extract.py index 562b0cdb..47b54429 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -4,11 +4,33 @@ # You can obtain one at http://mozilla.org/MPL/2.0/. import tarfile +import zipfile from unittest.mock import patch import pytest # pylint: disable=import-error -from fuzzfetch.extract import extract_tar +from fuzzfetch.extract import extract_tar, extract_zip + + +def test_zipfile_extract(tmp_path): + """basic extract_zip functions""" + (tmp_path / "empty").touch() + (tmp_path / "folder").mkdir() + with zipfile.ZipFile(tmp_path / "test.zip", "w") as zip_fp: + zip_fp.write(tmp_path / "empty", "./firefox/firefox") + zip_fp.write(tmp_path / "empty", "buildinfo.txt") + zip_fp.write(tmp_path / "folder", "folder") + + (tmp_path / "out").mkdir() + extract_zip(tmp_path / "test.zip", tmp_path / "out") + assert set((tmp_path / "out").glob("**/*")) == { + tmp_path / "out" / "firefox", + tmp_path / "out" / "folder", + tmp_path / "out" / "buildinfo.txt", + } + assert (tmp_path / "out" / "firefox").is_file() + assert (tmp_path / "out" / "folder").is_dir() + assert (tmp_path / "out" / "buildinfo.txt").is_file() @patch("fuzzfetch.extract.TAR_PATH", None)