xiaoshiai-hub 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

xiaoshiai_hub/client.py

@@ -4,10 +4,11 @@ XiaoShi AI Hub Client
 
 import base64
 import os
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 
 import requests
+from xiaoshiai_hub.envelope_crypto import DataKey
 
 try:
     from tqdm.auto import tqdm
@@ -25,10 +26,11 @@ from .types import  Repository, Ref, GitContent
 # 默认基础 URL，可通过环境变量 MOHA_ENDPOINT 覆盖
 DEFAULT_BASE_URL = os.environ.get(
     "MOHA_ENDPOINT",
-    "https://rune-api.develop.xiaoshiai.cn/"
+    "https://rune-api.develop.xiaoshiai.cn"
 )
 
 
+
 class HubClient:
     """Client for interacting with XiaoShi AI Hub API."""
 
@@ -85,6 +87,30 @@ class HubClient:
             return response
         except requests.RequestException as e:
             raise HTTPError(f"Request failed: {str(e)}")
+        
+    def cancel_repository_encrypted(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+    ) -> None:
+        """Cancel repository encrypted flag."""
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/encryption/cancel"
+        response = self._make_request("PUT", url)
+        if response.status_code != 200:
+            raise HTTPError(f"Failed to cancel repository encrypted flag: {response.text}")
+        
+    def set_repository_encrypted(
+        self,
+        organization: str,
+        repo_type: str,
+        repo_name: str,
+    ) -> None:
+        """Set repository encrypted flag."""
+        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}")    
     
     def get_repository_info(
         self,
@@ -103,17 +129,17 @@ class HubClient:
         Returns:
             Repository information
         """
-        url = f"{self.base_url}/v1/organizations/{organization}/{repo_type}/{repo_name}"
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}"
         response = self._make_request("GET", url)
         data = response.json()
 
         # Parse annotations if present
-        annotations = {}
+        annotations: Dict[str, str] = {}
         if 'annotations' in data and isinstance(data['annotations'], dict):
             annotations = data['annotations']
 
         # Parse metadata if present
-        metadata = {}
+        metadata: Dict[str, List[str]] = {}
         if 'metadata' in data and isinstance(data['metadata'], dict):
             metadata = data['metadata']
 
@@ -143,7 +169,7 @@ class HubClient:
         Returns:
             List of references
         """
-        url = f"{self.base_url}/v1/organizations/{organization}/{repo_type}/{repo_name}/refs"
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/refs"
         response = self._make_request("GET", url)
         data = response.json()
 
@@ -204,9 +230,9 @@ class HubClient:
             Git content information
         """
         if path:
-            url = f"{self.base_url}/v1/organizations/{organization}/{repo_type}/{repo_name}/contents/{branch}/{path}"
+            url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/contents/{branch}/{path}"
         else:
-            url = f"{self.base_url}/v1/organizations/{organization}/{repo_type}/{repo_name}/contents/{branch}"
+            url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/contents/{branch}"
         
         response = self._make_request("GET", url)
         data = response.json()
@@ -240,7 +266,7 @@ class HubClient:
             repo_type: Repository type ("models" or "datasets")
             repo_name: Repository name
         """
-        url = f"{self.base_url}/v1/organizations/{organization}/{repo_type}/{repo_name}/downloads"
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/downloads"
         response = self._make_request("POST", url)
         # 检查http code
         if response.status_code != 200:
@@ -268,7 +294,7 @@ class HubClient:
             local_path: Local path to save the file
             show_progress: Whether to show download progress bar
         """
-        url = f"{self.base_url}/v1/organizations/{organization}/{repo_type}/{repo_name}/resolve/{branch}/{file_path}"
+        url = f"{self.base_url}/moha/v1/organizations/{organization}/{repo_type}/{repo_name}/resolve/{branch}/{file_path}"
         response = self._make_request("GET", url, stream=True)
 
         # Get file size from headers
@@ -311,3 +337,24 @@ class HubClient:
             if progress_bar is not None:
                 progress_bar.close()
 
+
+    def generate_data_key(self, password: Optional[str] = None) -> DataKey:
+        url = f"{self.base_url}/api/kms/generate-data-key"
+        try:
+            resp = requests.post(
+                url,
+                json={"password": password},
+                timeout=30
+            )
+            resp.raise_for_status()
+            data = resp.json()
+            plaintext_key = base64.b64decode(data["plaintextKey"])
+            encrypted_key = data["encryptedKey"]
+            
+            return DataKey(
+                plaintext_key=plaintext_key,
+                encrypted_key=encrypted_key
+            )
+            
+        except requests.RequestException as e:
+            raise e

xiaoshiai_hub/download.py

@@ -7,14 +7,15 @@ import os
 from pathlib import Path
 from typing import List, Optional, Union
 
+from .exceptions import RepositoryNotFoundError
+
 try:
     from tqdm.auto import tqdm
 except ImportError:
     tqdm = None
 
-from .client import HubClient, DEFAULT_BASE_URL
-from .types import GitContent
-from .exceptions import RepositoryNotFoundError
+from .client import HubClient
+
 
 
 def _match_pattern(name: str, pattern: str) -> bool:
@@ -263,6 +264,16 @@ def moha_hub_download(
         password=password,
         token=token,
     )
+
+    try:
+        client.get_repository_info(organization, repo_type, repo_name)
+    except RepositoryNotFoundError:
+        raise RepositoryNotFoundError(
+            f"Repository not found: {organization}/{repo_type}/{repo_name}. "
+            f"Please create the repository first."
+        )
+
+
     # 获取默认分支
     if revision is None:
         revision = client.get_default_branch(organization, repo_type, repo_name)
@@ -348,6 +359,13 @@ def snapshot_download(
         password=password,
         token=token,
     )
+    try:
+        client.get_repository_info(organization, repo_type, repo_name)
+    except RepositoryNotFoundError:
+        raise RepositoryNotFoundError(
+            f"Repository not found: {organization}/{repo_type}/{repo_name}. "
+            f"Please create the repository first."
+        )
     if revision is None:
         revision = client.get_default_branch(organization, repo_type, repo_name)
 

xiaoshiai_hub/envelope_crypto.py

@@ -0,0 +1,237 @@
+"""
+信封加密核心模块
+
+实现基于 AES-256-CTR 的信封加密/解密逻辑。
+支持随机访问和流式解密，适合大文件部分读取场景。
+
+信封加密文件格式:
+- 前 4 字节: 元数据长度 (big-endian)
+- 元数据: JSON 格式，包含 encryptedKey 和 password
+- 16 字节: IV (用于 AES-256-CTR)
+- 剩余部分: 加密后的文件内容
+"""
+
+import io
+import os
+import json
+from dataclasses import dataclass
+from pathlib import Path
+import shutil
+import tempfile
+from typing import Optional, Union
+
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.backends import default_backend
+
+
+# IV 长度 (AES-256-CTR)
+IV_SIZE = 16
+# 元数据长度字段大小
+METADATA_LENGTH_SIZE = 4
+CHUNK_SIZE = 64 * 1024  # 64KB chunks
+
+@dataclass
+class DataKey:
+    """数据密钥对象，包含明文密钥和加密后的密钥"""
+    plaintext_key: bytes  # 明文密钥（32字节，用于 AES-256）
+    encrypted_key: str 
+
+def _raw_open(path, mode="rb"):
+    """
+    获取原始的文件句柄，绕过 decrypt_patch 的 hook。
+
+    使用 io.FileIO 直接访问文件系统，不经过 builtins.open。
+    """
+    # io.FileIO 是底层实现，不会被 builtins.open 的 hook 影响
+    file_io = io.FileIO(path, mode.replace("b", ""))
+    if "b" in mode:
+        return io.BufferedReader(file_io) if "r" in mode else io.BufferedWriter(file_io)
+    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()
+    )
+
+
+@dataclass
+class EnvelopeMetadata:
+    """加密文件元数据"""
+    encrypted_key: str   # 加密后的数据密钥（Base64）
+    password: str        # 用于解密 DEK 的密码
+
+    def to_bytes(self) -> bytes:
+        """序列化为 JSON 字节"""
+        return json.dumps({
+            "encryptedKey": self.encrypted_key,
+            "password": self.password
+        }).encode("utf-8")
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "EnvelopeMetadata":
+        """从 JSON 字节反序列化"""
+        obj = json.loads(data.decode("utf-8"))
+        return cls(
+            encrypted_key=obj["encryptedKey"],
+            password=obj["password"]
+        )
+
+
+def _envelope_encrypt_file(
+    input_path: Union[str, Path],
+    output_path: Union[str, Path],
+    password: str,
+    data_key: DataKey,
+) -> None:
+    """
+    使用信封加密对文件进行加密 (AES-256-CTR)
+
+    Args:
+        input_path: 明文文件路径
+        output_path: 加密文件输出路径
+        password: 用于 KMS 密钥派生的密码
+        kms_client: KMS 客户端实例（可选，不提供则使用默认客户端）
+    """
+    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)
+    encryptor = cipher.encryptor()
+    with _raw_open(input_path, "rb") as f:
+        plaintext = f.read()
+
+    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
+    metadata = EnvelopeMetadata(
+        encrypted_key=data_key.encrypted_key,
+        password=password
+    )
+    metadata_bytes = metadata.to_bytes()
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with _raw_open(output_path, "wb") as f:
+        f.write(len(metadata_bytes).to_bytes(METADATA_LENGTH_SIZE, "big"))
+        f.write(metadata_bytes)
+        f.write(iv)
+        f.write(ciphertext)
+    del data_key.plaintext_key
+
+
+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 一致
+) -> Path:
+    """使用信封加密模式加密单个文件 (AES-256-CTR)。
+
+    信封加密使用 KMS 服务生成数据密钥，数据密钥用于加密文件内容，
+    加密后的数据密钥存储在文件头部。
+
+    Args:
+        source: 源文件路径
+        password: 用于 KMS 密钥派生的密码
+        dest: 目标文件路径（默认添加 .encrypted 后缀）
+        replace: 是否原地加密（替换原文件）
+        chunked: 是否使用流式加密（适用于大文件，减少内存占用）
+        chunk_size: 流式加密时每次处理的块大小
+
+    Returns:
+        加密后的文件路径
+    """
+    src = Path(source)
+    if not src.exists() or not src.is_file():
+        raise FileNotFoundError(f"Source file not found: {src}")
+
+    # 确定输出路径
+    if replace:
+        dst = _make_temp_path(src.parent, ".encrypted.tmp")
+    elif dest:
+        dst = Path(dest)
+    else:
+        dst = src.with_suffix(src.suffix + ".encrypted")
+
+    dst.parent.mkdir(parents=True, exist_ok=True)
+
+    if chunked:
+        _envelope_encrypt_file_streaming(src, dst, password, data_key, chunk_size)
+    else:
+        _envelope_encrypt_file(src, dst, password, data_key)
+
+    # 替换模式：移动加密文件到原位置
+    if replace:
+        shutil.move(str(src), str(dst))
+        return src
+
+    return dst
+
+def _make_temp_path(parent: Path, suffix: str) -> Path:
+    """创建临时文件路径（使用 mkstemp 避免竞态条件）"""
+    fd, path = tempfile.mkstemp(suffix=suffix, dir=parent)
+    os.close(fd)
+    return Path(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,
+) -> None:
+    """
+    使用信封加密对大文件进行流式加密 (AES-256-CTR)
+
+    AES-CTR 模式天然支持流式加密，不需要将整个文件加载到内存。
+    加密后的文件格式与普通模式完全相同，可以用普通模式解密。
+
+    Args:
+        input_path: 明文文件路径
+        output_path: 加密文件输出路径
+        password: 用于 KMS 密钥派生的密码
+        kms_client: KMS 客户端实例
+        chunk_size: 每次读取的块大小
+    """
+    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
+    )
+    metadata_bytes = metadata.to_bytes()
+    cipher = create_aes_ctr_cipher(data_key.plaintext_key, iv)
+    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:
+        # 写入元数据
+        fout.write(len(metadata_bytes).to_bytes(METADATA_LENGTH_SIZE, "big"))
+        fout.write(metadata_bytes)
+        # 写入 IV
+        fout.write(iv)
+
+        # 流式加密
+        while True:
+            chunk = fin.read(chunk_size)
+            if not chunk:
+                break
+            encrypted_chunk = encryptor.update(chunk)
+            fout.write(encrypted_chunk)
+
+        # 完成加密
+        final = encryptor.finalize()
+        if final:
+            fout.write(final)
+
+    del data_key.plaintext_key
+

xiaoshiai_hub/upload.py

@@ -12,10 +12,10 @@ from typing import List, Optional, Union, Dict
 import requests
 
 from xiaoshiai_hub.client import DEFAULT_BASE_URL, HubClient
+from xiaoshiai_hub.envelope_crypto import DataKey, envelope_enc_file
 from .exceptions import HubException, AuthenticationError, RepositoryNotFoundError
 
-from xpai_enc import enc_file
-from key_manager import KeyManager
+
 
 try:
     from tqdm.auto import tqdm
@@ -40,7 +40,7 @@ def _build_api_url(
 ) -> str:
     """Build API upload URL."""
     base_url = (base_url or DEFAULT_BASE_URL).rstrip('/')
-    return f"{base_url}/{organization}/{repo_type}/{repo_name}/api/upload"
+    return f"{base_url}/moha/{organization}/{repo_type}/{repo_name}/api/upload"
 
 
 def _create_session(
@@ -96,7 +96,8 @@ class _TqdmUploadWrapper:
     def __enter__(self):
         return self
 
-    def __exit__(self, _exc_type, _exc_val, _exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore[no-untyped-def]
+        del exc_type, exc_val, exc_tb  # Unused but required by context manager protocol
         if self.pbar is not None:
             self.pbar.close()
         self.file_obj.close()
@@ -123,6 +124,7 @@ def _upload_file_with_progress(
 def _encrypt_file_if_needed(
     file_path: Path,
     encryption_password: Optional[str] = None,
+    data_key: Optional[DataKey] = None,
 ) -> tuple[Optional[Path], Optional[Path]]:
     """
     Encrypt file if encryption_password is provided, file is large enough, and file extension is encryptable.
@@ -130,12 +132,13 @@ def _encrypt_file_if_needed(
     Args:
         file_path: Path to the file to encrypt
         encryption_password: Password for encryption
+        data_key: DataKey for encryption (required if encryption_password is provided)
 
     Returns:
         Tuple of (encrypted_file_path, temp_dir_path)
         If no encryption, returns (None, None)
     """
-    if not encryption_password:
+    if not encryption_password or not data_key:
         return None, None
 
     # Check file size first (only encrypt files >= 5MB)
@@ -147,21 +150,18 @@ def _encrypt_file_if_needed(
     if file_path.suffix.lower() not in ENCRYPTABLE_EXTENSIONS:
         return None, None
 
-    # Generate encryption key from password
-    encryption_key = KeyManager.generate_key(encryption_password)
-
-    # Create temporary directory for encrypted file
     temp_dir = Path(tempfile.mkdtemp())
     encrypted_file = temp_dir / file_path.name
-    manifest_path = temp_dir / "xpai_encryption_manifest.enc"
 
     # Encrypt the file
-    enc_file(
+    envelope_enc_file(
         source=file_path,
+        password=encryption_password,
+        data_key=data_key,
         dest=encrypted_file,
-        encryption_key=encryption_key,
-        manifest_path=manifest_path,
+        chunked=True,
     )
+    print(f"Encrypted file path: {encrypted_file}")
 
     return encrypted_file, temp_dir
 
@@ -349,15 +349,10 @@ def upload_folder(
     if '.git' not in ignore_patterns and '.git/' not in ignore_patterns:
         ignore_patterns.append('.git')
     ignore_patterns.append('.gitattributes')
-
-    # Prepare encryption if needed
-    encryption_key = None
-    if encryption_password:
-        encryption_key = KeyManager.generate_key(encryption_password)
-
     # Create or use temporary directory for encrypted files
     temp_dir_path: Optional[Path] = None
-    if encryption_key:
+    data_key: Optional[DataKey] = None
+    if encryption_password:
         if temp_dir:
             # Use user-specified temp directory
             temp_dir_path = Path(temp_dir)
@@ -365,6 +360,8 @@ def upload_folder(
         else:
             # Auto-create temp directory
             temp_dir_path = Path(tempfile.mkdtemp())
+        # Generate data key for encryption
+        data_key = client.generate_data_key(encryption_password)
 
     try:
         # Create session and API URL
@@ -381,6 +378,7 @@ def upload_folder(
         files_to_upload = []
         large_files = []  # Files >= 5MB
         small_files = []  # Files < 5MB
+        has_encrypted_files = False  # Track if any file was encrypted
         for root, dirs, files in os.walk(folder_path):
             # Filter directories
             dirs[:] = [d for d in dirs if not any(
@@ -405,19 +403,22 @@ def upload_folder(
 
                 # Encrypt file if needed (only for large files with specific extensions)
                 actual_file = local_file
-                if (encryption_key and temp_dir_path and
+                if (encryption_password and temp_dir_path and data_key and
                     file_size >= 5 * 1024 * 1024 and  # Only encrypt files >= 5MB
                     local_file.suffix.lower() in ENCRYPTABLE_EXTENSIONS):
                     # Preserve directory structure in temp dir
                     encrypted_file = temp_dir_path / rel_file_path
                     encrypted_file.parent.mkdir(parents=True, exist_ok=True)
-                    enc_file(
+                    envelope_enc_file(
                         source=local_file,
+                        password=encryption_password,
+                        data_key=data_key,
                         dest=encrypted_file,
-                        encryption_key=encryption_key,
-                        manifest_path=temp_dir_path / "xpai_encryption_manifest.enc",
+                        chunked=True,
                     )
+                    print(f"Encrypted file path: {encrypted_file}")
                     actual_file = encrypted_file
+                    has_encrypted_files = True
 
                 # Determine upload method based on size
                 if file_size >= 5 * 1024 * 1024:  # 5MB
@@ -441,24 +442,6 @@ def upload_folder(
                     })
                     small_files.append((rel_file_path, file_size))
 
-        # Add encryption manifest file if it exists
-        if temp_dir_path:
-            manifest_file = temp_dir_path / "xpai_encryption_manifest.enc"
-            if manifest_file.exists():
-                manifest_size = manifest_file.stat().st_size
-                # Read manifest file content
-                with open(manifest_file, 'rb') as f:
-                    manifest_content = f.read()
-                manifest_b64 = base64.b64encode(manifest_content).decode('utf-8')
-
-                # Add to files to upload
-                files_to_upload.append({
-                    "path": "xpai_encryption_manifest.enc",
-                    "content": manifest_b64,
-                    "size": manifest_size,
-                })
-                small_files.append(("xpai_encryption_manifest.enc", manifest_size))
-
         print(f"Found {len(files_to_upload)} files to upload ({len(large_files)} large files, {len(small_files)} small files)")
         # Show progress for small files upload
         small_files_pbar = None
@@ -475,7 +458,6 @@ def upload_folder(
         # Upload files via API
         try:
             result = _upload_files_via_api(session, api_url, files_to_upload, commit_message, revision)
-
             # Update progress bar for small files
             if small_files_pbar is not None:
                 for _, size in small_files:
@@ -499,6 +481,14 @@ def upload_folder(
                     )
                     print(f"Upload completed: {remote_path}")
 
+        # Set repository encrypted flag if any file was encrypted
+        if has_encrypted_files:
+            try:
+                client.set_repository_encrypted(organization, repo_type, repo_name)
+                print(f"Set repository encrypted flag for {repo_id}")
+            except Exception as e:
+                print(f"Warning: Failed to set repository encrypted flag: {e}")
+
         print(f"Successfully uploaded to {repo_id}")
         return result
     finally:
@@ -565,9 +555,15 @@ def upload_file(
     if not path_file.is_file():
         raise ValueError(f"Path is not a file: {path_file}")
 
+    # Generate data key if encryption is needed
+    data_key: Optional[DataKey] = None
+    if encryption_password:
+        data_key = client.generate_data_key(encryption_password)
+
     # Encrypt file if needed
-    encrypted_file, temp_dir = _encrypt_file_if_needed(path_file, encryption_password)
+    encrypted_file, temp_dir = _encrypt_file_if_needed(path_file, encryption_password, data_key)
     actual_file = encrypted_file if encrypted_file else path_file
+    file_was_encrypted = encrypted_file is not None
 
     try:
         # Create session and API URL
@@ -582,15 +578,6 @@ def upload_file(
 
         # Check file size
         file_size = actual_file.stat().st_size
-
-        # Check if encryption manifest file exists
-        manifest_file = None
-        if temp_dir:
-            manifest_path = temp_dir / "xpai_encryption_manifest.enc"
-            if manifest_path.exists():
-                manifest_file = manifest_path
-                print(f"Found encryption manifest file: xpai_encryption_manifest.enc")
-
         # Upload main file
         if file_size >= 5 * 1024 * 1024:  # 5MB
             # Large file
@@ -599,16 +586,13 @@ def upload_file(
             # Small file
             result = _upload_small_file(session, api_url, actual_file, path_in_repo, commit_message, revision)
 
-        # Upload manifest file if it exists
-        if manifest_file:
-            _upload_small_file(
-                session,
-                api_url,
-                manifest_file,
-                "xpai_encryption_manifest.enc",
-                f"{commit_message} (manifest)",
-                revision
-            )
+        # Set repository encrypted flag if file was encrypted
+        if file_was_encrypted:
+            try:
+                client.set_repository_encrypted(organization, repo_type, repo_name)
+                print(f"Set repository encrypted flag for {repo_id}")
+            except Exception as e:
+                print(f"Warning: Failed to set repository encrypted flag: {e}")
 
         print(f"Successfully uploaded {path_in_repo} to {repo_id}")
         return result