xiaoshiai-hub 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

xiaoshiai_hub/client.py

@@ -8,7 +8,7 @@ from typing import Dict, List, Optional
 
 
 import requests
-from xiaoshiai_hub.envelope_crypto import DataKey
+from xiaoshiai_hub.envelope_crypto import DEFAULT_ALGORITHM, Algorithm, DataKey
 
 try:
     from tqdm.auto import tqdm
@@ -110,8 +110,186 @@ class HubClient:
         url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/encryption/set"
         response = self._make_request("PUT", url)
         if response.status_code != 200:
-            raise HTTPError(f"Failed to set repository encrypted flag: {response.text}")    
+            raise HTTPError(f"Failed to set repository encrypted flag: {response.text}")
+
+    # 创建分支
+    def create_branch(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+        branch_name: str,
+        from_branch: str,
+    ) -> None:
+        """
+        创建新分支。如果分支已存在，则直接返回。
+
+        Args:
+            organization: 组织名称
+            repo_type: 仓库类型 ("models" 或 "datasets")
+            repo_name: 仓库名称
+            branch_name: 要创建的分支名称
+            from_branch: 基于哪个分支创建
+        """
+        # 先检查分支是否已存在
+        refs = self.get_repository_refs(organization, repo_type, repo_name)
+        for ref in refs:
+            if ref.name == branch_name:
+                # 分支已存在，直接返回
+                return
+
+        # 分支不存在，创建新分支
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/refs/{branch_name}"
+        body = {
+            "base": from_branch,
+        }
+        response = self._make_request("POST", url, json=body)
+        if response.status_code != 200:
+            raise HTTPError(f"Failed to create branch: {response.text}")
+
+    # 删除分支
+    def delete_branch(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+        branch_name: str,
+    ) -> None:
+        """
+        删除分支。如果分支不存在，则直接返回。
+
+        Args:
+            organization: 组织名称
+            repo_type: 仓库类型 ("models" 或 "datasets")
+            repo_name: 仓库名称
+            branch_name: 要删除的分支名称
+        """
+        # 先检查分支是否存在
+        refs = self.get_repository_refs(organization, repo_type, repo_name)
+        branch_exists = False
+        for ref in refs:
+            if ref.name == branch_name:
+                branch_exists = True
+                break
+
+        if not branch_exists:
+            # 分支不存在，直接返回
+            return
+
+        # 分支存在，删除它
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/refs/{branch_name}"
+        response = self._make_request("DELETE", url)
+        if response.status_code != 200:
+            raise HTTPError(f"Failed to delete branch: {response.text}")
+
+
+    # 更新仓库,先获取仓库信息，然后更新
+    def update_repository(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+        description: Optional[str] = None,
+        visibility: str = "internal",
+        annotations: Optional[Dict[str, str]] = None,
+        metadata: Optional[Dict[str, List[str]]] = None,
+        base_model: Optional[List[str]] = None,
+        relationship: Optional[str] = None,
+    ) -> None:
+        """Update repository information."""
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}"
+        body: Dict = {
+            "name": repo_name,
+            "organization": organization,
+            "type": repo_type,
+        }
+        if annotations:
+            body["annotations"] = annotations
+        if description:
+            body["description"] = description
+        if visibility:
+            body["visibility"] = visibility
+        if metadata:
+            body["metadata"] = metadata
+        if base_model or relationship:
+            body["requestgenealogy"] = {}
+            if base_model:
+                body["requestgenealogy"]["baseModel"] = base_model
+            if relationship:
+                body["requestgenealogy"]["relationship"] = relationship
+        response = self._make_request("PUT", url, json=body)
+        if response.status_code != 200:
+            raise HTTPError(f"Failed to update repository: {response.text}")    
     
