cmshoppe/tests/test_update_installer.py

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

from app import release_manifest, update_installer


class FakeResponse(io.BytesIO):
    def __init__(self, payload, url):
        super().__init__(payload)
        self.url = url

    def geturl(self):
        return self.url

    def __enter__(self):
        return self

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


class FakeOpener:
    def __init__(self, payload, url="https://updates.example.test/app.zip"):
        self.payload = payload
        self.url = url
        self.calls = 0

    def open(self, _request, timeout=None):
        assert timeout
        self.calls += 1
        return FakeResponse(self.payload, self.url)


class UpdateInstallerTests(unittest.TestCase):
    @staticmethod
    def public_resolver(_host, _port):
        return ["93.184.216.34"]

    def make_zip(self, root, version="1.2.3"):
        package = Path(root) / "package"
        (package / "_internal").mkdir(parents=True)
        (package / "cmshopee.exe").write_bytes(b"new-exe")
        (package / "_internal" / "runtime.dll").write_bytes(b"runtime")
        (package / "_internal" / "empty.marker").write_bytes(b"")
        (package / "version.txt").write_text(version, encoding="ascii")
        (package / "README.txt").write_text("说明", encoding="utf-8")
        release_manifest.write_package_manifest(package, version)
        zip_path = Path(root) / "package.zip"
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as archive:
            for path in package.rglob("*"):
                if path.is_file():
                    archive.write(path, path.relative_to(package).as_posix())
        return zip_path.read_bytes()

    def metadata(self, payload, **changes):
        values = {
            "version": "1.2.3",
            "download_url": "https://updates.example.test/app.zip",
            "sha256": hashlib.sha256(payload).hexdigest(),
            "size_bytes": len(payload),
            "package_format": release_manifest.PACKAGE_FORMAT,
            "updater_protocol": release_manifest.UPDATER_PROTOCOL,
            "min_updater_protocol": 1,
        }
        values.update(changes)
        return update_installer.UpdatePackageMetadata(**values)

    def test_online_release_uses_internal_package_contract_defaults(self):
        info = SimpleNamespace(
            latest_version="0.1.6",
            download_url="http://downloads.external.test/cmshopee.zip",
            sha256="a" * 64,
            size_bytes=53318076,
        )

        metadata = update_installer.metadata_from_update_info(info)

        self.assertEqual(release_manifest.PACKAGE_FORMAT, metadata.package_format)
        self.assertEqual(release_manifest.UPDATER_PROTOCOL, metadata.updater_protocol)
        self.assertEqual(release_manifest.UPDATER_PROTOCOL, metadata.min_updater_protocol)

    def test_prepare_update_verifies_and_reuses_pending(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            payload = self.make_zip(temp_dir)
            install_root = Path(temp_dir) / "installed"
            install_root.mkdir()
            (install_root / "cmshopee.exe").write_bytes(b"old-exe")
            (install_root / "data").mkdir()
            opener = FakeOpener(payload)

            staged = update_installer.prepare_update(
                self.metadata(payload),
                install_root,
                resolver=self.public_resolver,
                opener=opener,
            )
            reused = update_installer.prepare_update(
                self.metadata(payload),
                install_root,
                resolver=self.public_resolver,
                opener=opener,
            )

            self.assertEqual(1, opener.calls)
            self.assertEqual(staged.staging_dir, reused.staging_dir)
            self.assertEqual(b"old-exe", (install_root / "cmshopee.exe").read_bytes())
            self.assertTrue((install_root / "data").is_dir())
            pending = json.loads(staged.pending_path.read_text(encoding="utf-8"))
            self.assertNotIn("download_url", pending)

    def test_prepare_update_accepts_public_http_and_writes_verified_pending(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            payload = self.make_zip(temp_dir)
            install_root = Path(temp_dir) / "installed"
            install_root.mkdir()
            metadata = self.metadata(
                payload,
                download_url="http://downloads.external.test/releases/app.zip",
            )

            staged = update_installer.prepare_update(
                metadata,
                install_root,
                resolver=self.public_resolver,
                opener=FakeOpener(payload, url=metadata.download_url),
            )

            pending = json.loads(staged.pending_path.read_text(encoding="utf-8"))
            self.assertEqual("verified", pending["stage"])
            self.assertEqual(metadata.sha256, pending["sha256"])

    def test_download_rejects_invalid_metadata_and_unsafe_urls(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            payload = b"zip"
            root = Path(temp_dir)
            cases = (
                self.metadata(payload, sha256=""),
                self.metadata(payload, size_bytes=0),
                self.metadata(payload, download_url="file:///tmp/app.zip"),
                self.metadata(payload, download_url="http://user:secret@example.test/app.zip"),
                self.metadata(payload, download_url="http://127.0.0.1/app.zip"),
                self.metadata(payload, download_url="http://10.0.0.8/app.zip"),
                self.metadata(payload, download_url="http://[::1]/app.zip"),
                self.metadata(payload, download_url="http://localhost/app.zip"),
                self.metadata(payload, download_url="http://[invalid/app.zip"),
            )
            for metadata in cases:
                with self.subTest(metadata=metadata):
                    with self.assertRaises(update_installer.UpdateInstallError):
                        update_installer.download_package(
                            metadata,
                            root,
                            resolver=self.public_resolver,
                            opener=FakeOpener(payload),
                        )

    def test_download_allows_public_http_https_domains_and_public_ip(self):
        payload = b"zip"
        urls = (
            "http://cdn-one.example.test:8080/releases/app.zip",
            "https://storage.other.test/app.zip?token=hidden",
            "http://8.8.8.8/app.zip",
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            for index, url in enumerate(urls):
                with self.subTest(url=url):
                    root = Path(temp_dir) / str(index)
                    downloaded = update_installer.download_package(
                        self.metadata(payload, download_url=url),
                        root,
                        resolver=self.public_resolver,
                        opener=FakeOpener(payload, url=url),
                    )
                    self.assertEqual(payload, downloaded.read_bytes())

    def test_download_rejects_domain_resolving_to_non_public_address(self):
        payload = b"zip"
        with tempfile.TemporaryDirectory() as temp_dir:
            with self.assertRaisesRegex(update_installer.UpdateInstallError, "本机或内网"):
                update_installer.download_package(
                    self.metadata(payload, download_url="http://files.example.test/app.zip"),
                    temp_dir,
                    resolver=lambda _host, _port: ["192.168.1.20"],
                    opener=FakeOpener(payload, url="http://files.example.test/app.zip"),
                )

    def test_redirect_handler_allows_public_cross_domain_and_protocol_redirects(self):
        handler = update_installer._SafeRedirectHandler(self.public_resolver)
        request = urllib.request.Request("https://origin.example.test/app.zip")
        for target in (
            "http://cdn.example.test:8080/releases/app.zip",
            "https://objects.other.test/app.zip",
        ):
            with self.subTest(target=target):
                redirected = handler.redirect_request(request, None, 302, "Found", {}, target)
                self.assertEqual(target, redirected.full_url)

        with self.assertRaisesRegex(update_installer.UpdateInstallError, "本机或内网"):
            handler.redirect_request(
                request,
                None,
                302,
                "Found",
                {},
                "http://169.254.169.254/latest/meta-data",
            )

    def test_download_revalidates_final_response_url(self):
        payload = b"zip"
        with tempfile.TemporaryDirectory() as temp_dir:
            with self.assertRaisesRegex(update_installer.UpdateInstallError, "本机或内网"):
                update_installer.download_package(
                    self.metadata(payload, download_url="http://cdn.example.test/app.zip"),
                    temp_dir,
                    resolver=self.public_resolver,
                    opener=FakeOpener(payload, url="http://127.0.0.1/app.zip"),
                )

    def test_http_download_still_requires_matching_hash(self):
        payload = b"zip"
        metadata = self.metadata(
            payload,
            download_url="http://cdn.example.test/app.zip",
            sha256="0" * 64,
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            with self.assertRaisesRegex(update_installer.UpdateInstallError, "完整性校验失败"):
                update_installer.download_package(
                    metadata,
                    temp_dir,
                    resolver=self.public_resolver,
                    opener=FakeOpener(payload, url=metadata.download_url),
                )
            self.assertFalse(
                (Path(temp_dir) / ".cmshopee-update/downloads/1.2.3.zip").exists()
            )

    def test_cancel_removes_partial_download(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            payload = b"zip-content"
            with self.assertRaises(update_installer.UpdateCancelled):
                update_installer.download_package(
                    self.metadata(payload),
                    temp_dir,
                    resolver=self.public_resolver,
                    opener=FakeOpener(payload),
                    cancelled=lambda: True,
                )
            self.assertFalse(
                (Path(temp_dir) / ".cmshopee-update/downloads/1.2.3.zip.part").exists()
            )

    def test_safe_extract_rejects_path_traversal_case_duplicates_and_symlink(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            for index, entries in enumerate(
                (
                    [("../escape.txt", b"x", None)],
                    [("same.txt", b"a", None), ("SAME.txt", b"b", None)],
                    [("link", b"target", (stat_mode := (0o120777 << 16)))],
                )
            ):
                zip_path = Path(temp_dir) / ("bad-%d.zip" % index)
                with zipfile.ZipFile(zip_path, "w") as archive:
                    for name, content, external_attr in entries:
                        info = zipfile.ZipInfo(name)
                        if external_attr is not None:
                            info.create_system = 3
                            info.external_attr = stat_mode
                        archive.writestr(info, content)
                with self.assertRaises(update_installer.UpdateInstallError):
                    update_installer.safe_extract(zip_path, Path(temp_dir) / ("out-%d" % index))

    def test_staging_rejects_manifest_tampering_and_extra_files(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            payload = self.make_zip(temp_dir)
            metadata = self.metadata(payload)
            package = Path(temp_dir) / "package"
            (package / "cmshopee.exe").write_bytes(b"tampered")
            with self.assertRaises(update_installer.UpdateInstallError):
                update_installer.validate_staging(package, metadata)

            (package / "cmshopee.exe").write_bytes(b"new-exe")
            (package / "extra.bin").write_bytes(b"extra")
            with self.assertRaises(update_installer.UpdateInstallError):
                update_installer.validate_staging(package, metadata)

    def test_zip_bomb_ratio_is_rejected(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            zip_path = Path(temp_dir) / "bomb.zip"
            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as archive:
                archive.writestr("large.bin", b"0" * 1024 * 1024)
            with self.assertRaises(update_installer.UpdateInstallError):
                update_installer.safe_extract(zip_path, Path(temp_dir) / "out")


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