cmautobuy/client/test/test_update_service.py

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

import base64
import hashlib
import io
import json
import tempfile
import unittest
import urllib.request
import zipfile
from pathlib import Path

from src.update_service import (
    DEFAULT_UPDATE_MANIFEST_URL,
    LEGACY_DEFAULT_UPDATE_MANIFEST_URL,
    UnsafeUpdateArchiveError,
    UpdateConfigurationError,
    UpdateIntegrityError,
    UpdateCredentials,
    UpdateService,
    _SameOriginRedirectHandler,
    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 = []
        self.requested_headers = []

    def __call__(self, request, timeout):
        self.requested_urls.append((request.full_url, timeout))
        self.requested_headers.append(dict(request.header_items()))
        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,
        )
        self.assertEqual(
            validate_manifest_url(DEFAULT_UPDATE_MANIFEST_URL),
            DEFAULT_UPDATE_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_basic_auth_is_used_for_manifest_and_update_without_repr_leak(self):
        update_content = make_update_zip()
        service, opener = self.service_with(update_content)
        credentials = UpdateCredentials("release-reader", "unit-test-password")

        result = service.check(
            self.manifest_url,
            "0.1.0",
            credentials=credentials,
        )
        service.download_and_stage(result.update, credentials=credentials)

        expected = "Basic " + base64.b64encode(
            b"release-reader:unit-test-password"
        ).decode("ascii")
        self.assertEqual(len(opener.requested_headers), 2)
        self.assertTrue(
            all(
                headers.get("Authorization") == expected
                for headers in opener.requested_headers
            )
        )
        self.assertNotIn("unit-test-password", repr(credentials))

    def test_legacy_unicode_manifest_path_is_encoded_for_http_request(self):
        request = UpdateService._make_request(
            LEGACY_DEFAULT_UPDATE_MANIFEST_URL,
            UpdateCredentials("release-reader", "unit-test-password"),
        )

        self.assertNotIn("——", request.full_url)
        self.assertIn("%E2%80%94%E2%80%94", request.full_url)

    def test_authenticated_redirect_cannot_change_origin(self):
        handler = _SameOriginRedirectHandler()
        request = urllib.request.Request(
            "https://updates.example.test/manifest.json",
            headers={"Authorization": "Basic test"},
        )

        with self.assertRaises(UpdateConfigurationError):
            handler.redirect_request(
                request,
                None,
                302,
                "Found",
                {},
                "https://other.example.test/manifest.json",
            )

    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.1.0")


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