1#
2# Copyright (C) 2018 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import os
18import os.path
19import zipfile
20
21import common
22import test_utils
23from add_img_to_target_files import (
24    AddPackRadioImages,
25    CheckAbOtaImages)
26from rangelib import RangeSet
27from common import AddCareMapForAbOta, GetCareMap
28
29
30OPTIONS = common.OPTIONS
31
32
33class AddImagesToTargetFilesTest(test_utils.ReleaseToolsTestCase):
34
35  def setUp(self):
36    OPTIONS.input_tmp = common.MakeTempDir()
37
38  @staticmethod
39  def _create_images(images, prefix):
40    """Creates images under OPTIONS.input_tmp/prefix."""
41    path = os.path.join(OPTIONS.input_tmp, prefix)
42    if not os.path.exists(path):
43      os.mkdir(path)
44
45    for image in images:
46      image_path = os.path.join(path, image + '.img')
47      with open(image_path, 'wb') as image_fp:
48        image_fp.write(image.encode())
49
50    images_path = os.path.join(OPTIONS.input_tmp, 'IMAGES')
51    if not os.path.exists(images_path):
52      os.mkdir(images_path)
53    return images, images_path
54
55  def test_CheckAbOtaImages_imageExistsUnderImages(self):
56    """Tests the case with existing images under IMAGES/."""
57    images, _ = self._create_images(['aboot', 'xbl'], 'IMAGES')
58    CheckAbOtaImages(None, images)
59
60  def test_CheckAbOtaImages_imageExistsUnderRadio(self):
61    """Tests the case with some image under RADIO/."""
62    images, _ = self._create_images(['system', 'vendor'], 'IMAGES')
63    radio_path = os.path.join(OPTIONS.input_tmp, 'RADIO')
64    if not os.path.exists(radio_path):
65      os.mkdir(radio_path)
66    with open(os.path.join(radio_path, 'modem.img'), 'wb') as image_fp:
67      image_fp.write('modem'.encode())
68    CheckAbOtaImages(None, images + ['modem'])
69
70  def test_CheckAbOtaImages_missingImages(self):
71    images, _ = self._create_images(['aboot', 'xbl'], 'RADIO')
72    self.assertRaises(
73        AssertionError, CheckAbOtaImages, None, images + ['baz'])
74
75  def test_AddPackRadioImages(self):
76    images, images_path = self._create_images(['foo', 'bar'], 'RADIO')
77    AddPackRadioImages(None, images)
78
79    for image in images:
80      self.assertTrue(
81          os.path.exists(os.path.join(images_path, image + '.img')))
82
83  def test_AddPackRadioImages_with_suffix(self):
84    images, images_path = self._create_images(['foo', 'bar'], 'RADIO')
85    images_with_suffix = [image + '.img' for image in images]
86    AddPackRadioImages(None, images_with_suffix)
87
88    for image in images:
89      self.assertTrue(
90          os.path.exists(os.path.join(images_path, image + '.img')))
91
92  def test_AddPackRadioImages_zipOutput(self):
93    images, _ = self._create_images(['foo', 'bar'], 'RADIO')
94
95    # Set up the output zip.
96    output_file = common.MakeTempFile(suffix='.zip')
97    with zipfile.ZipFile(output_file, 'w', allowZip64=True) as output_zip:
98      AddPackRadioImages(output_zip, images)
99
100    with zipfile.ZipFile(output_file, 'r', allowZip64=True) as verify_zip:
101      for image in images:
102        self.assertIn('IMAGES/' + image + '.img', verify_zip.namelist())
103
104  def test_AddPackRadioImages_imageExists(self):
105    images, images_path = self._create_images(['foo', 'bar'], 'RADIO')
106
107    # Additionally create images under IMAGES/ so that they should be skipped.
108    images, images_path = self._create_images(['foo', 'bar'], 'IMAGES')
109
110    AddPackRadioImages(None, images)
111
112    for image in images:
113      self.assertTrue(
114          os.path.exists(os.path.join(images_path, image + '.img')))
115
116  def test_AddPackRadioImages_missingImages(self):
117    images, _ = self._create_images(['foo', 'bar'], 'RADIO')
118    AddPackRadioImages(None, images)
119
120    self.assertRaises(AssertionError, AddPackRadioImages, None,
121                      images + ['baz'])
122
123  @staticmethod
124  def _test_AddCareMapForAbOta():
125    """Helper function to set up the test for test_AddCareMapForAbOta()."""
126    OPTIONS.info_dict = {
127        'extfs_sparse_flag' : '-s',
128        'system_image_size' : 65536,
129        'vendor_image_size' : 40960,
130        'system_verity_block_device': '/dev/block/system',
131        'vendor_verity_block_device': '/dev/block/vendor',
132        'system.build.prop': common.PartitionBuildProps.FromDictionary(
133            'system', {
134                'ro.system.build.fingerprint':
135                'google/sailfish/12345:user/dev-keys'}
136        ),
137        'vendor.build.prop': common.PartitionBuildProps.FromDictionary(
138            'vendor', {
139                'ro.vendor.build.fingerprint':
140                'google/sailfish/678:user/dev-keys'}
141        ),
142    }
143
144    # Prepare the META/ folder.
145    meta_path = os.path.join(OPTIONS.input_tmp, 'META')
146    if not os.path.exists(meta_path):
147      os.mkdir(meta_path)
148
149    system_image = test_utils.construct_sparse_image([
150        (0xCAC1, 6),
151        (0xCAC3, 4),
152        (0xCAC1, 8)])
153    vendor_image = test_utils.construct_sparse_image([
154        (0xCAC2, 12)])
155
156    image_paths = {
157        'system' : system_image,
158        'vendor' : vendor_image,
159    }
160    return image_paths
161
162  def _verifyCareMap(self, expected, file_name):
163    """Parses the care_map.pb; and checks the content in plain text."""
164    text_file = common.MakeTempFile(prefix="caremap-", suffix=".txt")
165
166    # Calls an external binary to convert the proto message.
167    cmd = ["care_map_generator", "--parse_proto", file_name, text_file]
168    common.RunAndCheckOutput(cmd)
169
170    with open(text_file) as verify_fp:
171      plain_text = verify_fp.read()
172    self.assertEqual('\n'.join(expected), plain_text)
173
174  @test_utils.SkipIfExternalToolsUnavailable()
175  def test_AddCareMapForAbOta(self):
176    image_paths = self._test_AddCareMapForAbOta()
177
178    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
179    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
180
181    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
182                "ro.system.build.fingerprint",
183                "google/sailfish/12345:user/dev-keys",
184                'vendor', RangeSet("0-9").to_string_raw(),
185                "ro.vendor.build.fingerprint",
186                "google/sailfish/678:user/dev-keys"]
187
188    self._verifyCareMap(expected, care_map_file)
189
190  @test_utils.SkipIfExternalToolsUnavailable()
191  def test_AddCareMapForAbOta_withNonCareMapPartitions(self):
192    """Partitions without care_map should be ignored."""
193    image_paths = self._test_AddCareMapForAbOta()
194
195    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
196    AddCareMapForAbOta(
197        care_map_file, ['boot', 'system', 'vendor', 'vbmeta'], image_paths)
198
199    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
200                "ro.system.build.fingerprint",
201                "google/sailfish/12345:user/dev-keys",
202                'vendor', RangeSet("0-9").to_string_raw(),
203                "ro.vendor.build.fingerprint",
204                "google/sailfish/678:user/dev-keys"]
205
206    self._verifyCareMap(expected, care_map_file)
207
208  @test_utils.SkipIfExternalToolsUnavailable()
209  def test_AddCareMapForAbOta_withAvb(self):
210    """Tests the case for device using AVB."""
211    image_paths = self._test_AddCareMapForAbOta()
212    OPTIONS.info_dict = {
213        'extfs_sparse_flag': '-s',
214        'system_image_size': 65536,
215        'vendor_image_size': 40960,
216        'avb_system_hashtree_enable': 'true',
217        'avb_vendor_hashtree_enable': 'true',
218        'system.build.prop': common.PartitionBuildProps.FromDictionary(
219            'system', {
220                'ro.system.build.fingerprint':
221                'google/sailfish/12345:user/dev-keys'}
222        ),
223        'vendor.build.prop': common.PartitionBuildProps.FromDictionary(
224            'vendor', {
225                'ro.vendor.build.fingerprint':
226                'google/sailfish/678:user/dev-keys'}
227        ),
228    }
229
230    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
231    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
232
233    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
234                "ro.system.build.fingerprint",
235                "google/sailfish/12345:user/dev-keys",
236                'vendor', RangeSet("0-9").to_string_raw(),
237                "ro.vendor.build.fingerprint",
238                "google/sailfish/678:user/dev-keys"]
239
240    self._verifyCareMap(expected, care_map_file)
241
242  @test_utils.SkipIfExternalToolsUnavailable()
243  def test_AddCareMapForAbOta_noFingerprint(self):
244    """Tests the case for partitions without fingerprint."""
245    image_paths = self._test_AddCareMapForAbOta()
246    OPTIONS.info_dict = {
247        'extfs_sparse_flag' : '-s',
248        'system_image_size' : 65536,
249        'vendor_image_size' : 40960,
250        'system_verity_block_device': '/dev/block/system',
251        'vendor_verity_block_device': '/dev/block/vendor',
252    }
253
254    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
255    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
256
257    expected = ['system', RangeSet("0-5 10-15").to_string_raw(), "unknown",
258                "unknown", 'vendor', RangeSet("0-9").to_string_raw(), "unknown",
259                "unknown"]
260
261    self._verifyCareMap(expected, care_map_file)
262
263  @test_utils.SkipIfExternalToolsUnavailable()
264  def test_AddCareMapForAbOta_withThumbprint(self):
265    """Tests the case for partitions with thumbprint."""
266    image_paths = self._test_AddCareMapForAbOta()
267    OPTIONS.info_dict = {
268        'extfs_sparse_flag': '-s',
269        'system_image_size': 65536,
270        'vendor_image_size': 40960,
271        'system_verity_block_device': '/dev/block/system',
272        'vendor_verity_block_device': '/dev/block/vendor',
273        'system.build.prop': common.PartitionBuildProps.FromDictionary(
274            'system', {
275                'ro.system.build.thumbprint':
276                'google/sailfish/123:user/dev-keys'}
277        ),
278        'vendor.build.prop': common.PartitionBuildProps.FromDictionary(
279            'vendor', {
280                'ro.vendor.build.thumbprint':
281                'google/sailfish/456:user/dev-keys'}
282        ),
283    }
284
285    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
286    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
287
288    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
289                "ro.system.build.thumbprint",
290                "google/sailfish/123:user/dev-keys",
291                'vendor', RangeSet("0-9").to_string_raw(),
292                "ro.vendor.build.thumbprint",
293                "google/sailfish/456:user/dev-keys"]
294
295    self._verifyCareMap(expected, care_map_file)
296
297  @test_utils.SkipIfExternalToolsUnavailable()
298  def test_AddCareMapForAbOta_skipPartition(self):
299    image_paths = self._test_AddCareMapForAbOta()
300
301    # Remove vendor_image_size to invalidate the care_map for vendor.img.
302    del OPTIONS.info_dict['vendor_image_size']
303
304    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
305    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
306
307    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
308                "ro.system.build.fingerprint",
309                "google/sailfish/12345:user/dev-keys"]
310
311    self._verifyCareMap(expected, care_map_file)
312
313  @test_utils.SkipIfExternalToolsUnavailable()
314  def test_AddCareMapForAbOta_skipAllPartitions(self):
315    image_paths = self._test_AddCareMapForAbOta()
316
317    # Remove the image_size properties for all the partitions.
318    del OPTIONS.info_dict['system_image_size']
319    del OPTIONS.info_dict['vendor_image_size']
320
321    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
322    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
323
324    self.assertFalse(os.path.exists(care_map_file))
325
326  def test_AddCareMapForAbOta_verityNotEnabled(self):
327    """No care_map.pb should be generated if verity not enabled."""
328    image_paths = self._test_AddCareMapForAbOta()
329    OPTIONS.info_dict = {}
330    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
331    AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
332
333    self.assertFalse(os.path.exists(care_map_file))
334
335  def test_AddCareMapForAbOta_missingImageFile(self):
336    """Missing image file should be considered fatal."""
337    image_paths = self._test_AddCareMapForAbOta()
338    image_paths['vendor'] = ''
339    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
340    self.assertRaises(common.ExternalError, AddCareMapForAbOta, care_map_file,
341                      ['system', 'vendor'], image_paths)
342
343  @test_utils.SkipIfExternalToolsUnavailable()
344  def test_AddCareMapForAbOta_zipOutput(self):
345    """Tests the case with ZIP output."""
346    image_paths = self._test_AddCareMapForAbOta()
347
348    output_file = common.MakeTempFile(suffix='.zip')
349    with zipfile.ZipFile(output_file, 'w', allowZip64=True) as output_zip:
350      AddCareMapForAbOta(output_zip, ['system', 'vendor'], image_paths)
351
352    care_map_name = "META/care_map.pb"
353    temp_dir = common.MakeTempDir()
354    with zipfile.ZipFile(output_file, 'r', allowZip64=True) as verify_zip:
355      self.assertTrue(care_map_name in verify_zip.namelist())
356      verify_zip.extract(care_map_name, path=temp_dir)
357
358    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
359                "ro.system.build.fingerprint",
360                "google/sailfish/12345:user/dev-keys",
361                'vendor', RangeSet("0-9").to_string_raw(),
362                "ro.vendor.build.fingerprint",
363                "google/sailfish/678:user/dev-keys"]
364    self._verifyCareMap(expected, os.path.join(temp_dir, care_map_name))
365
366  @test_utils.SkipIfExternalToolsUnavailable()
367  def test_AddCareMapForAbOta_zipOutput_careMapEntryExists(self):
368    """Tests the case with ZIP output which already has care_map entry."""
369    image_paths = self._test_AddCareMapForAbOta()
370
371    output_file = common.MakeTempFile(suffix='.zip')
372    with zipfile.ZipFile(output_file, 'w', allowZip64=True) as output_zip:
373      # Create an existing META/care_map.pb entry.
374      common.ZipWriteStr(output_zip, 'META/care_map.pb',
375                         'fake care_map.pb')
376
377      # Request to add META/care_map.pb again.
378      AddCareMapForAbOta(output_zip, ['system', 'vendor'], image_paths)
379
380    # The one under OPTIONS.input_tmp must have been replaced.
381    care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
382    expected = ['system', RangeSet("0-5 10-15").to_string_raw(),
383                "ro.system.build.fingerprint",
384                "google/sailfish/12345:user/dev-keys",
385                'vendor', RangeSet("0-9").to_string_raw(),
386                "ro.vendor.build.fingerprint",
387                "google/sailfish/678:user/dev-keys"]
388
389    self._verifyCareMap(expected, care_map_file)
390
391    # The existing entry should be scheduled to be replaced.
392    self.assertIn('META/care_map.pb', OPTIONS.replace_updated_files_list)
393
394  def test_GetCareMap(self):
395    sparse_image = test_utils.construct_sparse_image([
396        (0xCAC1, 6),
397        (0xCAC3, 4),
398        (0xCAC1, 6)])
399    OPTIONS.info_dict = {
400        'extfs_sparse_flag' : '-s',
401        'system_image_size' : 53248,
402    }
403    name, care_map = GetCareMap('system', sparse_image)
404    self.assertEqual('system', name)
405    self.assertEqual(RangeSet("0-5 10-12").to_string_raw(), care_map)
406
407  def test_GetCareMap_invalidPartition(self):
408    self.assertRaises(AssertionError, GetCareMap, 'oem', None)
409
410  def test_GetCareMap_invalidAdjustedPartitionSize(self):
411    sparse_image = test_utils.construct_sparse_image([
412        (0xCAC1, 6),
413        (0xCAC3, 4),
414        (0xCAC1, 6)])
415    OPTIONS.info_dict = {
416        'extfs_sparse_flag' : '-s',
417        'system_image_size' : -45056,
418    }
419    self.assertRaises(AssertionError, GetCareMap, 'system', sparse_image)
420
421  def test_GetCareMap_nonSparseImage(self):
422    OPTIONS.info_dict = {
423        'system_image_size' : 53248,
424    }
425    # 'foo' is the image filename, which is expected to be not used by
426    # GetCareMap().
427    name, care_map = GetCareMap('system', 'foo')
428    self.assertEqual('system', name)
429    self.assertEqual(RangeSet("0-12").to_string_raw(), care_map)
430