/* Copyright 2017 - 2026 R. Thomas
 * Copyright 2017 - 2026 Quarkslab
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <algorithm>
#include <iterator>

#ifndef MBEDTLS_DECLARE_PRIVATE_IDENTIFIERS
  #define MBEDTLS_DECLARE_PRIVATE_IDENTIFIERS
#endif

#include "logging.hpp"
#include "psa/crypto_builtin_primitives.h"

#include "LIEF/PE/utils.hpp"
#include "LIEF/PE/Binary.hpp"
#include "LIEF/PE/Import.hpp"
#include "LIEF/PE/Section.hpp"
#include "LIEF/PE/ImportEntry.hpp"
#include "LIEF/BinaryStream/SpanStream.hpp"
#include "LIEF/BinaryStream/FileStream.hpp"
#include "PE/Structures.hpp"

#include "LIEF/utils.hpp"

#include "utils/ordinals_lookup_tables/libraries_table.hpp"
#include "utils/ordinals_lookup_tables_std/libraries_table.hpp"

#include "hash_stream.hpp"
#include "internal_utils.hpp"


namespace LIEF::PE {

static constexpr size_t SIZEOF_OPT_HEADER_32 = 0xE0;
static constexpr size_t SIZEOF_OPT_HEADER_64 = 0xF0;

inline std::string to_lower(std::string str) {
  std::string lower = str;
  std::transform(str.begin(), str.end(), lower.begin(), ::tolower);
  return lower;
}

bool is_pe(BinaryStream& stream) {
  using signature_t = std::array<char, sizeof(details::PE_Magic)>;
  ScopedStream scoped(stream, 0);
  if (auto dos_header = scoped->read<details::pe_dos_header>()) {
    if (dos_header->Magic != /* MZ */ 0x5a4d) {
      return false;
    }
    scoped->setpos(dos_header->AddressOfNewExeHeader);
    if (auto res_sig = scoped->read<signature_t>()) {
      const auto signature = *res_sig;
      return std::equal(std::begin(signature), std::end(signature),
                        std::begin(details::PE_Magic));
    }
  }
  return false;
}

bool is_pe(const std::string& file) {
  if (auto stream = FileStream::from_file(file)) {
    return is_pe(*stream);
  }
  return false;
}


bool is_pe(const std::vector<uint8_t>& raw) {
  if (auto stream = SpanStream::from_vector(raw)) {
    return is_pe(*stream);
  }
  return false;
}


result<PE_TYPE> get_type_from_stream(BinaryStream& stream) {
  const uint64_t cpos = stream.pos();
  stream.setpos(0);
  auto dos_hdr = stream.read<details::pe_dos_header>();
  if (!dos_hdr) {
    LIEF_ERR("Failed to read the DOS header structure");
    return make_error_code(dos_hdr.error());
  }
  stream.setpos(dos_hdr->AddressOfNewExeHeader);
  if (!stream.can_read<details::pe_header>()) {
    LIEF_ERR("Failed to read the PE header");
    return make_error_code(lief_errors::read_error);
  }

  auto header = stream.read<details::pe_header>();
  if (!header) {
    LIEF_ERR("Failed to read the PE header");
    return make_error_code(header.error());
  }

  const size_t sizeof_opt_header = header->SizeOfOptionalHeader;
  if (sizeof_opt_header != SIZEOF_OPT_HEADER_32 &&
      sizeof_opt_header != SIZEOF_OPT_HEADER_64)
  {
    LIEF_WARN(
        "SizeOfOptionalHeader seems invalid: {:#x}. Expecting {:#x} or {:#x}",
        sizeof_opt_header, SIZEOF_OPT_HEADER_32, SIZEOF_OPT_HEADER_64
    );
  }

  auto opt_hdr = stream.read<details::pe32_optional_header>();
  if (!opt_hdr) {
    LIEF_ERR("Failed to read the PE optional header");
    return make_error_code(opt_hdr.error());
  }
  const auto type = static_cast<PE_TYPE>(opt_hdr->Magic);
  stream.setpos(cpos); // Restore the original position

  /*
   * We first try to determine PE's type from the OptionalHeader's Magic.
   * In some cases (cf. https://github.com/lief-project/LIEF/issues/644),
   * these magic bytes can be "corrupted" without afecting the executing.
   * So the second test to determine the PE's type is to check the value
   * of SizeOfOptionalHeader we should match sizeof(pe32_optional_header)
   * for a 32 bits PE or sizeof(pe64_optional_header) for a 64 bits PE file
   */
  if (type == PE_TYPE::PE32 || type == PE_TYPE::PE32_PLUS) {
    return type;
  }

  if (sizeof_opt_header == SIZEOF_OPT_HEADER_32) {
    return PE_TYPE::PE32;
  }

  if (sizeof_opt_header == SIZEOF_OPT_HEADER_64) {
    return PE_TYPE::PE32_PLUS;
  }

  LIEF_ERR("Failed to determine PE type (PE32 / PE32+)");
  return make_error_code(lief_errors::file_format_error);
}

