// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Mobius Forensic Toolkit
// Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017 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 <mobius/partition/partition_system_impl_gpt.h>
#include <mobius/partition/partition_system_impl_dos.h>
#include <mobius/decoder/data_decoder.h>
#include <mobius/io/bytearray_io.h>
#include <mobius/io/sector_reader_adaptor.h>
#include <mobius/charset.h>
#include <algorithm>

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// constants
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static const std::string GPT_UNUSED_ENTRY = "00000000-0000-0000-0000-000000000000";

namespace mobius
{
namespace partition
{
using sector_type = partition_system_impl_base::sector_type;

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief format GUID type
//! \param data GUID as bytearray
//! \return formatted GUID
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static const std::string
_format_guid_type (const mobius::bytearray& data)
{
  std::uint32_t g1 =
    (std::uint32_t (data[0])) |
    (std::uint32_t (data[1]) << 8) |
    (std::uint32_t (data[2]) << 16) |
    (std::uint32_t (data[3]) << 24);

  std::uint16_t g2 =
    (std::uint16_t (data[4])) |
    (std::uint16_t (data[5]) << 8);

  std::uint16_t g3 =
    (std::uint16_t (data[6])) |
    (std::uint16_t (data[7]) << 8);

  std::uint16_t g4 =
    (std::uint16_t (data[8]) << 8) |
    (std::uint16_t (data[9]));

  std::uint64_t g5 =
    (std::uint64_t (data[10]) << 40) |
    (std::uint64_t (data[11]) << 32) |
    (std::uint64_t (data[12]) << 24) |
    (std::uint64_t (data[13]) << 16) |
    (std::uint64_t (data[14]) << 8) |
    (std::uint64_t (data[15]));

  char buffer[128];
  sprintf (buffer, "%08X-%04X-%04X-%04X-%012lX", g1, g2, g3, g4, g5);

  return buffer;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief read GPT header
//! \param reader stream reader
//! \return data
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=        
static const mobius::bytearray
_read_gpt_header (mobius::io::sector_reader_adaptor reader)
{
  mobius::bytearray data;
  
  try
    {
      reader.seek (1);                  // first GPT header
      data = reader.read ();
    }
  catch (const std::runtime_error& e)
    {
      reader.seek (-1);                 // backup GPT header
      data = reader.read ();
    }

  return data;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief read GPT partition table
//! \param reader stream reader
//! \return data
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=        
static const mobius::bytearray
_read_gpt_partition_table (
        mobius::io::sector_reader_adaptor reader,
        sector_type sector,
        sector_type sectors)
{
  mobius::bytearray data;
  
  try
    {
      reader.seek (sector);             // main partition table
      data = reader.read (sectors);
    }
  catch (const std::runtime_error& e)
    {
      reader.seek (reader.get_sectors () - sectors - 1);  // backup partition table
      data = reader.read (sectors);
    }

  return data;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief check if stream has an instance of GPT partition system
//! \param disk disk object
//! \return true/false
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
bool
partition_system_impl_gpt::is_instance (const mobius::disk::disk disk)
{
  mobius::io::sector_reader_adaptor reader (disk.new_reader (), disk.get_sector_size ());
  auto data = _read_gpt_header (reader);
  const mobius::bytearray GPT_SIGNATURE = "EFI PART";

  return data.slice (0, 7) == GPT_SIGNATURE;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief constructor
//! \param disk disk object
//! \see UEFI Specification, version 2.7
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
partition_system_impl_gpt::partition_system_impl_gpt (const mobius::disk::disk disk)
{
  mobius::io::sector_reader_adaptor reader (disk.new_reader (), disk.get_sector_size ());
  entry e;

  // retrieve DOS partitions first
  if (partition_system_impl_dos::is_instance (disk))
    {
      partition_system_impl_dos psystem_dos (disk);
      auto partitions = psystem_dos.get_partitions ();

      bool protective_mbr = std::none_of (
              partitions.begin (),
              partitions.end (),
              [](partition p){ return p.get_type () != "0xee"; }
      );

      // if DOS partition system has only protective partitions, create protective MBR entry
      if (protective_mbr)
        {
          e.starting_sector = 0;
          e.ending_sector = 0;
          e.sectors = 1;
          e.type = "mbr";
          e.description = "Protective MBR";
          _add_entry (e);
        }

      else
        {
          for (auto e : psystem_dos.get_entries ())
            _add_entry (e);

          for (auto p : partitions)
            _add_partition (p);

          _add_freespaces (reader);
          return;
        }
    }
  
  // decode GPT header
  auto data = _read_gpt_header (reader);
  mobius::decoder::data_decoder decoder (mobius::io::new_bytearray_reader (data));
  decoder.skip (72);
  auto partition_entry_lba = decoder.get_uint64_le ();
  auto partitions = decoder.get_uint32_le ();
  auto partition_entry_size = decoder.get_uint32_le ();
  
  // add GPT header entry
  e.starting_sector = 1;
  e.ending_sector = 1;
  e.sectors = 1;
  e.type = "gpt.header";
  e.description = "GPT Header";
  _add_entry (e);
  
  // add GPT partition table entry
  auto sectors = (partitions * partition_entry_size + reader.get_sector_size () - 1) / reader.get_sector_size ();

  e.starting_sector = partition_entry_lba;
  e.ending_sector = partition_entry_lba + sectors - 1;
  e.sectors = sectors;
  e.type = "gpt.table";
  e.description = "GPT partition table";
  _add_entry (e);

  // read partition table
  data = _read_gpt_partition_table (reader, partition_entry_lba, sectors);
  mobius::decoder::data_decoder decoder_table (mobius::io::new_bytearray_reader (data));

  // add partitions
  for (std::uint32_t i = 0;i < partitions;i++)
    {
      // UEFI 2.7 - 5.3.3
      auto type_guid = _format_guid_type (decoder_table.get_bytearray_by_size (16));
      auto unique_partition_guid = decoder_table.get_bytearray_by_size (16);
      auto starting_sector = decoder_table.get_uint64_le ();
      auto ending_sector = decoder_table.get_uint64_le ();
      auto attributes = decoder_table.get_uint64_le ();
      auto name = mobius::conv_charset_to_utf8 (decoder_table.get_bytearray_by_size (72), "utf-16");
      
      if (type_guid != GPT_UNUSED_ENTRY)
        {
          // create partition
          partition p;
          p.set_starting_sector (starting_sector);
          p.set_ending_sector (ending_sector);
          p.set_starting_address (starting_sector * reader.get_sector_size ());
          p.set_ending_address ((p.get_ending_sector () + 1) * reader.get_sector_size () - 1);
          p.set_name (name);
          p.set_type (type_guid);
          p.is_bootable (attributes & 0x00000004);
          p.is_readable (true);
          _add_partition (p);

          // create partition table entry
          entry e;
          e.starting_sector = starting_sector;
          e.ending_sector = ending_sector;
          e.sectors = ending_sector - starting_sector + 1;
          e.type = "partition";
          e.description = name + " partition";
          _add_entry (e);
        }
    }
  
  // add backup partition table
  e.starting_sector = reader.get_sectors () - sectors - 1;
  e.ending_sector = reader.get_sectors () - 2;
  e.sectors = sectors;
  e.type = "gpt.table";
  e.description = "Backup GPT partition table";
  _add_entry (e);

  // add backup GPT header
  e.starting_sector = reader.get_sectors () - 1;
  e.ending_sector = reader.get_sectors () - 1;
  e.sectors = 1;
  e.type = "gpt.header";
  e.description = "Backup GPT Header";
  _add_entry (e);
  
  // add freespaces between non-contiguous partitions
  _add_freespaces (reader);
}

} // namespace partition
} // namespace mobius
