cmshoppe/tests/test_update_installer.py

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

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):
    trusted_hosts = {"updates.example.test"}

    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_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,
                trusted_hosts=self.trusted_hosts,
                opener=opener,
            )
            reused = update_installer.prepare_update(
                self.metadata(payload),
                install_root,
                trusted_hosts=self.trusted_hosts,
                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_download_rejects_hash_size_http_and_untrusted_host(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="http://updates.example.test/app.zip"),
                self.metadata(payload, download_url="https://other.example.test/app.zip"),
            )
            for metadata in cases:
                with self.subTest(metadata=metadata):
                    with self.assertRaises(update_installer.UpdateInstallError):
                        update_installer.download_package(
                            metadata,
                            root,
                            trusted_hosts=self.trusted_hosts,
                            opener=FakeOpener(payload),
                        )

    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,
                    trusted_hosts=self.trusted_hosts,
                    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()