+    # 删除仓库
+    def delete_repository(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+    ) -> None:
+        """Delete repository."""
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}"
+        response = self._make_request("DELETE", url)
+        if response.status_code != 200:
+            raise HTTPError(f"Failed to delete repository: {response.text}")
+
+    # 创建仓库
+    def create_repository(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+        description: Optional[str] = None,
+        visibility: str = "internal",
+        metadata: Optional[Dict[str, List[str]]] = None,
+        base_model: Optional[List[str]] = None,
+        relationship: Optional[str] = None,
+    ) -> None:
+        """
+        创建新仓库。
+
+        Args:
+            organization: 组织名称
+            repo_type: 仓库类型 ("models" 或 "datasets")
+            repo_name: 仓库名称
+            description: 仓库描述
+            visibility: 可见性 ("public", "internal", "private")
+            metadata: 元数据，包含 license, tasks, languages, tags, frameworks 等
+            annotations: 注解
+            base_model: 基础模型列表 (如 ["demo/yyyy"])
+            relationship: 与基础模型的关系 ("adapter", "finetune", "quantized", "merge", "repackage" 等)
+
+        Returns:
+            创建的仓库信息
+        """
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}"
+
+        # 构建请求体
+        body: Dict = {
+            "name": repo_name,
+            "organization": organization,
+            "visibility": visibility,
+        }
+
+        if description:
+            body["description"] = description
+        if metadata:
+            body["metadata"] = metadata
+
+        # 构建 genealogy（模型谱系）
+        if base_model or relationship:
+            body["requestgenealogy"] = {}
+            if base_model:
+                body["requestgenealogy"]["baseModel"] = base_model
+            if relationship:
+                body["requestgenealogy"]["relationship"] = relationship
+
+        response = self._make_request("POST", url, json=body)
+        if response.status_code != 200:
+            raise HTTPError(f"Failed to create repository: {response.text}")
+        
+    # 获取仓库信息
     def get_repository_info(
         self,
         organization: str,
@@ -143,15 +321,25 @@ class HubClient:
         if 'metadata' in data and isinstance(data['metadata'], dict):
             metadata = data['metadata']
 
+        # Parse genealogy if present
+        genealogy: Optional[Dict] = None
+        if 'genealogy' in data and isinstance(data['genealogy'], dict):
+            genealogy = data['genealogy']
+
         return Repository(
             name=data.get('name', repo_name),
-            organization=organization,
-            type=repo_type,
+            organization=data.get('organization', organization),
+            owner=data.get('owner', ''),
+            creator=data.get('creator', ''),
+            type=data.get('type', repo_type),
+            visibility=data.get('visibility', 'internal'),
+            genealogy=genealogy,
             description=data.get('description'),
             metadata=metadata,
             annotations=annotations,
         )
     
+    # 获取仓库的所有分支
     def get_repository_refs(
         self,
         organization: str,
@@ -338,12 +526,18 @@ class HubClient:
                 progress_bar.close()
 
 
-    def generate_data_key(self, password: Optional[str] = None) -> DataKey:
+    def generate_data_key(
+            self, 
+            algorithm : Optional[str]  = DEFAULT_ALGORITHM,
+            password: Optional[str] = None) -> DataKey:
         url = f"{self.base_url}/api/kms/generate-data-key"
         try:
             resp = requests.post(
                 url,
-                json={"password": password},
+                json={
+                    "algorithm": algorithm,
+                    "password": password,
+                },
                 timeout=30
             )
             resp.raise_for_status()

xiaoshiai_hub/envelope_crypto.py

@@ -1,20 +1,25 @@
 """
 信封加密核心模块
 
-实现基于 AES-256-CTR 的信封加密/解密逻辑。
+实现基于 AES-256-CTR 和 SM4-CTR 的信封加密/解密逻辑。
 支持随机访问和流式解密，适合大文件部分读取场景。
 
 信封加密文件格式:
 - 前 4 字节: 元数据长度 (big-endian)
-- 元数据: JSON 格式，包含 encryptedKey 和 password
-- 16 字节: IV (用于 AES-256-CTR)
+- 元数据: JSON 格式，包含 encryptedKey 和 algorithm
+- 16 字节: IV (用于 CTR 模式)
 - 剩余部分: 加密后的文件内容
+
+支持的算法:
+- AES: 使用 32 字节密钥
+- SM4: 使用 16 字节密钥
 """
 
 import io
 import os
 import json
 from dataclasses import dataclass
+from enum import Enum
 from pathlib import Path
 import shutil
 import tempfile
@@ -24,17 +29,29 @@ from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.hazmat.backends import default_backend
 
 
-# IV 长度 (AES-256-CTR)
+# IV 长度 (AES-256-CTR 和 SM4-CTR 都使用 16 字节 IV)
 IV_SIZE = 16
 # 元数据长度字段大小
 METADATA_LENGTH_SIZE = 4
 CHUNK_SIZE = 64 * 1024  # 64KB chunks
 
+
+class Algorithm(str, Enum):
+    """支持的加密算法"""
+    AES = "AES"
+    SM4 = "SM4"
+
+
+# 默认算法
+DEFAULT_ALGORITHM = Algorithm.AES
+
+
 @dataclass
 class DataKey:
     """数据密钥对象，包含明文密钥和加密后的密钥"""
-    plaintext_key: bytes  # 明文密钥（32字节，用于 AES-256）
-    encrypted_key: str 
+    plaintext_key: bytes  # 明文密钥（AES-256: 32字节，SM4: 16字节）
+    encrypted_key: str
+
 
 def _raw_open(path, mode="rb"):
     """
@@ -49,26 +66,51 @@ def _raw_open(path, mode="rb"):
     return file_io
 
 
-def create_aes_ctr_cipher(key: bytes, iv: bytes):
-    """创建 AES-256-CTR cipher"""
-    return Cipher(
-        algorithms.AES(key),
-        modes.CTR(iv),
-        backend=default_backend()
-    )
+def create_cipher(key: bytes, iv: bytes, algorithm: Algorithm = DEFAULT_ALGORITHM):
+    """
+    创建加密 cipher
+
+    Args:
+        key: 密钥（AES-256: 32字节，SM4: 16字节）
+        iv: 初始化向量（16字节）
+        algorithm: 加密算法
+
+    Returns:
+        Cipher 对象
+    """
+    if algorithm == Algorithm.AES:
+        if len(key) != 32:
+            raise ValueError(f"AES requires 32-byte key, got {len(key)} bytes")
+        return Cipher(
+            algorithms.AES(key),
+            modes.CTR(iv),
+            backend=default_backend()
+        )
+    elif algorithm == Algorithm.SM4:
+        if len(key) != 16:
+            raise ValueError(f"SM4 requires 16-byte key, got {len(key)} bytes")
+        return Cipher(
+            algorithms.SM4(key),
+            modes.CTR(iv),
+            backend=default_backend()
+        )
+    else:
+        raise ValueError(f"Unsupported algorithm: {algorithm}")
+
+
 
 
 @dataclass
 class EnvelopeMetadata:
     """加密文件元数据"""
     encrypted_key: str   # 加密后的数据密钥（Base64）
-    password: str        # 用于解密 DEK 的密码
+    algorithm: str       # 加密算法名称
 
     def to_bytes(self) -> bytes:
         """序列化为 JSON 字节"""
         return json.dumps({
             "encryptedKey": self.encrypted_key,
-            "password": self.password
+            "algorithm": self.algorithm
         }).encode("utf-8")
 
     @classmethod
@@ -77,30 +119,30 @@ class EnvelopeMetadata:
         obj = json.loads(data.decode("utf-8"))
         return cls(
             encrypted_key=obj["encryptedKey"],
-            password=obj["password"]
+            algorithm=obj.get("algorithm", Algorithm.AES.value)  # 兼容旧格式
         )
 
 
 def _envelope_encrypt_file(
     input_path: Union[str, Path],
     output_path: Union[str, Path],
-    password: str,
     data_key: DataKey,
+    algorithm: Algorithm = DEFAULT_ALGORITHM,
 ) -> None:
     """
-    使用信封加密对文件进行加密 (AES-256-CTR)
+    使用信封加密对文件进行加密
 
     Args:
         input_path: 明文文件路径
         output_path: 加密文件输出路径
-        password: 用于 KMS 密钥派生的密码
-        kms_client: KMS 客户端实例（可选，不提供则使用默认客户端）
+        data_key: 数据密钥对象
+        algorithm: 加密算法（默认 AES）
     """
     input_path = Path(input_path)
     output_path = Path(output_path)
 
     iv = os.urandom(IV_SIZE)
-    cipher = create_aes_ctr_cipher(data_key.plaintext_key, iv)
+    cipher = create_cipher(data_key.plaintext_key, iv, algorithm)
     encryptor = cipher.encryptor()
     with _raw_open(input_path, "rb") as f:
         plaintext = f.read()
@@ -108,7 +150,7 @@ def _envelope_encrypt_file(
     ciphertext = encryptor.update(plaintext) + encryptor.finalize()
     metadata = EnvelopeMetadata(
         encrypted_key=data_key.encrypted_key,
-        password=password
+        algorithm=algorithm.value
     )
     metadata_bytes = metadata.to_bytes()
     output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -123,25 +165,26 @@ def _envelope_encrypt_file(
 def envelope_enc_file(
     source: Union[Path, str],
     *,
-    password: str,
     data_key: DataKey,
     dest: Optional[Union[Path, str]] = None,
     replace: bool = False,
     chunked: bool = True,
-    chunk_size: int = CHUNK_SIZE,  # 默认 64KB，与 envelope_crypto 一致
+    chunk_size: int = CHUNK_SIZE,
+    algorithm: Algorithm = DEFAULT_ALGORITHM,
 ) -> Path:
-    """使用信封加密模式加密单个文件 (AES-256-CTR)。
+    """使用信封加密模式加密单个文件。
 
     信封加密使用 KMS 服务生成数据密钥，数据密钥用于加密文件内容，
     加密后的数据密钥存储在文件头部。
 
     Args:
         source: 源文件路径
-        password: 用于 KMS 密钥派生的密码
+        data_key: 数据密钥对象
         dest: 目标文件路径（默认添加 .encrypted 后缀）
         replace: 是否原地加密（替换原文件）
         chunked: 是否使用流式加密（适用于大文件，减少内存占用）
         chunk_size: 流式加密时每次处理的块大小
+        algorithm: 加密算法（默认 AES，可选 SM4）
 
     Returns:
         加密后的文件路径
@@ -161,9 +204,9 @@ def envelope_enc_file(
     dst.parent.mkdir(parents=True, exist_ok=True)
 
     if chunked:
-        _envelope_encrypt_file_streaming(src, dst, password, data_key, chunk_size)
+        _envelope_encrypt_file_streaming(src, dst, data_key, chunk_size, algorithm)
     else:
-        _envelope_encrypt_file(src, dst, password, data_key)
+        _envelope_encrypt_file(src, dst, data_key, algorithm)
 
     # 替换模式：移动加密文件到原位置
     if replace:
@@ -185,32 +228,32 @@ def _make_temp_path(parent: Path, suffix: str) -> Path:
 def _envelope_encrypt_file_streaming(
     input_path: Union[str, Path],
     output_path: Union[str, Path],
-    password: str,
     data_key: DataKey,
     chunk_size: int = CHUNK_SIZE,
+    algorithm: Algorithm = DEFAULT_ALGORITHM,
 ) -> None:
     """
-    使用信封加密对大文件进行流式加密 (AES-256-CTR)
+    使用信封加密对大文件进行流式加密
 
-    AES-CTR 模式天然支持流式加密，不需要将整个文件加载到内存。
+    CTR 模式天然支持流式加密，不需要将整个文件加载到内存。
     加密后的文件格式与普通模式完全相同，可以用普通模式解密。
 
     Args:
         input_path: 明文文件路径
         output_path: 加密文件输出路径
-        password: 用于 KMS 密钥派生的密码
-        kms_client: KMS 客户端实例
+        data_key: 数据密钥对象
         chunk_size: 每次读取的块大小
+        algorithm: 加密算法（默认 AES）
     """
     input_path = Path(input_path)
     output_path = Path(output_path)
     iv = os.urandom(IV_SIZE)
     metadata = EnvelopeMetadata(
         encrypted_key=data_key.encrypted_key,
-        password=password
+        algorithm=algorithm.value
     )
     metadata_bytes = metadata.to_bytes()
-    cipher = create_aes_ctr_cipher(data_key.plaintext_key, iv)
+    cipher = create_cipher(data_key.plaintext_key, iv, algorithm)
     encryptor = cipher.encryptor()
     output_path.parent.mkdir(parents=True, exist_ok=True)
     with _raw_open(input_path, "rb") as fin, _raw_open(output_path, "wb") as fout:

xiaoshiai_hub/types.py

@@ -12,7 +12,11 @@ class Repository:
     """Repository information."""
     name: str
     organization: str
+    owner: str
+    creator: str
     type: str  # "models" or "datasets"
+    visibility: str
+    genealogy: Optional[Dict] = None
     description: Optional[str] = None
     metadata: Dict[str, List[str]] = field(default_factory=dict)  # Repository metadata
     annotations: Dict[str, str] = field(default_factory=dict)  # Repository annotations/metadata