// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// 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 "turing.h"
#include <mobius/core/application.h>
#include <mobius/database/connection_pool.h>
#include <mobius/string_functions.h>
#include <string>
#include <tuple>
#include <set>

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Usage:
//
// mobius::turing::turing turing;
// auto connection = turing.new_connection ()    // only when running in a new thread!!!
// auto transaction = turing.new_transaction ()  // only when modifying data
// ...
// operations (set_hash, remove_hashes, get_hash_password, get_hashes, ...)
// ...
// transaction.commit ();                        // only when have modified data
// connection.release ();                        // only when running in a new thread!!!
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

namespace mobius
{
namespace turing
{
namespace
{
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
// Constants
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static const std::string LM_NULL = "aad3b435b51404ee";
static const std::string NT_NULL = "31d6cfe0d16ae931b73c59d7e0c089c0";
static constexpr int SCHEMA_VERSION = 2;

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Turing data structure
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
static class data
{
public:
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  //! \brief Get connection pool
  //! \return Connection pool object
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  mobius::database::connection_pool
  get_pool ()
  {
    if (!data_loaded_)
      {
        _load_data ();
        data_loaded_ = true;
      }

    return pool_;
  }

private:
  mobius::database::connection_pool pool_;
  bool data_loaded_ = false;

  void _load_data ();
} data_;

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Load data structure
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
data::_load_data ()
{
  mobius::core::application app;
  pool_.set_path (app.get_config_path ("turing.sqlite"));

  auto db = pool_.get_database ();
  db.execute ("PRAGMA foreign_keys = OFF;");

  auto transaction = db.new_transaction ();

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // create table 'meta'
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  db.execute (
    "CREATE TABLE IF NOT EXISTS meta ("
                 "key TEXT PRIMARY KEY,"
               "value TEXT NOT NULL"
    ");");

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // set schema version
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  int current_version = 1;

  auto stmt = db.new_statement (
                "SELECT value "
                  "FROM meta "
                 "WHERE key = 'version'");

  if (stmt.fetch_row ())
    {
      current_version = stmt.get_column_int (0);

      if (current_version < SCHEMA_VERSION)
        {
          // update schema version
          stmt = db.new_statement (
                   "UPDATE meta "
                      "SET value = ? "
                    "WHERE key = 'version'");

          stmt.bind (1, SCHEMA_VERSION);
          stmt.execute ();
        }
    }

  else
    {
      // insert 'version' metadata
      stmt = db.new_statement (
               "INSERT INTO meta "
                    "VALUES ('version', ?)");

      stmt.bind (1, SCHEMA_VERSION);
      stmt.execute ();
    }

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // if 'hash' table is deprecated, save data and drop table
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  std::set <std::tuple <std::string, std::string, std::string>> data;

  if (db.table_has_column ("hash", "password_status"))
    {
      db.execute ("ALTER TABLE hash RENAME TO hash_old");

      stmt = db.new_statement (
                  "SELECT type, value, password "
                    "FROM hash_old "
                   "WHERE password IS NOT NULL");

      while (stmt.fetch_row ())
        {
          const std::string hash_type = stmt.get_column_string (0);
          const std::string hash_value = stmt.get_column_string (1);
          const std::string password = stmt.get_column_string (2);

          if (hash_type == "lm")
            {
              data.emplace (hash_type,
                            hash_value.substr (0, 16),
                            mobius::string::toupper (password.substr (0, 7)));

              if (password.length () > 7)
                data.emplace (hash_type,
                              hash_value.substr (16),
                              mobius::string::toupper (password.substr (7)));
            }

          else
            data.emplace (hash_type, hash_value, password);
        }

      db.execute ("DROP TABLE hash_args");
      db.execute ("DROP TABLE hash_old");
    }

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // create 'hash' table
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  db.execute ("CREATE TABLE IF NOT EXISTS hash"
                             "(type TEXT NOT NULL,"
                             "value TEXT NOT NULL,"
                          "password TEXT NOT NULL,"
                                        "PRIMARY KEY (type, value))");

  stmt = db.new_statement ("INSERT INTO hash VALUES (?, ?, ?)");

  for (const auto& row : data)
    {
      stmt.bind (1, std::get<0> (row));
      stmt.bind (2, std::get<1> (row));
      stmt.bind (3, std::get<2> (row));
      stmt.execute ();
    }

  transaction.commit ();
  db.execute ("PRAGMA foreign_keys = ON;");
}

} // namespace

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Get database object for current thread
//! \return database object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::database::database
get_database ()
{
  auto pool = data_.get_pool ();
  return pool.get_database ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new connection for Turing database
//! \return new connection object
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::database::connection
turing::new_connection ()
{
  auto pool = data_.get_pool ();
  return pool.acquire ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Create new transaction for Turing database
//! \return new database transaction
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
mobius::database::transaction
turing::new_transaction ()
{
  auto db = get_database ();
  return db.new_transaction ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Check if hash is set
//! \param hash_type Hash type
//! \param hash_value Hash value
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
bool
turing::has_hash (
    const std::string& hash_type,
    const std::string& hash_value)
{
  auto db = get_database ();

  auto stmt = db.new_statement (
                "SELECT 1 "
                  "FROM hash "
                 "WHERE type = ? "
                   "AND value = ?");

  stmt.bind (1, hash_type);
  stmt.bind (2, hash_value);

  return bool (stmt.fetch_row ());
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Set hash
//! \param hash_type Hash type
//! \param hash_value Hash value
//! \param password Password as UTF-8 string
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
turing::set_hash (
    const std::string& hash_type,
    const std::string& hash_value,
    const std::string& password)
{
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // handle LM hashes as two separated halves
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  if (hash_type == "lm" && hash_value.length () > 16)
    {
      set_hash (hash_type,
                hash_value.substr (0, 16),
                mobius::string::toupper (password.substr (0, 7)));

      if (password.length () > 7)
        set_hash (hash_type,
                  hash_value.substr (16),
                  mobius::string::toupper (password.substr (7)));
    }

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // if hash already exists then return
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  if (has_hash (hash_type, hash_value))
    return;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // insert hash
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  auto db = get_database ();

  auto stmt = db.new_statement (
                "INSERT INTO hash "
                     "VALUES (?, ?, ?)");

  stmt.bind (1, hash_type);
  stmt.bind (2, hash_value);
  stmt.bind (3, password);
  stmt.execute ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Get hash password
//! \param hash_type Hash type
//! \param hash_value Hash value
//! \return status, password
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
std::pair <turing::pwd_status, std::string>
turing::get_hash_password (
    const std::string& hash_type,
    const std::string& hash_value) const
{
  std::string password;
  pwd_status status = pwd_status::not_found;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // test empty values
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  if (hash_type == "lm" && hash_value == LM_NULL)
    status = pwd_status::found;

  else if (hash_type == "nt" && hash_value == NT_NULL)
    status = pwd_status::found;

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // handle LM hashes as two separated halves
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  else if (hash_type == "lm" && hash_value.length () > 16)
    {
      auto lm_1 = get_hash_password (hash_type, hash_value.substr (0, 16));
      auto lm_2 = get_hash_password (hash_type, hash_value.substr (16));
      
      if (lm_1.first == pwd_status::found && lm_2.first == pwd_status::found)
        {
          password = lm_1.second + lm_2.second;
          status = pwd_status::found;
        }

      else if (lm_1.first == pwd_status::found)
        {
          password = lm_1.second + "???????";
          status = pwd_status::lm_1_found;
        }

      else if (lm_2.first == pwd_status::found)
        {
          password = "???????" + lm_2.second;
          status = pwd_status::lm_2_found;
        }
    }

  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  // select hash from table
  // =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
  else
    {
      auto db = get_database ();

      auto stmt = db.new_statement (
                "SELECT password "
                  "FROM hash "
                 "WHERE type = ? "
                   "AND value = ?");

      stmt.bind (1, hash_type);
      stmt.bind (2, hash_value);
      
      if (stmt.fetch_row ())
        {
          password = stmt.get_column_string (0);
          status = pwd_status::found;
        }
    }

  return std::make_pair (status, password);
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Remove hashes
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
void
turing::remove_hashes ()
{
  auto db = get_database ();
  auto stmt = db.new_statement ("DELETE FROM hash");

  stmt.execute ();
}

// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
//! \brief Get hashes
//! \return Tuples <type, value, password>
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
turing::hashlist_type
turing::get_hashes () const
{
  auto db = get_database ();

  auto stmt = db.new_statement (
                "SELECT type, value, password "
                  "FROM hash "
              "ORDER BY type, value");

  hashlist_type hashes;

  while (stmt.fetch_row ())
    {
      const std::string hash_type = stmt.get_column_string (0);
      const std::string hash_value = stmt.get_column_string (1);
      const std::string password = stmt.get_column_string (2);

      hashes.emplace_back (hash_type, hash_value, password);
    }

  return hashes;
}

} // namespace turing
} // namespace mobius
