// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// 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/imagefile/ewf/writer_impl.h>
#include <mobius/io/file.h>
#include <mobius/exception.inc>
#include <mobius/uri.h>
#include <stdexcept>

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

namespace mobius
{
namespace imagefile
{
namespace ewf
{

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief constructor
//! \param imagefile_impl imagefile implementation object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
writer_impl::writer_impl (const imagefile_impl& imagefile_impl)
 : segment_size_ (imagefile_impl.get_segment_size ()),
   chunk_size_ (imagefile_impl.get_sector_size () * CHUNK_SECTORS),
   sector_size_ (imagefile_impl.get_sector_size ()),
   compression_level_ (imagefile_impl.get_compression_level ())
{
  // 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_EXCEPTION_MSG ("segment size too small"));

  // parse URL
  const std::string url_string = imagefile_impl.get_url ();
  mobius::uri url (url_string);
  url_extension_ = url.get_extension ();
  url_prefix_ = url_string.substr (0, url_string.length () - url_extension_.length ());

  // create first segment
  _new_segment_writer ();
  segment_writer_.set_drive_vendor (imagefile_impl.get_drive_vendor ());
  segment_writer_.set_drive_model (imagefile_impl.get_drive_model ());
  segment_writer_.set_drive_serial_number (imagefile_impl.get_drive_serial_number ());
  segment_writer_.set_acquisition_user (imagefile_impl.get_acquisition_user ());
  segment_writer_.create ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief destructor
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
writer_impl::~writer_impl ()
{
  segment_writer_.set_md5_hash (mobius::bytearray (16));
  mobius::bytearray guid (16);
  guid.random ();

  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_number_);
      sw.set_guid (guid);
      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_EXCEPTION_MSG ("writer is not seekable"));
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief write bytes to writer
//! \param data a bytearray
//! \return number of bytes written
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
writer_impl::size_type
writer_impl::write (const mobius::bytearray& data)
{
  mobius::bytearray buffer = data;
  
  while (!buffer.empty ())
    {
      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
          url_extension_ = _get_next_extension (url_extension_);
          segment_number_++;
          _new_segment_writer ();
          segment_writer_.create ();
        }
    }

  hash_.update (data.begin (), data.end ());
  return data.size ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief flush data to file
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
writer_impl::flush ()
{
  segment_writer_.flush ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief create new segment file
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
writer_impl::_new_segment_writer ()
{
  mobius::io::file f (url_prefix_ + url_extension_);

  segment_writer_ = segment_writer (f.new_writer (), segment_number_);
  segment_writer_.set_segment_size (segment_size_);
  segment_writer_.set_chunk_size (chunk_size_);
  segment_writer_.set_compression_level (compression_level_);

  segment_writer_list_.push_back (segment_writer_);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief return next extension
//! \param extension
//! \return next extension
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
const 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;
}

} // namespace ewf
} // namespace imagefile
} // namespace mobius
