blob: db3f83b1cd8017e8f84bf89e7b956829feee98f6 [file] [log] [blame]
Xiaochu Liudeed0232018-06-26 10:25:34 -07001# -*- coding: utf-8 -*-
2# Copyright 2018 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Script to generate a DLC (Downloadable Content) artifact."""
7
8from __future__ import print_function
9
10import hashlib
11import json
12import math
13import os
14
15from chromite.lib import commandline
16from chromite.lib import cros_build_lib
17from chromite.lib import osutils
18
19
Amin Hassani22a25eb2019-01-11 14:25:02 -080020_SQUASHFS_TYPE = 'squashfs'
21_EXT4_TYPE = 'ext4'
22
Xiaochu Liudeed0232018-06-26 10:25:34 -070023def HashFile(file_path):
24 """Calculate the sha256 hash of a file.
25
26 Args:
27 file_path: (str) path to the file.
28
29 Returns:
30 [str]: The sha256 hash of the file.
31 """
32 sha256 = hashlib.sha256()
33 with open(file_path, 'rb') as f:
34 for b in iter(lambda: f.read(2048), b''):
35 sha256.update(b)
36 return sha256.hexdigest()
37
38
Amin Hassani174eb7e2019-01-18 11:11:24 -080039class DlcGenerator(object):
Xiaochu Liudeed0232018-06-26 10:25:34 -070040 """Object to generate DLC artifacts."""
41 # Block size for the DLC image.
42 # We use 4K for various reasons:
43 # 1. it's what imageloader (linux kernel) supports.
44 # 2. it's what verity supports.
45 _BLOCK_SIZE = 4096
46 # Blocks in the initial sparse image.
47 _BLOCKS = 500000
48 # Version of manifest file.
49 _MANIFEST_VERSION = 1
50
Amin Hassanicc7ffce2019-01-11 14:57:52 -080051 # The DLC root path inside the DLC module.
52 _DLC_ROOT_DIR = 'root'
53
Colin Howesfc607682018-09-28 16:46:32 -070054 def __init__(self, img_dir, meta_dir, src_dir, fs_type, pre_allocated_blocks,
55 version, dlc_id, name):
Xiaochu Liudeed0232018-06-26 10:25:34 -070056 """Object initializer.
57
58 Args:
Colin Howesfc607682018-09-28 16:46:32 -070059 img_dir: (str) path to the DLC image dest root directory.
60 meta_dir: (str) path to the DLC metadata dest root directory.
Xiaochu Liudeed0232018-06-26 10:25:34 -070061 src_dir: (str) path to the DLC source root directory.
62 fs_type: (str) file system type.
63 pre_allocated_blocks: (int) number of blocks pre-allocated on device.
64 version: (str) DLC version.
65 dlc_id: (str) DLC ID.
66 name: (str) DLC name.
67 """
68 self.src_dir = src_dir
69 self.fs_type = fs_type
70 self.pre_allocated_blocks = pre_allocated_blocks
71 self.version = version
72 self.dlc_id = dlc_id
73 self.name = name
74 # Create path for all final artifacts.
Colin Howesfc607682018-09-28 16:46:32 -070075 self.dest_image = os.path.join(img_dir, 'dlc.img')
76 self.dest_table = os.path.join(meta_dir, 'table')
77 self.dest_imageloader_json = os.path.join(meta_dir, 'imageloader.json')
Xiaochu Liudeed0232018-06-26 10:25:34 -070078
79 def SquashOwnerships(self, path):
80 """Squash the owernships & permissions for files.
81
82 Args:
83 path: (str) path that contains all files to be processed.
84 """
85 cros_build_lib.SudoRunCommand(['chown', '-R', '0:0', path])
86 cros_build_lib.SudoRunCommand(
87 ['find', path, '-exec', 'touch', '-h', '-t', '197001010000.00', '{}',
88 '+'])
89
90 def CreateExt4Image(self):
91 """Create an ext4 image."""
92 with osutils.TempDir(prefix='dlc_') as temp_dir:
93 mount_point = os.path.join(temp_dir, 'mount_point')
94 # Create a raw image file.
95 with open(self.dest_image, 'w') as f:
96 f.truncate(self._BLOCKS * self._BLOCK_SIZE)
97 # Create an ext4 file system on the raw image.
98 cros_build_lib.RunCommand(
99 ['/sbin/mkfs.ext4', '-b', str(self._BLOCK_SIZE), '-O',
100 '^has_journal', self.dest_image], capture_output=True)
101 # Create the mount_point directory.
102 osutils.SafeMakedirs(mount_point)
103 # Mount the ext4 image.
104 osutils.MountDir(self.dest_image, mount_point, mount_opts=('loop', 'rw'))
Amin Hassanicc7ffce2019-01-11 14:57:52 -0800105
106 dlc_root_dir = os.path.join(mount_point, self._DLC_ROOT_DIR)
107 osutils.SafeMakedirs(dlc_root_dir)
Xiaochu Liudeed0232018-06-26 10:25:34 -0700108 try:
109 # Copy DLC files over to the image.
Colin Howes839c0042018-12-03 14:50:43 -0800110 osutils.CopyDirContents(self.src_dir, dlc_root_dir)
111
Xiaochu Liudeed0232018-06-26 10:25:34 -0700112 self.SquashOwnerships(mount_point)
113 finally:
114 # Unmount the ext4 image.
115 osutils.UmountDir(mount_point)
116 # Shrink to minimum size.
117 cros_build_lib.RunCommand(
118 ['/sbin/e2fsck', '-y', '-f', self.dest_image], capture_output=True)
119 cros_build_lib.RunCommand(
120 ['/sbin/resize2fs', '-M', self.dest_image], capture_output=True)
121
122 def CreateSquashfsImage(self):
123 """Create a squashfs image."""
124 with osutils.TempDir(prefix='dlc_') as temp_dir:
Amin Hassani22a25eb2019-01-11 14:25:02 -0800125 squashfs_root = os.path.join(temp_dir, 'squashfs-root')
Amin Hassanicc7ffce2019-01-11 14:57:52 -0800126 dlc_root_dir = os.path.join(squashfs_root, self._DLC_ROOT_DIR)
127 osutils.SafeMakedirs(dlc_root_dir)
Colin Howes839c0042018-12-03 14:50:43 -0800128 osutils.CopyDirContents(self.src_dir, dlc_root_dir)
Amin Hassanicc7ffce2019-01-11 14:57:52 -0800129
Amin Hassani22a25eb2019-01-11 14:25:02 -0800130 self.SquashOwnerships(squashfs_root)
131
132 cros_build_lib.RunCommand(['mksquashfs', squashfs_root, self.dest_image,
133 '-4k-align', '-noappend'],
134 capture_output=True)
135
136 # We changed the ownership and permissions of the squashfs_root
137 # directory. Now we need to remove it manually.
138 osutils.RmDir(squashfs_root, sudo=True)
Xiaochu Liudeed0232018-06-26 10:25:34 -0700139
140 def CreateImage(self):
141 """Create the image and copy the DLC files to it."""
Amin Hassani22a25eb2019-01-11 14:25:02 -0800142 if self.fs_type == _EXT4_TYPE:
Xiaochu Liudeed0232018-06-26 10:25:34 -0700143 self.CreateExt4Image()
Amin Hassani22a25eb2019-01-11 14:25:02 -0800144 elif self.fs_type == _SQUASHFS_TYPE:
Xiaochu Liudeed0232018-06-26 10:25:34 -0700145 self.CreateSquashfsImage()
146 else:
147 raise ValueError('Wrong fs type: %s used:' % self.fs_type)
148
149 def GetImageloaderJsonContent(self, image_hash, table_hash, blocks):
150 """Return the content of imageloader.json file.
151
152 Args:
153 image_hash: (str) sha256 hash of the DLC image.
154 table_hash: (str) sha256 hash of the DLC table file.
155 blocks: (int) number of blocks in the DLC image.
156
157 Returns:
158 [str]: content of imageloader.json file.
159 """
160 return {
161 'fs-type': self.fs_type,
162 'id': self.dlc_id,
163 'image-sha256-hash': image_hash,
164 'image-type': 'dlc',
165 'is-removable': True,
166 'manifest-version': self._MANIFEST_VERSION,
167 'name': self.name,
168 'pre-allocated-size': self.pre_allocated_blocks * self._BLOCK_SIZE,
169 'size': blocks * self._BLOCK_SIZE,
170 'table-sha256-hash': table_hash,
171 'version': self.version,
172 }
173
174 def GenerateVerity(self):
175 """Generate verity parameters and hashes for the image."""
176 with osutils.TempDir(prefix='dlc_') as temp_dir:
177 hash_tree = os.path.join(temp_dir, 'hash_tree')
178 # Get blocks in the image.
179 blocks = math.ceil(
180 os.path.getsize(self.dest_image) / self._BLOCK_SIZE)
181 result = cros_build_lib.RunCommand(
182 ['verity', 'mode=create', 'alg=sha256', 'payload=' + self.dest_image,
183 'payload_blocks=' + str(blocks), 'hashtree=' + hash_tree,
184 'salt=random'], capture_output=True)
185 table = result.output
186
187 # Append the merkle tree to the image.
188 osutils.WriteFile(self.dest_image, osutils.ReadFile(hash_tree), 'a+')
189
190 # Write verity parameter to table file.
191 osutils.WriteFile(self.dest_table, table)
192
193 # Compute image hash.
194 image_hash = HashFile(self.dest_image)
195 table_hash = HashFile(self.dest_table)
196 # Write image hash to imageloader.json file.
197 blocks = math.ceil(
198 os.path.getsize(self.dest_image) / self._BLOCK_SIZE)
199 imageloader_json_content = self.GetImageloaderJsonContent(
200 image_hash, table_hash, int(blocks))
201 with open(self.dest_imageloader_json, 'w') as f:
202 json.dump(imageloader_json_content, f)
203
204 def GenerateDLC(self):
205 """Generate a DLC artifact."""
206 # Create the image and copy the DLC files to it.
207 self.CreateImage()
208 # Generate hash tree and other metadata.
209 self.GenerateVerity()
210
211
212def GetParser():
213 parser = commandline.ArgumentParser(description=__doc__)
214 # Required arguments:
215 required = parser.add_argument_group('Required Arguments')
216 required.add_argument('--src-dir', type='path', metavar='SRC_DIR_PATH',
217 required=True,
218 help='Root directory path that contains all DLC files '
219 'to be packed.')
Colin Howesfc607682018-09-28 16:46:32 -0700220 required.add_argument('--img-dir', type='path', metavar='IMG_DIR_PATH',
Xiaochu Liudeed0232018-06-26 10:25:34 -0700221 required=True,
Colin Howesfc607682018-09-28 16:46:32 -0700222 help='Root directory path that contains DLC image file '
223 'output.')
224 required.add_argument('--meta-dir', type='path', metavar='META_DIR_PATH',
225 required=True,
226 help='Root directory path that contains DLC metadata '
227 'output.')
Xiaochu Liudeed0232018-06-26 10:25:34 -0700228 required.add_argument('--pre-allocated-blocks', type=int,
229 metavar='PREALLOCATEDBLOCKS', required=True,
230 help='Number of blocks (block size is 4k) that need to'
231 'be pre-allocated on device.')
232 required.add_argument('--version', metavar='VERSION', required=True,
233 help='DLC Version.')
234 required.add_argument('--id', metavar='ID', required=True,
235 help='DLC ID (unique per DLC).')
236 required.add_argument('--name', metavar='NAME', required=True,
237 help='A human-readable name for the DLC.')
Amin Hassani22a25eb2019-01-11 14:25:02 -0800238
239 args = parser.add_argument_group('Arguments')
240 args.add_argument('--fs-type', metavar='FS_TYPE', default=_SQUASHFS_TYPE,
241 choices=(_SQUASHFS_TYPE, _EXT4_TYPE),
242 help='File system type of the image.')
243
Xiaochu Liudeed0232018-06-26 10:25:34 -0700244 return parser
245
246
247def main(argv):
248 opts = GetParser().parse_args(argv)
249 opts.Freeze()
250
251 # Generate final DLC files.
Amin Hassani174eb7e2019-01-18 11:11:24 -0800252 dlc_generator = DlcGenerator(opts.img_dir, opts.meta_dir, opts.src_dir,
Colin Howesfc607682018-09-28 16:46:32 -0700253 opts.fs_type, opts.pre_allocated_blocks,
254 opts.version, opts.id, opts.name)
Xiaochu Liudeed0232018-06-26 10:25:34 -0700255 dlc_generator.GenerateDLC()