#!/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)