525 lines
20 KiB
Python
525 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import io
|
||
import json
|
||
import os
|
||
import sys
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import urllib3
|
||
from minio import Minio
|
||
from minio.error import S3Error
|
||
|
||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
if PROJECT_ROOT not in sys.path:
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
from config import (
|
||
MINIO_ENABLED,
|
||
MINIO_ENDPOINT,
|
||
MINIO_ACCESS_KEY,
|
||
MINIO_SECRET_KEY,
|
||
MINIO_BUCKET,
|
||
MINIO_SECURE,
|
||
APK_MINIO_MAX_BYTES,
|
||
)
|
||
|
||
DOWNLOAD_QUEUE_KEY = "download-queue/current_batch.json"
|
||
RESULTS_PREFIX = "download-results/"
|
||
APKS_PREFIX = "apks/"
|
||
|
||
|
||
class MinioStorage:
|
||
"""
|
||
Minio 对象存储客户端
|
||
|
||
注意：此类依赖 Minio 服务，仅在 MINIO_ENABLED=True 时可用。
|
||
跨团队部署时，如果没有 Minio 服务，请在配置中设置 MINIO_ENABLED=False。
|
||
|
||
降级方案：
|
||
- 当 MINIO_ENABLED=False 时，直接实例化会抛出 RuntimeError
|
||
- 推荐使用 create() 工厂方法，它会在禁用时返回 None
|
||
- 调用方应检查返回值并使用直接下载模式（APK_DOWNLOAD_MODE='direct'）
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
endpoint: str = MINIO_ENDPOINT,
|
||
access_key: str = MINIO_ACCESS_KEY,
|
||
secret_key: str = MINIO_SECRET_KEY,
|
||
bucket: str = MINIO_BUCKET,
|
||
secure: bool = MINIO_SECURE,
|
||
):
|
||
if not MINIO_ENABLED:
|
||
raise RuntimeError(
|
||
"Minio 功能已禁用（MINIO_ENABLED=False）。\n"
|
||
"如需使用 Minio，请在配置文件中设置 MINIO_ENABLED=True 并配置 Minio 连接信息。\n"
|
||
"跨团队部署时，推荐使用 create() 工厂方法和 APK_DOWNLOAD_MODE='direct' 直接下载模式。"
|
||
)
|
||
|
||
if not endpoint or not access_key or not secret_key:
|
||
raise ValueError(
|
||
"Minio 配置不完整。请在配置文件中设置：\n"
|
||
" - MINIO_ENDPOINT\n"
|
||
" - MINIO_ACCESS_KEY\n"
|
||
" - MINIO_SECRET_KEY"
|
||
)
|
||
|
||
self.endpoint = endpoint
|
||
self.bucket = bucket
|
||
http_client = urllib3.PoolManager(
|
||
timeout=urllib3.Timeout(connect=10, read=60),
|
||
maxsize=30,
|
||
retries=urllib3.Retry(total=3, backoff_factor=0.3),
|
||
)
|
||
self.client = Minio(endpoint, access_key=access_key, secret_key=secret_key,
|
||
secure=secure, http_client=http_client)
|
||
self._ensure_bucket()
|
||
|
||
@classmethod
|
||
def create(
|
||
cls,
|
||
endpoint: str = MINIO_ENDPOINT,
|
||
access_key: str = MINIO_ACCESS_KEY,
|
||
secret_key: str = MINIO_SECRET_KEY,
|
||
bucket: str = MINIO_BUCKET,
|
||
secure: bool = MINIO_SECURE,
|
||
) -> Optional["MinioStorage"]:
|
||
"""
|
||
工厂方法：创建 MinioStorage 实例，如果 MINIO_ENABLED=False 则返回 None。
|
||
|
||
用法示例：
|
||
storage = MinioStorage.create()
|
||
if storage is None:
|
||
# 降级到直接下载模式
|
||
print("Minio 已禁用，使用直接下载模式")
|
||
return
|
||
|
||
返回:
|
||
MinioStorage 实例或 None（当 MINIO_ENABLED=False 时）
|
||
"""
|
||
if not MINIO_ENABLED:
|
||
return None
|
||
return cls(endpoint, access_key, secret_key, bucket, secure)
|
||
|
||
@staticmethod
|
||
def is_enabled() -> bool:
|
||
"""
|
||
检查 Minio 功能是否启用。
|
||
|
||
返回:
|
||
bool: True 表示 Minio 已启用并可用
|
||
"""
|
||
return MINIO_ENABLED
|
||
|
||
def _ensure_bucket(self) -> None:
|
||
if not self.client.bucket_exists(self.bucket):
|
||
self.client.make_bucket(self.bucket)
|
||
|
||
# ── low-level ops ──────────────────────────────────────────────
|
||
|
||
def upload_file(self, local_path: str, object_name: str) -> str:
|
||
with open(local_path, "rb") as fh:
|
||
fh.seek(0, os.SEEK_END)
|
||
size = fh.tell()
|
||
fh.seek(0)
|
||
result = self.client.put_object(
|
||
self.bucket, object_name, fh, length=size,
|
||
)
|
||
return result.etag or ""
|
||
|
||
def upload_json(self, data: Any, object_name: str) -> str:
|
||
payload = json.dumps(data, ensure_ascii=False, indent=2).encode("utf-8")
|
||
self.client.put_object(self.bucket, object_name, io.BytesIO(payload),
|
||
length=len(payload), content_type="application/json")
|
||
return object_name
|
||
|
||
def download_file(self, object_name: str, local_path: str,
|
||
*, overwrite: bool = True) -> str:
|
||
os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
|
||
self.client.fget_object(self.bucket, object_name, local_path)
|
||
return local_path
|
||
|
||
def download_json(self, object_name: str, retries: int = 3) -> Optional[Any]:
|
||
import time as _time
|
||
for attempt in range(retries):
|
||
try:
|
||
response = self.client.get_object(self.bucket, object_name)
|
||
data = response.read()
|
||
response.close()
|
||
response.release_conn()
|
||
return json.loads(data.decode("utf-8"))
|
||
except S3Error as exc:
|
||
if exc.code == "NoSuchKey":
|
||
return None
|
||
if attempt < retries - 1:
|
||
_time.sleep(1 * (attempt + 1))
|
||
else:
|
||
raise
|
||
except Exception as exc:
|
||
if attempt < retries - 1:
|
||
_time.sleep(1 * (attempt + 1))
|
||
else:
|
||
raise
|
||
return None
|
||
|
||
def list_objects(self, prefix: str = "") -> List[Dict[str, Any]]:
|
||
objects = self.client.list_objects(self.bucket, prefix=prefix, recursive=True)
|
||
return [
|
||
{"name": obj.object_name, "size": obj.size or 0,
|
||
"last_modified": obj.last_modified.isoformat() if obj.last_modified else ""}
|
||
for obj in objects
|
||
]
|
||
|
||
def delete_object(self, object_name: str) -> None:
|
||
self.client.remove_object(self.bucket, object_name)
|
||
|
||
def object_exists(self, object_name: str) -> bool:
|
||
try:
|
||
self.client.stat_object(self.bucket, object_name)
|
||
return True
|
||
except S3Error as exc:
|
||
if exc.code == "NoSuchKey":
|
||
return False
|
||
raise
|
||
|
||
# ── download queue ─────────────────────────────────────────────
|
||
|
||
def push_download_queue(self, tasks: List[Dict[str, str]]) -> None:
|
||
tasks = [t for t in tasks if self._is_valid_queue_entry(t)]
|
||
self.upload_json(tasks, DOWNLOAD_QUEUE_KEY)
|
||
|
||
def prepend_to_download_queue(self, tasks: List[Dict[str, str]]) -> None:
|
||
existing = self.pull_download_queue()
|
||
existing = [t for t in existing if self._is_valid_queue_entry(t)]
|
||
existing_packages = {t.get("package_name", "") for t in existing}
|
||
for task in reversed(tasks):
|
||
pkg = task.get("package_name", "")
|
||
if pkg in existing_packages:
|
||
existing = [t for t in existing if t.get("package_name") != pkg]
|
||
existing.insert(0, task)
|
||
self.upload_json(existing, DOWNLOAD_QUEUE_KEY)
|
||
|
||
def pull_download_queue(self) -> List[Dict[str, str]]:
|
||
import time as _time
|
||
for attempt in range(3):
|
||
try:
|
||
data = self.download_json(DOWNLOAD_QUEUE_KEY)
|
||
if not isinstance(data, list):
|
||
return []
|
||
return [t for t in data if self._is_valid_queue_entry(t)]
|
||
except Exception:
|
||
if attempt == 2:
|
||
return []
|
||
_time.sleep(1 * (attempt + 1))
|
||
|
||
@staticmethod
|
||
def _is_valid_queue_entry(task: Dict[str, str]) -> bool:
|
||
pkg = str(task.get("package_name") or "").strip()
|
||
if not pkg:
|
||
return False
|
||
import re
|
||
if not re.match(r'^[a-zA-Z][a-zA-Z0-9_.]{2,}$', pkg):
|
||
return False
|
||
return True
|
||
|
||
# ── APK upload / download ──────────────────────────────────────
|
||
|
||
def upload_apk(self, package_name: str, local_dir: str,
|
||
download_date: str, version_code: str = "",
|
||
delete_after_upload: bool = False) -> Dict[str, Any]:
|
||
manifest: Dict[str, Any] = {
|
||
"package_name": package_name,
|
||
"download_date": download_date,
|
||
"version_code": version_code,
|
||
"files": [],
|
||
}
|
||
pkg_dir = os.path.join(local_dir, package_name)
|
||
if not os.path.isdir(pkg_dir):
|
||
raise FileNotFoundError(f"APK directory not found: {pkg_dir}")
|
||
|
||
for filename in sorted(os.listdir(pkg_dir)):
|
||
filepath = os.path.join(pkg_dir, filename)
|
||
if not os.path.isfile(filepath):
|
||
continue
|
||
if not filename.endswith((".apk", ".xapk", ".apkm")):
|
||
continue
|
||
|
||
md5 = _file_md5(filepath)
|
||
filesize = os.path.getsize(filepath)
|
||
remote_name = f"{APKS_PREFIX}{package_name}/{download_date}_{version_code}/{filename}"
|
||
self.upload_file(filepath, remote_name)
|
||
|
||
manifest["files"].append({
|
||
"remote_name": remote_name,
|
||
"filename": filename,
|
||
"size": filesize,
|
||
"md5": md5,
|
||
})
|
||
|
||
return manifest
|
||
|
||
def download_apk(self, package_name: str, manifest: Dict[str, Any],
|
||
dest_dir: str) -> List[str]:
|
||
local_paths: List[str] = []
|
||
pkg_dir = os.path.join(dest_dir, package_name)
|
||
|
||
for old_file in os.listdir(pkg_dir) if os.path.isdir(pkg_dir) else []:
|
||
if old_file.endswith(('.apk', '.xapk', '.apkm')):
|
||
try:
|
||
os.remove(os.path.join(pkg_dir, old_file))
|
||
except OSError:
|
||
pass
|
||
os.makedirs(pkg_dir, exist_ok=True)
|
||
|
||
# 优先使用 uploaded_files（包含 remote_name），回退到 files
|
||
file_list = manifest.get("uploaded_files") or manifest.get("files") or []
|
||
for file_info in file_list:
|
||
remote_name = file_info["remote_name"]
|
||
filename = file_info.get("filename", os.path.basename(remote_name))
|
||
local_path = os.path.join(pkg_dir, filename)
|
||
|
||
self.download_file(remote_name, local_path)
|
||
|
||
actual_md5 = _file_md5(local_path)
|
||
expected_md5 = file_info.get("md5", "")
|
||
if expected_md5 and actual_md5 != expected_md5:
|
||
os.unlink(local_path)
|
||
raise IOError(
|
||
f"MD5 mismatch for {filename}: expected {expected_md5}, got {actual_md5}"
|
||
)
|
||
|
||
local_paths.append(local_path)
|
||
|
||
return local_paths
|
||
|
||
# ── results ────────────────────────────────────────────────────
|
||
|
||
def write_download_result(self, package_name: str,
|
||
result: Dict[str, Any]) -> None:
|
||
key = f"{RESULTS_PREFIX}{package_name}.json"
|
||
self.upload_json(result, key)
|
||
|
||
def read_download_result(self, package_name: str) -> Optional[Dict[str, Any]]:
|
||
key = f"{RESULTS_PREFIX}{package_name}.json"
|
||
return self.download_json(key)
|
||
|
||
def list_download_results(self) -> Dict[str, Dict[str, Any]]:
|
||
results: Dict[str, Dict[str, Any]] = {}
|
||
for obj in self.list_objects(RESULTS_PREFIX):
|
||
if not obj["name"].endswith(".json"):
|
||
continue
|
||
package_name = obj["name"][len(RESULTS_PREFIX):-len(".json")]
|
||
if not package_name:
|
||
continue
|
||
data = self.download_json(obj["name"])
|
||
if isinstance(data, dict):
|
||
results[package_name] = data
|
||
return results
|
||
|
||
def delete_download_result(self, package_name: str) -> None:
|
||
key = f"{RESULTS_PREFIX}{package_name}.json"
|
||
self.delete_object(key)
|
||
|
||
# ── old APK cleanup ────────────────────────────────────────────
|
||
|
||
def list_apk_versions(self, package_name: str) -> List[str]:
|
||
prefix = f"{APKS_PREFIX}{package_name}/"
|
||
objects = self.list_objects(prefix)
|
||
dirs = set()
|
||
for obj in objects:
|
||
rel = obj["name"][len(prefix):]
|
||
parts = rel.split("/")
|
||
if len(parts) >= 2:
|
||
dirs.add(parts[0])
|
||
return sorted(dirs)
|
||
|
||
def delete_apk_version(self, package_name: str, version_dir: str) -> None:
|
||
prefix = f"{APKS_PREFIX}{package_name}/{version_dir}/"
|
||
for obj in self.list_objects(prefix):
|
||
self.delete_object(obj["name"])
|
||
|
||
def cleanup_old_versions(self, package_name: str, keep: int = 1) -> int:
|
||
versions = self.list_apk_versions(package_name)
|
||
if len(versions) <= keep:
|
||
return 0
|
||
deleted = 0
|
||
for version in versions[:-keep]:
|
||
self.delete_apk_version(package_name, version)
|
||
deleted += 1
|
||
return deleted
|
||
|
||
# ── storage limit enforcement ───────────────────────────────────
|
||
|
||
def list_all_versions_by_age(self) -> List[Dict[str, Any]]:
|
||
"""按版本聚合列出所有 APK 版本，按最早修改时间排序。
|
||
|
||
每个版本条目包含该版本下所有文件的总大小和最晚 last_modified。
|
||
这确保了 enforce_storage_limit 能正确统计每个版本的完整大小。
|
||
"""
|
||
# 先按文件遍历，按 (package_name, version_dir) 聚合
|
||
version_map: Dict[str, Dict[str, Any]] = {}
|
||
for obj in self.list_objects(APKS_PREFIX):
|
||
name = obj["name"]
|
||
if not name.endswith((".apk", ".xapk", ".apkm")):
|
||
continue
|
||
rel = name[len(APKS_PREFIX):]
|
||
parts = rel.split("/", 2)
|
||
if len(parts) < 3:
|
||
continue
|
||
package_name = parts[0]
|
||
version_dir = parts[1]
|
||
key = f"{package_name}/{version_dir}"
|
||
file_size = obj.get("size", 0)
|
||
file_modified = obj.get("last_modified", "")
|
||
|
||
if key not in version_map:
|
||
version_map[key] = {
|
||
"package_name": package_name,
|
||
"version_dir": version_dir,
|
||
"size": file_size,
|
||
"last_modified": file_modified,
|
||
}
|
||
else:
|
||
# 累加文件大小，取最晚的 last_modified
|
||
version_map[key]["size"] += file_size
|
||
if file_modified > version_map[key]["last_modified"]:
|
||
version_map[key]["last_modified"] = file_modified
|
||
|
||
versions = list(version_map.values())
|
||
versions.sort(key=lambda v: v.get("last_modified", ""))
|
||
return versions
|
||
|
||
def enforce_storage_limit(self, max_bytes: int = 0) -> dict:
|
||
limit = max_bytes or APK_MINIO_MAX_BYTES
|
||
# list_all_versions_by_age 已按版本聚合，size 为该版本所有文件的总大小
|
||
all_versions = self.list_all_versions_by_age()
|
||
current_size = sum(v["size"] for v in all_versions)
|
||
if current_size <= limit:
|
||
return {"before_bytes": current_size, "after_bytes": current_size,
|
||
"deleted_versions": 0, "deleted_bytes": 0}
|
||
if not all_versions:
|
||
return {"before_bytes": current_size, "after_bytes": current_size,
|
||
"deleted_versions": 0, "deleted_bytes": 0}
|
||
|
||
# 清理到 limit 的 80%
|
||
target = int(limit * 0.8)
|
||
deleted_bytes = 0
|
||
to_delete: List[Dict[str, Any]] = []
|
||
|
||
for ver in all_versions:
|
||
if current_size - deleted_bytes <= target:
|
||
break
|
||
to_delete.append(ver)
|
||
# ver["size"] 已是该版本所有文件的总大小，无需额外聚合
|
||
deleted_bytes += ver["size"]
|
||
|
||
for ver in to_delete:
|
||
self.delete_apk_version(ver["package_name"], ver["version_dir"])
|
||
|
||
after_versions = self.list_all_versions_by_age()
|
||
after_size = sum(v["size"] for v in after_versions)
|
||
return {
|
||
"before_bytes": current_size,
|
||
"after_bytes": after_size,
|
||
"deleted_versions": len(to_delete),
|
||
"deleted_bytes": deleted_bytes,
|
||
}
|
||
|
||
|
||
# ── helpers ────────────────────────────────────────────────────────
|
||
|
||
def _file_md5(filepath: str) -> str:
|
||
h = hashlib.md5()
|
||
with open(filepath, "rb") as fh:
|
||
for chunk in iter(lambda: fh.read(65536), b""):
|
||
h.update(chunk)
|
||
return h.hexdigest()
|
||
|
||
|
||
# ── test ───────────────────────────────────────────────────────────
|
||
|
||
def main() -> int:
|
||
print(f"MinIO enabled: {MINIO_ENABLED}")
|
||
if not MINIO_ENABLED:
|
||
print("MinIO is disabled. Set MINIO_ENABLED=True to run this test.")
|
||
return 0
|
||
|
||
print(f"MinIO endpoint: {MINIO_ENDPOINT}")
|
||
print(f"MinIO bucket: {MINIO_BUCKET}")
|
||
print(f"MinIO secure: {MINIO_SECURE}")
|
||
|
||
# 使用工厂方法创建实例
|
||
storage = MinioStorage.create()
|
||
if storage is None:
|
||
print("Failed to create MinioStorage instance (MINIO_ENABLED=False)")
|
||
return 1
|
||
|
||
# 1. Test push/pull download queue
|
||
test_tasks = [
|
||
{"package_name": "com.test.app1", "last_updated": "2026-05-01", "app_name": "Test App 1"},
|
||
{"package_name": "com.test.app2", "last_updated": "2026-05-15", "app_name": "Test App 2"},
|
||
]
|
||
print("\n[1] push_download_queue ...")
|
||
storage.push_download_queue(test_tasks)
|
||
print(" pushed OK")
|
||
|
||
pulled = storage.pull_download_queue()
|
||
print(f" pulled: {len(pulled)} tasks")
|
||
for t in pulled:
|
||
print(f" {t['package_name']}")
|
||
|
||
# 2. Test upload/download a dummy file
|
||
print("\n[2] upload / download file ...")
|
||
dummy_dir = "/tmp/apk_cloud_test/com.test.app1"
|
||
os.makedirs(dummy_dir, exist_ok=True)
|
||
dummy_path = os.path.join(dummy_dir, "base.apk")
|
||
with open(dummy_path, "w") as fh:
|
||
fh.write("fake apk content for testing\n")
|
||
manifest = storage.upload_apk(
|
||
"com.test.app1", "/tmp/apk_cloud_test", download_date="2026-05-19", version_code="1",
|
||
)
|
||
print(f" uploaded manifest: {json.dumps(manifest, indent=2)}")
|
||
|
||
# 3. Test write / read download result
|
||
result = {
|
||
"status": "ok",
|
||
"download_date": "2026-05-19T10:00:00",
|
||
"version_code": "1",
|
||
"files": manifest["files"],
|
||
}
|
||
storage.write_download_result("com.test.app1", result)
|
||
read_back = storage.read_download_result("com.test.app1")
|
||
print(f"\n[3] write/read result: status={read_back.get('status') if read_back else 'NONE'}")
|
||
|
||
# 4. Test download APK
|
||
print("\n[4] download APK ...")
|
||
os.makedirs("/tmp/apk_cloud_test_dl", exist_ok=True)
|
||
local_paths = storage.download_apk("com.test.app1", manifest, "/tmp/apk_cloud_test_dl")
|
||
print(f" downloaded: {local_paths}")
|
||
|
||
# 5. List results
|
||
print("\n[5] list_download_results ...")
|
||
all_results = storage.list_download_results()
|
||
for pkg, res in all_results.items():
|
||
print(f" {pkg}: {res.get('status')}")
|
||
|
||
# Cleanup test data
|
||
print("\n[6] cleanup ...")
|
||
storage.delete_download_result("com.test.app1")
|
||
storage.delete_object(DOWNLOAD_QUEUE_KEY)
|
||
storage.delete_apk_version("com.test.app1", "2026-05-19_1")
|
||
import shutil
|
||
shutil.rmtree("/tmp/apk_cloud_test", ignore_errors=True)
|
||
shutil.rmtree("/tmp/apk_cloud_test_dl", ignore_errors=True)
|
||
print(" done")
|
||
|
||
print("\nAll tests passed!")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|