parttool.py
1 #!/usr/bin/env python 2 # 3 # parttool is used to perform partition level operations - reading, 4 # writing, erasing and getting info about the partition. 5 # 6 # Copyright 2018 Espressif Systems (Shanghai) PTE LTD 7 # 8 # Licensed under the Apache License, Version 2.0 (the "License"); 9 # you may not use this file except in compliance with the License. 10 # You may obtain a copy of the License at 11 # 12 # http:#www.apache.org/licenses/LICENSE-2.0 13 # 14 # Unless required by applicable law or agreed to in writing, software 15 # distributed under the License is distributed on an "AS IS" BASIS, 16 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 17 # See the License for the specific language governing permissions and 18 # limitations under the License. 19 from __future__ import print_function, division 20 import argparse 21 import os 22 import sys 23 import subprocess 24 import tempfile 25 import re 26 import gen_esp32part as gen 27 28 29 __version__ = '2.0' 30 31 COMPONENTS_PATH = os.path.expandvars(os.path.join("$IDF_PATH", "components")) 32 ESPTOOL_PY = os.path.join(COMPONENTS_PATH, "esptool_py", "esptool", "esptool.py") 33 34 PARTITION_TABLE_OFFSET = 0x8000 35 36 37 quiet = False 38 39 40 def status(msg): 41 if not quiet: 42 print(msg) 43 44 45 class _PartitionId(): 46 47 def __init__(self, name=None, p_type=None, subtype=None, part_list=None): 48 self.name = name 49 self.type = p_type 50 self.subtype = subtype 51 self.part_list = part_list 52 53 54 class PartitionName(_PartitionId): 55 56 def __init__(self, name): 57 _PartitionId.__init__(self, name=name) 58 59 60 class PartitionType(_PartitionId): 61 62 def __init__(self, p_type, subtype, part_list=None): 63 _PartitionId.__init__(self, p_type=p_type, subtype=subtype, part_list=part_list) 64 65 66 PARTITION_BOOT_DEFAULT = _PartitionId() 67 68 69 class ParttoolTarget(): 70 71 def __init__(self, port=None, baud=None, partition_table_offset=PARTITION_TABLE_OFFSET, partition_table_file=None, 72 esptool_args=[], esptool_write_args=[], esptool_read_args=[], esptool_erase_args=[]): 73 self.port = port 74 self.baud = baud 75 76 gen.offset_part_table = partition_table_offset 77 78 def parse_esptool_args(esptool_args): 79 results = list() 80 for arg in esptool_args: 81 pattern = re.compile(r"(.+)=(.+)") 82 result = pattern.match(arg) 83 try: 84 key = result.group(1) 85 value = result.group(2) 86 results.extend(["--" + key, value]) 87 except AttributeError: 88 results.extend(["--" + arg]) 89 return results 90 91 self.esptool_args = parse_esptool_args(esptool_args) 92 self.esptool_write_args = parse_esptool_args(esptool_write_args) 93 self.esptool_read_args = parse_esptool_args(esptool_read_args) 94 self.esptool_erase_args = parse_esptool_args(esptool_erase_args) 95 96 if partition_table_file: 97 partition_table = None 98 with open(partition_table_file, "rb") as f: 99 input_is_binary = (f.read(2) == gen.PartitionDefinition.MAGIC_BYTES) 100 f.seek(0) 101 if input_is_binary: 102 partition_table = gen.PartitionTable.from_binary(f.read()) 103 104 if partition_table is None: 105 with open(partition_table_file, "r") as f: 106 f.seek(0) 107 partition_table = gen.PartitionTable.from_csv(f.read()) 108 else: 109 temp_file = tempfile.NamedTemporaryFile(delete=False) 110 temp_file.close() 111 112 try: 113 self._call_esptool(["read_flash", str(partition_table_offset), str(gen.MAX_PARTITION_LENGTH), temp_file.name]) 114 with open(temp_file.name, "rb") as f: 115 partition_table = gen.PartitionTable.from_binary(f.read()) 116 finally: 117 os.unlink(temp_file.name) 118 119 self.partition_table = partition_table 120 121 # set `out` to None to redirect the output to the STDOUT 122 # otherwise set `out` to file descriptor 123 # beware that the method does not close the file descriptor 124 def _call_esptool(self, args, out=None): 125 esptool_args = [sys.executable, ESPTOOL_PY] + self.esptool_args 126 127 if self.port: 128 esptool_args += ["--port", self.port] 129 130 if self.baud: 131 esptool_args += ["--baud", str(self.baud)] 132 133 esptool_args += args 134 135 print("Running %s..." % (" ".join(esptool_args))) 136 try: 137 subprocess.check_call(esptool_args, stdout=out, stderr=subprocess.STDOUT) 138 except subprocess.CalledProcessError as e: 139 print("An exception: **", str(e), "** occurred in _call_esptool.", file=out) 140 raise e 141 142 def get_partition_info(self, partition_id): 143 partition = None 144 145 if partition_id.name: 146 partition = self.partition_table.find_by_name(partition_id.name) 147 elif partition_id.type and partition_id.subtype: 148 partition = list(self.partition_table.find_by_type(partition_id.type, partition_id.subtype)) 149 if not partition_id.part_list: 150 partition = partition[0] 151 else: # default boot partition 152 search = ["factory"] + ["ota_{}".format(d) for d in range(16)] 153 for subtype in search: 154 partition = next(self.partition_table.find_by_type("app", subtype), None) 155 if partition: 156 break 157 158 if not partition: 159 raise Exception("Partition does not exist") 160 161 return partition 162 163 def erase_partition(self, partition_id): 164 partition = self.get_partition_info(partition_id) 165 self._call_esptool(["erase_region", str(partition.offset), str(partition.size)] + self.esptool_erase_args) 166 167 def read_partition(self, partition_id, output): 168 partition = self.get_partition_info(partition_id) 169 self._call_esptool(["read_flash", str(partition.offset), str(partition.size), output] + self.esptool_read_args) 170 171 def write_partition(self, partition_id, input): 172 self.erase_partition(partition_id) 173 174 partition = self.get_partition_info(partition_id) 175 176 with open(input, "rb") as input_file: 177 content_len = len(input_file.read()) 178 179 if content_len > partition.size: 180 raise Exception("Input file size exceeds partition size") 181 182 self._call_esptool(["write_flash", str(partition.offset), input] + self.esptool_write_args) 183 184 185 def _write_partition(target, partition_id, input): 186 target.write_partition(partition_id, input) 187 partition = target.get_partition_info(partition_id) 188 status("Written contents of file '{}' at offset 0x{:x}".format(input, partition.offset)) 189 190 191 def _read_partition(target, partition_id, output): 192 target.read_partition(partition_id, output) 193 partition = target.get_partition_info(partition_id) 194 status("Read partition '{}' contents from device at offset 0x{:x} to file '{}'" 195 .format(partition.name, partition.offset, output)) 196 197 198 def _erase_partition(target, partition_id): 199 target.erase_partition(partition_id) 200 partition = target.get_partition_info(partition_id) 201 status("Erased partition '{}' at offset 0x{:x}".format(partition.name, partition.offset)) 202 203 204 def _get_partition_info(target, partition_id, info): 205 try: 206 partitions = target.get_partition_info(partition_id) 207 if not isinstance(partitions, list): 208 partitions = [partitions] 209 except Exception: 210 return 211 212 infos = [] 213 214 try: 215 for p in partitions: 216 info_dict = { 217 "name": '{}'.format(p.name), 218 "offset": '0x{:x}'.format(p.offset), 219 "size": '0x{:x}'.format(p.size), 220 "encrypted": '{}'.format(p.encrypted) 221 } 222 for i in info: 223 infos += [info_dict[i]] 224 except KeyError: 225 raise RuntimeError("Request for unknown partition info {}".format(i)) 226 227 print(" ".join(infos)) 228 229 230 def main(): 231 global quiet 232 233 parser = argparse.ArgumentParser("ESP-IDF Partitions Tool") 234 235 parser.add_argument("--quiet", "-q", help="suppress stderr messages", action="store_true") 236 parser.add_argument("--esptool-args", help="additional main arguments for esptool", nargs="+") 237 parser.add_argument("--esptool-write-args", help="additional subcommand arguments when writing to flash", nargs="+") 238 parser.add_argument("--esptool-read-args", help="additional subcommand arguments when reading flash", nargs="+") 239 parser.add_argument("--esptool-erase-args", help="additional subcommand arguments when erasing regions of flash", nargs="+") 240 241 # By default the device attached to the specified port is queried for the partition table. If a partition table file 242 # is specified, that is used instead. 243 parser.add_argument("--port", "-p", help="port where the target device of the command is connected to; the partition table is sourced from this device \ 244 when the partition table file is not defined") 245 parser.add_argument("--baud", "-b", help="baudrate to use", type=int) 246 247 parser.add_argument("--partition-table-offset", "-o", help="offset to read the partition table from", type=str) 248 parser.add_argument("--partition-table-file", "-f", help="file (CSV/binary) to read the partition table from; \ 249 overrides device attached to specified port as the partition table source when defined") 250 251 partition_selection_parser = argparse.ArgumentParser(add_help=False) 252 253 # Specify what partition to perform the operation on. This can either be specified using the 254 # partition name or the first partition that matches the specified type/subtype 255 partition_selection_args = partition_selection_parser.add_mutually_exclusive_group() 256 257 partition_selection_args.add_argument("--partition-name", "-n", help="name of the partition") 258 partition_selection_args.add_argument("--partition-type", "-t", help="type of the partition") 259 partition_selection_args.add_argument('--partition-boot-default', "-d", help='select the default boot partition \ 260 using the same fallback logic as the IDF bootloader', action="store_true") 261 262 partition_selection_parser.add_argument("--partition-subtype", "-s", help="subtype of the partition") 263 264 subparsers = parser.add_subparsers(dest="operation", help="run parttool -h for additional help") 265 266 # Specify the supported operations 267 read_part_subparser = subparsers.add_parser("read_partition", help="read partition from device and dump contents into a file", 268 parents=[partition_selection_parser]) 269 read_part_subparser.add_argument("--output", help="file to dump the read partition contents to") 270 271 write_part_subparser = subparsers.add_parser("write_partition", help="write contents of a binary file to partition on device", 272 parents=[partition_selection_parser]) 273 write_part_subparser.add_argument("--input", help="file whose contents are to be written to the partition offset") 274 275 subparsers.add_parser("erase_partition", help="erase the contents of a partition on the device", parents=[partition_selection_parser]) 276 277 print_partition_info_subparser = subparsers.add_parser("get_partition_info", help="get partition information", parents=[partition_selection_parser]) 278 print_partition_info_subparser.add_argument("--info", help="type of partition information to get", 279 choices=["name", "offset", "size", "encrypted"], default=["offset", "size"], nargs="+") 280 print_partition_info_subparser.add_argument('--part_list', help="Get a list of partitions suitable for a given type", action='store_true') 281 282 args = parser.parse_args() 283 quiet = args.quiet 284 285 # No operation specified, display help and exit 286 if args.operation is None: 287 if not quiet: 288 parser.print_help() 289 sys.exit(1) 290 291 # Prepare the partition to perform operation on 292 if args.partition_name: 293 partition_id = PartitionName(args.partition_name) 294 elif args.partition_type: 295 if not args.partition_subtype: 296 raise RuntimeError("--partition-subtype should be defined when --partition-type is defined") 297 partition_id = PartitionType(args.partition_type, args.partition_subtype, args.part_list) 298 elif args.partition_boot_default: 299 partition_id = PARTITION_BOOT_DEFAULT 300 else: 301 raise RuntimeError("Partition to operate on should be defined using --partition-name OR \ 302 partition-type,--partition-subtype OR partition-boot-default") 303 304 # Prepare the device to perform operation on 305 target_args = {} 306 307 if args.port: 308 target_args["port"] = args.port 309 310 if args.baud: 311 target_args["baud"] = args.baud 312 313 if args.partition_table_file: 314 target_args["partition_table_file"] = args.partition_table_file 315 316 if args.partition_table_offset: 317 target_args["partition_table_offset"] = int(args.partition_table_offset, 0) 318 319 if args.esptool_args: 320 target_args["esptool_args"] = args.esptool_args 321 322 if args.esptool_write_args: 323 target_args["esptool_write_args"] = args.esptool_write_args 324 325 if args.esptool_read_args: 326 target_args["esptool_read_args"] = args.esptool_read_args 327 328 if args.esptool_erase_args: 329 target_args["esptool_erase_args"] = args.esptool_erase_args 330 331 target = ParttoolTarget(**target_args) 332 333 # Create the operation table and execute the operation 334 common_args = {'target':target, 'partition_id':partition_id} 335 parttool_ops = { 336 'erase_partition':(_erase_partition, []), 337 'read_partition':(_read_partition, ["output"]), 338 'write_partition':(_write_partition, ["input"]), 339 'get_partition_info':(_get_partition_info, ["info"]) 340 } 341 342 (op, op_args) = parttool_ops[args.operation] 343 344 for op_arg in op_args: 345 common_args.update({op_arg:vars(args)[op_arg]}) 346 347 if quiet: 348 # If exceptions occur, suppress and exit quietly 349 try: 350 op(**common_args) 351 except Exception: 352 sys.exit(2) 353 else: 354 try: 355 op(**common_args) 356 except gen.InputError as e: 357 print(e, file=sys.stderr) 358 sys.exit(2) 359 360 361 if __name__ == '__main__': 362 main()