// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Mobius Forensic Toolkit
// Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018,2019,2020,2021,2022,2023,2024 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 "writer_impl.h"
#include <mobius/io/file.h>
#include <mobius/io/uri.h>
#include <mobius/exception.inc>
#include <stdexcept>

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief number of sectors per chunk
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static constexpr int SECTOR_SIZE = 512;
static constexpr int CHUNK_SECTORS = 64;

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Constructor
//! \param imagefile_impl imagefile implementation object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
writer_impl::writer_impl (const imagefile_impl& imagefile_impl)
  : segment_size_ (std::int64_t (imagefile_impl.get_attribute ("segment_size"))),
    chunk_size_ (imagefile_impl.get_sector_size () * CHUNK_SECTORS),
    sector_size_ (imagefile_impl.get_sector_size ()),
    compression_level_ (std::int64_t (imagefile_impl.get_attribute ("compression_level"))),
    segments_ (imagefile_impl.get_segment_array ()),
    hash_ ("md5")
{
  // validate segment size. The segment size must be at least large enough
  // to store a chunk of data.
  constexpr size_type SECTION_HEADER_SIZE = 76;
  constexpr size_type HEADER_SECTION_SIZE = 500;
  constexpr size_type VOLUME_SECTION_SIZE = 1128;
  constexpr size_type HASH_SECTION_SIZE = 112;
  constexpr size_type DONE_SECTION_SIZE = 76;
  constexpr size_type TABLE_SECTION_SIZE = 76 + 24 + 4 + 4;

  const size_type min_segment_size =
    HEADER_SECTION_SIZE * 3 +                 // header2, header2 and header
    VOLUME_SECTION_SIZE +                     // volume section
    SECTION_HEADER_SIZE + chunk_size_ +       // sectors section
    TABLE_SECTION_SIZE * 2 +                  // table and table2 sections
    VOLUME_SECTION_SIZE +                     // data section (equal to volume section)
    HASH_SECTION_SIZE +                       // hash section
    DONE_SECTION_SIZE;                        // done section

  if (segment_size_ < min_segment_size)
    throw std::runtime_error (mobius::MOBIUS_EXCEPTION_MSG ("segment size too small"));

  // create imagefile GUID
  guid_ = mobius::bytearray (16);
  guid_.random ();

  // create first segment
  auto& segment_writer = _new_segment_writer ();

  auto acquisition_user = imagefile_impl.get_attribute ("get_acquisition_user");
  auto drive_vendor = imagefile_impl.get_attribute ("drive_vendor");
  auto drive_model = imagefile_impl.get_attribute ("drive_model");
  auto drive_serial_number = imagefile_impl.get_attribute ("drive_serial_number");
  
  if (drive_vendor.is_string ())
    segment_writer.set_drive_vendor (std::string (drive_vendor));

  if (drive_model.is_string ())
    segment_writer.set_drive_model (std::string (drive_model));

  if (drive_serial_number.is_string ())
    segment_writer.set_drive_serial_number (std::string (drive_serial_number));
  
  if (acquisition_user.is_string ())
    segment_writer.set_acquisition_user (std::string (acquisition_user));

  segment_writer.create ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Destructor
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
writer_impl::~writer_impl ()
{
  // pad image to SECTOR_SIZE, if necessary
  if (size_ % SECTOR_SIZE)
    {
      auto pad_size = SECTOR_SIZE - (size_ % SECTOR_SIZE);
      write (mobius::bytearray (pad_size));
    }

  // close segments
  for (auto& sw : segment_writer_list_)
    {
      sw.set_total_size (size_);
      sw.set_chunk_sectors (CHUNK_SECTORS);
      sw.set_sector_size (sector_size_);
      sw.set_segment_count (segment_writer_list_.size ());
      sw.set_md5_hash (hash_.get_digest ());
      sw.close ();
    }
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Set write position
//! \param offset offset in bytes
//! \param w either beginning, current or end
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
writer_impl::seek (offset_type, whence_type)
{
  throw std::runtime_error (mobius::MOBIUS_EXCEPTION_MSG ("writer is not seekable"));
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Write bytes to stream
//! \param data Data
//! \return Number of bytes written
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
writer_impl::size_type
writer_impl::write (const mobius::bytearray& data)
{
  mobius::bytearray buffer = data;

  while (!buffer.empty ())
    {
      auto& segment_writer = _get_current_segment_writer ();

      auto bytes_written = segment_writer.write (buffer);
      size_ += bytes_written;

      if (bytes_written == buffer.size ())
        buffer.clear ();

      else
        {
          // discard bytes already written
          buffer = buffer.slice (bytes_written, buffer.size () - 1);

          // create next segment file
          auto& new_segment_writer = _new_segment_writer ();
          new_segment_writer.create ();
        }
    }

  hash_.update (data);
  return data.size ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Flush data to file
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
writer_impl::flush ()
{
  auto& segment_writer = _get_current_segment_writer ();
  segment_writer.flush ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new segment file
//! \return Reference to current segment writer
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
segment_writer&
writer_impl::_new_segment_writer ()
{
  auto segment_number = segment_writer_list_.size () + 1;
  auto writer = segments_.new_writer (segment_number - 1);
  segment_writer_list_.emplace_back (writer, segment_number);

  auto& segment_writer = segment_writer_list_[segment_writer_list_.size () - 1];
  segment_writer.set_segment_size (segment_size_);
  segment_writer.set_chunk_size (chunk_size_);
  segment_writer.set_compression_level (compression_level_);
  segment_writer.set_guid (guid_);

  return segment_writer;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Get current segment writer
//! \return Reference to current segment writer
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
segment_writer&
writer_impl::_get_current_segment_writer ()
{
  if (segment_writer_list_.size () == 0)
    throw std::runtime_error (mobius::MOBIUS_EXCEPTION_MSG ("no segment writer found"));

  return segment_writer_list_[segment_writer_list_.size () - 1];
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Return next extension
//! \param extension
//! \return next extension
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
std::string
writer_impl::_get_next_extension (const std::string& extension) const
{
  std::string tmp_extension;

  if (extension == "E99")
    tmp_extension = "EAA";

  else
    {
      tmp_extension = extension;
      int pos = extension.length () - 1;
      bool carry = true;

      while (pos >= 0 && carry)
        {
          if (tmp_extension[pos] == '9')
            tmp_extension[pos] = '0';

          else if (tmp_extension[pos] == 'Z')
            tmp_extension[pos] = 'A';

          else
            {
              ++tmp_extension[pos];
              carry = false;
            }

          if (carry)
            --pos;
        }
    }

  return tmp_extension;
}
