Skip to content

Commit

Permalink
gptsum: expose tool functionality as a library
Browse files Browse the repository at this point in the history
This commit extends the `gptsum` module to now expose an API so one can
access the tool functionality from other Python code.
  • Loading branch information
NicolasT committed Mar 25, 2021
1 parent d9f6134 commit 5366cc3
Show file tree
Hide file tree
Showing 7 changed files with 358 additions and 39 deletions.
3 changes: 3 additions & 0 deletions .flake8
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,11 @@ strictness=short
docstring_style=sphinx

# Configure flake8-rst-docstrings
rst-directives =
versionadded,
rst-roles =
class,
exc,
func,
meth,
mod,
Expand Down
9 changes: 9 additions & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
Python API
==========
.. automodule:: gptsum
:members:

.. autodata:: gptsum.__author__
.. autodata:: gptsum.__contact__
.. autodata:: gptsum.__license__
.. autodata:: gptsum.__version__
1 change: 1 addition & 0 deletions docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,5 +13,6 @@
html_theme = "furo"

extensions = [
"sphinx.ext.autodoc",
"sphinxarg.ext",
]
1 change: 1 addition & 0 deletions docs/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,5 @@ Command Line Interface
:hidden:
:maxdepth: 1

api
contributing
185 changes: 179 additions & 6 deletions src/gptsum/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,187 @@
gptsum: a tool to make disk images using GPT partitions self-verifiable.
Like `isomd5sum </~https://github.com/rhinstaller/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)
50 changes: 17 additions & 33 deletions src/gptsum/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand All @@ -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:
Expand Down
Loading

0 comments on commit 5366cc3

Please sign in to comment.