lou
159 lines
6.8 KiB
Python
159 lines
6.8 KiB
Python
"""服务端集成测试: 真实 FastAPI + 真实 PG + 真实加密/分卷 (非 mock)
|
|
|
|
流程: 客户端 crypto 加密 -> splitter 切卷 -> metadata 组装 init json
|
|
-> 通过 HTTP 上传 -> complete -> 验证明文还原 + 入库
|
|
|
|
运行: cd /home/lou/文档/server && .venv/bin/python -m unittest discover -s tests -v
|
|
注意: 需要本地 PG (7zencrypt_test 库) 与客户端项目 (7z-encrypt) 的 crypto/splitter/metadata
|
|
"""
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# 测试用配置 (先于 api 导入)
|
|
os.environ["SZ_DB_PATH"] = "/tmp/sz-test.db"
|
|
os.environ["SZ_TMP_ROOT"] = "/tmp/sz-test-tmp"
|
|
os.environ["SZ_STORAGE_ROOT"] = "/tmp/sz-test-storage"
|
|
os.environ["SZ_KEYRING"] = "/tmp/sz-test-keyring.json"
|
|
|
|
from fastapi.testclient import TestClient # noqa: E402
|
|
|
|
from db import ServerDB # noqa: E402
|
|
# 注意: api 在 setUpClass 里延迟 import (模块级 Keyring 单例需在密钥文件就绪后加载)
|
|
|
|
# 客户端模块 (7z-encrypt 项目)
|
|
CLIENT_DIR = "/home/lou/文档/7z-encrypt"
|
|
sys.path.insert(0, CLIENT_DIR)
|
|
from crypto import CryptoEngine # type: ignore[import-not-found] # noqa: E402
|
|
from metadata import build_init_json # type: ignore[import-not-found] # noqa: E402
|
|
from splitter import split_stream # type: ignore[import-not-found] # noqa: E402
|
|
|
|
import shutil # noqa: E402
|
|
|
|
TEST_KEY_ID = "test_key"
|
|
|
|
|
|
def _prepare(keyring_path: str, data: bytes, chunk_size: int = 1 << 20):
|
|
"""客户端管线: 加密 -> 分卷 -> init json, 返回 (init_json, chunk_files)"""
|
|
engine = CryptoEngine(keyring_path)
|
|
enc = io.BytesIO()
|
|
params = engine.encrypt_stream(io.BytesIO(data), enc, TEST_KEY_ID)
|
|
enc_bytes = enc.getvalue()
|
|
total_sha256 = hashlib.sha256(enc_bytes).hexdigest()
|
|
|
|
tmpdir = tempfile.mkdtemp(prefix="sz-chunks-")
|
|
manifest = split_stream(io.BytesIO(enc_bytes), chunk_size, os.path.join(tmpdir, "c"))
|
|
chunk_files = {ch["index"]: ch["filename"] for ch in manifest}
|
|
init_json = build_init_json("test.bin", len(data), total_sha256, params, manifest)
|
|
return init_json, chunk_files
|
|
|
|
|
|
class ServerPipelineTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
# 清空测试库 (SQLite: 删文件重建, 代替 TRUNCATE)
|
|
for f in ("/tmp/sz-test.db", "/tmp/sz-test.db-wal", "/tmp/sz-test.db-shm"):
|
|
os.path.exists(f) and os.unlink(f)
|
|
db = ServerDB()
|
|
# 清空运行时目录 + 旧密钥文件
|
|
for d in ("/tmp/sz-test-tmp", "/tmp/sz-test-storage"):
|
|
shutil.rmtree(d, ignore_errors=True)
|
|
os.path.exists(os.environ["SZ_KEYRING"]) and os.unlink(os.environ["SZ_KEYRING"])
|
|
# 客户端密钥 -> 服务端 keyring (预共享, 须在 import api 之前就绪)
|
|
client_kr = os.path.join(tempfile.mkdtemp(prefix="sz-kr-"), "keyring.json")
|
|
engine = CryptoEngine(client_kr)
|
|
engine.generate_key(TEST_KEY_ID)
|
|
os.makedirs(os.path.dirname(os.environ["SZ_KEYRING"]), exist_ok=True)
|
|
shutil.copy(client_kr, os.environ["SZ_KEYRING"])
|
|
cls.client_kr = client_kr
|
|
# 延迟 import: Keyring 单例此时才能读到正确密钥
|
|
from api import app
|
|
cls.client = TestClient(app)
|
|
|
|
def _upload(self, init_json, chunk_files):
|
|
r = self.client.post("/api/transfer/init", json=init_json)
|
|
self.assertEqual(r.status_code, 200, r.text)
|
|
tid = r.json()["transfer_id"]
|
|
for idx in sorted(chunk_files):
|
|
with open(chunk_files[idx], "rb") as f:
|
|
r = self.client.put(f"/api/transfer/{tid}/chunk/{idx}", content=f.read())
|
|
self.assertEqual(r.status_code, 200, r.text)
|
|
r = self.client.get(f"/api/transfer/{tid}/chunks")
|
|
self.assertEqual(r.status_code, 200)
|
|
self.assertEqual(len(r.json()["received"]), len(chunk_files))
|
|
return tid
|
|
|
|
def test_full_pipeline_roundtrip(self):
|
|
data = os.urandom(3 << 20) # 3MB
|
|
init_json, chunk_files = _prepare(self.client_kr, data)
|
|
tid = self._upload(init_json, chunk_files)
|
|
|
|
r = self.client.post(f"/api/transfer/{tid}/complete")
|
|
self.assertEqual(r.status_code, 200, r.text)
|
|
self.assertEqual(r.json()["status"], "done")
|
|
file_id = r.json()["file_id"]
|
|
self.assertTrue(file_id)
|
|
|
|
# 存储目录明文 == 原始数据
|
|
db = ServerDB()
|
|
row = None
|
|
with db._conn() as conn:
|
|
row = conn.execute("SELECT path, size FROM files WHERE file_id = ?", (file_id,)).fetchone()
|
|
self.assertIsNotNone(row)
|
|
with open(row["path"], "rb") as f:
|
|
restored = f.read()
|
|
self.assertEqual(restored, data)
|
|
self.assertEqual(row["size"], len(data))
|
|
|
|
# transfer 状态 done
|
|
with db._conn() as conn:
|
|
st = conn.execute("SELECT status FROM transfers WHERE transfer_id = ?", (tid,)).fetchone()
|
|
self.assertEqual(st["status"], "done")
|
|
|
|
def test_bad_chunk_rejected(self):
|
|
data = os.urandom(1 << 20)
|
|
init_json, chunk_files = _prepare(self.client_kr, data, chunk_size=512 * 1024)
|
|
r = self.client.post("/api/transfer/init", json=init_json)
|
|
tid = r.json()["transfer_id"]
|
|
# 篡改第一卷
|
|
bad = bytearray(open(chunk_files[1], "rb").read())
|
|
bad[10] ^= 0xFF
|
|
r = self.client.put(f"/api/transfer/{tid}/chunk/1", content=bytes(bad))
|
|
self.assertEqual(r.status_code, 409)
|
|
|
|
def test_incomplete_complete_409(self):
|
|
data = os.urandom(1 << 20)
|
|
init_json, chunk_files = _prepare(self.client_kr, data, chunk_size=512 * 1024)
|
|
r = self.client.post("/api/transfer/init", json=init_json)
|
|
tid = r.json()["transfer_id"]
|
|
# 只传 1 卷就 complete
|
|
with open(chunk_files[1], "rb") as f:
|
|
self.client.put(f"/api/transfer/{tid}/chunk/1", content=f.read())
|
|
r = self.client.post(f"/api/transfer/{tid}/complete")
|
|
self.assertEqual(r.status_code, 409)
|
|
|
|
def test_duplicate_chunk_idempotent(self):
|
|
data = os.urandom(512 * 1024)
|
|
init_json, chunk_files = _prepare(self.client_kr, data, chunk_size=256 * 1024)
|
|
r = self.client.post("/api/transfer/init", json=init_json)
|
|
tid = r.json()["transfer_id"]
|
|
with open(chunk_files[1], "rb") as f:
|
|
payload = f.read()
|
|
self.assertEqual(self.client.put(f"/api/transfer/{tid}/chunk/1", content=payload).status_code, 200)
|
|
# 重复传同一卷 (断点续传重发) -> 幂等 ok
|
|
self.assertEqual(self.client.put(f"/api/transfer/{tid}/chunk/1", content=payload).status_code, 200)
|
|
|
|
def test_unknown_transfer_404(self):
|
|
self.assertEqual(self.client.get("/api/transfer/nope/chunks").status_code, 404)
|
|
self.assertEqual(self.client.put("/api/transfer/nope/chunk/1", content=b"x").status_code, 404)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|