"""Tool to create/modify/inspect fastboot packed images.

This tool provides commands to:
-   `info`: Display information about a packed image file.
-   `packversion`: Print the pack version of a packed image file.
-   `create`: Create a packed image file from a configuration (XML/YAML).
-   `unpack`: Extract files from a packed image.
"""

import argparse
import logging
import os
#
# Copyright 2021 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys

from google3.third_party.devsite.androidsource.en.docs.core.architecture.bootloader.tools.pixel.fw_unpack import fbpack
from google3.third_party.devsite.androidsource.en.docs.core.architecture.bootloader.tools.pixel.fw_unpack import xmlparser
from google3.third_party.devsite.androidsource.en.docs.core.architecture.bootloader.tools.pixel.fw_unpack import xmlserializer
from google3.third_party.devsite.androidsource.en.docs.core.architecture.bootloader.tools.pixel.fw_unpack import yamlparser


# !/usr/bin/env python3


def bytes_to_str(bstr):
  return bstr.decode().rstrip('\x00')


def print_pack_header(pack):
  """Prints the fields of a PackHeader object.

  Args:
      pack: The PackHeader object to print.
  """
  print('magic:              {:#x}'.format(pack.magic))
  print('version:            {}'.format(pack.version))
  print('header size:        {}'.format(pack.header_size))
  print('entry header size:  {}'.format(pack.entry_header_size))
  platform = bytes_to_str(pack.platform)
  print('platform:           {}'.format(platform))
  pack_version = bytes_to_str(pack.pack_version)
  print('pack version:       {}'.format(pack_version))
  print('slock type:         {}'.format(pack.slot_type))
  print('data align:         {}'.format(pack.data_align))
  print('total entries:      {}'.format(pack.total_entries))
  print('total size:         {}'.format(pack.total_size))


def print_pack_entry(entry, prefix):
  """Prints the fields of a PackEntry object.

  Args:
      entry: The PackEntry object to print.
      prefix: A string to prepend to each line of output.
  """
  name = bytes_to_str(entry.name)
  print('{}name:       {}'.format(prefix, name))
  etype = 'unknown'
  if entry.type == fbpack.FBPACK_PARTITION_TABLE:
    etype = 'partiton table'
  elif entry.type == fbpack.FBPACK_PARTITION_DATA:
    etype = 'partition'
  elif entry.type == fbpack.FBPACK_SIDELOAD_DATA:
    etype = 'sideload'
  else:
    print('entry else')
  print('{}type:       {}'.format(prefix, etype))
  product = bytes_to_str(entry.product)
  print('{}product:    {}'.format(prefix, product))
  print('{}offset:     {:#x} ({})'.format(prefix, entry.offset, entry.offset))
  print('{}size:       {:#x} ({})'.format(prefix, entry.size, entry.size))
  print('{}slotted:    {}'.format(entry.size, bool(entry.slotted)))
  print('{}crc32:      {:#08x}'.format(prefix, entry.crc32))


def cmd_info(args):
  """Displays information about a packed image file.

  Args:
      args: The parsed command-line arguments, including the 'file' to inspect.
  """
  with open(args.file, 'rb') as f:
    pack = fbpack.PackHeader.from_bytes(f.read(len(fbpack.PackHeader())))

    if pack.version != fbpack.FBPACK_VERSION:
      raise NotImplementedError('unsupported version {}'.format(pack.version))

    print('Header:')
    print_pack_header(pack)

    print('\nEntries:')
    for i in range(1, pack.total_entries + 1):
      entry = fbpack.PackEntry.from_bytes(f.read(len(fbpack.PackEntry())))
      print('Entry {}: {{'.format(i))
      print_pack_entry(entry, '    ')
      print('}')


def cmd_packversion(args):
  with open(args.file, 'rb') as f:
    pack = fbpack.PackHeader.from_bytes(f.read(len(fbpack.PackHeader())))

    if pack.version != fbpack.FBPACK_VERSION:
      raise NotImplementedError('unsupported version {}'.format(pack.version))

    pack_version = bytes_to_str(pack.pack_version)
    print(pack_version)


def align_up(val, align):
  return (val + align - 1) & ~(align - 1)


def create_pack_file(file_name, in_dir_name, pack):
  """Creates a packed fastboot image file.

  Args:
      file_name: The path to the output packed image file.
      in_dir_name: The directory containing the input files to pack.
      pack: The PackHeader object containing the entries to pack.
  """
  pack.total_entries = len(pack.entries)
  offset = pack.header_size + pack.total_entries * pack.entry_header_size
  with open(file_name, 'wb') as f:
    # write entries data
    for entry in pack.entries:
      # align data
      offset = align_up(offset, pack.data_align)
      entry.offset = offset
      f.seek(offset)
      fin_name = os.path.join(in_dir_name, entry.filepath)
      with open(fin_name, 'rb') as fin:
        data = fin.read()
        entry.size = len(data)
        f.write(data)
        offset += len(data)

    pack.total_size = offset
    f.seek(0)
    # write pack header
    f.write(bytes(pack))
    # iterate over entries again to write entry header
    for entry in pack.entries:
      f.write(bytes(entry))


def cmd_create(args):
  """Creates a packed fastboot image file from a configuration file.

  Args:
      args: The parsed command-line arguments, including 'file', 'pack_version',
        'out_dir', and 'in_dir'.

  Raises:
      NotImplementedError: If the configuration file type is not supported (only
        .xml and .yaml are supported).
  """
  if args.file.lower().endswith('.xml'):
    pack = xmlparser.parse(args.file)
  elif args.file.lower().endswith('.yaml'):
    pack = yamlparser.parse(args.file)
  else:
    raise NotImplementedError('{} type not supported'.format(args.file))

  pack.pack_version = bytes(str(args.pack_version).encode('ascii'))
  pack.header_size = len(pack)

  # create output directory if missing
  if not os.path.isdir(args.out_dir):
    os.makedirs(args.out_dir, 0o755)

  file_name = os.path.join(args.out_dir, pack.name + '.img')

  create_pack_file(file_name, args.in_dir, pack)


