Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 1 | # Copyright 2022 The ChromiumOS Authors |
| 2 | # Use of this source code is governed by a BSD-style license that can be |
| 3 | # found in the LICENSE file. |
| 4 | |
| 5 | """Parse and operate based on disk layout files. |
| 6 | |
| 7 | For information on the JSON format, see: |
| 8 | http://dev.chromium.org/chromium-os/developer-guide/disk-layout-format |
| 9 | |
| 10 | The --adjust_part flag takes arguments like: |
| 11 | <label>:<op><size> |
| 12 | Where: |
| 13 | <label> is a label name as found in the disk layout file |
| 14 | <op> is one of the three: + - = |
| 15 | <size> is a number followed by an optional size qualifier: |
| 16 | B, KiB, MiB, GiB, TiB: bytes, kibi-, mebi-, gibi-, tebi- (base 1024) |
| 17 | B, K, M, G, T: short hand for above |
| 18 | B, KB, MB, GB, TB: bytes, kilo-, mega-, giga-, tera- (base 1000) |
| 19 | |
| 20 | This will set the ROOT-A partition size to 1 gibibytes (1024 * 1024 * 1024 * 1): |
| 21 | --adjust_part ROOT-A:=1GiB |
| 22 | This will grow the ROOT-A partition size by 500 mebibytes (1024 * 1024 * 500): |
| 23 | --adjust_part ROOT-A:+500MiB |
| 24 | This will shrink the ROOT-A partition size by 10 mebibytes (1024 * 1024 * 10): |
| 25 | --adjust_part ROOT-A:-20MiB |
| 26 | """ |
| 27 | |
| 28 | import argparse |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 29 | import inspect |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 30 | import os |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 31 | import sys |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 32 | from typing import List, Union |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 33 | |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 34 | from chromite.lib import disk_layout |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 35 | |
| 36 | |
| 37 | def WritePartitionScript( |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 38 | options: List[str], |
| 39 | image_type: str, |
| 40 | layout_filename: Union[str, os.PathLike], |
| 41 | sfilename: Union[str, os.PathLike], |
| 42 | vfilename: Union[str, os.PathLike], |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 43 | ): |
| 44 | """Writes a shell script with functions for the base and requested layouts. |
| 45 | |
| 46 | Args: |
| 47 | options: Flags passed to the script |
| 48 | image_type: Type of image eg base/test/dev/factory_install |
| 49 | layout_filename: Path to partition configuration file |
| 50 | sfilename: Filename to write the finished script to |
| 51 | vfilename: Filename to write the partition variables json data to |
| 52 | """ |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 53 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 54 | return disk_layout.DiskLayout( |
| 55 | layout_filename, options.adjust_part.split() |
| 56 | ).WritePartitionScript(image_type, sfilename, vfilename) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 57 | |
| 58 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 59 | def GetBlockSize( |
| 60 | _options: List[str], layout_filename: Union[str, os.PathLike] |
| 61 | ) -> str: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 62 | """Returns the partition table block size. |
| 63 | |
| 64 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 65 | options: Flags passed to the script. |
| 66 | layout_filename: Path to partition configuration file. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 67 | |
| 68 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 69 | Block size of all partitions in the layout. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 70 | """ |
| 71 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 72 | return disk_layout.DiskLayout(layout_filename).GetBlockSize() |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 73 | |
| 74 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 75 | def GetFilesystemBlockSize( |
| 76 | _options: List[str], layout_filename: Union[str, os.PathLike] |
| 77 | ) -> str: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 78 | """Returns the filesystem block size. |
| 79 | |
| 80 | This is used for all partitions in the table that have filesystems. |
| 81 | |
| 82 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 83 | options: Flags passed to the script. |
| 84 | layout_filename: Path to partition configuration file. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 85 | |
| 86 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 87 | Block size of all filesystems in the layout. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 88 | """ |
| 89 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 90 | return disk_layout.DiskLayout(layout_filename).GetFilesystemBlockSize() |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 91 | |
| 92 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 93 | def GetPartitionSize( |
| 94 | options: List[str], |
| 95 | image_type: str, |
| 96 | layout_filename: Union[str, os.PathLike], |
| 97 | num: str, |
| 98 | ) -> str: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 99 | """Returns the partition size of a given partition for a given layout type. |
| 100 | |
| 101 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 102 | options: Flags passed to the script. |
| 103 | image_type: Type of image eg base/test/dev/factory_install. |
| 104 | layout_filename: Path to partition configuration file. |
| 105 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 106 | |
| 107 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 108 | Size of selected partition in bytes. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 109 | """ |
| 110 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 111 | return disk_layout.DiskLayout( |
| 112 | layout_filename, options.adjust_part.split() |
| 113 | ).GetPartitionSize(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 114 | |
| 115 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 116 | def GetFormat( |
| 117 | options: List[str], |
| 118 | image_type: str, |
| 119 | layout_filename: Union[str, os.PathLike], |
| 120 | num: str, |
| 121 | ) -> str: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 122 | """Returns the format of a given partition for a given layout type. |
| 123 | |
| 124 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 125 | options: Flags passed to the script. |
| 126 | image_type: Type of image eg base/test/dev/factory_install. |
| 127 | layout_filename: Path to partition configuration file. |
| 128 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 129 | |
| 130 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 131 | Format of the selected partition's filesystem. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 132 | """ |
| 133 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 134 | return disk_layout.DiskLayout( |
| 135 | layout_filename, options.adjust_part.split() |
| 136 | ).GetFormat(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 137 | |
| 138 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 139 | def GetFilesystemFormat( |
| 140 | options: List[str], |
| 141 | image_type: str, |
| 142 | layout_filename: Union[str, os.PathLike], |
| 143 | num: str, |
| 144 | ) -> str: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 145 | """Returns the filesystem format of a partition for a given layout type. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 146 | |
| 147 | Args: |
| 148 | options: Flags passed to the script |
| 149 | image_type: Type of image eg base/test/dev/factory_install |
| 150 | layout_filename: Path to partition configuration file |
| 151 | num: Number of the partition you want to read from |
| 152 | |
| 153 | Returns: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 154 | Format of the selected partition's filesystem |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 155 | """ |
| 156 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 157 | return disk_layout.DiskLayout( |
| 158 | layout_filename, options.adjust_part.split() |
| 159 | ).GetFilesystemFormat(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 160 | |
| 161 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 162 | def GetFilesystemSize( |
| 163 | options: List[str], |
| 164 | image_type: str, |
| 165 | layout_filename: Union[str, os.PathLike], |
| 166 | num: str, |
| 167 | ) -> int: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 168 | """Returns the filesystem size of a given partition for a given layout type. |
| 169 | |
| 170 | If no filesystem size is specified, returns the partition size. |
| 171 | |
| 172 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 173 | options: Flags passed to the script. |
| 174 | image_type: Type of image eg base/test/dev/factory_install. |
| 175 | layout_filename: Path to partition configuration file. |
| 176 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 177 | |
| 178 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 179 | Size of selected partition filesystem in bytes. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 180 | """ |
| 181 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 182 | return disk_layout.DiskLayout( |
| 183 | layout_filename, options.adjust_part.split() |
| 184 | ).GetFilesystemSize(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 185 | |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 186 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 187 | def GetImageTypes( |
| 188 | _options: List[str], layout_filename: Union[str, os.PathLike] |
| 189 | ) -> str: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 190 | """Returns a list of all the image types in the layout. |
| 191 | |
| 192 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 193 | options: Flags passed to the script. |
| 194 | layout_filename: Path to partition configuration file. |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 195 | |
| 196 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 197 | List of all image types. |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 198 | """ |
| 199 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 200 | image_types = disk_layout.DiskLayout(layout_filename).GetImageTypes() |
| 201 | return " ".join(image_types) |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 202 | |
| 203 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 204 | def GetFilesystemOptions( |
| 205 | options: List[str], |
| 206 | image_type: str, |
| 207 | layout_filename: Union[str, os.PathLike], |
| 208 | num: str, |
| 209 | ) -> str: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 210 | """Returns the filesystem options of a given partition and layout type. |
| 211 | |
| 212 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 213 | options: Flags passed to the script. |
| 214 | image_type: Type of image eg base/test/dev/factory_install. |
| 215 | layout_filename: Path to partition configuration file. |
| 216 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 217 | |
| 218 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 219 | The selected partition's filesystem options. |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 220 | """ |
| 221 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 222 | return disk_layout.DiskLayout( |
| 223 | layout_filename, options.adjust_part.split() |
| 224 | ).GetFilesystemOptions(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 225 | |
| 226 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 227 | def GetLabel( |
| 228 | options: List[str], |
| 229 | image_type: str, |
| 230 | layout_filename: Union[str, os.PathLike], |
| 231 | num: str, |
| 232 | ) -> str: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 233 | """Returns the label for a given partition. |
| 234 | |
| 235 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 236 | options: Flags passed to the script. |
| 237 | image_type: Type of image eg base/test/dev/factory_install. |
| 238 | layout_filename: Path to partition configuration file. |
| 239 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 240 | |
| 241 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 242 | Label of selected partition, or 'UNTITLED' if none specified. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 243 | """ |
| 244 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 245 | return disk_layout.DiskLayout( |
| 246 | layout_filename, options.adjust_part.split() |
| 247 | ).GetLabel(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 248 | |
| 249 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 250 | def GetNumber( |
| 251 | options: List[str], |
| 252 | image_type: str, |
| 253 | layout_filename: Union[str, os.PathLike], |
| 254 | label, |
| 255 | ) -> int: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 256 | """Returns the partition number of a given label. |
| 257 | |
| 258 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 259 | options: Flags passed to the script. |
| 260 | image_type: Type of image eg base/test/dev/factory_install. |
| 261 | layout_filename: Path to partition configuration file. |
| 262 | label: Number of the partition you want to read from. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 263 | |
| 264 | Returns: |
| 265 | The number of the partition corresponding to the label. |
| 266 | """ |
| 267 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 268 | return disk_layout.DiskLayout( |
| 269 | layout_filename, options.adjust_part.split() |
| 270 | ).GetNumber(image_type, label) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 271 | |
| 272 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 273 | def GetReservedEraseBlocks( |
| 274 | options: List[str], |
| 275 | image_type: str, |
| 276 | layout_filename: Union[str, os.PathLike], |
| 277 | num: str, |
| 278 | ) -> int: |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 279 | """Returns the number of erase blocks reserved in the partition. |
| 280 | |
| 281 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 282 | options: Flags passed to the script. |
| 283 | image_type: Type of image eg base/test/dev/factory_install. |
| 284 | layout_filename: Path to partition configuration file. |
| 285 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 286 | |
| 287 | Returns: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 288 | Number of reserved erase blocks. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 289 | """ |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 290 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 291 | return disk_layout.DiskLayout( |
| 292 | layout_filename, options.adjust_part.split() |
| 293 | ).GetReservedEraseBlocks(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 294 | |
| 295 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 296 | def GetType( |
| 297 | options: List[str], |
| 298 | image_type: str, |
| 299 | layout_filename: Union[str, os.PathLike], |
| 300 | num: str, |
| 301 | ) -> str: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 302 | """Returns the type of a given partition for a given layout. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 303 | |
| 304 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 305 | options: Flags passed to the script. |
| 306 | image_type: Type of image eg base/test/dev/factory_install. |
| 307 | layout_filename: Path to partition configuration file. |
| 308 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 309 | |
| 310 | Returns: |
| 311 | Type of the specified partition. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 312 | """ |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 313 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 314 | return disk_layout.DiskLayout( |
| 315 | layout_filename, options.adjust_part.split() |
| 316 | ).GetType(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 317 | |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 318 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 319 | def GetPartitions( |
| 320 | options: List[str], |
| 321 | image_type: str, |
| 322 | layout_filename: Union[str, os.PathLike], |
| 323 | ) -> str: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 324 | """Returns the partition numbers for the image_type. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 325 | |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 326 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 327 | options: Flags passed to the script. |
| 328 | image_type: Type of image eg base/test/dev/factory_install. |
| 329 | layout_filename: Path to partition configuration file. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 330 | |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 331 | Returns: |
| 332 | A space delimited string of partition numbers. |
| 333 | """ |
| 334 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 335 | return disk_layout.DiskLayout( |
| 336 | layout_filename, options.adjust_part.split() |
| 337 | ).GetPartitions(image_type) |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 338 | |
| 339 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 340 | def GetUUID( |
| 341 | options: List[str], |
| 342 | image_type: str, |
| 343 | layout_filename: Union[str, os.PathLike], |
| 344 | num: str, |
| 345 | ) -> str: |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 346 | """Returns the filesystem UUID of a given partition for a given layout type. |
| 347 | |
| 348 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 349 | options: Flags passed to the script. |
| 350 | image_type: Type of image eg base/test/dev/factory_install. |
| 351 | layout_filename: Path to partition configuration file. |
| 352 | num: Number of the partition you want to read from. |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 353 | |
| 354 | Returns: |
| 355 | UUID of specified partition. Defaults to random if not set. |
| 356 | """ |
| 357 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 358 | return disk_layout.DiskLayout( |
| 359 | layout_filename, options.adjust_part.split() |
| 360 | ).GetUUID(image_type, int(num)) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 361 | |
| 362 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 363 | def DoDebugOutput( |
| 364 | options: List[str], |
| 365 | layout_filename: Union[str, os.PathLike], |
| 366 | image_type: str, |
| 367 | ): |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 368 | """Prints out a human readable disk layout in on-disk order. |
| 369 | |
| 370 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 371 | options: Flags passed to the script. |
| 372 | layout_filename: Path to partition configuration file. |
| 373 | image_type: Type of image e.g. ALL/LIST/base/test/dev/factory_install. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 374 | """ |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 375 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 376 | return disk_layout.DiskLayout( |
| 377 | layout_filename, options.adjust_part.split() |
| 378 | ).DoDebugOutput(image_type) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 379 | |
| 380 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 381 | def Validate( |
| 382 | options: List[str], |
| 383 | image_type: str, |
| 384 | layout_filename: Union[str, os.PathLike], |
| 385 | ): |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 386 | """Validates a layout file, used before reading sizes to check for errors. |
| 387 | |
| 388 | Args: |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 389 | options: Flags passed to the script. |
| 390 | image_type: Type of image eg base/test/dev/factory_install. |
| 391 | layout_filename: Path to partition configuration file. |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 392 | """ |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 393 | |
Ram Chandrasekar | 02bfe44 | 2022-10-18 21:55:46 +0000 | [diff] [blame^] | 394 | return disk_layout.DiskLayout( |
| 395 | layout_filename, options.adjust_part.split() |
| 396 | ).Validate(image_type) |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 397 | |
| 398 | |
| 399 | class ArgsAction(argparse.Action): # pylint: disable=no-init |
| 400 | """Helper to add all arguments to an args array. |
| 401 | |
| 402 | ArgumentParser does not let you specify the same dest for multiple args. |
| 403 | We take care of appending to the 'args' array ourselves here. |
| 404 | """ |
| 405 | |
| 406 | def __call__(self, parser, namespace, values, option_string=None): |
| 407 | args = getattr(namespace, "args", []) |
| 408 | args.append(values) |
| 409 | setattr(namespace, "args", args) |
| 410 | |
| 411 | |
| 412 | class HelpAllAction(argparse.Action): |
| 413 | """Display all subcommands help in one go.""" |
| 414 | |
| 415 | def __init__(self, *args, **kwargs): |
| 416 | if "nargs" in kwargs: |
| 417 | raise ValueError("nargs not allowed") |
| 418 | kwargs["nargs"] = 0 |
| 419 | argparse.Action.__init__(self, *args, **kwargs) |
| 420 | |
| 421 | def __call__(self, parser, namespace, values, option_string=None): |
| 422 | print("%s\nCommands:" % (parser.description,), end="") |
| 423 | subparser = getattr(namespace, "help_all") |
| 424 | for key, subparser in namespace.help_all.choices.items(): |
| 425 | # Should we include the desc of each arg too ? |
| 426 | print( |
| 427 | "\n %s %s\n %s" |
| 428 | % ( |
| 429 | key, |
| 430 | subparser.get_default("help_all"), |
| 431 | subparser.description, |
| 432 | ) |
| 433 | ) |
| 434 | sys.exit(0) |
| 435 | |
| 436 | |
| 437 | def GetParser(): |
Ram Chandrasekar | 633e1bb | 2022-09-20 20:02:31 +0000 | [diff] [blame] | 438 | """Return a parser for the CLI. |
| 439 | |
| 440 | We use the function docstring to build the cli argument, help text |
| 441 | and their arguments. |
| 442 | """ |
Ram Chandrasekar | 613e1f4 | 2022-08-25 00:02:12 +0000 | [diff] [blame] | 443 | parser = argparse.ArgumentParser( |
| 444 | description=__doc__, |
| 445 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 446 | ) |
| 447 | parser.add_argument( |
| 448 | "--adjust_part", |
| 449 | metavar="SPEC", |
| 450 | default="", |
| 451 | help="adjust partition sizes", |
| 452 | ) |
| 453 | |
| 454 | action_map = { |
| 455 | "write": WritePartitionScript, |
| 456 | "readblocksize": GetBlockSize, |
| 457 | "readfsblocksize": GetFilesystemBlockSize, |
| 458 | "readpartsize": GetPartitionSize, |
| 459 | "readformat": GetFormat, |
| 460 | "readfsformat": GetFilesystemFormat, |
| 461 | "readfssize": GetFilesystemSize, |
| 462 | "readimagetypes": GetImageTypes, |
| 463 | "readfsoptions": GetFilesystemOptions, |
| 464 | "readlabel": GetLabel, |
| 465 | "readnumber": GetNumber, |
| 466 | "readreservederaseblocks": GetReservedEraseBlocks, |
| 467 | "readtype": GetType, |
| 468 | "readpartitionnums": GetPartitions, |
| 469 | "readuuid": GetUUID, |
| 470 | "debug": DoDebugOutput, |
| 471 | "validate": Validate, |
| 472 | } |
| 473 | |
| 474 | # Subparsers are required by default under Python 2. Python 3 changed to |
| 475 | # not required, but didn't include a required option until 3.7. Setting |
| 476 | # the required member works in all versions (and setting dest name). |
| 477 | subparsers = parser.add_subparsers(title="Commands", dest="command") |
| 478 | subparsers.required = True |
| 479 | |
| 480 | for name, func in sorted(action_map.items()): |
| 481 | # Turn the func's docstring into something we can show the user. |
| 482 | desc, doc = func.__doc__.split("\n", 1) |
| 483 | # Extract the help for each argument. |
| 484 | args_help = {} |
| 485 | for line in doc.splitlines(): |
| 486 | if ":" in line: |
| 487 | arg, text = line.split(":", 1) |
| 488 | args_help[arg.strip()] = text.strip() |
| 489 | |
| 490 | argspec = inspect.getfullargspec(func) |
| 491 | # Skip the first argument as that'll be the options field. |
| 492 | args = argspec.args[1:] |
| 493 | |
| 494 | subparser = subparsers.add_parser(name, description=desc, help=desc) |
| 495 | subparser.set_defaults( |
| 496 | callback=func, help_all=" ".join("<%s>" % x for x in args) |
| 497 | ) |
| 498 | for arg in args: |
| 499 | subparser.add_argument(arg, action=ArgsAction, help=args_help[arg]) |
| 500 | |
| 501 | parser.add_argument( |
| 502 | "--help-all", |
| 503 | action=HelpAllAction, |
| 504 | default=subparsers, |
| 505 | help="show all commands and their help in one screen", |
| 506 | ) |
| 507 | |
| 508 | return parser |
| 509 | |
| 510 | |
| 511 | def main(argv): |
| 512 | parser = GetParser() |
| 513 | opts = parser.parse_args(argv) |
| 514 | |
| 515 | ret = opts.callback(opts, *opts.args) |
| 516 | if ret is not None: |
| 517 | print(ret) |
| 518 | |
| 519 | |
| 520 | if __name__ == "__main__": |
| 521 | sys.exit(main(sys.argv[1:])) |