update_engine/scripts/update_payload/common.py

# Copyright 2013 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Utilities for update payload processing."""

import base64

from update_payload import update_metadata_pb2
from update_payload.error import PayloadError


#
# Constants.
#
SIG_ASN1_HEADER = (
    b"\x30\x31\x30\x0d\x06\x09\x60\x86"
    b"\x48\x01\x65\x03\x04\x02\x01\x05"
    b"\x00\x04\x20"
)

BRILLO_MAJOR_PAYLOAD_VERSION = 2

SOURCE_MINOR_PAYLOAD_VERSION = 2
OPSRCHASH_MINOR_PAYLOAD_VERSION = 3
BROTLI_BSDIFF_MINOR_PAYLOAD_VERSION = 4
PUFFDIFF_MINOR_PAYLOAD_VERSION = 5

KERNEL = "kernel"
ROOTFS = "root"
# Tuple of (name in system, name in protobuf).
CROS_PARTITIONS = ((KERNEL, KERNEL), (ROOTFS, "rootfs"))


#
# Payload operation types.
#
class OpType(object):
    """Container for operation type constants."""

    _CLASS = update_metadata_pb2.InstallOperation
    REPLACE = _CLASS.REPLACE
    REPLACE_BZ = _CLASS.REPLACE_BZ
    SOURCE_COPY = _CLASS.SOURCE_COPY
    SOURCE_BSDIFF = _CLASS.SOURCE_BSDIFF
    ZERO = _CLASS.ZERO
    DISCARD = _CLASS.DISCARD
    REPLACE_XZ = _CLASS.REPLACE_XZ
    REPLACE_ZSTD = _CLASS.REPLACE_ZSTD
    REPLACE_ZSTD_INCREASED_WINDOW = _CLASS.REPLACE_ZSTD_INCREASED_WINDOW
    PUFFDIFF = _CLASS.PUFFDIFF
    BROTLI_BSDIFF = _CLASS.BROTLI_BSDIFF
    ALL = (
        REPLACE,
        REPLACE_BZ,
        SOURCE_COPY,
        SOURCE_BSDIFF,
        ZERO,
        DISCARD,
        REPLACE_XZ,
        REPLACE_ZSTD,
        REPLACE_ZSTD_INCREASED_WINDOW,
        PUFFDIFF,
        BROTLI_BSDIFF,
    )
    NAMES = {
        REPLACE: "REPLACE",
        REPLACE_BZ: "REPLACE_BZ",
        SOURCE_COPY: "SOURCE_COPY",
        SOURCE_BSDIFF: "SOURCE_BSDIFF",
        ZERO: "ZERO",
        DISCARD: "DISCARD",
        REPLACE_XZ: "REPLACE_XZ",
        REPLACE_ZSTD: "REPLACE_ZSTD",
        REPLACE_ZSTD_INCREASED_WINDOW: "REPLACE_ZSTD_INCREASED_WINDOW",
        PUFFDIFF: "PUFFDIFF",
        BROTLI_BSDIFF: "BROTLI_BSDIFF",
    }

    def __init__(self):
        pass


#
# Checked and hashed reading of data.
#
def IntPackingFmtStr(size, is_unsigned):
    """Returns an integer format string for use by the struct module.

    Args:
      size: the integer size in bytes (2, 4 or 8)
      is_unsigned: whether it is signed or not

    Returns:
      A format string for packing/unpacking integer values; assumes network byte
      order (big-endian).

    Raises:
      PayloadError if something is wrong with the arguments.
    """
    # Determine the base conversion format.
    if size == 2:
        fmt = "h"
    elif size == 4:
        fmt = "i"
    elif size == 8:
        fmt = "q"
    else:
        raise PayloadError("unsupport numeric field size (%s)" % size)

    # Signed or unsigned?
    if is_unsigned:
        fmt = fmt.upper()

    # Make it network byte order (big-endian).
    fmt = "!" + fmt

    return fmt


def Read(file_obj, length, offset=None, hasher=None):
    """Reads binary data from a file.

    Args:
      file_obj: an open file object
      length: the length of the data to read
      offset: an offset to seek to prior to reading; this is an absolute offset
              from either the beginning (non-negative) or end (negative) of the
              file.  (optional)
      hasher: a hashing object to pass the read data through (optional)

    Returns:
      A string containing the read data.

    Raises:
      PayloadError if a read error occurred or not enough data was read.
    """
    if offset is not None:
        if offset >= 0:
            file_obj.seek(offset)
        else:
            file_obj.seek(offset, 2)

    try:
        data = file_obj.read(length)
    except IOError as e:
        raise PayloadError(
            "error reading from file (%s): %s" % (file_obj.name, e)
        )

    if len(data) != length:
        raise PayloadError(
            "reading from file (%s) too short (%d instead of %d bytes)"
            % (file_obj.name, len(data), length)
        )

    if hasher:
        hasher.update(data)

    return data


#
# Formatting functions.
#
def FormatExtent(ex, block_size=0):
    end_block = ex.start_block + ex.num_blocks
    if block_size:
        return "%d->%d * %d" % (ex.start_block, end_block, block_size)
    return "%d->%d" % (ex.start_block, end_block)


def FormatSha256(digest):
    """Returns a canonical string representation of a SHA256 digest."""
    return base64.b64encode(digest).decode("utf-8")


#
# Useful iterators.
#
def _ObjNameIter(items, base_name, reverse=False, name_format_func=None):
    """A generic (item, name) tuple iterators.

    Args:
      items: the sequence of objects to iterate on
      base_name: the base name for all objects
      reverse: whether iteration should be in reverse order
      name_format_func: a function to apply to the name string

    Yields:
      An iterator whose i-th invocation returns (items[i], name), where name ==
      base_name + '[i]' (with a formatting function optionally applied to it).
    """
    idx, inc = (len(items), -1) if reverse else (1, 1)
    if reverse:
        items = reversed(items)
    for item in items:
        item_name = "%s[%d]" % (base_name, idx)
        if name_format_func:
            item_name = name_format_func(item, item_name)
        yield (item, item_name)
        idx += inc


def _OperationNameFormatter(op, op_name):
    return "%s(%s)" % (op_name, OpType.NAMES.get(op.type, "?"))


def OperationIter(operations, base_name, reverse=False):
    """An (item, name) iterator for update operations."""
    return _ObjNameIter(
        operations,
        base_name,
        reverse=reverse,
        name_format_func=_OperationNameFormatter,
    )


def ExtentIter(extents, base_name, reverse=False):
    """An (item, name) iterator for operation extents."""
    return _ObjNameIter(extents, base_name, reverse=reverse)


def SignatureIter(sigs, base_name, reverse=False):
    """An (item, name) iterator for signatures."""
    return _ObjNameIter(sigs, base_name, reverse=reverse)