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