diff --git a/hfutils/archive/base.py b/hfutils/archive/base.py index 2b644f7dac..be2cfd3be2 100644 --- a/hfutils/archive/base.py +++ b/hfutils/archive/base.py @@ -1,6 +1,6 @@ import os.path import warnings -from typing import List, Dict, Tuple, Callable +from typing import List, Dict, Tuple, Callable, Optional _KNOWN_ARCHIVE_TYPES: Dict[str, Tuple[List[str], Callable, Callable]] = {} @@ -85,7 +85,7 @@ def get_archive_type(archive_file: str) -> str: raise ValueError(f'Unknown type of archive file {archive_file!r}.') -def archive_unpack(archive_file: str, directory: str, silent: bool = False): +def archive_unpack(archive_file: str, directory: str, silent: bool = False, password: Optional[str] = None): """ Unpack an archive file into a directory using the specified archive type. @@ -95,10 +95,12 @@ def archive_unpack(archive_file: str, directory: str, silent: bool = False): :type directory: str :param silent: If True, suppress warnings during the unpacking process. :type silent: bool + :param password: The password to extract the archive file. + :type password: str, optional :return: The path to the unpacked directory. :rtype: str """ type_name = get_archive_type(archive_file) _, _, fn_unpack = _KNOWN_ARCHIVE_TYPES[type_name] - return fn_unpack(archive_file, directory, silent=silent) + return fn_unpack(archive_file, directory, silent=silent, password=password) diff --git a/hfutils/archive/rar.py b/hfutils/archive/rar.py index fb4dba254e..674a0186b0 100644 --- a/hfutils/archive/rar.py +++ b/hfutils/archive/rar.py @@ -1,4 +1,5 @@ import os +from typing import Optional from .base import register_archive_type from ..utils import tqdm @@ -14,10 +15,12 @@ def _rar_pack(directory, zip_file, silent: bool = False): raise RuntimeError('RAR format packing is not supported.') -def _rar_unpack(rar_file, directory, silent: bool = False): +def _rar_unpack(rar_file, directory, silent: bool = False, password: Optional[str] = None): directory = os.fspath(directory) os.makedirs(directory, exist_ok=True) with rarfile.RarFile(rar_file, 'r') as zf: + if password is not None: + zf.setpassword(password) progress = tqdm(zf.namelist(), silent=silent, desc=f'Unpacking {directory!r} ...') for rarinfo in progress: progress.set_description(rarinfo) diff --git a/hfutils/archive/sevenz.py b/hfutils/archive/sevenz.py index 573287b731..6910c5b378 100644 --- a/hfutils/archive/sevenz.py +++ b/hfutils/archive/sevenz.py @@ -1,4 +1,5 @@ import os +from typing import Optional from .base import register_archive_type from ..utils import tqdm, walk_files @@ -17,10 +18,10 @@ def _7z_pack(directory, sz_file, silent: bool = False): zf.write(os.path.join(directory, file), file) -def _7z_unpack(sz_file, directory, silent: bool = False): +def _7z_unpack(sz_file, directory, silent: bool = False, password: Optional[str] = None): directory = os.fspath(directory) os.makedirs(directory, exist_ok=True) - with py7zr.SevenZipFile(sz_file, 'r') as zf: + with py7zr.SevenZipFile(sz_file, 'r', password=password) as zf: progress = tqdm(zf.getnames(), silent=silent, desc=f'Unpacking {directory!r} ...') for name in progress: progress.set_description(name) diff --git a/hfutils/archive/tar.py b/hfutils/archive/tar.py index 5e4bcf52f4..1400fc53c3 100644 --- a/hfutils/archive/tar.py +++ b/hfutils/archive/tar.py @@ -1,8 +1,9 @@ import copy import os.path import tarfile +import warnings from functools import partial -from typing import Literal +from typing import Literal, Optional from .base import register_archive_type from .zip import _ZLIB_SUPPORTED @@ -47,7 +48,10 @@ def _tarfile_pack(directory, tar_file, compress: CompressTyping = "gzip", silent tar.add(os.path.join(directory, file), file) -def _tarfile_unpack(tar_file, directory, silent: bool = False, numeric_owner=False): +def _tarfile_unpack(tar_file, directory, silent: bool = False, numeric_owner=False, password: Optional[str] = None): + if password is not None: + warnings.warn('Password is not supported in tar archive files.\n' + 'So assigned password will be ignored.') with tarfile.open(tar_file) as tar: directories = [] progress = tqdm(tar, silent=silent, desc=f'Unpacking {directory!r} ...') diff --git a/hfutils/archive/zip.py b/hfutils/archive/zip.py index 48c679007b..9b190c3654 100644 --- a/hfutils/archive/zip.py +++ b/hfutils/archive/zip.py @@ -1,5 +1,6 @@ import os.path import zipfile +from typing import Optional from .base import register_archive_type from ..utils import tqdm, walk_files @@ -21,10 +22,12 @@ def _zip_pack(directory, zip_file, silent: bool = False): zf.write(os.path.join(directory, file), file) -def _zip_unpack(zip_file, directory, silent: bool = False): +def _zip_unpack(zip_file, directory, silent: bool = False, password: Optional[str] = None): directory = os.fspath(directory) os.makedirs(directory, exist_ok=True) with zipfile.ZipFile(zip_file, 'r') as zf: + if password is not None: + zf.setpassword(password.encode(encoding='utf-8')) progress = tqdm(zf.namelist(), silent=silent, desc=f'Unpacking {directory!r} ...') for zipinfo in progress: progress.set_description(zipinfo) diff --git a/test/archive/test_rar.py b/test/archive/test_rar.py index 002ddeed1b..00348d6c92 100644 --- a/test/archive/test_rar.py +++ b/test/archive/test_rar.py @@ -18,6 +18,11 @@ def raw_rar(): return get_testfile('raw.rar') +@pytest.fixture() +def raw_password_rar(): + return get_testfile('raw-password.rar') + + @pytest.mark.unittest class TestArchiveRar: @skipUnless(rarfile, 'rarfile module required.') @@ -53,3 +58,16 @@ def test_archive_unpack(self, raw_rar, check_unpack_dir): with disable_output(): archive_unpack(raw_rar, '.') check_unpack_dir('.') + + @skipUnless(rarfile, 'rarfile module required.') + def test_archive_unpack_with_password_failed(self, raw_password_rar): + with isolated_directory(): + with disable_output(), pytest.raises(rarfile.PasswordRequired): + archive_unpack(raw_password_rar, '.') + + @skipUnless(rarfile, 'rarfile module required.') + def test_archive_unpack_with_password(self, raw_password_rar, check_unpack_dir): + with isolated_directory(): + with disable_output(): + archive_unpack(raw_password_rar, '.', password='password') + check_unpack_dir('.') diff --git a/test/archive/test_sevenz.py b/test/archive/test_sevenz.py index 3922b8aafd..4c3e1d0035 100644 --- a/test/archive/test_sevenz.py +++ b/test/archive/test_sevenz.py @@ -18,6 +18,11 @@ def raw_7z(): return get_testfile('raw.7z') +@pytest.fixture() +def raw_password_7z(): + return get_testfile('raw-password.7z') + + @pytest.mark.unittest class TestArchive7z: @skipUnless(py7zr, 'py7zr module required.') @@ -58,3 +63,16 @@ def test_archive_unpack(self, raw_7z, check_unpack_dir): with disable_output(): archive_unpack(raw_7z, '.') check_unpack_dir('.') + + @skipUnless(py7zr, 'py7zr module required.') + def test_archive_unpack_with_password_failed(self, raw_password_7z): + with isolated_directory(): + with disable_output(), pytest.raises(py7zr.exceptions.PasswordRequired): + archive_unpack(raw_password_7z, '.') + + @skipUnless(py7zr, 'py7zr module required.') + def test_archive_unpack_with_password(self, raw_password_7z, check_unpack_dir): + with isolated_directory(): + with disable_output(): + archive_unpack(raw_password_7z, '.', password='password') + check_unpack_dir('.') diff --git a/test/archive/test_tar.py b/test/archive/test_tar.py index a1098ff925..db521fa460 100644 --- a/test/archive/test_tar.py +++ b/test/archive/test_tar.py @@ -49,3 +49,15 @@ def test_archive_unpack(self, check_unpack_dir, type_, ext): with disable_output(): archive_unpack(get_testfile(f'raw{ext}'), '.') check_unpack_dir('.') + + @pytest.mark.parametrize(['type_', 'ext'], [ + ('tar', '.tar'), + ('gztar', '.tar.gz'), + ('bztar', '.tar.bz2'), + ('xztar', '.tar.xz'), + ]) + def test_archive_unpack_password_ignore(self, check_unpack_dir, type_, ext): + with isolated_directory(): + with disable_output(), pytest.warns(UserWarning): + archive_unpack(get_testfile(f'raw{ext}'), '.', password='password') + check_unpack_dir('.') diff --git a/test/archive/test_zip.py b/test/archive/test_zip.py index 87876ec4b0..a6b50a9dcd 100644 --- a/test/archive/test_zip.py +++ b/test/archive/test_zip.py @@ -13,6 +13,11 @@ def raw_zip(): return get_testfile('raw.zip') +@pytest.fixture() +def raw_password_zip(): + return get_testfile('raw-password.zip') + + @pytest.mark.unittest class TestArchiveZip: def test_get_archive_type(self): @@ -37,3 +42,14 @@ def test_archive_unpack(self, raw_zip, check_unpack_dir): with disable_output(): archive_unpack(raw_zip, '.') check_unpack_dir('.') + + def test_archive_unpack_with_password_failed(self, raw_password_zip): + with isolated_directory(): + with disable_output(), pytest.raises(RuntimeError): + archive_unpack(raw_password_zip, '.') + + def test_archive_unpack_with_password(self, raw_password_zip, check_unpack_dir): + with isolated_directory(): + with disable_output(): + archive_unpack(raw_password_zip, '.', password='password') + check_unpack_dir('.') diff --git a/test/testfile/raw-password.7z b/test/testfile/raw-password.7z new file mode 100755 index 0000000000..1521625acb Binary files /dev/null and b/test/testfile/raw-password.7z differ diff --git a/test/testfile/raw-password.rar b/test/testfile/raw-password.rar new file mode 100755 index 0000000000..0c329b73b8 Binary files /dev/null and b/test/testfile/raw-password.rar differ diff --git a/test/testfile/raw-password.zip b/test/testfile/raw-password.zip new file mode 100755 index 0000000000..83527a4d02 Binary files /dev/null and b/test/testfile/raw-password.zip differ