| #!/usr/bin/python |
| # Copyright (c) 2012 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| import copy |
| import json |
| import os |
| import sys |
| |
| # First sector we can use. |
| START_SECTOR = 64 |
| |
| class ConfigNotFound(Exception): |
| pass |
| class PartitionNotFound(Exception): |
| pass |
| class InvalidLayout(Exception): |
| pass |
| |
| |
| def LoadPartitionConfig(filename): |
| """Loads a partition tables configuration file into a Python object. |
| |
| Args: |
| filename: Filename to load into object |
| Returns: |
| Object containing disk layout configuration |
| """ |
| |
| if not os.path.exists(filename): |
| raise ConfigNotFound("Partition config %s was not found!" % filename) |
| with open(filename) as f: |
| config = json.load(f) |
| f.close() |
| |
| metadata = config["metadata"] |
| metadata["block_size"] = int(metadata["block_size"]) |
| |
| for layout_name, layout in config["layouts"].items(): |
| for part in layout: |
| part["blocks"] = int(part["blocks"]) |
| part["bytes"] = part["blocks"] * metadata["block_size"] |
| |
| if "fs_blocks" in part: |
| part["fs_blocks"] = int(part["fs_blocks"]) |
| part["fs_bytes"] = part["fs_blocks"] * metadata["fs_block_size"] |
| |
| if part["fs_bytes"] > part["bytes"]: |
| raise InvalidLayout("Filesystem may not be larger than partition") |
| |
| return config |
| |
| |
| def GetTableTotals(config, partitions): |
| """Calculates total sizes/counts for a partition table. |
| |
| Args: |
| config: Partition configuration file object |
| partitions: List of partitions to process |
| Returns: |
| Dict containing totals data |
| """ |
| |
| ret = { |
| "expand_count": 0, |
| "expand_min": 0, |
| "block_count": START_SECTOR * config["metadata"]["block_size"] |
| } |
| |
| # Total up the size of all non-expanding partitions to get the minimum |
| # required disk size. |
| for partition in partitions: |
| if "features" in partition and "expand" in partition["features"]: |
| ret["expand_count"] += 1 |
| ret["expand_min"] += partition["blocks"] |
| else: |
| ret["block_count"] += partition["blocks"] |
| |
| # At present, only one expanding partition is permitted. |
| # Whilst it'd be possible to have two, we don't need this yet |
| # and it complicates things, so it's been left out for now. |
| if ret["expand_count"] > 1: |
| raise InvalidLayout("1 expand partition allowed, %d requested" |
| % ret["expand_count"]) |
| |
| ret["min_disk_size"] = ret["block_count"] + ret["expand_min"] |
| |
| return ret |
| |
| |
| def GetPartitionTable(config, image_type): |
| """Generates requested image_type layout from a layout configuration. |
| |
| This loads the base table and then overlays the requested layout over |
| the base layout. |
| |
| Args: |
| config: Partition configuration file object |
| image_type: Type of image eg base/test/dev/factory_install |
| Returns: |
| Object representing a selected partition table |
| """ |
| |
| partitions = config["layouts"]["base"] |
| |
| if image_type != "base": |
| for partition_t in config["layouts"][image_type]: |
| for partition in partitions: |
| if partition["type"] == "blank" or partition_t["type"] == "blank": |
| continue |
| if partition_t["num"] == partition["num"]: |
| for k, v in partition_t.items(): |
| partition[k] = v |
| |
| return partitions |
| |
| |
| def GetScriptShell(): |
| """Loads and returns the skeleton script for our output script. |
| |
| Returns: |
| A string containg the skeleton script |
| """ |
| |
| script_shell_path = os.path.join(os.path.dirname(__file__), "cgpt_shell.sh") |
| with open(script_shell_path, "r") as f: |
| script_shell = "".join(f.readlines()) |
| f.close() |
| |
| # Before we return, insert the path to this tool so somebody reading the |
| # script later can tell where it was generated. |
| script_shell = script_shell.replace("@SCRIPT_GENERATOR@", script_shell_path) |
| |
| return script_shell |
| |
| |
| def WriteLayoutFunction(sfile, func_name, image_type, config): |
| """Writes a shell script function to write out a given partition table. |
| |
| Args: |
| sfile: File handle we're writing to |
| func_name: Function name to write out for specified layout |
| image_type: Type of image eg base/test/dev/factory_install |
| config: Partition configuration file object |
| """ |
| |
| partitions = GetPartitionTable(config, image_type) |
| partition_totals = GetTableTotals(config, partitions) |
| |
| sfile.write("%s() {\ncreate_image $1 %d %s\n" % ( |
| func_name, partition_totals["min_disk_size"], |
| config["metadata"]["block_size"])) |
| |
| sfile.write("CURR=%d\n" % START_SECTOR) |
| sfile.write("$GPT create $1\n") |
| |
| # Pass 1: Set up the expanding partition size. |
| for partition in partitions: |
| partition["var"] = partition["blocks"] |
| if partition["type"] != "blank": |
| |
| if partition["num"] == 1: |
| if "features" in partition and "expand" in partition["features"]: |
| sfile.write("if [ -b $1 ]; then\n") |
| sfile.write("STATEFUL_SIZE=$(( $(numsectors $1) - %d))\n" % |
| partition_totals["block_count"]) |
| sfile.write("else\n") |
| sfile.write("STATEFUL_SIZE=%s\n" % partition["blocks"]) |
| sfile.write("fi\n") |
| partition["var"] = "$STATEFUL_SIZE" |
| |
| |
| sfile.write("STATEFUL_SIZE=$((STATEFUL_SIZE-(STATEFUL_SIZE %% %d)))\n" % |
| config["metadata"]["fs_block_size"]) |
| |
| # Pass 2: Write out all the cgpt add commands. |
| for partition in partitions: |
| if partition["type"] != "blank": |
| sfile.write("$GPT add -i %d -b $CURR -s %s -t %s -l %s $1 && " % ( |
| partition["num"], str(partition["var"]), partition["type"], |
| partition["label"])) |
| |
| # Increment the CURR counter ready for the next partition. |
| sfile.write("CURR=$(( $CURR + %s ))\n" % partition["var"]) |
| |
| # Set default priorities on kernel partitions |
| sfile.write("$GPT add -i 2 -S 0 -T 15 -P 15 $1\n") |
| sfile.write("$GPT add -i 4 -S 0 -T 15 -P 0 $1\n") |
| sfile.write("$GPT add -i 6 -S 0 -T 15 -P 0 $1\n") |
| |
| sfile.write("$GPT boot -p -b $2 -i 12 $1\n") |
| sfile.write("$GPT show $1\n") |
| sfile.write("}\n") |
| |
| |
| def GetPartitionByNumber(partitions, num): |
| """Given a partition table and number returns the partition object. |
| |
| Args: |
| partitions: List of partitions to search in |
| num: Number of partition to find |
| Returns: |
| An object for the selected partition |
| """ |
| for partition in partitions: |
| if partition["type"] == "blank": |
| continue |
| if partition["num"] == int(num): |
| return partition |
| |
| raise PartitionNotFound("Partition not found") |
| |
| |
| def WritePartitionScript(image_type, layout_filename, sfilename): |
| """Writes a shell script with functions for the base and requested layouts. |
| |
| Args: |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| sfilename: Filename to write the finished script to |
| """ |
| |
| config = LoadPartitionConfig(layout_filename) |
| |
| sfile = open(sfilename, "w") |
| script_shell = GetScriptShell() |
| sfile.write(script_shell) |
| |
| WriteLayoutFunction(sfile, "write_base_table", "base", config) |
| WriteLayoutFunction(sfile, "write_partition_table", image_type, config) |
| |
| sfile.close() |
| |
| |
| def GetBlockSize(layout_filename): |
| """Returns the partition table block size. |
| |
| Args: |
| layout_filename: Path to partition configuration file |
| Returns: |
| Block size of all partitions in the layout |
| """ |
| |
| config = LoadPartitionConfig(layout_filename) |
| return config["metadata"]["block_size"] |
| |
| |
| def GetFilesystemBlockSize(layout_filename): |
| """Returns the filesystem block size. |
| |
| This is used for all partitions in the table that have filesystems. |
| |
| Args: |
| layout_filename: Path to partition configuration file |
| Returns: |
| Block size of all filesystems in the layout |
| """ |
| |
| config = LoadPartitionConfig(layout_filename) |
| return config["metadata"]["fs_block_size"] |
| |
| |
| def GetPartitionSize(image_type, layout_filename, num): |
| """Returns the partition size of a given partition for a given layout type. |
| |
| Args: |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| num: Number of the partition you want to read from |
| Returns: |
| Size of selected partition in bytes |
| """ |
| |
| config = LoadPartitionConfig(layout_filename) |
| partitions = GetPartitionTable(config, image_type) |
| partition = GetPartitionByNumber(partitions, num) |
| |
| return partition["bytes"] |
| |
| |
| def GetFilesystemSize(image_type, layout_filename, num): |
| """Returns the filesystem size of a given partition for a given layout type. |
| |
| If no filesystem size is specified, returns the partition size. |
| |
| Args: |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| num: Number of the partition you want to read from |
| Returns: |
| Size of selected partition filesystem in bytes |
| """ |
| |
| config = LoadPartitionConfig(layout_filename) |
| partitions = GetPartitionTable(config, image_type) |
| partition = GetPartitionByNumber(partitions, num) |
| |
| if "fs_bytes" in partition: |
| return partition["fs_bytes"] |
| else: |
| return partition["bytes"] |
| |
| |
| def GetLabel(image_type, layout_filename, num): |
| """Returns the label for a given partition. |
| |
| Args: |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| num: Number of the partition you want to read from |
| Returns: |
| Label of selected partition, or "UNTITLED" if none specified |
| """ |
| |
| config = LoadPartitionConfig(layout_filename) |
| partitions = GetPartitionTable(config, image_type) |
| partition = GetPartitionByNumber(partitions, num) |
| |
| if "label" in partition: |
| return partition["label"] |
| else: |
| return "UNTITLED" |
| |
| |
| def DoDebugOutput(image_type, layout_filename): |
| """Prints out a human readable disk layout in on-disk order. |
| |
| This will round values larger than 1MB, it's exists to quickly |
| visually verify a layout looks correct. |
| |
| Args: |
| image_type: Type of image eg base/test/dev/factory_install |
| layout_filename: Path to partition configuration file |
| """ |
| config = LoadPartitionConfig(layout_filename) |
| partitions = GetPartitionTable(config, image_type) |
| |
| for partition in partitions: |
| if partition["bytes"] < 1024 * 1024: |
| size = "%d bytes" % partition["bytes"] |
| else: |
| size = "%d MB" % (partition["bytes"] / 1024 / 1024) |
| if "label" in partition: |
| if "fs_bytes" in partition: |
| if partition["fs_bytes"] < 1024 * 1024: |
| fs_size = "%d bytes" % partition["fs_bytes"] |
| else: |
| fs_size = "%d MB" % (partition["fs_bytes"] / 1024 / 1024) |
| print "%s - %s/%s" % (partition["label"], fs_size, size) |
| else: |
| print "%s - %s" % (partition["label"], size) |
| else: |
| print "blank - %s" % size |
| |
| |
| def main(argv): |
| action_map = { |
| "write": { |
| "argc": 4, |
| "usage": "<image_type> <partition_config_file> <script_file>", |
| "func": WritePartitionScript |
| }, |
| "readblocksize": { |
| "argc": 2, |
| "usage": "<partition_config_file>", |
| "func": GetBlockSize |
| }, |
| "readfsblocksize": { |
| "argc": 2, |
| "usage": "<partition_config_file>", |
| "func": GetFilesystemBlockSize |
| }, |
| "readpartsize": { |
| "argc": 4, |
| "usage": "<image_type> <partition_config_file> <partition_num>", |
| "func": GetPartitionSize |
| }, |
| "readfssize": { |
| "argc": 4, |
| "usage": "<image_type> <partition_config_file> <partition_num>", |
| "func": GetFilesystemSize |
| }, |
| "readlabel": { |
| "argc": 4, |
| "usage": "<image_type> <partition_config_file> <partition_num>", |
| "func": GetLabel |
| }, |
| "debug": { |
| "argc": 3, |
| "usage": "<image_type> <partition_config_file>", |
| "func": DoDebugOutput |
| } |
| } |
| |
| if len(sys.argv) < 2 or sys.argv[1] not in action_map: |
| print "Usage: %s <action>\n" % sys.argv[0] |
| print "Valid actions are:" |
| for action in action_map: |
| print " %s %s" % (action, action_map[action]["usage"]) |
| sys.exit(1) |
| else: |
| action_name = sys.argv[1] |
| action = action_map[action_name] |
| if action["argc"] == len(sys.argv) - 1: |
| print action["func"](*sys.argv[2:]) |
| else: |
| sys.exit("Usage: %s %s %s" % (sys.argv[0], sys.argv[1], action["usage"])) |
| |
| if __name__ == "__main__": |
| main(sys.argv) |