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