images.py

# Copyright (C) 2019 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific

import os
import threading
from hashlib import sha1

from rangelib import RangeSet

__all__ = ["EmptyImage", "DataImage", "FileImage"]


class Image(object):
    def RangeSha1(self, ranges):
        raise NotImplementedError

    def ReadRangeSet(self, ranges):
        raise NotImplementedError

    def TotalSha1(self, include_clobbered_blocks=False):
        raise NotImplementedError

    def WriteRangeDataToFd(self, ranges, fd):
        raise NotImplementedError


class EmptyImage(Image):
    """A zero-length image."""

    def __init__(self):
        self.blocksize = 4096
        self.care_map = RangeSet()
        self.clobbered_blocks = RangeSet()
        self.extended = RangeSet()
        self.total_blocks = 0
        self.file_map = {}
        self.hashtree_info = None

    def RangeSha1(self, ranges):
        return sha1().hexdigest()

    def ReadRangeSet(self, ranges):
        return ()

    def TotalSha1(self, include_clobbered_blocks=False):
        # EmptyImage always carries empty clobbered_blocks, so
        # include_clobbered_blocks can be ignored.
        assert self.clobbered_blocks.size() == 0
        return sha1().hexdigest()

    def WriteRangeDataToFd(self, ranges, fd):
        raise ValueError("Can't write data from EmptyImage to file")


class DataImage(Image):
    """An image wrapped around a single string of data."""

    def __init__(self, data, trim=False, pad=False):
        self.data = data
        self.blocksize = 4096

        assert not (trim and pad)

        partial = len(self.data) % self.blocksize
        padded = False
        if partial > 0:
            if trim:
                self.data = self.data[:-partial]
            elif pad:
                self.data += "\0" * (self.blocksize - partial)
                padded = True
            else:
                raise ValueError(
                    (
                        "data for DataImage must be multiple of %d bytes "
                        "unless trim or pad is specified"
                    )
                    % (self.blocksize,)
                )

        assert len(self.data) % self.blocksize == 0

        self.total_blocks = len(self.data) // self.blocksize
        self.care_map = RangeSet(data=(0, self.total_blocks))
        # When the last block is padded, we always write the whole block even for
        # incremental OTAs. Because otherwise the last block may get skipped if
        # unchanged for an incremental, but would fail the post-install
        # verification if it has non-zero contents in the padding bytes.
        # Bug: 23828506
        if padded:
            clobbered_blocks = [self.total_blocks - 1, self.total_blocks]
        else:
            clobbered_blocks = []
        self.clobbered_blocks = clobbered_blocks
        self.extended = RangeSet()

        zero_blocks = []
        nonzero_blocks = []
        reference = "\0" * self.blocksize

        for i in range(self.total_blocks - 1 if padded else self.total_blocks):
            d = self.data[i * self.blocksize : (i + 1) * self.blocksize]
            if d == reference:
                zero_blocks.append(i)
                zero_blocks.append(i + 1)
            else:
                nonzero_blocks.append(i)
                nonzero_blocks.append(i + 1)

        assert zero_blocks or nonzero_blocks or clobbered_blocks

        self.file_map = dict()
        if zero_blocks:
            self.file_map["__ZERO"] = RangeSet(data=zero_blocks)
        if nonzero_blocks:
            self.file_map["__NONZERO"] = RangeSet(data=nonzero_blocks)
        if clobbered_blocks:
            self.file_map["__COPY"] = RangeSet(data=clobbered_blocks)

    def _GetRangeData(self, ranges):
        for s, e in ranges:
            yield self.data[s * self.blocksize : e * self.blocksize]

    def RangeSha1(self, ranges):
        h = sha1()
        for data in self._GetRangeData(ranges):  # pylint: disable=not-an-iterable
            h.update(data)
        return h.hexdigest()

    def ReadRangeSet(self, ranges):
        return list(self._GetRangeData(ranges))

    def TotalSha1(self, include_clobbered_blocks=False):
        if not include_clobbered_blocks:
            return self.RangeSha1(self.care_map.subtract(self.clobbered_blocks))
        return sha1(self.data).hexdigest()

    def WriteRangeDataToFd(self, ranges, fd):
        for data in self._GetRangeData(ranges):  # pylint: disable=not-an-iterable
            fd.write(data)


class FileImage(Image):
    """An image wrapped around a raw image file."""

    def __init__(self, path, hashtree_info_generator=None):
        self.path = path
        self.blocksize = 4096
        self._file_size = os.path.getsize(self.path)
        self._file = open(self.path, "rb")

        if self._file_size % self.blocksize != 0:
            raise ValueError(
                "Size of file %s must be multiple of %d bytes, but is %d" % self.path,
                self.blocksize,
                self._file_size,
            )

        self.total_blocks = self._file_size // self.blocksize
        self.care_map = RangeSet(data=(0, self.total_blocks))
        self.clobbered_blocks = RangeSet()
        self.extended = RangeSet()

        self.generator_lock = threading.Lock()

        self.hashtree_info = None
        if hashtree_info_generator:
            self.hashtree_info = hashtree_info_generator.Generate(self)

        zero_blocks = []
        nonzero_blocks = []
        reference = "\0" * self.blocksize

        for i in range(self.total_blocks):
            d = self._file.read(self.blocksize)
            if d == reference:
                zero_blocks.append(i)
                zero_blocks.append(i + 1)
            else:
                nonzero_blocks.append(i)
                nonzero_blocks.append(i + 1)

        assert zero_blocks or nonzero_blocks

        self.file_map = {}
        if zero_blocks:
            self.file_map["__ZERO"] = RangeSet(data=zero_blocks)
        if nonzero_blocks:
            self.file_map["__NONZERO"] = RangeSet(data=nonzero_blocks)
        if self.hashtree_info:
            self.file_map["__HASHTREE"] = self.hashtree_info.hashtree_range

    def __del__(self):
        self._file.close()

    def _GetRangeData(self, ranges):
        # Use a lock to protect the generator so that we will not run two
        # instances of this generator on the same object simultaneously.
        with self.generator_lock:
            for s, e in ranges:
                self._file.seek(s * self.blocksize)
                for _ in range(s, e):
                    yield self._file.read(self.blocksize)

    def RangeSha1(self, ranges):
        h = sha1()
        for data in self._GetRangeData(ranges):  # pylint: disable=not-an-iterable
            h.update(data)
        return h.hexdigest()

    def ReadRangeSet(self, ranges):
        return list(self._GetRangeData(ranges))

    def TotalSha1(self, include_clobbered_blocks=False):
        assert not self.clobbered_blocks
        return self.RangeSha1(self.care_map)

    def WriteRangeDataToFd(self, ranges, fd):
        for data in self._GetRangeData(ranges):  # pylint: disable=not-an-iterable
            fd.write(data)