mirror of
https://github.com/flatcar/scripts.git
synced 2025-09-23 14:41:31 +02:00
Merge GetPartitionTable and partition alignment from WritePartitionTable into LoadPartitionConfig so that all this config manipulation code is in one place and inheritance from the 'base' layout is more predictable.
444 lines
13 KiB
Python
Executable File
444 lines
13 KiB
Python
Executable File
#!/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 argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import uuid
|
|
|
|
# First sector we can use.
|
|
GPT_RESERVED_SECTORS = 34
|
|
|
|
class ConfigNotFound(Exception):
|
|
pass
|
|
class PartitionNotFound(Exception):
|
|
pass
|
|
class InvalidLayout(Exception):
|
|
pass
|
|
class InvalidAdjustment(Exception):
|
|
pass
|
|
|
|
|
|
def LoadPartitionConfig(options):
|
|
"""Loads a partition tables configuration file into a Python object.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Returns:
|
|
Object containing disk layout configuration
|
|
"""
|
|
|
|
valid_keys = set(('_comment', 'metadata', 'layouts'))
|
|
valid_layout_keys = set((
|
|
'_comment', 'type', 'num', 'label', 'blocks', 'block_size', 'fs_blocks',
|
|
'fs_block_size', 'features', 'uuid', 'alignment'))
|
|
integer_layout_keys = set((
|
|
'blocks', 'block_size', 'fs_blocks', 'fs_block_size', 'alignment'))
|
|
required_layout_keys = set(('type', 'num', 'label', 'blocks'))
|
|
|
|
filename = options.disk_layout_file
|
|
if not os.path.exists(filename):
|
|
raise ConfigNotFound('Partition config %s was not found!' % filename)
|
|
with open(filename) as f:
|
|
config = json.load(f)
|
|
|
|
unknown_keys = set(config.keys()) - valid_keys
|
|
if unknown_keys:
|
|
raise InvalidLayout('Unknown items: %s' % ' '.join(unknown_keys))
|
|
|
|
try:
|
|
metadata = config['metadata']
|
|
base = config['layouts']['base']
|
|
for key in ('alignment', 'block_size', 'fs_block_size'):
|
|
metadata[key] = int(metadata[key])
|
|
except KeyError as e:
|
|
raise InvalidLayout('Metadata is missing required entries: %s' % e)
|
|
|
|
# Sometimes qemu-img expects disks sizes aligned to 64k
|
|
align_bytes = metadata['alignment'] * metadata['block_size']
|
|
if align_bytes < 65536 or align_bytes % 65536 != 0:
|
|
raise InvalidLayout('Invalid alignment, 64KB or better required')
|
|
|
|
|
|
def VerifyLayout(layout_name, layout, base=None):
|
|
for part_num, part in layout.iteritems():
|
|
part['num'] = int(part_num)
|
|
part_keys = set(part.iterkeys())
|
|
unknown_keys = part_keys - valid_layout_keys
|
|
if unknown_keys:
|
|
raise InvalidLayout('Unknown items in partition %s %s: %r' %
|
|
(layout_name, part_num, ' '.join(unknown_keys)))
|
|
|
|
for int_key in integer_layout_keys.intersection(part_keys):
|
|
part[int_key] = int(part[int_key])
|
|
|
|
if part.get('type', None) == 'blank':
|
|
continue
|
|
|
|
if base:
|
|
part_keys.update(base.iterkeys())
|
|
|
|
missing_keys = required_layout_keys - part_keys
|
|
if missing_keys:
|
|
raise InvalidLayout('Missing items in partition %s %s: %s' %
|
|
(layout_name, part_num, ' '.join(missing_keys)))
|
|
|
|
if 'uuid' in part:
|
|
try:
|
|
# double check the string formatting
|
|
part['uuid'] = str(uuid.UUID(part['uuid']))
|
|
except ValueError as e:
|
|
raise InvalidLayout('Invalid uuid %r: %s' % (part['uuid'], e))
|
|
|
|
|
|
def Align(count, alignment):
|
|
offset = count % alignment
|
|
if offset:
|
|
count += alignment - offset
|
|
return count
|
|
|
|
|
|
def FillExtraValues(layout_name, layout, base=None):
|
|
# Reserved size for first GPT
|
|
disk_block_count = GPT_RESERVED_SECTORS
|
|
|
|
# Fill in default values from base,
|
|
# dict doesn't have a update()+setdefault() method so this looks tedious
|
|
if base:
|
|
for part_num, base_part in base.iteritems():
|
|
part = layout.setdefault(part_num, {})
|
|
for base_key, base_value in base_part.iteritems():
|
|
part.setdefault(base_key, base_value)
|
|
|
|
for part_num, part in layout.iteritems():
|
|
if part['type'] == 'blank':
|
|
continue
|
|
|
|
part.setdefault('alignment', metadata['alignment'])
|
|
part['bytes'] = part['blocks'] * metadata['block_size']
|
|
part.setdefault('fs_block_size', metadata['fs_block_size'])
|
|
part.setdefault('fs_blocks', part['bytes'] // part['fs_block_size'])
|
|
part['fs_bytes'] = part['fs_blocks'] * part['fs_block_size']
|
|
|
|
if part['fs_bytes'] > part['bytes']:
|
|
raise InvalidLayout(
|
|
'Filesystem may not be larger than partition: %s %s: %d > %d' %
|
|
(layout_name, part_num, part['fs_bytes'], part['bytes']))
|
|
|
|
disk_block_count = Align(disk_block_count, part['alignment'])
|
|
part['first_block'] = disk_block_count
|
|
disk_block_count += part['blocks']
|
|
|
|
part.setdefault('uuid', str(uuid.uuid4()))
|
|
|
|
# Reserved size for second GPT plus align disk image size
|
|
disk_block_count += GPT_RESERVED_SECTORS
|
|
disk_block_count = Align(disk_block_count, metadata['alignment'])
|
|
|
|
# If this is the requested layout stash the disk size into the global
|
|
# metadata. Kinda odd but the best place I've got with this data structure.
|
|
if layout_name == options.disk_layout:
|
|
metadata['blocks'] = disk_block_count
|
|
|
|
|
|
# Verify 'base' before other layouts because it is inherited by the others
|
|
# Fill in extra/default values in base last so they aren't inherited
|
|
VerifyLayout('base', base)
|
|
for layout_name, layout in config['layouts'].iteritems():
|
|
if layout_name == 'base':
|
|
continue
|
|
VerifyLayout(layout_name, layout, base)
|
|
FillExtraValues(layout_name, layout, base)
|
|
FillExtraValues('base', base)
|
|
|
|
return config, config['layouts'][options.disk_layout]
|
|
|
|
|
|
def GetPartitionTableFromConfig(options):
|
|
"""Loads a partition table and returns a given partition table type
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Returns:
|
|
A list defining all known partitions.
|
|
"""
|
|
|
|
config, partitions = LoadPartitionConfig(options)
|
|
return partitions
|
|
|
|
|
|
def WritePartitionTable(options):
|
|
"""Writes the given partition table to a disk image or device.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
"""
|
|
|
|
def Cgpt(*args):
|
|
subprocess.check_call(['cgpt'] + [str(a) for a in args])
|
|
|
|
config, partitions = LoadPartitionConfig(options)
|
|
|
|
Cgpt('create', '-c', '-s', config['metadata']['blocks'], options.disk_image)
|
|
|
|
esp_number = None
|
|
for partition in partitions.itervalues():
|
|
if partition['type'] != 'blank':
|
|
Cgpt('add', '-i', partition['num'],
|
|
'-b', partition['first_block'],
|
|
'-s', partition['blocks'],
|
|
'-t', partition['type'],
|
|
'-l', partition['label'],
|
|
'-u', partition['uuid'],
|
|
options.disk_image)
|
|
|
|
if partition['type'] == 'efi':
|
|
esp_number = partition['num']
|
|
|
|
if esp_number is None:
|
|
raise InvalidLayout('Table does not include an EFI partition.')
|
|
|
|
if options.mbr_boot_code:
|
|
Cgpt('boot', '-p',
|
|
'-b', options.mbr_boot_code,
|
|
'-i', esp_number,
|
|
options.disk_image)
|
|
|
|
Cgpt('show', options.disk_image)
|
|
|
|
|
|
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
|
|
"""
|
|
|
|
partition = partitions.get(str(num), None)
|
|
if not partition or partition['type'] == 'blank':
|
|
raise PartitionNotFound('Partition not found')
|
|
|
|
return partition
|
|
|
|
|
|
def GetPartitionByLabel(partitions, label):
|
|
"""Given a partition table and label returns the partition object.
|
|
|
|
Args:
|
|
partitions: List of partitions to search in
|
|
label: Label of partition to find
|
|
Returns:
|
|
An object for the selected partition
|
|
"""
|
|
for partition in partitions.itervalues():
|
|
if partition['type'] == 'blank':
|
|
continue
|
|
elif partition['label'] == label:
|
|
return partition
|
|
|
|
raise PartitionNotFound('Partition not found')
|
|
|
|
|
|
def GetBlockSize(options):
|
|
"""Returns the partition table block size.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
Block size of all partitions in the layout
|
|
"""
|
|
|
|
config, partitions = LoadPartitionConfig(options)
|
|
print config['metadata']['block_size']
|
|
|
|
|
|
def GetFilesystemBlockSize(options):
|
|
"""Returns the filesystem block size.
|
|
|
|
This is used for all partitions in the table that have filesystems.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
Block size of all filesystems in the layout
|
|
"""
|
|
|
|
config, partitions = LoadPartitionConfig(options)
|
|
print config['metadata']['fs_block_size']
|
|
|
|
|
|
def GetPartitionSize(options):
|
|
"""Returns the partition size of a given partition for a given layout type.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
Size of selected partition in bytes
|
|
"""
|
|
|
|
partitions = GetPartitionTableFromConfig(options)
|
|
partition = GetPartitionByNumber(partitions, options.partition_num)
|
|
print partition['bytes']
|
|
|
|
|
|
def GetFilesystemSize(options):
|
|
"""Returns the filesystem size of a given partition for a given layout type.
|
|
|
|
If no filesystem size is specified, returns the partition size.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
Size of selected partition filesystem in bytes
|
|
"""
|
|
|
|
partitions = GetPartitionTableFromConfig(options)
|
|
partition = GetPartitionByNumber(partitions, options.partition_num)
|
|
print partition.get('fs_bytes', partition['bytes'])
|
|
|
|
|
|
def GetLabel(options):
|
|
"""Returns the label for a given partition.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
Label of selected partition, or 'UNTITLED' if none specified
|
|
"""
|
|
|
|
partitions = GetPartitionTableFromConfig(options)
|
|
partition = GetPartitionByNumber(partitions, options.partition_num)
|
|
print partition.get('label', 'UNTITLED')
|
|
|
|
|
|
def GetNum(options):
|
|
"""Returns the number for a given label.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
Number of selected partition, or '-1' if there is no number
|
|
"""
|
|
|
|
partitions = GetPartitionTableFromConfig(options)
|
|
partition = GetPartitionByLabel(partitions, options.label)
|
|
print partition.get('num', '-1')
|
|
|
|
|
|
def GetUuid(options):
|
|
"""Returns the unique partition UUID for a given label.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
Prints:
|
|
String containing the requested UUID
|
|
"""
|
|
|
|
partitions = GetPartitionTableFromConfig(options)
|
|
partition = GetPartitionByLabel(partitions, options.label)
|
|
print partition.get('uuid', '')
|
|
|
|
|
|
def DoDebugOutput(options):
|
|
"""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:
|
|
options: Flags passed to the script
|
|
"""
|
|
partitions = GetPartitionTableFromConfig(options)
|
|
|
|
for num, partition in sorted(partitions.iteritems()):
|
|
if partition['type'] != 'blank':
|
|
if partition['bytes'] < 1024 * 1024:
|
|
size = '%d bytes' % partition['bytes']
|
|
else:
|
|
size = '%d MB' % (partition['bytes'] / 1024 / 1024)
|
|
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/%s' % (num, partition['label'], fs_size, size)
|
|
else:
|
|
print '%s: %s - %s' % (num, partition['label'], size)
|
|
else:
|
|
print '%s: blank' % num
|
|
|
|
|
|
def DoParseOnly(options):
|
|
"""Parses a layout file only, used before reading sizes to check for errors.
|
|
|
|
Args:
|
|
options: Flags passed to the script
|
|
"""
|
|
GetPartitionTableFromConfig(options)
|
|
|
|
|
|
def main(argv):
|
|
default_layout_file = os.environ.get('DISK_LAYOUT_FILE',
|
|
os.path.join(os.path.dirname(__file__), 'legacy_disk_layout.json'))
|
|
default_layout_type = os.environ.get('DISK_LAYOUT_TYPE', 'base')
|
|
|
|
parser = argparse.ArgumentParser(
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument('--disk_layout_file', default=default_layout_file,
|
|
help='path to disk layout json file')
|
|
parser.add_argument('--disk_layout', default=default_layout_type,
|
|
help='disk layout type from the json file')
|
|
actions = parser.add_subparsers(title='actions')
|
|
|
|
a = actions.add_parser('write_gpt', help='write gpt to new image')
|
|
a.add_argument('--mbr_boot_code',
|
|
help='path to mbr boot block, such as syslinux/gptmbr.bin')
|
|
a.add_argument('disk_image', help='path to disk image file')
|
|
a.set_defaults(func=WritePartitionTable)
|
|
|
|
a = actions.add_parser('readblocksize', help='get device block size')
|
|
a.set_defaults(func=GetBlockSize)
|
|
|
|
a = actions.add_parser('readfsblocksize', help='get filesystem block size')
|
|
a.set_defaults(func=GetFilesystemBlockSize)
|
|
|
|
a = actions.add_parser('readpartsize', help='get partition size')
|
|
a.add_argument('partition_num', type=int, help='partition number')
|
|
a.set_defaults(func=GetPartitionSize)
|
|
|
|
a = actions.add_parser('readfssize', help='get filesystem size')
|
|
a.add_argument('partition_num', type=int, help='partition number')
|
|
a.set_defaults(func=GetFilesystemSize)
|
|
|
|
a = actions.add_parser('readlabel', help='get partition label')
|
|
a.add_argument('partition_num', type=int, help='partition number')
|
|
a.set_defaults(func=GetLabel)
|
|
|
|
a = actions.add_parser('readnum', help='get partition number')
|
|
a.add_argument('label', help='partition label')
|
|
a.set_defaults(func=GetNum)
|
|
|
|
a = actions.add_parser('readuuid', help='get partition uuid')
|
|
a.add_argument('label', help='partition label')
|
|
a.set_defaults(func=GetUuid)
|
|
|
|
a = actions.add_parser('debug', help='dump debug output')
|
|
a.set_defaults(func=DoDebugOutput)
|
|
|
|
a = actions.add_parser('parseonly', help='validate config')
|
|
a.set_defaults(func=DoParseOnly)
|
|
|
|
options = parser.parse_args(argv[1:])
|
|
options.func(options)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main(sys.argv)
|