1# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Utilities for update payload processing."""
6
7from __future__ import print_function
8
9from error import PayloadError
10import update_metadata_pb2
11
12
13#
14# Constants.
15#
16PSEUDO_EXTENT_MARKER = (1L << 64) - 1  # UINT64_MAX
17
18SIG_ASN1_HEADER = (
19    '\x30\x31\x30\x0d\x06\x09\x60\x86'
20    '\x48\x01\x65\x03\x04\x02\x01\x05'
21    '\x00\x04\x20'
22)
23
24CHROMEOS_MAJOR_PAYLOAD_VERSION = 1
25BRILLO_MAJOR_PAYLOAD_VERSION = 2
26
27INPLACE_MINOR_PAYLOAD_VERSION = 1
28SOURCE_MINOR_PAYLOAD_VERSION = 2
29OPSRCHASH_MINOR_PAYLOAD_VERSION = 3
30IMGDIFF_MINOR_PAYLOAD_VERSION = 4
31
32#
33# Payload operation types.
34#
35class OpType(object):
36  """Container for operation type constants."""
37  _CLASS = update_metadata_pb2.InstallOperation
38  # pylint: disable=E1101
39  REPLACE = _CLASS.REPLACE
40  REPLACE_BZ = _CLASS.REPLACE_BZ
41  MOVE = _CLASS.MOVE
42  BSDIFF = _CLASS.BSDIFF
43  SOURCE_COPY = _CLASS.SOURCE_COPY
44  SOURCE_BSDIFF = _CLASS.SOURCE_BSDIFF
45  ZERO = _CLASS.ZERO
46  DISCARD = _CLASS.DISCARD
47  REPLACE_XZ = _CLASS.REPLACE_XZ
48  IMGDIFF = _CLASS.IMGDIFF
49  ALL = (REPLACE, REPLACE_BZ, MOVE, BSDIFF, SOURCE_COPY, SOURCE_BSDIFF, ZERO,
50         DISCARD, REPLACE_XZ, IMGDIFF)
51  NAMES = {
52      REPLACE: 'REPLACE',
53      REPLACE_BZ: 'REPLACE_BZ',
54      MOVE: 'MOVE',
55      BSDIFF: 'BSDIFF',
56      SOURCE_COPY: 'SOURCE_COPY',
57      SOURCE_BSDIFF: 'SOURCE_BSDIFF',
58      ZERO: 'ZERO',
59      DISCARD: 'DISCARD',
60      REPLACE_XZ: 'REPLACE_XZ',
61      IMGDIFF: 'IMGDIFF',
62  }
63
64  def __init__(self):
65    pass
66
67
68#
69# Checked and hashed reading of data.
70#
71def IntPackingFmtStr(size, is_unsigned):
72  """Returns an integer format string for use by the struct module.
73
74  Args:
75    size: the integer size in bytes (2, 4 or 8)
76    is_unsigned: whether it is signed or not
77
78  Returns:
79    A format string for packing/unpacking integer values; assumes network byte
80    order (big-endian).
81
82  Raises:
83    PayloadError if something is wrong with the arguments.
84  """
85  # Determine the base conversion format.
86  if size == 2:
87    fmt = 'h'
88  elif size == 4:
89    fmt = 'i'
90  elif size == 8:
91    fmt = 'q'
92  else:
93    raise PayloadError('unsupport numeric field size (%s)' % size)
94
95  # Signed or unsigned?
96  if is_unsigned:
97    fmt = fmt.upper()
98
99  # Make it network byte order (big-endian).
100  fmt = '!' + fmt
101
102  return fmt
103
104
105def Read(file_obj, length, offset=None, hasher=None):
106  """Reads binary data from a file.
107
108  Args:
109    file_obj: an open file object
110    length: the length of the data to read
111    offset: an offset to seek to prior to reading; this is an absolute offset
112            from either the beginning (non-negative) or end (negative) of the
113            file.  (optional)
114    hasher: a hashing object to pass the read data through (optional)
115
116  Returns:
117    A string containing the read data.
118
119  Raises:
120    PayloadError if a read error occurred or not enough data was read.
121  """
122  if offset is not None:
123    if offset >= 0:
124      file_obj.seek(offset)
125    else:
126      file_obj.seek(offset, 2)
127
128  try:
129    data = file_obj.read(length)
130  except IOError, e:
131    raise PayloadError('error reading from file (%s): %s' % (file_obj.name, e))
132
133  if len(data) != length:
134    raise PayloadError(
135        'reading from file (%s) too short (%d instead of %d bytes)' %
136        (file_obj.name, len(data), length))
137
138  if hasher:
139    hasher.update(data)
140
141  return data
142
143
144#
145# Formatting functions.
146#
147def FormatExtent(ex, block_size=0):
148  end_block = ex.start_block + ex.num_blocks
149  if block_size:
150    return '%d->%d * %d' % (ex.start_block, end_block, block_size)
151  else:
152    return '%d->%d' % (ex.start_block, end_block)
153
154
155def FormatSha256(digest):
156  """Returns a canonical string representation of a SHA256 digest."""
157  return digest.encode('base64').strip()
158
159
160#
161# Useful iterators.
162#
163def _ObjNameIter(items, base_name, reverse=False, name_format_func=None):
164  """A generic (item, name) tuple iterators.
165
166  Args:
167    items: the sequence of objects to iterate on
168    base_name: the base name for all objects
169    reverse: whether iteration should be in reverse order
170    name_format_func: a function to apply to the name string
171
172  Yields:
173    An iterator whose i-th invocation returns (items[i], name), where name ==
174    base_name + '[i]' (with a formatting function optionally applied to it).
175  """
176  idx, inc = (len(items), -1) if reverse else (1, 1)
177  if reverse:
178    items = reversed(items)
179  for item in items:
180    item_name = '%s[%d]' % (base_name, idx)
181    if name_format_func:
182      item_name = name_format_func(item, item_name)
183    yield (item, item_name)
184    idx += inc
185
186
187def _OperationNameFormatter(op, op_name):
188  return '%s(%s)' % (op_name, OpType.NAMES.get(op.type, '?'))
189
190
191def OperationIter(operations, base_name, reverse=False):
192  """An (item, name) iterator for update operations."""
193  return _ObjNameIter(operations, base_name, reverse=reverse,
194                      name_format_func=_OperationNameFormatter)
195
196
197def ExtentIter(extents, base_name, reverse=False):
198  """An (item, name) iterator for operation extents."""
199  return _ObjNameIter(extents, base_name, reverse=reverse)
200
201
202def SignatureIter(sigs, base_name, reverse=False):
203  """An (item, name) iterator for signatures."""
204  return _ObjNameIter(sigs, base_name, reverse=reverse)
205