dev(narugo): add support for password

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)

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)

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)

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} ...')

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)

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('.')

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('.')

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('.')

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('.')