// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Mobius Forensic Toolkit
// Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019,2020,2021,2022,2023 Eduardo Aguiar
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2, or (at your option) any later
// version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
// Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#include "partition_system_impl.h"
#include <mobius/charset.h>
#include <mobius/exception.inc>
#include <mobius/core/log.h>
#include <mobius/decoder/data_decoder.h>
#include <mobius/io/bytearray_io.h>
#include <mobius/io/sector_reader_adaptor.h>
#include <stdexcept>

namespace
{
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Datatypes
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
using sector_type = mobius::vfs::partition_system_impl_base::sector_type;

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Constants
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static const std::string GPT_UNUSED_ENTRY = "00000000-0000-0000-0000-000000000000";
static const std::string GPT_MS_BASIC_DATA = "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7";
static const std::string GPT_SIGNATURE = "EFI PART";

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief GPT Header data
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
struct gpt_header
{
  bool is_valid = false;
  sector_type sector;
  sector_type sectors;
  std::string signature;
  std::string revision;
  std::uint32_t size;
  std::uint32_t crc32;
  std::uint64_t lba;
  std::uint64_t alternate_lba;
  std::uint64_t first_usable_lba;
  std::uint64_t last_usable_lba;
  std::string disk_guid;
  std::uint64_t partition_entry_lba;
  std::uint32_t number_of_partition_entries;
  std::uint32_t size_of_partition_entry;
  std::uint32_t partition_entry_array_crc32;
};

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief GPT Partition data
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
struct gpt_partition
{
  std::string type_guid;
  std::string unique_guid;
  std::uint64_t starting_lba;
  std::uint64_t ending_lba;
  std::uint64_t attributes;
  std::string name;
};

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Decode GPT header
//! \param reader Reader object
//! \param sector_size Sector size
//! \param sector Header sector
//! \return Entry object
//! \see UEFI 2.9 - section 5.3.2
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static gpt_header
_decode_gpt_header (
  mobius::io::reader reader,
  std::uint32_t sector_size,
  sector_type sector
)
{
  gpt_header header;
  header.sector = sector;
  header.sectors = 1;

  try
    {
      reader.seek (sector_size * sector);
      mobius::decoder::data_decoder decoder (reader);
      header.signature = decoder.get_string_by_size (8);

      auto revision_l = decoder.get_uint16_le ();
      auto revision_h = decoder.get_uint16_le ();
      header.revision = std::to_string (revision_h) + '.' + std::to_string (revision_l);

      header.size = decoder.get_uint32_le ();
      header.crc32 = decoder.get_uint32_le ();
      decoder.skip (4);     // reserved
      header.lba = decoder.get_uint64_le ();
      header.alternate_lba = decoder.get_uint64_le ();
      header.first_usable_lba = decoder.get_uint64_le ();
      header.last_usable_lba = decoder.get_uint64_le ();
      header.disk_guid = decoder.get_guid ();
      header.partition_entry_lba = decoder.get_uint64_le ();
      header.number_of_partition_entries = decoder.get_uint32_le ();
      header.size_of_partition_entry = decoder.get_uint32_le ();
      header.partition_entry_array_crc32 = decoder.get_uint32_le ();
      header.is_valid = (header.signature == GPT_SIGNATURE);
    }
  catch (const std::exception& e)
    {
    }

  return header;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Decode GPT partition table
//! \param reader Reader object
//! \param sector_size Sector size
//! \param sector Partition table sector
//! \param number_of_partition_entries Number of partition entries
//! \return Array of valid partitions
//! \see UEFI 2.9 - section 5.3.3
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static std::vector <gpt_partition>
_decode_gpt_partition_table (
  mobius::io::reader reader,
  std::uint32_t sector_size,
  sector_type sector,
  std::uint32_t number_of_partition_entries
)
{
  std::vector <gpt_partition> partitions;

  try
    {
      reader.seek (sector_size * sector);
      mobius::decoder::data_decoder decoder (reader);

      for (std::uint32_t i = 0; i < number_of_partition_entries; i++)
        {
          auto type_guid = decoder.get_guid ();

          if (type_guid == GPT_UNUSED_ENTRY)
            decoder.skip (112);

          else
            {
              gpt_partition p;
              p.type_guid = type_guid;
              p.unique_guid = decoder.get_guid ();
              p.starting_lba = decoder.get_uint64_le ();
              p.ending_lba = decoder.get_uint64_le ();
              p.attributes = decoder.get_uint64_le ();
              p.name = mobius::conv_charset_to_utf8 (decoder.get_bytearray_by_size (72), "utf-16");
              partitions.push_back (p);
            }
        }
    }
  catch (const std::exception& e)
    {
      partitions.clear ();
    }

  return partitions;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new protective MBR entry
//! \return Entry
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static mobius::vfs::partition_system_entry
_new_protective_mbr_entry ()
{
  mobius::vfs::partition_system_entry e;
  e.set_type ("mbr");
  e.set_starting_sector (0);
  e.set_ending_sector (0);
  e.set_sectors (1);
  e.set_description ("Protective MBR");

  return e;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create gpt.header entry from GPT Header
//! \param header GPT header object
//! \return Entry
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static mobius::vfs::partition_system_entry
_new_header_entry (const gpt_header& header)
{
  mobius::vfs::partition_system_entry e;
  e.set_type ("gpt.header");
  e.set_starting_sector (header.sector);
  e.set_ending_sector (header.sector);
  e.set_sectors (header.sectors);
  e.set_attribute ("signature", header.signature);
  e.set_attribute ("revision", header.revision);
  e.set_attribute ("header_size", header.size);
  e.set_attribute ("header_crc32", header.crc32);
  e.set_attribute ("lba", header.lba);
  e.set_attribute ("alternate_lba", header.alternate_lba);
  e.set_attribute ("first_usable_lba", header.first_usable_lba);
  e.set_attribute ("last_usable_lba", header.last_usable_lba);
  e.set_attribute ("disk_guid", header.disk_guid);
  e.set_attribute ("partition_entry_lba", header.partition_entry_lba);
  e.set_attribute ("number_of_partition_entries", header.number_of_partition_entries);
  e.set_attribute ("size_of_partition_entry", header.size_of_partition_entry);
  e.set_attribute ("partition_entry_array_crc32", header.partition_entry_array_crc32);

  return e;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create gpt.table entry from GPT Header
//! \param header GPT header object
//! \param sector_size Sector size
//! \return Entry
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static mobius::vfs::partition_system_entry
_new_table_entry (const gpt_header& header, std::uint32_t sector_size)
{
  auto partition_table_sectors = (header.number_of_partition_entries * header.size_of_partition_entry + sector_size - 1) / sector_size;

  mobius::vfs::partition_system_entry e;
  e.set_type ("gpt.table");
  e.set_starting_sector (header.partition_entry_lba);
  e.set_ending_sector (e.get_starting_sector () + partition_table_sectors - 1);
  e.set_sectors (e.get_ending_sector () - e.get_starting_sector () + 1);
  e.set_attribute ("entry_size", header.size_of_partition_entry);
  e.set_attribute ("entries", header.number_of_partition_entries);

  return e;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create gpt.partition entry from GPT partition
//! \param partition GPT partition object
//! \return Entry
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static mobius::vfs::partition_system_entry
_new_partition_entry (const mobius::vfs::partition& p)
{
  mobius::vfs::partition_system_entry e;
  e.set_type ("gpt.partition");
  e.set_starting_sector (p.get_starting_sector ());
  e.set_ending_sector (p.get_ending_sector ());
  e.set_sectors (e.get_ending_sector () - e.get_starting_sector () + 1);
  
  if (p.get_name ().empty ())
    e.set_description ("Unnamed partition");

  else
    e.set_description (p.get_name ());

  std::string flags;

  if (p.is_bootable ()) flags += 'B';
  if (p.is_hidden ()) flags += 'H';
  if (p.is_primary ()) flags += 'P';
  if (p.is_extended ()) flags += 'E';
  if (p.is_logical ()) flags += 'L';

  e.set_flags (flags);

  e.set_attribute ("is_bootable", p.is_bootable ());
  e.set_attribute ("is_primary", p.is_primary ());
  e.set_attribute ("is_extended", p.is_extended ());
  e.set_attribute ("is_logical", p.is_logical ());
  e.set_attribute ("is_hidden", p.is_hidden ());
  e.set_attribute ("is_readable", p.is_readable ());
  e.set_attribute ("is_writeable", p.is_writeable ());

  for (const auto& [name, value] : p.get_attributes ())
    e.set_attribute (name, value);

  return e;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create partition from GPT partition object
//! \param partition GPT partition object
//! \param sector_size Sector size
//! \return Entry
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static mobius::vfs::partition
_new_partition (const gpt_partition& partition, std::uint32_t sector_size)
{
  mobius::vfs::partition p;
  p.set_starting_sector (partition.starting_lba);
  p.set_ending_sector (partition.ending_lba);
  p.set_sectors (partition.ending_lba - partition.starting_lba + 1);
  p.set_starting_address (partition.starting_lba * sector_size);
  p.set_ending_address ((partition.ending_lba + 1) * sector_size - 1);
  p.set_size (p.get_sectors () * sector_size);
  p.set_name (partition.name);
  p.set_type (partition.type_guid);
  p.is_bootable (partition.attributes & 0x00000004);
  p.is_readable (true);
  p.is_writeable (true);

  if (partition.type_guid == GPT_MS_BASIC_DATA)
    {
      p.is_writeable (partition.attributes ^ 0x10000000);
      p.is_hidden (partition.attributes & 0x40000000);
    }

  p.set_attribute ("type_guid", partition.type_guid);
  p.set_attribute ("unique_guid", partition.unique_guid);
  p.set_attribute ("attributes", partition.attributes);

  return p;
}

} // namespace

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Check if stream has an instance of GPT partition system
//! \param reader Reader object
//! \param sector_size Sector size in bytes
//! \return true/false
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
bool
partition_system_impl::is_instance (
  const mobius::io::reader& reader,
  std::uint32_t sector_size
)
{
  // try to read primary GPT header
  auto header = _decode_gpt_header (reader, sector_size, 1);
  if (header.is_valid)
    return true;

  // try to read backup GPT header
  sector_type sectors = (reader.get_size () + sector_size - 1) / sector_size;
  header = _decode_gpt_header (reader, sector_size, sectors - 1);
  if (header.is_valid)
    return true;

  return false;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Constructor
//! \param reader Reader object
//! \param sector_size Sector size in bytes
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
partition_system_impl::partition_system_impl (
  const mobius::io::reader& reader,
  std::uint32_t sector_size
)
  : reader_ (reader),
    sector_size_ (sector_size)
{
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Load data
//! \see UEFI 2.9 - section 5.3.1
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
partition_system_impl::_load () const
{
  if (is_loaded_)
    return;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // Read GPT headers
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  sector_type sectors = (reader_.get_size () + sector_size_ - 1) / sector_size_;
  auto primary_gpt_header = _decode_gpt_header (reader_, sector_size_, 1);
  auto backup_gpt_header = _decode_gpt_header (reader_, sector_size_, sectors - 1);
  gpt_header header;
  
  if (primary_gpt_header.is_valid)
    header = primary_gpt_header;

  else if (backup_gpt_header.is_valid)
    header = backup_gpt_header;

  else
    throw std::runtime_error (mobius::MOBIUS_EXCEPTION_MSG ("no valid GPT header found"));

  attributes_.set ("revision", header.revision);
  attributes_.set ("disk_guid", header.disk_guid);
  attributes_.set ("first_usable_lba", header.first_usable_lba);
  attributes_.set ("last_usable_lba", header.last_usable_lba);

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // Read partitions
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  std::vector <gpt_partition> partitions;

  // try to read primary partition table
  if (primary_gpt_header.is_valid)
    partitions = _decode_gpt_partition_table (
       reader_,
       sector_size_,
       primary_gpt_header.partition_entry_lba,
       primary_gpt_header.number_of_partition_entries
    );
  
  // if there is no partition, try to read backup partition table
  if (partitions.empty () && backup_gpt_header.is_valid)
    partitions = _decode_gpt_partition_table (
       reader_,
       sector_size_,
       backup_gpt_header.partition_entry_lba,
       backup_gpt_header.number_of_partition_entries
    );

  attributes_.set ("number_of_partitions", partitions.size ());

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // Create entries
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  auto e = _new_protective_mbr_entry ();
  entries_.push_back (e);
  
  // Create primary GPT header and primary GPT partition table entries
  if (primary_gpt_header.is_valid)
    {
      auto e = _new_header_entry (primary_gpt_header);
      e.set_description ("Primary GPT Header");
      entries_.push_back (e);

      e = _new_table_entry (primary_gpt_header, sector_size_);
      e.set_description ("Primary GPT Partition Table");
      entries_.push_back (e);
    }

  // Add partitions
  for (const auto& partition : partitions)
    {
      // create partition
      auto p = _new_partition (partition, sector_size_);
      partitions_.push_back (p);

      // create partition table entry
      auto e = _new_partition_entry (p);

      std::string flags;
      if (p.is_bootable ()) flags += 'B';
      if (p.is_hidden ()) flags += 'H';
      if (p.is_primary ()) flags += 'P';
      if (p.is_extended ()) flags += 'E';
      if (p.is_logical ()) flags += 'L';
      if (!p.is_writeable ()) flags += 'R';
      e.set_flags (flags);

      entries_.push_back (e);
    }

  // Create backup GPT header and backup GPT partition table entries
  if (backup_gpt_header.is_valid)
    {
      auto e = _new_header_entry (backup_gpt_header);
      e.set_description ("Backup GPT Header");
      entries_.push_back (e);

      e = _new_table_entry (backup_gpt_header, sector_size_);
      e.set_description ("Backup GPT Partition Table");
      entries_.push_back (e);
    }

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // Add freespaces between non-contiguous partitions
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  _add_freespaces (entries_, sectors);
  _set_addresses (entries_, sector_size_);
  is_loaded_ = true;
}