def product_match(products, product):
  return product in bytes_to_str(products).split('|')


def copyfileobj(src, dst, file_size):
  while file_size > 0:
    buf = src.read(min(128 * 1024, file_size))
    dst.write(buf)
    file_size -= len(buf)


def cmd_unpack(args):
  """Unpacks files from a packed fastboot image.

  Extracts specified partitions or generates configuration files from a packed
  image file.

  Args:
      args: The parsed command-line arguments, including 'file', 'out_dir',
        'product', 'unpack_ver', 'unpack_config_xml', and 'partitions'.

  Raises:
      NotImplementedError: If the pack file version is not supported.
      RuntimeError: If no images are selected for unpacking.
  """
  with open(args.file, 'rb') as f:
    pack = fbpack.PackHeader.from_bytes(f.read(len(fbpack.PackHeader())))

    if pack.version != fbpack.FBPACK_VERSION:
      raise NotImplementedError('unsupported version {}'.format(pack.version))

    pack.name = os.path.splitext(os.path.basename(args.file))[0]

    pack.entries = []
    # create list of entries we want to extact
    for _ in range(pack.total_entries):
      entry = fbpack.PackEntry.from_bytes(f.read(len(fbpack.PackEntry())))
      name = bytes_to_str(entry.name)
      if not args.partitions or name in args.partitions:
        # if both product are valid then match product name too
        if (
            not args.product
            or not entry.product
            or product_match(entry.product, args.product)
        ):
          pack.entries.append(entry)

    if not pack.entries and not args.unpack_ver and not args.unpack_config_xml:
      raise RuntimeError('no images to unpack')

    # create output directory if it does not exist
    if not os.path.isdir(args.out_dir):
      os.makedirs(args.out_dir, 0o755)

    out_files = {}
    # write file per entry
    for entry in pack.entries:
      name = bytes_to_str(entry.name)
      logging.info(
          'Unpacking %s (size: %s, offset: %s)',
          name,
          entry.size,
          entry.offset,
      )
      f.seek(entry.offset)
      entry_filename = os.path.join(args.out_dir, name + '.img')
      instance = out_files.get(entry_filename, 0) + 1
      out_files[entry_filename] = instance
      if instance > 1:
        entry_filename = os.path.join(
            args.out_dir, name + '({}).img'.format(instance - 1)
        )
      with open(entry_filename, 'wb') as entry_file:
        copyfileobj(f, entry_file, entry.size)
      entry.filename = entry_filename

    if args.unpack_ver:
      ver_file_path = os.path.join(args.out_dir, 'version.txt')
      with open(ver_file_path, 'w') as ver_file:
        ver_file.write(bytes_to_str(pack.pack_version))

    if args.unpack_config_xml:
      xml_file_path = os.path.join(args.out_dir, 'repack.xml')
      xmlserializer.serialize(pack, xml_file_path)

  logging.info('Done')


def parse_args():
  """Parses command-line arguments for the fbpacktool.

  Returns:
      The parsed arguments object.
  """
  parser = argparse.ArgumentParser(
      description='Tool to create/modify/inspect fastboot packed images'
  )
  parser.add_argument(
      '-v',
      '--verbosity',
      action='count',
      default=0,
      help='increase output verbosity',
  )

  subparsers = parser.add_subparsers()

  # info command
  info = subparsers.add_parser('info')
  info.add_argument('file', help='packed image file')
  info.set_defaults(func=cmd_info)

  # packversion command
  packversion = subparsers.add_parser('packversion')
  packversion.add_argument('file', help='packed image file')
  packversion.set_defaults(func=cmd_packversion)

  # create command
  create = subparsers.add_parser('create')
  create.add_argument(
      '-d', '--in_dir', help='directory to search for data files', default='.'
  )
  create.add_argument(
      '-o',
      '--out_dir',
      help='output directory for the packed image',
      default='.',
  )
  create.add_argument(
      '-v', '--pack_version', help='Packed image version ', default=''
  )
  create.add_argument(
      'file', help='config file describing packed image (yaml/xml)'
  )
  create.set_defaults(func=cmd_create)

  # unpack command
  unpack = subparsers.add_parser('unpack')
  unpack.add_argument(
      '-o', '--out_dir', help='directory to store unpacked images', default='.'
  )
  unpack.add_argument(
      '-p', '--product', help='filter images by product', default=''
  )
  unpack.add_argument(
      '-v', '--unpack_ver', help='Unpack version to a file', action='store_true'
  )
  unpack.add_argument(
      '-c',
      '--unpack_config_xml',
      help='Generate XML config file from packed file',
      action='store_true',
  )
  unpack.add_argument('file', help='packed image file')
  unpack.add_argument(
      'partitions',
      metavar='PART',
      type=str,
      nargs='*',
      help='Partition names to extract (default all).',
  )
  unpack.set_defaults(func=cmd_unpack)

  args = parser.parse_args()
  # make sure a command was passed
  if not hasattr(args, 'func'):
    parser.print_usage()
    print('fbpacktool.py: error: no command was passed')
    sys.exit(2)

  return args


def main():
  args = parse_args()

  if args.verbosity >= 2:
    log_level = logging.DEBUG
  elif args.verbosity == 1:
    log_level = logging.INFO
  else:
    log_level = logging.WARNING

  logging.basicConfig(level=log_level)

  # execute command
  args.func(args)


if __name__ == '__main__':
  main()