mirror of
https://github.com/flatcar/scripts.git
synced 2025-09-24 07:01:13 +02:00
508 lines
15 KiB
Python
Executable File
508 lines
15 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 re
|
|
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'))
|
|
|
|
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)
|
|
|
|
try:
|
|
metadata = config['metadata']
|
|
for key in ('alignment', 'block_size', 'fs_block_size'):
|
|
metadata[key] = int(metadata[key])
|
|
|
|
# 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')
|
|
|
|
unknown_keys = set(config.keys()) - valid_keys
|
|
if unknown_keys:
|
|
raise InvalidLayout('Unknown items: %r' % unknown_keys)
|
|
|
|
if len(config['layouts']) <= 0:
|
|
raise InvalidLayout('Missing "layouts" entries')
|
|
|
|
for layout_name, layout in config['layouts'].items():
|
|
for part in layout:
|
|
unknown_keys = set(part.keys()) - valid_layout_keys
|
|
if unknown_keys:
|
|
raise InvalidLayout('Unknown items in layout %s: %r' %
|
|
(layout_name, unknown_keys))
|
|
|
|
if part['type'] != 'blank':
|
|
for s in ('num', 'label'):
|
|
if not s in part:
|
|
raise InvalidLayout('Layout "%s" missing "%s"' % (layout_name, s))
|
|
|
|
part['alignment'] = int(part.get('alignment', metadata['alignment']))
|
|
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: %s %s: %d > %d' %
|
|
(layout_name, part['label'], part['fs_bytes'], part['bytes']))
|
|
|
|
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))
|
|
else:
|
|
part['uuid'] = str(uuid.uuid4())
|
|
except KeyError as e:
|
|
raise InvalidLayout('Layout is missing required entries: %s' % e)
|
|
|
|
return config
|
|
|
|
|
|
def GetPartitionTable(options, config):
|
|
"""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:
|
|
options: Flags passed to the script
|
|
config: Partition configuration file object
|
|
Returns:
|
|
Object representing a selected partition table
|
|
"""
|
|
|
|
partitions = config['layouts']['base']
|
|
metadata = config['metadata']
|
|
image_type = options.disk_layout
|
|
|
|
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
|
|
|
|
for adjustment_str in options.adjust_part.split():
|
|
adjustment = adjustment_str.split(':')
|
|
if len(adjustment) < 2:
|
|
raise InvalidAdjustment('Adjustment specified was incomplete')
|
|
|
|
label = adjustment[0]
|
|
operator = adjustment[1][0]
|
|
operand = adjustment[1][1:]
|
|
ApplyPartitionAdjustment(partitions, metadata, label, operator, operand)
|
|
|
|
return partitions
|
|
|
|
|
|
def ApplyPartitionAdjustment(partitions, metadata, label, operator, operand):
|
|
"""Applies an adjustment to a partition specified by label
|
|
|
|
Args:
|
|
partitions: Partition table to modify
|
|
metadata: Partition table metadata
|
|
label: The label of the partition to adjust
|
|
operator: Type of adjustment (+/-/=)
|
|
operand: How much to adjust by
|
|
"""
|
|
|
|
partition = GetPartitionByLabel(partitions, label)
|
|
|
|
operand_digits = re.sub('\D', '', operand)
|
|
size_factor = block_factor = 1
|
|
suffix = operand[len(operand_digits):]
|
|
if suffix:
|
|
size_factors = { 'B': 0, 'K': 1, 'M': 2, 'G': 3, 'T': 4, }
|
|
try:
|
|
size_factor = size_factors[suffix[0].upper()]
|
|
except KeyError:
|
|
raise InvalidAdjustment('Unknown size type %s' % suffix)
|
|
if size_factor == 0 and len(suffix) > 1:
|
|
raise InvalidAdjustment('Unknown size type %s' % suffix)
|
|
block_factors = { '': 1024, 'B': 1000, 'IB': 1024, }
|
|
try:
|
|
block_factor = block_factors[suffix[1:].upper()]
|
|
except KeyError:
|
|
raise InvalidAdjustment('Unknown size type %s' % suffix)
|
|
|
|
operand_bytes = int(operand_digits) * pow(block_factor, size_factor)
|
|
|
|
if operand_bytes % metadata['block_size'] == 0:
|
|
operand_blocks = operand_bytes / metadata['block_size']
|
|
else:
|
|
raise InvalidAdjustment('Adjustment size not divisible by block size')
|
|
|
|
if operator == '+':
|
|
partition['blocks'] += operand_blocks
|
|
partition['bytes'] += operand_bytes
|
|
elif operator == '-':
|
|
partition['blocks'] -= operand_blocks
|
|
partition['bytes'] -= operand_bytes
|
|
elif operator == '=':
|
|
partition['blocks'] = operand_blocks
|
|
partition['bytes'] = operand_bytes
|
|
else:
|
|
raise ValueError('unknown operator %s' % operator)
|
|
|
|
if partition['type'] == 'rootfs':
|
|
# If we're adjusting a rootFS partition, we assume the full partition size
|
|
# specified is being used for the filesytem, minus the space reserved for
|
|
# the hashpad.
|
|
partition['fs_bytes'] = partition['bytes']
|
|
partition['fs_blocks'] = partition['fs_bytes'] / metadata['fs_block_size']
|
|
partition['blocks'] = int(partition['blocks'] * 1.15)
|
|
partition['bytes'] = partition['blocks'] * metadata['block_size']
|
|
|
|
|
|
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 = LoadPartitionConfig(options)
|
|
partitions = GetPartitionTable(options, config)
|
|
|
|
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])
|
|
|
|
def Align(count, alignment):
|
|
offset = count % alignment
|
|
if offset:
|
|
count += alignment - offset
|
|
return count
|
|
|
|
config = LoadPartitionConfig(options)
|
|
partitions = GetPartitionTable(options, config)
|
|
disk_block_count = GPT_RESERVED_SECTORS
|
|
|
|
for partition in partitions:
|
|
disk_block_count = Align(disk_block_count, partition['alignment'])
|
|
disk_block_count += partition['blocks']
|
|
|
|
disk_block_count += GPT_RESERVED_SECTORS
|
|
# Sometimes qemu-img expects disks sizes aligned to 64k
|
|
disk_block_count = Align(disk_block_count, config['metadata']['alignment'])
|
|
|
|
Cgpt('create', '-c', '-s', disk_block_count, options.disk_image)
|
|
|
|
sector = GPT_RESERVED_SECTORS
|
|
esp_number = None
|
|
for partition in partitions:
|
|
sector = Align(sector, partition['alignment'])
|
|
if partition['type'] != 'blank':
|
|
Cgpt('add', '-i', partition['num'],
|
|
'-b', sector,
|
|
'-s', partition['blocks'],
|
|
'-t', partition['type'],
|
|
'-l', partition['label'],
|
|
'-u', partition['uuid'],
|
|
options.disk_image)
|
|
|
|
if partition['type'] == 'efi':
|
|
esp_number = partition['num']
|
|
|
|
sector += partition['blocks']
|
|
|
|
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
|
|
"""
|
|
for partition in partitions:
|
|
if partition['type'] == 'blank':
|
|
continue
|
|
if partition['num'] == int(num):
|
|
return partition
|
|
|
|
raise PartitionNotFound('Partition not found')
|
|
|
|
|
|
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:
|
|
if 'label' not in partition:
|
|
continue
|
|
if 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 = 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 = 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 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 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('--adjust_part', default='',
|
|
help='adjust partition sizes')
|
|
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)
|