// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Mobius Forensic Toolkit
// Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016,2017,2018 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 "imagefile_impl.h"
#include "reader_impl_aes.h"
#include "reader_impl_blowfish.h"
#include "reader_impl_plaintext.h"
#include <mobius/decoder/data_decoder.h>
#include <mobius/io/file.h>
#include <mobius/io/reader.h>
#include <mobius/bytearray.h>
#include <mobius/exception.inc>
#include <mobius/string_functions.h>
#include <stdexcept>

#include <iostream>

namespace mobius
{
namespace imagefile
{
namespace msr
{
namespace
{
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// constants
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

//! \brief header size in bytes
constexpr int HEADER_SIZE = 16384;

//! \brief header signatures
const mobius::bytearray HEADER_SIGNATURE_V0 = {0x98, 0x92, 0x04, 0x71};
const mobius::bytearray HEADER_SIGNATURE_V1 = {0x12, 0xa1, 0x58, 0x32};
const mobius::bytearray HEADER_SIGNATURE_V2 = {0xa7, 0xb2, 0x62, 0x5a};

//! \brief header encryption key
const mobius::bytearray HEADER_ENCRYPTION_KEY =
{
  0x06, 0x42, 0x21, 0x98, 0x03, 0x69, 0x5e, 0xb1,
  0x5f, 0x40, 0x60, 0x8c, 0x2e, 0x36, 0x00, 0x06
};

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief read and decrypt file header
//! \param reader reader
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::bytearray
read_header_data (mobius::io::reader reader)
{
  const mobius::bytearray encrypted_data = reader.read (HEADER_SIZE);

  mobius::crypt::cipher_aes aes (HEADER_ENCRYPTION_KEY, "cbc");
  return aes.decrypt (encrypted_data);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief describe encryption algorithm
//! \param encryption_algorithm encryption algorithm
//! \return string
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
std::string
encryption_algorithm_to_string (int encryption_algorithm)
{
  std::string text = std::to_string (encryption_algorithm);

  if (encryption_algorithm == 0)
    text += " (no encryption)";

  else if (encryption_algorithm == 1)
    text += " (AES-128)";

  else if (encryption_algorithm == 2)
    text += " (AES-256)";

  else if (encryption_algorithm == 3)
    text += " (Blowfish-448)";

  else
    text += " (unknown)";

  return text;
}

} // namespace

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief check if URL is an instance of imagefile msr
//! \return true/false
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
bool
imagefile_impl::is_instance (const std::string& url)
{
  bool is_instance = false;
  mobius::io::file f (url);

  if (f && f.exists ())
    {
      auto reader = f.new_reader ();
      const mobius::bytearray data = read_header_data (reader);

      is_instance = data && (data.slice (8192, 8195) == HEADER_SIGNATURE_V0 ||
                             data.slice (8192, 8195) == HEADER_SIGNATURE_V1 ||
                             data.slice (8192, 8195) == HEADER_SIGNATURE_V2);
    }

  return is_instance;
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief construct object
//! \param url imagefile URL
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
imagefile_impl::imagefile_impl (const std::string& url)
  : url_ (url)
{
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief create new reader for imagefile
//! \return reader
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::io::reader
imagefile_impl::new_reader () const
{
  _load_metadata ();

  if (encryption_algorithm_ == 0)
    return mobius::io::reader (std::make_shared <reader_impl_plaintext> (*this));

  else if (encryption_algorithm_ == 1 || encryption_algorithm_ == 2)
    return mobius::io::reader (std::make_shared <reader_impl_aes> (*this));

  else if (encryption_algorithm_ == 3)
    throw std::runtime_error (MOBIUS_EXCEPTION_MSG ("Blowfish-448 encryption not supported"));
  //return mobius::io::reader (std::make_shared <reader_impl_blowfish> (*this));

  else
    throw std::runtime_error (MOBIUS_EXCEPTION_MSG ("unknown/invalid encryption algorithm"));
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief create new writer for imagefile
//! \return writer
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::io::writer
imagefile_impl::new_writer () const
{
  throw std::runtime_error (MOBIUS_EXCEPTION_MSG ("writer not implemented"));
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief get metadata
//! \return imagefile metadata
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::metadata
imagefile_impl::get_metadata () const
{
  return mobius::metadata
  {
    {
      "url",
      "URL",
      "std::string",
      url_
    },
    {
      "type",
      "type",
      "std::string",
      get_type ()
    },
    {
      "size",
      "size",
      "size_type",
      std::to_string (get_size ()) + " bytes"
    },
    {
      "sectors",
      "number of sectors",
      "size_type",
      std::to_string (get_sectors ())
    },
    {
      "sector_size",
      "sector size",
      "size_type",
      std::to_string (get_sector_size ()) + " bytes"
    },
    {
      "version",
      "driver version",
      "std::uint32_t",
      std::to_string (get_version ())
    },
    {
      "signature",
      "file signature",
      "std::uint32_t",
      "0x" + mobius::string::to_hex (get_signature (), 8)
    },
    {
      "encryption_algorithm",
      "encryption algorithm",
      "std::uint32_t",
      encryption_algorithm_to_string (get_encryption_algorithm ())
    },
    {
      "encryption_key",
      "encryption key",
      "mobius::bytearray",
      get_encryption_key ().to_hexstring ()
    },
    {
      "device_id",
      "device ID",
      "std::string",
      get_device_id ()
    },
    {
      "last_metadata_time",
      "last metadata modification date/time",
      "mobius::datetime::datetime",
      to_string (get_last_metadata_time ())
    },
    {
      "last_modification_time",
      "last modification date/time",
      "mobius::datetime::datetime",
      to_string (get_last_modification_time ())
    },
    {
      "last_access_time",
      "last access date/time",
      "mobius::datetime::datetime",
      to_string (get_last_access_time ())
    },
  };
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief load metadata
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
imagefile_impl::_load_metadata () const
{
  if (metadata_loaded_)
    return;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // parse header data
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  mobius::io::file f (url_);

  if (!f || !f.exists ())
    return;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // file metadata
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  last_metadata_time_ = f.get_last_metadata_time ();
  last_modification_time_ = f.get_last_modification_time ();
  last_access_time_ = f.get_last_access_time ();

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // header metadata
  // 0x00      uint32_t        signature
  // *0x04      uint32_t        ??
  // 0x08      uint64_t        size
  // 0x10      uint32_t        flag: password checked (0x00 = no, 0x01 = yes)
  // 0x14      s (16)          device ID
  // *0x24      s (16)          some hash
  // 0x34      uint32_t        sector size
  // 0x38      uint32_t        encryption algorithm
  // 0x3c      var             encryption key
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  auto reader = f.new_reader ();
  const mobius::bytearray data = read_header_data (reader);
  mobius::decoder::data_decoder decoder (data);

  decoder.skip (8192);
  signature_ = decoder.get_uint32_le ();
  decoder.skip (4);                                             // unknown1
  size_ = decoder.get_uint64_le ();
  decoder.skip (4);                                             // flag: password checked (0x00 = no, 0x01 = yes)
  device_id_ = decoder.get_string_by_size (16);
  auto password_hash = decoder.get_bytearray_by_size (16);
  sector_size_ = decoder.get_uint32_le ();
  encryption_algorithm_ = decoder.get_uint32_le ();

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // version
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  if (signature_ == 0x71049298)
    {
      version_ = 0;
      sector_size_ = 512;
    }

  else if (signature_ == 0x3258a112)
    version_ = 1;

  else if (signature_ == 0x5a62b2a7)
    version_ = 2;

  else
    version_ = 0xffffffff;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // sectors
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  if (sector_size_ > 0)
    sectors_ = (size_ + sector_size_ - 1) / sector_size_;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // encryption key
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  if (version_ == 0)
    {
      //! \todo V0 keys are hardcoded
    }

  else
    {
      if (encryption_algorithm_ == 1)               // AES-128
        encryption_key_ = decoder.get_bytearray_by_size (16);

      else if (encryption_algorithm_ == 2)          // AES-256
        encryption_key_ = decoder.get_bytearray_by_size (32);

      else if (encryption_algorithm_ == 3)          // Blowfish-448
        encryption_key_ = decoder.get_bytearray_by_size (16);

      else if (encryption_algorithm_)
        throw std::runtime_error (MOBIUS_EXCEPTION_MSG ("unknown/invalid encryption algorithm"));
    }

  metadata_loaded_ = true;
}

} // namespace msr
} // namespace imagefile
} // namespace mobius
