// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// 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/decoder/data_decoder.h>
#include <mobius/io/bytearray_io.h>
#include <mobius/io/sector_reader_adaptor.h>
#include <map>
#include <cstdint>
#include <string>
#include <iomanip>
#include <sstream>

namespace
{
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Partition description by type
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static const std::map <std::uint8_t, const std::string> PARTITION_DESCRIPTION =
{
  {0x01, "FAT-12"},
  {0x04, "FAT-16 (<32 MB)"},
  {0x05, "Extended (CHS)"},
  {0x06, "FAT-16B"},
  {0x07, "NTFS/HPFS"},
  {0x0b, "FAT-32 (CHS)"},
  {0x0c, "FAT-32 (LBA)"},
  {0x0e, "FAT-16 (LBA)"},
  {0x0f, "Extended (LBA)"},
  {0x11, "Hidden FAT-12"},
  {0x12, "Hibernation/firmware"},
  {0x14, "Hidden FAT-16 (<32 MB)"},
  {0x15, "Hidden extended (CHS)"},
  {0x16, "Hidden FAT-16B"},
  {0x17, "Hidden NTFS/HPFS"},
  {0x1b, "Hidden FAT-32 (CHS)"},
  {0x1c, "Hidden FAT-32 (LBA)"},
  {0x1e, "Hidden FAT-16 (LBA)"},
  {0x1f, "Hidden extended (LBA)"},
  {0x27, "Windows Recovery Environment"},
  {0x82, "Linux swap space"},
  {0x83, "Linux"},
  {0x85, "Linux extended"},
  {0x86, "Linux RAID"},
  {0x8e, "Linux LVM"},
  {0xee, "GPT protective MBR"},
  {0xef, "EFI system"}
};

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Convert uint8_t to string
//! \param value uint8_t value
//! \return string
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static const std::string
to_hexstring (std::uint8_t value)
{
  std::stringstream stream;
  stream << std::setfill ('0') << std::setw (2) << std::hex << int (value);

  return stream.str ();
}

} // namespace

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Check if stream has an instance of DOS partition system
//! \param reader Reader object
//! \param sector_size Sector size in bytes
//! \return true/false
//! \see Linux source code: block/partitions/msdos.c
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
bool
partition_system_impl::is_instance (
  const mobius::io::reader& reader,
  std::uint32_t sector_size
)
{
  constexpr int MBR_PARTITION_OFFSET = 0x1be;
  constexpr int MBR_SIGNATURE_OFFSET = 0x1fe;
  constexpr std::uint8_t EFI_GPT_TYPE = 0xee;

  auto r = reader;
  r.seek (0);
  auto data = r.read (sector_size);

  return
    data.size () >= 512 &&

    // MBR signature
    data[MBR_SIGNATURE_OFFSET] == 0x55 &&
    data[MBR_SIGNATURE_OFFSET + 1] == 0xaa &&
    
    // Boot indicators must be either 0x00 or 0x80
    (data[MBR_PARTITION_OFFSET] & 0x7f) == 0x00 &&
    (data[MBR_PARTITION_OFFSET+16] & 0x7f) == 0x00 &&
    (data[MBR_PARTITION_OFFSET+32] & 0x7f) == 0x00 &&
    (data[MBR_PARTITION_OFFSET+48] & 0x7f) == 0x00 &&

    // It must have at least one valid non protective partition
    (
      (data[MBR_PARTITION_OFFSET+4] != EFI_GPT_TYPE &&
       data[MBR_PARTITION_OFFSET+4] != 0x00) ||
      (data[MBR_PARTITION_OFFSET+16+4] != EFI_GPT_TYPE &&
       data[MBR_PARTITION_OFFSET+16+4] != 0x00) ||
      (data[MBR_PARTITION_OFFSET+32+4] != EFI_GPT_TYPE &&
       data[MBR_PARTITION_OFFSET+32+4] != 0x00) ||
      (data[MBR_PARTITION_OFFSET+48+4] != EFI_GPT_TYPE &&
       data[MBR_PARTITION_OFFSET+48+4] != 0x00)
    );
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \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
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
partition_system_impl::_load () const
{
  if (is_loaded_)
    return;

  sector_type sectors = (reader_.get_size () + sector_size_ - 1) / sector_size_;

  _scan_partitions ();
  _add_freespaces (entries_, sectors);
  _set_addresses (entries_, sector_size_);

  attributes_.set ("number_of_partitions", partitions_.size ());
  is_loaded_ = true;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Scan entries recursively through the partition system
//! \param reader stream reader
//! \param sector starting sector (default = 0)
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
partition_system_impl::_scan_partitions (sector_type sector) const
{
  // create MBR entry
  mobius::vfs::partition_system_entry e_mbr;
  e_mbr.set_starting_sector (sector);
  e_mbr.set_ending_sector (sector);
  e_mbr.set_sectors (1);
  e_mbr.set_type ("mbr");
  e_mbr.set_description ((sector == 0) ? "Master Boot Record (MBR)" : "Extended Master Boot Record");
  entries_.push_back (e_mbr);

  // read MBR
  auto reader = reader_;
  reader.seek (sector * sector_size_);
  mobius::decoder::data_decoder decoder (reader);
  decoder.skip (440);   // bootcode
  
  if (sector == 0)
    {
      attributes_.set ("disk_id", decoder.get_uint32_le ());
      attributes_.set ("is_copy_protected", decoder.get_uint16_le () == 0x5a5a);
    }

  else
    decoder.skip (6);

  // create partition entries
  std::vector <sector_type> extended_list;

  for (int i = 0; i < 4; i++)
    {
      auto drive_index = decoder.get_uint8 ();
      auto start_chs = decoder.get_bytearray_by_size (3);
      auto type = decoder.get_uint8 ();
      auto end_chs = decoder.get_bytearray_by_size (3);
      auto starting_sector = decoder.get_uint32_le () + sector;
      auto sectors = decoder.get_uint32_le ();

      if (sectors)
        {
          mobius::vfs::partition p;
          p.set_starting_sector (starting_sector);
          p.set_ending_sector (starting_sector + sectors - 1);
          p.set_sectors (sectors);
          p.set_starting_address (starting_sector * sector_size_);
          p.set_ending_address ((p.get_ending_sector () + 1) * sector_size_ - 1);
          p.set_size (sectors * sector_size_);
          p.is_bootable (drive_index & 0x80);
          p.is_primary (sector == 0);
          p.is_extended (type == 0x05 || type == 0x0f || type == 0x15 || type == 0x1f || type == 0x85);
          p.is_logical (!p.is_primary () && !p.is_extended ());
          p.is_hidden (type == 0x14 || type == 0x15 || type == 0x16 || type == 0x17 || type == 0x1b || type == 0x1c || type == 0x1e || type == 0x1f);
          p.is_readable (true);
          p.is_writeable (true);
          p.set_type ("0x" + to_hexstring (type));
          p.set_attribute ("type", type);
          p.set_attribute ("drive_index", drive_index);
          p.set_attribute ("start_chs", start_chs);
          p.set_attribute ("end_chs", end_chs);
          partitions_.push_back (p);

          // create partition table entry
          mobius::vfs::partition_system_entry e;
          e.set_starting_sector (starting_sector);
          e.set_ending_sector (starting_sector + sectors - 1);
          e.set_sectors (sectors);
          e.set_type ("partition");

          auto iterator = PARTITION_DESCRIPTION.find (type);

          if (iterator != PARTITION_DESCRIPTION.end ())
            e.set_description (iterator->second + " partition");
          else
            e.set_description ("Partition (type = " + to_hexstring (type) + ')');

          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 ("partition_type", type);
          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);

          entries_.push_back (e);

          // if partition is extended, add to extended_list
          if (p.is_extended ())
            extended_list.push_back (starting_sector);
        }
    }

  // scan extended partitions
  for (auto sector : extended_list)
    _scan_partitions (sector);
}
