// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Mobius Forensic Toolkit
// Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019,2020,2021 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_dos.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>

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \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 mobius
{
namespace partition
{
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief check if stream has an instance of DOS partition system
//! \param reader stream reader
//! \return true/false
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
bool
partition_system_impl_dos::is_instance (const mobius::disk::disk disk)
{
  constexpr int MBR_SIGNATURE_OFFSET = 0x1fe;

  auto reader = disk.new_reader ();
  reader.seek (0);
  auto data = reader.read (disk.get_sector_size ());

  return
    data[MBR_SIGNATURE_OFFSET] == 0x55 &&
    data[MBR_SIGNATURE_OFFSET + 1] == 0xaa &&
    data.slice (3, 6) != "NTFS" &&
    data.slice (54, 58) != "FAT12" &&
    data.slice (54, 58) != "FAT16" &&
    data.slice (82, 86) != "FAT32" &&
    data.slice (11, 18) != "\x00\x02\x01\x01\x00\x02\xe0\x00";
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief constructor
//! \param reader stream reader
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
partition_system_impl_dos::partition_system_impl_dos (const mobius::disk::disk disk)
{
  mobius::io::sector_reader_adaptor reader (disk.new_reader (), disk.get_sector_size ());
  _scan_partitions (reader);
  _add_freespaces (reader);

}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief scan entries recursively through the partition system
//! \param reader stream reader
//! \param sector starting sector (default = 0)
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
partition_system_impl_dos::_scan_partitions (
  mobius::io::sector_reader_adaptor reader,
  sector_type sector)
{
  // create MBR entry
  entry e;
  e.starting_sector = sector;
  e.ending_sector = sector;
  e.sectors = 1;
  e.type = "mbr";
  e.description = (sector == 0) ? "Master Boot Record (MBR)" : "Extended Master Boot Record";
  _add_entry (e);

  // read MBR
  reader.seek (sector);
  auto data = reader.read ();
  mobius::decoder::data_decoder decoder (mobius::io::new_bytearray_reader (data));
  decoder.skip (446);   // bootcode + disk_id + unknown

  // 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 ();
      auto ending_sector = starting_sector + sectors - 1;

      if (type)
        {
          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.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));
          _add_partition (p);

          // create partition table entry
          entry e;
          e.starting_sector = starting_sector;
          e.ending_sector = starting_sector + sectors - 1;
          e.sectors = sectors;
          e.type = "partition";

          auto iterator = PARTITION_DESCRIPTION.find (type);

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

          if (p.is_bootable ()) e.flags += 'B';
          if (p.is_hidden ()) e.flags += 'H';
          if (p.is_primary ()) e.flags += 'P';
          if (p.is_extended ()) e.flags += 'E';
          if (p.is_logical ()) e.flags += 'L';
          _add_entry (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 (reader, sector);
}

} // namespace partition
} // namespace mobius
