xiaoshiai-hub 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

xiaoshiai_hub/download.py

@@ -14,6 +14,30 @@ except ImportError:
 
 from .client import HubClient, DEFAULT_BASE_URL
 from .types import GitContent
+from .exceptions import EncryptionError, RepositoryNotFoundError
+
+
+def _should_decrypt_file(filename: str, decryption_exclude: Optional[List[str]] = None) -> bool:
+    """
+    Check if a file should be decrypted based on exclude patterns.
+
+    Args:
+        filename: Name of the file
+        decryption_exclude: List of patterns to exclude from decryption
+
+    Returns:
+        True if the file should be decrypted
+    """
+    import fnmatch
+
+    if not decryption_exclude:
+        return True
+
+    for pattern in decryption_exclude:
+        if fnmatch.fnmatch(filename, pattern):
+            return False
+
+    return True
 
 
 def _match_pattern(name: str, pattern: str) -> bool:
@@ -138,6 +162,9 @@ def _download_repository_recursively(
     ignore_patterns: Optional[List[str]] = None,
     verbose: bool = True,
     progress_bar = None,
+    encryption_metadata = None,
+    decryption_key: Optional[Union[str, bytes]] = None,
+    decryption_algorithm: Optional[str] = None,
 ) -> None:
     """
     Recursively download repository contents.
@@ -154,8 +181,11 @@ def _download_repository_recursively(
         ignore_patterns: Patterns to ignore
         verbose: Print progress messages
         progress_bar: Optional tqdm progress bar for overall progress
+        encryption_metadata: Encryption metadata from .moha_encryption file
+        decryption_key: Key to decrypt files
+        decryption_exclude: Patterns to exclude from decryption
+        decryption_algorithm: Algorithm to use for decryption
     """
