cmautobuy/client/test/test_update_service.py

"""在线更新清单、下载校验和安全解压测试。"""

import hashlib
import io
import json
import tempfile
import unittest
import zipfile
from pathlib import Path

from src.update_service import (
    UnsafeUpdateArchiveError,
    UpdateConfigurationError,
    UpdateIntegrityError,
    UpdateService,
    mark_current_version_healthy,
    parse_version,
    validate_manifest_url,
)


class FakeResponse(io.BytesIO):
    def __init__(self, content: bytes, url: str, declared_size=None):
        super().__init__(content)
        self._url = url
        self.headers = {}
        if declared_size is not None:
            self.headers["Content-Length"] = str(declared_size)

    def geturl(self):
        return self._url

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        self.close()


class FakeUrlOpen:
    def __init__(self, responses):
        self.responses = responses
        self.requested_urls = []

    def __call__(self, request, timeout):
        self.requested_urls.append((request.full_url, timeout))
        content, final_url = self.responses[request.full_url]
        return FakeResponse(content, final_url, len(content))


def make_update_zip(version="0.2.0", extra_entries=None):
    output = io.BytesIO()
    with zipfile.ZipFile(output, "w") as archive:
        archive.writestr("app/CMAutoBuy.exe", b"exe")
        archive.writestr("app/version.txt", version.encode("utf-8"))
        for name, content in extra_entries or []:
            archive.writestr(name, content)
    return output.getvalue()


def make_manifest(update_content, version="0.2.0"):
    return json.dumps(
        {
            "schema_version": 1,
            "product": "CMAutoBuy",
            "version": version,
            "update": {
                "file": f"CMAutoBuy-{version}-update.zip",
                "size": len(update_content),
                "sha256": hashlib.sha256(update_content).hexdigest(),
            },
        }
    ).encode("utf-8")


class UpdateServiceTest(unittest.TestCase):
    def setUp(self):
        self.temporary_directory = tempfile.TemporaryDirectory()
        self.update_directory = Path(self.temporary_directory.name) / "update"
        self.manifest_url = "https://updates.example.test/releases/autobuy%E2%80%94%E2%80%94manifest.json"

    def tearDown(self):
        self.temporary_directory.cleanup()

    def service_with(self, update_content, manifest_content=None):
        file_name = "CMAutoBuy-0.2.0-update.zip"
        update_url = f"https://updates.example.test/releases/{file_name}"
        opener = FakeUrlOpen(
            {
                self.manifest_url: (
                    manifest_content or make_manifest(update_content),
                    self.manifest_url,
                ),
                update_url: (update_content, update_url),
            }
        )
        return UpdateService(self.update_directory, opener), opener

    def test_strict_version_comparison(self):
        self.assertLess(parse_version("1.2.3"), parse_version("1.2.4"))
        self.assertEqual(parse_version("1.2.3"), parse_version("1.2.3.0"))
        with self.assertRaises(UpdateConfigurationError):
            parse_version("v1.2")

    def test_only_https_without_credentials_is_allowed(self):
        self.assertEqual(
            validate_manifest_url(self.manifest_url),
            self.manifest_url,
        )
        for invalid in (
            "http://updates.example.test/manifest.json",
            "https://user:pass@updates.example.test/manifest.json",
            "https://updates.example.test/manifest.json?token=secret",
            "not-a-url",
        ):
            with self.subTest(invalid=invalid), self.assertRaises(
                UpdateConfigurationError
            ):
                validate_manifest_url(invalid)

    def test_check_reports_newer_and_current_versions(self):
        update_content = make_update_zip()
        service, _opener = self.service_with(update_content)

        newer = service.check(self.manifest_url, "0.1.0")
        current = service.check(self.manifest_url, "0.2.0")

        self.assertTrue(newer.available)
        self.assertEqual(newer.latest_version, "0.2.0")
        self.assertIsNotNone(newer.update)
        self.assertFalse(current.available)
        self.assertIsNone(current.update)

    def test_download_stages_verified_app_and_pending_state(self):
        update_content = make_update_zip()
        service, _opener = self.service_with(update_content)
        result = service.check(self.manifest_url, "0.1.0")
        progress = []

        staged = service.download_and_stage(result.update, on_progress=progress.append)

        self.assertEqual(staged, self.update_directory / "app.new")
        self.assertEqual((staged / "version.txt").read_text(), "0.2.0")
        pending = json.loads(
            (self.update_directory / "pending.json").read_text(encoding="utf-8")
        )
        self.assertEqual(pending["state"], "ready")
        self.assertEqual(pending["version"], "0.2.0")
        self.assertEqual(progress[-1], 100)

    def test_hash_mismatch_never_creates_pending_state(self):
        update_content = make_update_zip()
        manifest = json.loads(make_manifest(update_content))
        manifest["update"]["sha256"] = "0" * 64
        service, _opener = self.service_with(
            update_content,
            json.dumps(manifest).encode("utf-8"),
        )
        result = service.check(self.manifest_url, "0.1.0")

        with self.assertRaises(UpdateIntegrityError):
            service.download_and_stage(result.update)

        self.assertFalse((self.update_directory / "pending.json").exists())
        self.assertFalse((self.update_directory / "app.new").exists())

    def test_path_traversal_archive_is_rejected(self):
        update_content = make_update_zip(
            extra_entries=[("app/../../outside.txt", b"unsafe")]
        )
        service, _opener = self.service_with(update_content)
        result = service.check(self.manifest_url, "0.1.0")

        with self.assertRaises(UnsafeUpdateArchiveError):
            service.download_and_stage(result.update)

        self.assertFalse((self.update_directory.parent / "outside.txt").exists())
        self.assertFalse((self.update_directory / "pending.json").exists())

    def test_packaged_version_must_match_manifest(self):
        update_content = make_update_zip(version="9.9.9")
        service, _opener = self.service_with(update_content)
        result = service.check(self.manifest_url, "0.1.0")

        with self.assertRaises(UpdateIntegrityError):
            service.download_and_stage(result.update)

    def test_health_marker_is_only_written_for_pending_update(self):
        mark_current_version_healthy(self.update_directory)
        self.assertFalse((self.update_directory / "healthy.json").exists())
        self.update_directory.mkdir(parents=True)
        (self.update_directory / "pending.json").write_text("{}", encoding="utf-8")

        mark_current_version_healthy(self.update_directory)

        health = json.loads(
            (self.update_directory / "healthy.json").read_text(encoding="utf-8")
        )
        self.assertEqual(health["version"], "0.2.0")


if __name__ == "__main__":
    unittest.main()