From 5366cc31f18efcbc3f98dd181918be812889919e Mon Sep 17 00:00:00 2001 From: Nicolas Trangez Date: Thu, 25 Mar 2021 23:43:04 +0100 Subject: [PATCH] gptsum: expose tool functionality as a library This commit extends the `gptsum` module to now expose an API so one can access the tool functionality from other Python code. --- .flake8 | 3 + docs/api.rst | 9 ++ docs/conf.py | 1 + docs/index.rst | 1 + src/gptsum/__init__.py | 185 +++++++++++++++++++++++++++++++++++++++-- src/gptsum/cli.py | 50 ++++------- tests/test_init.py | 148 +++++++++++++++++++++++++++++++++ 7 files changed, 358 insertions(+), 39 deletions(-) create mode 100644 docs/api.rst create mode 100644 tests/test_init.py diff --git a/.flake8 b/.flake8 index 5f4126de..f17372c9 100644 --- a/.flake8 +++ b/.flake8 @@ -11,8 +11,11 @@ strictness=short docstring_style=sphinx # Configure flake8-rst-docstrings +rst-directives = + versionadded, rst-roles = class, + exc, func, meth, mod, diff --git a/docs/api.rst b/docs/api.rst new file mode 100644 index 00000000..928581c0 --- /dev/null +++ b/docs/api.rst @@ -0,0 +1,9 @@ +Python API +========== +.. automodule:: gptsum + :members: + +.. autodata:: gptsum.__author__ +.. autodata:: gptsum.__contact__ +.. autodata:: gptsum.__license__ +.. autodata:: gptsum.__version__ diff --git a/docs/conf.py b/docs/conf.py index aa50a780..8848bde7 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -13,5 +13,6 @@ html_theme = "furo" extensions = [ + "sphinx.ext.autodoc", "sphinxarg.ext", ] diff --git a/docs/index.rst b/docs/index.rst index 662933cd..31ec4eba 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -11,4 +11,5 @@ Command Line Interface :hidden: :maxdepth: 1 + api contributing diff --git a/src/gptsum/__init__.py b/src/gptsum/__init__.py index fd6b7249..4be81b8b 100644 --- a/src/gptsum/__init__.py +++ b/src/gptsum/__init__.py @@ -2,14 +2,187 @@ gptsum: a tool to make disk images using GPT partitions self-verifiable. Like `isomd5sum `_. + +This module exposes the functionality provided by the tool through a Python API. +""" + +import dataclasses +import os +import sys +import uuid +from pathlib import Path +from typing import Callable, Optional, cast + +if sys.version_info >= (3, 8): + from importlib.metadata import Distribution, distribution +else: + from importlib_metadata import Distribution, distribution + +import gptsum +import gptsum.checksum +import gptsum.gpt + +__metadata__ = cast(Callable[[str], Distribution], distribution)(__name__).metadata +__author__ = __metadata__["Author"] +"""Package author name. + +.. versionadded:: 0.0.9 """ +__contact__ = __metadata__["Author-email"] +"""Package author contact e-mail address. + +.. versionadded:: 0.0.9 +""" +__license__ = __metadata__["License"] +"""Package license identifier. + +.. versionadded:: 0.0.9 +""" +__version__ = __metadata__["Version"] +"""Package version identifier. + +.. versionadded:: 0.0.1 +""" +del __metadata__ + + +def _check_fd_or_path(fd: Optional[int], path: Optional[Path]) -> None: + """Ensure either ``fd`` or ``path`` is given. + + >>> _check_fd_or_path(None, None) + Traceback (most recent call last): + ... + ValueError: Either fd or path must be given + + >>> _check_fd_or_path(0, Path("/")) + Traceback (most recent call last): + ... + ValueError: Both fd and path can't be given + + :param fd: Possible non-``None`` FD value + :param path: Possiblel non-``None`` path value + + :raises ValueError: Both ``fd`` and ``path`` are ``None`` + :raises ValueError: Both ``fd`` and ``path`` are not ``None`` + """ + if fd is None and path is None: + raise ValueError("Either fd or path must be given") + if fd is not None and path is not None: + raise ValueError("Both fd and path can't be given") + + +def get_guid(*, fd: Optional[int] = None, path: Optional[Path] = None) -> uuid.UUID: + """Get the GUID of a disk image. + + One of ``fd`` or ``path`` must be given. + + :param fd: Readable file-descriptor to an image file + :param path: Path of a readable image file + + :returns: GUID of the disk image + + .. versionadded:: 0.0.9 + """ + _check_fd_or_path(fd, path) + + with gptsum.gpt.GPTImage(fd=fd, path=path, open_mode=os.O_RDONLY) as image: + image.validate() + primary = image.read_primary_gpt_header() + return primary.disk_guid + + +def set_guid( + guid: uuid.UUID, *, fd: Optional[int] = None, path: Optional[Path] = None +) -> None: + """Set the GUID of a disk image. + + One of ``fd`` or ``path`` must be given. + + :param guid: GUID to write to the disk image + :param fd: Readable and writable file-descriptor to an image file + :param path: Path of a readable and writable image file + + .. versionadded:: 0.0.9 + """ + _check_fd_or_path(fd, path) + + with gptsum.gpt.GPTImage(fd=fd, path=path, open_mode=os.O_RDWR) as image: + image.validate() + image.update_guid(guid) + + +def calculate_expected_guid( + *, fd: Optional[int] = None, path: Optional[Path] = None +) -> uuid.UUID: + """Calculate the expected checksum GUID of a disk image. + + One of ``fd`` or ``path`` must be given. + + :param fd: Readable file-descriptor to an image file + :param path: Path of a readable image file + + :returns: Expected checksum GUID of the disk image + + .. versionadded:: 0.0.9 + """ + _check_fd_or_path(fd, path) + + with gptsum.gpt.GPTImage(fd=fd, path=path, open_mode=os.O_RDONLY) as image: + image.validate() + digest = gptsum.checksum.calculate(image) + return gptsum.checksum.digest_to_guid(digest) + + +def embed(*, fd: Optional[int] = None, path: Optional[Path] = None) -> None: + """Embed the calculated checksum GUID in an image file. + + In essence a combination of :func:`calculate_expected_guid` and :func:`set_guid`. + + One of ``fd`` or ``path`` must be given. + + :param fd: Readable and writable file-descriptor to an image file + :param path: Path to a readable and writable image file + + .. versionadded:: 0.0.9 + """ + _check_fd_or_path(fd, path) + + current_guid = get_guid(fd=fd, path=path) + checksum_guid = calculate_expected_guid(fd=fd, path=path) + + if current_guid != checksum_guid: + set_guid(checksum_guid, fd=fd, path=path) + + +@dataclasses.dataclass(frozen=True) +class VerificationFailure(Exception): + """Exception raised when :func:`verify` fails. + + .. versionadded:: 0.0.9 + """ + + expected: uuid.UUID + actual: uuid.UUID + + +def verify(*, fd: Optional[int] = None, path: Optional[Path] = None) -> None: + """Verify a GPT disk image GUID against the calculated checksum. + + In essence a combination of :func:`calculate_expected_guid` and :func:`get_guid`. + + One of ``fd`` or ``path`` must be given. + + :param fd: Readable file-descriptor to an image file + :param path: Path to a readable image file -from typing import cast + :raises VerificationFailure: Current GUID and checksum mismatch -try: - import importlib.metadata as importlib_metadata -except ImportError: # pragma: no cover - import importlib_metadata # type: ignore + .. versionadded:: 0.0.9 + """ + _check_fd_or_path(fd, path) + current_guid = get_guid(fd=fd, path=path) + checksum_guid = calculate_expected_guid(fd=fd, path=path) -__version__: str = cast(str, importlib_metadata.version(__name__)) # type: ignore + if current_guid != checksum_guid: + raise VerificationFailure(checksum_guid, current_guid) diff --git a/src/gptsum/cli.py b/src/gptsum/cli.py index 10075306..82322e48 100644 --- a/src/gptsum/cli.py +++ b/src/gptsum/cli.py @@ -91,7 +91,7 @@ def build_parser() -> argparse.ArgumentParser: "without writing it to the file." ), ) - calculate_guid_parser.set_defaults(func=calculate_guid) + calculate_guid_parser.set_defaults(func=calculate_expected_guid) calculate_guid_parser.add_argument( "image", type=argparse.FileType("rb"), @@ -104,53 +104,37 @@ def build_parser() -> argparse.ArgumentParser: def get_guid(ns: argparse.Namespace) -> None: """Execute the 'get-guid' subcommand.""" - with gptsum.gpt.GPTImage(fd=ns.image.fileno()) as image: - image.validate() - primary = image.read_primary_gpt_header() - print("{}".format(primary.disk_guid)) + guid = gptsum.get_guid(fd=ns.image.fileno()) + print("{}".format(guid)) def set_guid(ns: argparse.Namespace) -> None: """Execute the 'set-guid' subcommand.""" - with gptsum.gpt.GPTImage(fd=ns.image.fileno()) as image: - image.validate() - image.update_guid(ns.guid) + gptsum.set_guid(ns.guid, fd=ns.image.fileno()) -def calculate_guid(ns: argparse.Namespace) -> None: +def calculate_expected_guid(ns: argparse.Namespace) -> None: """Execute the 'calculate-expected-guid' subcommand.""" - with gptsum.gpt.GPTImage(fd=ns.image.fileno()) as image: - image.validate() - digest = gptsum.checksum.calculate(image) - checksum_guid = gptsum.checksum.digest_to_guid(digest) - print("{}".format(checksum_guid)) + guid = gptsum.calculate_expected_guid(fd=ns.image.fileno()) + print("{}".format(guid)) def embed(ns: argparse.Namespace) -> None: """Execute the 'embed' subcommand.""" - with gptsum.gpt.GPTImage(fd=ns.image.fileno()) as image: - image.validate() - digest = gptsum.checksum.calculate(image) - checksum_guid = gptsum.checksum.digest_to_guid(digest) - if checksum_guid != image.read_primary_gpt_header().disk_guid: - image.update_guid(checksum_guid) + gptsum.embed(fd=ns.image.fileno()) def verify(ns: argparse.Namespace) -> None: """Execute the 'verify' subcommand.""" - with gptsum.gpt.GPTImage(fd=ns.image.fileno()) as image: - image.validate() - primary = image.read_primary_gpt_header() - digest = gptsum.checksum.calculate(image) - checksum_guid = gptsum.checksum.digest_to_guid(digest) - - if primary.disk_guid != checksum_guid: - sys.stderr.write( - "Disk GUID doesn't match expected checksum, " - "got {}, expected {}\n".format(primary.disk_guid, checksum_guid) - ) - sys.stderr.flush() - sys.exit(1) + try: + gptsum.verify(fd=ns.image.fileno()) + except gptsum.VerificationFailure as exn: + sys.stderr.write( + "Disk GUID doesn't match expected checksum, " + "got {}, expected {}\n".format(exn.actual, exn.expected) + ) + sys.stderr.flush() + sys.exit(1) def main(args: Optional[List[str]] = None) -> None: diff --git a/tests/test_init.py b/tests/test_init.py new file mode 100644 index 00000000..289d8e1b --- /dev/null +++ b/tests/test_init.py @@ -0,0 +1,148 @@ +"""Tests for :mod:`gptsum`.""" + +import hashlib +import os +import uuid +from pathlib import Path +from typing import Any, Callable, List + +import pytest + +import gptsum +from gptsum import checksum +from tests import conftest + + +@pytest.mark.parametrize(("attr"), ["author", "contact", "license", "version"]) +def test_attribute(attr: str) -> None: + """Test expected metadata attributes on the package.""" + real_attr = "__{}__".format(attr) + assert hasattr(gptsum, real_attr) + assert type(getattr(gptsum, real_attr)) is str + + +@pytest.mark.parametrize( + ("func", "args"), + [ + (gptsum.get_guid, []), + (gptsum.set_guid, [uuid.UUID("656156c4-2456-4038-b7f3-dfa0ae263e95")]), + (gptsum.calculate_expected_guid, []), + (gptsum.embed, []), + (gptsum.verify, []), + ], +) +def test_argument_conflicts(func: Callable[..., Any], args: List[Any]) -> None: + """Ensure ``fd``/``path`` can't be passed at the same time.""" + with pytest.raises(ValueError, match="Either fd or path must be given"): + func(*args, fd=None, path=None) + + with pytest.raises(ValueError, match="Both fd and path can't be given"): + func(*args, fd=0, path=Path("/")) + + +@pytest.mark.parametrize( + ("disk_file", "expected_guid"), + [ + (conftest.TESTDATA_DISK, conftest.TESTDATA_DISK_GUID), + (conftest.TESTDATA_EMBEDDED_DISK, conftest.TESTDATA_EMBEDDED_DISK_GUID), + ], +) +def test_get_guid(disk_file: Path, expected_guid: uuid.UUID) -> None: + """Test :func:`gptsum.get_guid`.""" + assert gptsum.get_guid(path=disk_file) == expected_guid + + with open(disk_file, "rb") as fd: + assert gptsum.get_guid(fd=fd.fileno()) == expected_guid + + +@pytest.mark.parametrize(("method"), ["by_fd", "by_path"]) +def test_set_guid(method: str, disk_image: Path) -> None: + """Test :func:`gptsum.set_guid`.""" + new_guid = uuid.UUID("e81f41f2-1807-49aa-876b-57a7e727b3f8") + + assert gptsum.get_guid(path=disk_image) != new_guid + + if method == "by_fd": + with open(disk_image, "rb+") as fd: + gptsum.set_guid(new_guid, fd=fd.fileno()) + else: + assert method == "by_path" + gptsum.set_guid(new_guid, path=disk_image) + + assert gptsum.get_guid(path=disk_image) == new_guid + + +@pytest.mark.parametrize( + ("disk_file", "expected_guid"), + [ + (conftest.TESTDATA_DISK, conftest.TESTDATA_EMBEDDED_DISK_GUID), + (conftest.TESTDATA_EMBEDDED_DISK, conftest.TESTDATA_EMBEDDED_DISK_GUID), + ], +) +def test_calculate_expected_guid(disk_file: Path, expected_guid: uuid.UUID) -> None: + """Test :func:`gptsum.calculate_expected_guid`.""" + assert gptsum.calculate_expected_guid(path=disk_file) == expected_guid + + with open(disk_file, "rb") as fd: + assert gptsum.calculate_expected_guid(fd=fd.fileno()) == expected_guid + + +@pytest.mark.parametrize(("method"), ["by_fd", "by_path"]) +def test_embed(method: str, disk_image: Path) -> None: + """Test :func:`gptsum.embed`.""" + if method == "by_fd": + with open(disk_image, "rb+") as fd: + gptsum.embed(fd=fd.fileno()) + else: + assert method == "by_path" + gptsum.embed(path=disk_image) + + guid = gptsum.get_guid(path=disk_image) + assert guid == conftest.TESTDATA_EMBEDDED_DISK_GUID + + hash1 = hashlib.sha256() + with open(conftest.TESTDATA_EMBEDDED_DISK, "rb") as fd: + size = os.fstat(fd.fileno()).st_size + done = checksum.hash_file(hash1.update, fd.fileno(), size, 0) + assert done == size + + hash2 = hashlib.sha256() + with open(disk_image, "rb") as fd: + size = os.fstat(fd.fileno()).st_size + done = checksum.hash_file(hash2.update, fd.fileno(), size, 0) + assert done == size + + assert hash1.hexdigest() == hash2.hexdigest() + + +def test_embed_noop(disk_image: Path) -> None: + """Test to ensure issuing :option:`embed` on an already-prepared file is a no-op.""" + gptsum.embed(path=disk_image) + stat1 = os.stat(disk_image) + gptsum.embed(path=disk_image) + stat2 = os.stat(disk_image) + + assert stat2.st_mtime == stat1.st_mtime + assert stat2.st_ctime == stat1.st_ctime + + +@pytest.mark.parametrize(("method"), ["by_fd", "by_path"]) +def test_verify(method: str, disk_image: Path) -> None: + """Test :func:`gptsum.verify`.""" + if method == "by_fd": + with open(disk_image, "rb") as fd: + with pytest.raises(gptsum.VerificationFailure): + gptsum.verify(fd=fd.fileno()) + else: + assert method == "by_path" + with pytest.raises(gptsum.VerificationFailure): + gptsum.verify(path=disk_image) + + gptsum.embed(path=disk_image) + + if method == "by_fd": + with open(disk_image, "rb") as fd: + gptsum.verify(fd=fd.fileno()) + else: + assert method == "by_path" + gptsum.verify(path=disk_image)