result<PE_TYPE> get_type(const std::string& file) {
  if (auto stream = FileStream::from_file(file)) {
    return get_type_from_stream(*stream);
  }
  return make_error_code(lief_errors::file_format_error);
}

result<PE_TYPE> get_type(const std::vector<uint8_t>& raw) {
  if (auto stream = SpanStream::from_vector(raw)) {
    return get_type_from_stream(*stream);
  }
  return make_error_code(lief_errors::file_format_error);
}


std::string get_imphash_std(const Binary& binary) {
  static const std::set<std::string> ALLOWED_EXT = {"dll", "ocx", "sys"};
  std::vector<uint8_t> md5_buffer(16);
  if (!binary.has_imports()) {
    return "";
  }
  std::string lstr;
  bool first_entry = true;
  hashstream hs(hashstream::HASH::MD5);
  for (const Import& imp : binary.imports()) {
    Import resolved = imp;
    if (auto resolution =
            resolve_ordinals(imp, /* strict */ false, /* use_std */ true))
    {
      resolved = std::move(*resolution);
    }

    size_t ext_idx = resolved.name().find_last_of('.');
    std::string name = resolved.name();
    std::string ext;
    if (ext_idx != std::string::npos) {
      ext = to_lower(resolved.name().substr(ext_idx + 1));
    }
    if (ALLOWED_EXT.find(ext) != ALLOWED_EXT.end()) {
      name = name.substr(0, ext_idx);
    }

    std::string entries_string;
    for (const ImportEntry& e : resolved.entries()) {
      std::string funcname;
      if (e.is_ordinal()) {
        funcname = "ord" + std::to_string(e.ordinal());
      } else {
        funcname = e.name();
      }

      if (!entries_string.empty()) {
        entries_string += ',';
      }
      entries_string.append(name).append(".").append(funcname);
    }
    if (!first_entry) {
      lstr += ',';
    } else {
      first_entry = false;
    }
    lstr += to_lower(entries_string);

    // Use write(uint8_t*, size_t) instead of write(const std::string&) to avoid
    // including the null terminator
    hs.write(reinterpret_cast<const uint8_t*>(lstr.data()), lstr.size());
    lstr.clear();
  }

  return hex_dump(hs.raw(), "");
}


std::string get_imphash_lief(const Binary& binary) {
  std::vector<uint8_t> md5_buffer(16);
  if (!binary.has_imports()) {
    return std::to_string(0);
  }

  Binary::it_const_imports imports = binary.imports();

  std::string import_list;
  for (const Import& imp : imports) {
    Import resolved = imp;
    if (auto resolution = resolve_ordinals(imp)) {
      resolved = std::move(*resolution);
    }

    size_t ext_idx = resolved.name().find_last_of('.');
    std::string name_without_ext = resolved.name();

    if (ext_idx != std::string::npos) {
      name_without_ext = name_without_ext.substr(0, ext_idx);
    }

    std::string entries_string;
    for (const ImportEntry& e : resolved.entries()) {
      if (e.is_ordinal()) {
        entries_string.append(name_without_ext)
            .append(".#")
            .append(std::to_string(e.ordinal()));
      } else {
        entries_string.append(name_without_ext).append(".").append(e.name());
      }
    }
    import_list += to_lower(entries_string);
  }

  std::sort(import_list.begin(), import_list.end(), std::less<>());

  mbedtls_md5(reinterpret_cast<const uint8_t*>(import_list.data()),
              import_list.size(), md5_buffer.data());
  return hex_dump(md5_buffer, "");
}

std::string get_imphash(const Binary& binary, IMPHASH_MODE mode) {
  switch (mode) {
    case IMPHASH_MODE::LIEF:
    {
      return get_imphash_lief(binary);
    }
    case IMPHASH_MODE::PEFILE:
    {
      return get_imphash_std(binary);
    }
  }
  return "";
}

