def test_gptimage_constructor(small_file: Path) -> None:
"""Test the :class:`gpt.GPTImage` constructor."""
with pytest.raises(ValueError, match="Either fd or path must be given"):
gpt.GPTImage()
with pytest.raises(ValueError, match="Both fd and path can't be given"):
gpt.GPTImage(fd=0, path=small_file)
def test_gptimage_validate() -> None:
"""Test :meth:`gpt.GPTImage.validate`."""
with gpt.GPTImage(path=conftest.TESTDATA_DISK, open_mode=os.O_RDONLY) as image:
image.validate()
with open(conftest.TESTDATA_DISK, "rb") as fd:
with gpt.GPTImage(fd=fd.fileno()) as image:
image.validate()
def test_gptimage_too_small(small_file: Path) -> None:
"""Test :class:`gpt.GPTImage` with a (too) small file."""
with open(small_file, "rb") as fd:
with pytest.raises(gpt.InvalidImageError): # noqa: PT012
with gpt.GPTImage(fd=fd.fileno()):
pytest.fail("This code should not be reached") # pragma: no cover
with pytest.raises(gpt.InvalidImageError): # noqa: PT012
with gpt.GPTImage(path=small_file):
pytest.fail("This code should not be reached") # pragma: no cover
def test_gptimage_update_guid(disk_image: Path) -> None:
"""Test :meth:`gpt.GPTImage.update_guid`."""
new_guid = uuid.UUID("910fbf4b-f2ac-4d2c-8995-9dff42b9efde")
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDWR) as image:
image.update_guid(new_guid)
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
image.validate()
assert image.read_primary_gpt_header().disk_guid == new_guid
assert image.read_backup_gpt_header().disk_guid == new_guid
def test_calculate() -> None:
"""Test checksum calculation of an image twice."""
with gpt.GPTImage(path=conftest.TESTDATA_DISK,
open_mode=os.O_RDONLY) as image:
digest1 = checksum.calculate(image)
digest2 = checksum.calculate(image)
assert digest2 == digest1
assert digest1 == conftest.TESTDATA_EMBEDDED_DISK_GUID.bytes
with gpt.GPTImage(path=conftest.TESTDATA_EMBEDDED_DISK,
open_mode=os.O_RDONLY) as image:
assert checksum.calculate(
image) == conftest.TESTDATA_EMBEDDED_DISK_GUID.bytes
def test_calculate_inplace(disk_image: Path) -> None:
"""Test :func:`checksum.calculate` by modifying in image in-place."""
with open(disk_image, "rb") as fd:
real_hash = blake2b(fd)
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
digest = checksum.calculate(image)
assert digest != real_hash
# Overwrite CRCs and GUIDs with zeros, in-place
with open(disk_image, "rb+") as fd:
header_crc32_offset = 8 + 4 + 4
guid_offset = header_crc32_offset + 4 + 4 + 8 + 8 + 8 + 8
gpt.pwrite_all(fd.fileno(), b"\0" * 4,
gpt.MBR_SIZE + header_crc32_offset)
gpt.pwrite_all(fd.fileno(), b"\0" * 16, gpt.MBR_SIZE + guid_offset)
size = os.fstat(fd.fileno()).st_size
gpt.pwrite_all(fd.fileno(), b"\0" * 4,
size - gpt.LBA_SIZE + header_crc32_offset)
gpt.pwrite_all(fd.fileno(), b"\0" * 16,
size - gpt.LBA_SIZE + guid_offset)
with open(disk_image, "rb") as fd:
new_hash = blake2b(fd)
assert new_hash == digest
def test_gptimage_read_backup_gpt_header() -> None:
"""Test :meth:`gpt.GPTImage.read_backup_gpt_header`."""
with gpt.GPTImage(path=conftest.TESTDATA_DISK, open_mode=os.O_RDONLY) as image:
header = image.read_backup_gpt_header()
assert header.disk_guid == conftest.TESTDATA_DISK_GUID
assert header.first_usable_lba == 2048
assert header.last_usable_lba == 4062
assert header.current_lba == 4095
assert header.backup_lba == 1
def test_calculate_benchmark(
use_preadv: bool,
monkeypatch: pytest.MonkeyPatch,
benchmark: pytest_benchmark.fixture.BenchmarkFixture,
) -> None:
"""Benchmark :func:`checksum.calculate`."""
if not use_preadv and hasattr(os, "preadv"): # pragma: py-lt-37
monkeypatch.delattr(os, "preadv")
with gpt.GPTImage(path=conftest.TESTDATA_DISK,
open_mode=os.O_RDONLY) as image:
benchmark(checksum.calculate, image)
def test_gptimage_write_gpt_headers(disk_image: Path) -> None:
"""Test :meth:`gpt.GPTImage.write_gpt_headers`."""
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
primary = image.read_primary_gpt_header()
backup = image.read_backup_gpt_header()
new_guid = uuid.UUID("086991f8-75bd-4560-8943-e11c0c59422b")
new_primary = gpt.GPTHeader(
primary.current_lba,
primary.backup_lba,
primary.first_usable_lba,
primary.last_usable_lba,
new_guid,
primary.entries_starting_lba,
primary.num_entries,
primary.entry_size,
primary.entries_crc32,
)
new_backup = gpt.GPTHeader(
backup.current_lba,
backup.backup_lba,
backup.first_usable_lba,
backup.last_usable_lba,
new_guid,
backup.entries_starting_lba,
backup.num_entries,
backup.entry_size,
backup.entries_crc32,
)
with gpt.GPTImage(path=disk_image, open_mode=os.O_WRONLY) as image:
image.write_gpt_headers(new_primary, new_backup)
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
image.validate()
assert image.read_primary_gpt_header().disk_guid == new_guid
def test_gptimage_validate_invalid_image(disk_image: Path) -> None:
"""Test :meth:`gpt.GPTImage.validate` with an invalid image."""
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
primary = image.read_primary_gpt_header()
new_guid = uuid.UUID("21da705c-fec8-4857-8379-449d823a7155")
new_primary = gpt.GPTHeader(
primary.current_lba,
primary.backup_lba,
primary.first_usable_lba,
primary.last_usable_lba,
new_guid,
primary.entries_starting_lba,
primary.num_entries,
primary.entry_size,
primary.entries_crc32,
)
with open(disk_image, "rb+") as fd:
gpt.pwrite_all(fd.fileno(), new_primary.pack(), gpt.MBR_SIZE)
with pytest.raises(gpt.InvalidImageError): # noqa: PT012
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
image.validate()
def test_gptimage_write_gpt_headers_incorrect_current_lba(disk_image: Path) -> None:
"""Test :meth:`gpt.GPTImage.write_gpt_headers` with incorrect headers."""
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
primary = image.read_primary_gpt_header()
backup = image.read_backup_gpt_header()
with pytest.raises(
ValueError, match="Primary header has invalid 'current_lba', expected 1"
):
image.write_gpt_headers(backup, primary)
with open(disk_image, "rb+") as fd:
size = os.fstat(fd.fileno()).st_size
fd.truncate(size + gpt.LBA_SIZE)
with pytest.raises(
ValueError, match=r"Backup header has invalid 'current_lba', expected \d+"
):
image.write_gpt_headers(primary, backup)
def test_gptimage_write_gpt_headers_incompatible_headers(disk_image: Path) -> None:
"""Test :meth:`gpt.GPTImage.write_gpt_headers` with incompatible headers."""
with gpt.GPTImage(path=disk_image, open_mode=os.O_RDONLY) as image:
primary = image.read_primary_gpt_header()
backup = image.read_backup_gpt_header()
new_guid = uuid.UUID("913ef501-b10c-4fa4-8919-0d47f7d4d4bd")
new_primary = gpt.GPTHeader(
primary.current_lba,
primary.backup_lba,
primary.first_usable_lba,
primary.last_usable_lba,
new_guid,
primary.entries_starting_lba,
primary.num_entries,
primary.entry_size,
primary.entries_crc32,
)
with pytest.raises(
ValueError, match="Given headers are not backups of each other"
):
image.write_gpt_headers(new_primary, backup)
