224 lines
6.7 KiB
Python
224 lines
6.7 KiB
Python
|
# 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)
|