result<Import> resolve_ordinals(const Import& import, bool strict, bool use_std) {
  static const std::unordered_map<std::string, const char* (*)(uint32_t)>
      ordinals_library_tables = {
          {"kernel32.dll", &kernel32_dll_lookup},
          {"ntdll.dll", &ntdll_dll_lookup},
          {"advapi32.dll", &advapi32_dll_lookup},
          {"msvcp110.dll", &msvcp110_dll_lookup},
          {"msvcp120.dll", &msvcp120_dll_lookup},
          {"msvcr100.dll", &msvcr100_dll_lookup},
          {"msvcr110.dll", &msvcr110_dll_lookup},
          {"msvcr120.dll", &msvcr120_dll_lookup},
          {"user32.dll", &user32_dll_lookup},
          {"comctl32.dll", &comctl32_dll_lookup},
          {"ws2_32.dll", &ws2_32_dll_lookup},
          {"shcore.dll", &shcore_dll_lookup},
          {"oleaut32.dll", &oleaut32_dll_lookup},
          {"mfc42u.dll", &mfc42u_dll_lookup},
          {"shlwapi.dll", &shlwapi_dll_lookup},
          {"gdi32.dll", &gdi32_dll_lookup},
          {"shell32.dll", &shell32_dll_lookup},
  };

  static const std::unordered_map<std::string, const char* (*)(uint32_t)>
      imphashstd_ordinals_library_tables = {
          {"ws2_32.dll", &ws2_32_dll_lookup},
          {"wsock32.dll", &ws2_32_dll_lookup},
          {"oleaut32.dll", &oleaut32_dll_lookup},
  };

  using ordinal_resolver_t = const char* (*)(uint32_t);

  Import::it_const_entries entries = import.entries();

  if (std::all_of(entries.begin(), entries.end(),
                  [](const ImportEntry& entry) { return !entry.is_ordinal(); }))
  {
    // LIEF_DEBUG("All imports use name. No ordinal!");
    return import;
  }

  std::string name = to_lower(import.name());

  ordinal_resolver_t ordinal_resolver = nullptr;
  if (use_std) {
    auto it = imphashstd_ordinals_library_tables.find(name);
    if (it != imphashstd_ordinals_library_tables.end()) {
      ordinal_resolver = it->second;
    }
  } else {
    auto it = ordinals_library_tables.find(name);
    if (it != ordinals_library_tables.end()) {
      ordinal_resolver = it->second;
    }
  }

  if (ordinal_resolver == nullptr) {
    if (strict) {
      return make_error_code(lief_errors::not_implemented);
    }
    LIEF_DEBUG("Ordinal lookup table for '{}' not implemented", name);
    return import;
  }

  Import resolved_import = import;
  for (ImportEntry& entry : resolved_import.entries()) {
    if (entry.is_ordinal()) {
      LIEF_DEBUG("Dealing with: {}", to_string(entry));
      const char* entry_name =
          ordinal_resolver(static_cast<uint32_t>(entry.ordinal()));
      if (entry_name == nullptr) {
        if (strict) {
          return make_error_code(lief_errors::not_supported);
        }
        LIEF_DEBUG("Failed to resolve ordinal: #{:d}", entry.ordinal());
        continue;
      }
      entry.data(0);
      entry.name(entry_name);
    }
  }

  return resolved_import;
}
ALGORITHMS algo_from_oid(const std::string& oid) {
  static const std::unordered_map<std::string, ALGORITHMS> OID_MAP = {
      {"2.16.840.1.101.3.4.2.3", ALGORITHMS::SHA_512},
      {"2.16.840.1.101.3.4.2.2", ALGORITHMS::SHA_384},
      {"2.16.840.1.101.3.4.2.1", ALGORITHMS::SHA_256},
      {"1.3.14.3.2.26", ALGORITHMS::SHA_1},

      {"1.2.840.113549.2.5", ALGORITHMS::MD5},
      {"1.2.840.113549.2.4", ALGORITHMS::MD4},
      {"1.2.840.113549.2.2", ALGORITHMS::MD2},

      {"1.2.840.113549.1.1.1", ALGORITHMS::RSA},
      {"1.2.840.10045.2.1", ALGORITHMS::EC},

      {"1.2.840.113549.1.1.4", ALGORITHMS::MD5_RSA},
      {"1.2.840.10040.4.3", ALGORITHMS::SHA1_DSA},
      {"1.2.840.113549.1.1.5", ALGORITHMS::SHA1_RSA},
      {"1.2.840.113549.1.1.11", ALGORITHMS::SHA_256_RSA},
      {"1.2.840.113549.1.1.12", ALGORITHMS::SHA_384_RSA},
      {"1.2.840.113549.1.1.13", ALGORITHMS::SHA_512_RSA},
      {"1.2.840.10045.4.1", ALGORITHMS::SHA1_ECDSA},
      {"1.2.840.10045.4.3.2", ALGORITHMS::SHA_256_ECDSA},
      {"1.2.840.10045.4.3.3", ALGORITHMS::SHA_384_ECDSA},
      {"1.2.840.10045.4.3.4", ALGORITHMS::SHA_512_ECDSA},
  };


  const auto it = OID_MAP.find(oid.c_str());
  if (it == OID_MAP.end()) {
    return ALGORITHMS::UNKNOWN;
  }
  return it->second;
}

}