-    # Get content at current path
     content = client.get_repository_content(
         organization=organization,
         repo_type=repo_type,
@@ -168,7 +198,7 @@ def _download_repository_recursively(
     if content.entries:
         for entry in content.entries:
             if entry.type == "file":
-                # Check if file should be downloaded
+                # 检查文件是否应该被下载
                 if _should_download_file(entry.path, allow_patterns, ignore_patterns):
                     if verbose and progress_bar is None:
                         print(f"Downloading file: {entry.path}")
@@ -188,8 +218,45 @@ def _download_repository_recursively(
                         local_path=local_path,
                         show_progress=progress_bar is None,  # Show individual progress only if no overall progress
                     )
-
-                    # Update overall progress
+                    file_is_encrypted = False
+                    file_encryption_algorithm = None
+                    if encryption_metadata and encryption_metadata.files:
+                        for file_meta in encryption_metadata.files:
+                            if file_meta.path == entry.path:
+                                file_is_encrypted = True
+                                file_encryption_algorithm = file_meta.algorithm
+                                break
+
+                    if file_is_encrypted:
+                        if not decryption_key:
+                            raise EncryptionError(
+                                f"File '{entry.path}' is encrypted, but no decryption_key was provided. "
+                                "Please provide a decryption_key parameter."
+                            )
+                        if not decryption_algorithm:
+                            raise EncryptionError(
+                                f"File '{entry.path}' is encrypted, but no decryption_algorithm was provided. "
+                                f"The file was encrypted with '{file_encryption_algorithm}'. "
+                                "Please provide a decryption_algorithm parameter."
+                            )
+                        if decryption_algorithm != file_encryption_algorithm:
+                            raise EncryptionError(
+                                f"File '{entry.path}' is encrypted with '{file_encryption_algorithm}', "
+                                f"but decryption_algorithm '{decryption_algorithm}' was provided. "
+                                "Please use the correct decryption algorithm."
+                            )
+                        if verbose and progress_bar is None:
+                            print(f"Decrypting file: {entry.path}")
+                        from .encryption import EncryptionAlgorithm, decrypt_file as decrypt_file_func
+                        try:
+                            algorithm = EncryptionAlgorithm(decryption_algorithm)
+                        except ValueError:
+                            raise EncryptionError(
+                                f"Invalid decryption algorithm: {decryption_algorithm}. "
+                                f"Supported algorithms: {', '.join([a.value for a in EncryptionAlgorithm])}"
+                            )
+                        decrypt_file_func(Path(local_path), decryption_key, algorithm)
+                    
                     if progress_bar is not None:
                         progress_bar.update(1)
                 else:
@@ -199,8 +266,7 @@ def _download_repository_recursively(
             elif entry.type == "dir":
                 if verbose and progress_bar is None:
                     print(f"Entering directory: {entry.path}")
-
-                # Recursively download directory contents
+                # 递归下载
                 _download_repository_recursively(
                     client=client,
                     organization=organization,
@@ -213,26 +279,29 @@ def _download_repository_recursively(
                     ignore_patterns=ignore_patterns,
                     verbose=verbose,
                     progress_bar=progress_bar,
+                    encryption_metadata=encryption_metadata,
+                    decryption_key=decryption_key,
+                    decryption_algorithm=decryption_algorithm,
                 )
-
             else:
                 if verbose and progress_bar is None:
                     print(f"Skipping {entry.type}: {entry.path}")
 
 
-def hf_hub_download(
+def moha_hub_download(
     repo_id: str,
     filename: str,
     *,
     repo_type: str = "models",
     revision: Optional[str] = None,
-    cache_dir: Optional[Union[str, Path]] = None,
     local_dir: Optional[Union[str, Path]] = None,
     base_url: Optional[str] = None,
     username: Optional[str] = None,
     password: Optional[str] = None,
     token: Optional[str] = None,
     show_progress: bool = True,
+    decryption_key: Optional[Union[str, bytes]] = None,
+    decryption_algorithm: Optional[str] = None,
 ) -> str:
     """
     Download a single file from a repository.
@@ -251,6 +320,10 @@ def hf_hub_download(
         password: Password for authentication
         token: Token for authentication
         show_progress: Whether to show download progress bar
+        decryption_key: Key to decrypt the file if repository is encrypted (string for symmetric, PEM for asymmetric)
+        decryption_algorithm: Decryption algorithm to use (default: 'aes-256-cbc')
+            - Symmetric: 'aes-256-cbc', 'aes-256-gcm'
+            - Asymmetric: 'rsa-oaep', 'rsa-pkcs1v15' (requires RSA private key in PEM format)
 
     Returns:
         Path to the downloaded file
@@ -258,49 +331,79 @@ def hf_hub_download(
     Example:
         >>> file_path = hf_hub_download(
         ...     repo_id="demo/demo",
-        ...     filename="config.yaml",
+        ...     filename="data/config.yaml",
         ...     username="your-username",
         ...     password="your-password",
         ... )
+
+        >>> # Download from encrypted repository
+        >>> file_path = hf_hub_download(
+        ...     repo_id="demo/encrypted-model",
+        ...     filename="model.bin",
+        ...     decryption_key="my-secret-key",
+        ...     token="your-token",
+        ... )
     """
-    # Parse repo_id
     parts = repo_id.split('/')
     if len(parts) != 2:
         raise ValueError(f"Invalid repo_id format: {repo_id}. Expected 'organization/repo_name'")
-
     organization, repo_name = parts
-
-    # Create client
     client = HubClient(
         base_url=base_url,
         username=username,
         password=password,
         token=token,
     )
-
-    # Get repository info to determine branch
+    # 获取默认分支
     if revision is None:
-        repo_info = client.get_repository_info(organization, repo_type, repo_name)
-        revision = repo_info.default_branch or "main"
+        revision = client.get_default_branch(organization, repo_type, repo_name)
+
+    # 获取加密元数据
+    encryption_metadata = client.get_moha_encryption(
+        organization=organization,
+        repo_type=repo_type,
+        repo_name=repo_name,
+        reference=revision,
+    )
+
+    # 文件加密标识
+    file_is_encrypted = False
+    file_encryption_algorithm = None
+    # 查询文件是否加密
+    if encryption_metadata and encryption_metadata.files:
+        for file_meta in encryption_metadata.files:
+            if file_meta.path == filename:
+                file_is_encrypted = True
+                file_encryption_algorithm = file_meta.algorithm
+                break
+
+    # 如果该文件加密了，检查解密参数
+    if file_is_encrypted:
+        if not decryption_key:
+            raise EncryptionError(
+                f"File '{filename}' is encrypted, but no decryption_key was provided. "
+                "Please provide a decryption_key parameter."
+            )
+
+        if not decryption_algorithm:
+            raise EncryptionError(
+                f"File '{filename}' is encrypted, but no decryption_algorithm was provided. "
+                f"The file was encrypted with '{file_encryption_algorithm}'. "
+                "Please provide a decryption_algorithm parameter."
+            )
 
-    # Determine local path
+        if decryption_algorithm != file_encryption_algorithm:
+            raise EncryptionError(
+                f"File '{filename}' is encrypted with '{file_encryption_algorithm}', "
+                f"but decryption_algorithm '{decryption_algorithm}' was provided. "
+                "Please use the correct decryption algorithm."
+            )
     if local_dir:
         local_path = os.path.join(local_dir, filename)
-    elif cache_dir:
-        # Create cache structure similar to huggingface_hub
-        cache_path = os.path.join(
-            cache_dir,
-            f"{repo_type}--{organization}--{repo_name}",
-            "snapshots",
-            revision,
-            filename,
-        )
-        local_path = cache_path
     else:
-        # Default to current directory
         local_path = filename
 
-    # Download file
+    # 下载文件
     client.download_file(
         organization=organization,
         repo_type=repo_type,
@@ -311,15 +414,26 @@ def hf_hub_download(
         show_progress=show_progress,
     )
 
+    # 解密文件
+    if file_is_encrypted and decryption_key:
+        from .encryption import EncryptionAlgorithm, decrypt_file as decrypt_file_func
+        try:
+            algorithm = EncryptionAlgorithm(decryption_algorithm)
+        except ValueError:
+            raise EncryptionError(
+                f"Invalid decryption algorithm: {decryption_algorithm}. "
+                f"Supported algorithms: {', '.join([a.value for a in EncryptionAlgorithm])}"
+            )
+
+        decrypt_file_func(Path(local_path), decryption_key, algorithm)
+
     return local_path
 
 
 def snapshot_download(
     repo_id: str,
-    *,
     repo_type: str = "models",
     revision: Optional[str] = None,
-    cache_dir: Optional[Union[str, Path]] = None,
     local_dir: Optional[Union[str, Path]] = None,
     allow_patterns: Optional[Union[List[str], str]] = None,
     ignore_patterns: Optional[Union[List[str], str]] = None,
@@ -329,6 +443,8 @@ def snapshot_download(
     token: Optional[str] = None,
     verbose: bool = True,
     show_progress: bool = True,
+    decryption_key: Optional[Union[str, bytes]] = None,
+    decryption_algorithm: Optional[str] = None,
 ) -> str:
     """
     Download an entire repository snapshot.
@@ -339,7 +455,6 @@ def snapshot_download(
         repo_id: Repository ID in the format "organization/repo_name"
         repo_type: Type of repository ("models" or "datasets")
         revision: Branch/tag/commit to download from (default: main branch)
-        cache_dir: Directory to cache downloaded files
         local_dir: Directory to save files (if not using cache)
         allow_patterns: Pattern or list of patterns to allow (e.g., "*.yaml", "*.yml")
         ignore_patterns: Pattern or list of patterns to ignore (e.g., ".git*")
@@ -349,6 +464,10 @@ def snapshot_download(
         token: Token for authentication
         verbose: Print progress messages
         show_progress: Whether to show overall progress bar
+        decryption_key: Key to decrypt files if repository is encrypted (string for symmetric, PEM for asymmetric)
+        decryption_algorithm: Decryption algorithm to use (default: 'aes-256-cbc')
+            - Symmetric: 'aes-256-cbc', 'aes-256-gcm'
+            - Asymmetric: 'rsa-oaep', 'rsa-pkcs1v15' (requires RSA private key in PEM format)
 
     Returns:
         Path to the downloaded repository
@@ -362,46 +481,48 @@ def snapshot_download(
         ...     username="your-username",
         ...     password="your-password",
         ... )
+
+        >>> # Download from encrypted repository
+        >>> repo_path = snapshot_download(
+        ...     repo_id="demo/encrypted-model",
+        ...     repo_type="models",
+        ...     decryption_key="my-secret-key",
+        ...     decryption_exclude=["README.md", "*.txt"],  # Don't decrypt these files
+        ...     token="your-token",
+        ... )
     """
-    # Parse repo_id
     parts = repo_id.split('/')
     if len(parts) != 2:
         raise ValueError(f"Invalid repo_id format: {repo_id}. Expected 'organization/repo_name'")
 
     organization, repo_name = parts
 
-    # Normalize patterns to lists
     if isinstance(allow_patterns, str):
         allow_patterns = [allow_patterns]
     if isinstance(ignore_patterns, str):
         ignore_patterns = [ignore_patterns]
-
-    # Create client
     client = HubClient(
         base_url=base_url,
         username=username,
         password=password,
         token=token,
     )
-
-    # Get repository info
-    repo_info = client.get_repository_info(organization, repo_type, repo_name)
-
-    # Determine revision
     if revision is None:
-        revision = repo_info.default_branch or "main"
+        revision = client.get_default_branch(organization, repo_type, repo_name)
+    encryption_metadata = client.get_moha_encryption(
+        organization=organization,
+        repo_type=repo_type,
+        repo_name=repo_name,
+        reference=revision,
+    )
+
+    if encryption_metadata and encryption_metadata.files:
+        if verbose:
+            print(f"Repository has {len(encryption_metadata.files)} encrypted file(s). Files will be decrypted after download.")
 
     # Determine local directory
     if local_dir:
         download_dir = str(local_dir)
-    elif cache_dir:
-        # Create cache structure
-        download_dir = os.path.join(
-            cache_dir,
-            f"{repo_type}--{organization}--{repo_name}",
-            "snapshots",
-            revision,
-        )
     else:
         # Default to downloads directory
         download_dir = f"./downloads/{organization}_{repo_type}_{repo_name}"
@@ -412,13 +533,11 @@ def snapshot_download(
         print(f"Revision: {revision}")
         print(f"Destination: {download_dir}")
 
-    # Create progress bar if requested
     progress_bar = None
     if show_progress and tqdm is not None:
-        # Count total files to download
         if verbose:
             print(f"Fetching repository info...")
-
+        # 计算需要下载的文件总数
         total_files = _count_files_to_download(
             client=client,
             organization=organization,
@@ -438,7 +557,7 @@ def snapshot_download(
                 leave=True,
             )
 
-    # Download recursively
+    # 递归下载，不是使用git的方式
     try:
         _download_repository_recursively(
             client=client,
@@ -452,6 +571,9 @@ def snapshot_download(
             ignore_patterns=ignore_patterns,
             verbose=verbose,
             progress_bar=progress_bar,
+            encryption_metadata=encryption_metadata,
+            decryption_key=decryption_key,
+            decryption_algorithm=decryption_algorithm,
         )
     finally:
         if progress_bar is not None:
@@ -460,5 +582,17 @@ def snapshot_download(
     if verbose and not show_progress:
         print(f"Download completed to: {download_dir}")
 
+    # Add download count
+    try:
+        client.add_download_count(
+            organization=organization,
+            repo_type=repo_type,
+            repo_name=repo_name,
+        )
+    except Exception as e:
+        # Don't fail the download if adding count fails
+        if verbose:
+            print(f"Warning: Failed to add download count: {e}")
+
     return download_dir