140 lines
6.0 KiB
Python
140 lines
6.0 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
SOURCE_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(SOURCE_ROOT / "src"))
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
|
|
from fixture_image import build_fixture, first_fat_offset # noqa: E402
|
|
from test_config import sample_config # noqa: E402
|
|
from matrix_image_editor.config import create_config_file # noqa: E402
|
|
from matrix_image_editor.cli import run_command # noqa: E402
|
|
from matrix_image_editor.fat16 import Fat16Disk, Fat16Error # noqa: E402
|
|
from matrix_image_editor.operations import ( # noqa: E402
|
|
ImageOperationError,
|
|
copy_write_and_verify,
|
|
read_config,
|
|
write_and_verify,
|
|
)
|
|
from matrix_image_editor.paths import ImagePathError, cleanup_new_destination, validate_new_destination # noqa: E402
|
|
|
|
|
|
class ImageOperationTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temporary = tempfile.TemporaryDirectory(prefix="镜像编辑器-")
|
|
self.root = Path(self.temporary.name)
|
|
self.image = build_fixture(self.root / "原 镜像.img", create_config_file(sample_config()))
|
|
|
|
def tearDown(self) -> None:
|
|
self.temporary.cleanup()
|
|
|
|
def test_read_write_sidecar_and_reopen(self) -> None:
|
|
config = read_config(self.image)
|
|
config["wifi"]["ssid"] = "中文 新 WiFi"
|
|
progress: list[int] = []
|
|
verified = write_and_verify(self.image, config, progress.append)
|
|
self.assertEqual(verified["wifi"]["ssid"], "中文 新 WiFi")
|
|
sidecar = Path(str(self.image) + ".sha256")
|
|
raw = sidecar.read_bytes()
|
|
self.assertFalse(raw.startswith(b"\xef\xbb\xbf"))
|
|
self.assertEqual(raw, f"{hashlib.sha256(self.image.read_bytes()).hexdigest()} {self.image.name}\n".encode("utf-8"))
|
|
self.assertEqual(progress[-1], 100)
|
|
|
|
def test_rejects_bad_sidecar_forms_and_content(self) -> None:
|
|
sidecar = Path(str(self.image) + ".sha256")
|
|
digest = hashlib.sha256(self.image.read_bytes()).hexdigest()
|
|
invalid = [
|
|
b"\xef\xbb\xbf" + f"{digest} {self.image.name}\n".encode("utf-8"),
|
|
f"{digest.upper()} {self.image.name}\n".encode("utf-8"),
|
|
f"{digest} {self.image.name}\n".encode("utf-8"),
|
|
f"{digest} wrong.img\n".encode("utf-8"),
|
|
f"{'0' * 64} {self.image.name}\n".encode("utf-8"),
|
|
f"{digest} {self.image.name}\r\n".encode("utf-8"),
|
|
]
|
|
for payload in invalid:
|
|
with self.subTest(payload=payload[:12]):
|
|
sidecar.write_bytes(payload)
|
|
with self.assertRaises(ImageOperationError):
|
|
read_config(self.image)
|
|
|
|
def test_copy_apply_preserves_source_and_cleans_failure(self) -> None:
|
|
original = self.image.read_bytes()
|
|
config = sample_config()
|
|
config["account"]["username"] = "new_user"
|
|
target, verified = copy_write_and_verify(self.image, self.root / "新 镜像.img", config)
|
|
self.assertEqual(verified["account"]["username"], "new_user")
|
|
self.assertEqual(self.image.read_bytes(), original)
|
|
self.assertTrue(Path(str(target) + ".sha256").is_file())
|
|
|
|
failed = self.root / "失败.img"
|
|
wrong_version = sample_config()
|
|
wrong_version["software_version"] = "9.9.9"
|
|
with self.assertRaises(ImageOperationError):
|
|
copy_write_and_verify(self.image, failed, wrong_version)
|
|
self.assertFalse(failed.exists())
|
|
self.assertFalse(Path(str(failed) + ".sha256").exists())
|
|
|
|
def test_path_rejections(self) -> None:
|
|
with self.assertRaises(ImagePathError):
|
|
validate_new_destination(self.image, self.image)
|
|
existing = self.root / "existing.img"
|
|
existing.write_bytes(b"")
|
|
with self.assertRaises(ImagePathError):
|
|
validate_new_destination(self.image, existing)
|
|
orphan = self.root / "orphan.img"
|
|
Path(str(orphan) + ".sha256").write_text("orphan", encoding="utf-8")
|
|
with self.assertRaises(ImagePathError):
|
|
validate_new_destination(self.image, orphan)
|
|
with self.assertRaises(ImagePathError):
|
|
validate_new_destination(self.image, self.root / "wrong.txt")
|
|
if os.name == "nt":
|
|
with self.assertRaises(ImagePathError):
|
|
validate_new_destination(self.image, self.root / "CON.img")
|
|
|
|
def test_fat_cycle_and_partition_bounds_are_rejected(self) -> None:
|
|
cyclic = self.root / "cycle.img"
|
|
cyclic.write_bytes(self.image.read_bytes())
|
|
with cyclic.open("r+b") as stream:
|
|
stream.seek(first_fat_offset() + 2 * 2)
|
|
stream.write(struct.pack("<H", 2))
|
|
with Fat16Disk(cyclic, writable=False) as disk:
|
|
with self.assertRaises(Fat16Error):
|
|
disk.read_file("MSCCFG.BIN")
|
|
|
|
truncated = self.root / "truncated.img"
|
|
truncated.write_bytes(self.image.read_bytes()[:-512])
|
|
with self.assertRaises(Fat16Error):
|
|
Fat16Disk(truncated, writable=False)
|
|
|
|
def test_cli_contract_and_apply(self) -> None:
|
|
import contextlib
|
|
import io
|
|
|
|
output = io.StringIO()
|
|
with contextlib.redirect_stdout(output):
|
|
self.assertEqual(run_command(["--validate", str(self.image)]), 0)
|
|
document = json.loads(output.getvalue())
|
|
self.assertEqual(
|
|
set(document),
|
|
{"valid", "SoftwareVersion", "Username", "Ssid", "Mode", "Address"},
|
|
)
|
|
config_path = self.root / "config.json"
|
|
config_path.write_text(json.dumps(sample_config(), ensure_ascii=False), encoding="utf-8-sig")
|
|
target = self.root / "cli-new.img"
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
self.assertEqual(run_command(["--apply", str(self.image), str(config_path), str(target)]), 0)
|
|
self.assertTrue(target.is_file())
|
|
self.assertTrue(Path(str(target) + ".sha256").is_file())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|