#include "../../pro/registry.cpp"

#define PYTHON3_DLL "python3.dll"

//-------------------------------------------------------------------------
#define IDA_ADDLIB_VALUE "Python3TargetDLL"

#define PYTHON_INSTALLS_KEY L"Software\\Python"
#define PYTHON_INSTALL_PATH_SUBKEY L"InstallPath"
#define PYTHON_DISPLAY_NAME_SUBKEY L"DisplayName"
#define PYTHON_SYSVER_SUBKEY L"SysVersion"
#define PYTHON_SYSARCH_SUBKEY L"SysArchitecture"
#define PYTHON_INSTALL_PATH_DEFAULT_VALUE L""

//-------------------------------------------------------------------------
static bool read_string(
        qstring *out,
        HKEY key,
        const wchar16_t *value,
        qstring *errbuf=nullptr)
{
  DWORD type;
  DWORD size;
  if ( RegQueryValueExW(key, value, nullptr, &type, nullptr, &size) == ERROR_SUCCESS && type == REG_SZ )
  {
    bytevec_t buf;
    buf.resize(size);
    if ( RegQueryValueExW(key, value, nullptr, &type, buf.begin(), &size) == ERROR_SUCCESS && type == REG_SZ )
      return utf16_utf8(out, (wchar16_t *) buf.begin(), buf.size() / sizeof(wchar16_t));
  }
  if ( errbuf != nullptr )
    errbuf->sprnt("Couldn't query value \"%ls\"\n", value);
  return false;
}

//-------------------------------------------------------------------------
static bool extract_version_from_path(
        pylib_version_t *out,
        const char *_path)
{
  qstring qpath(_path);
  qpath.rtrim('\\');
  const char *path = qpath.c_str();
  const char *p = qbasename(path);
  if ( p == nullptr )
    return false;
  if ( !strnieq(p, "Python", 6) )
    return false;
  p += 6;
  if ( !qisdigit(p[0]) || !qisdigit(p[1]) )
    return false;
  out->raw = p;
  int major = p[0] - '0';
  int minor = p[1] - '0';
  int revision = 0;
  p += 2;
  if ( qisdigit(p[0]) )
  {
    revision = p[0] - '0';
    ++p;
  }
  out->major = major;
  out->minor = minor;
  out->revision = revision;
  out->modifiers = p;
  return true;
}


#pragma comment(lib, "version.lib")
//-------------------------------------------------------------------------
DWORD GetFileVersionNumber(const char *filename, DWORD *pdwMSVer, DWORD *pdwLSVer)
{
  DWORD             dwResult = NOERROR;
  unsigned          uiSize;
  DWORD             dwVerInfoSize;
  DWORD             dwHandle;
  PBYTE             prgbVersionInfo = nullptr;
  VS_FIXEDFILEINFO  *lpVSFixedFileInfo = nullptr;

  DWORD dwMSVer = 0xffffffff;
  DWORD dwLSVer = 0xffffffff;

  qwstring wfilename;
  if ( !utf8_utf16(&wfilename, filename) )
  {
    dwResult = ERROR_INVALID_PARAMETER;
    goto Finish;
  }

  dwVerInfoSize = GetFileVersionInfoSizeW(wfilename.c_str(), &dwHandle);
  if ( dwVerInfoSize != 0 )
  {
    prgbVersionInfo = (PBYTE)qalloc(dwVerInfoSize);
    if ( prgbVersionInfo == nullptr )
    {
      dwResult = ERROR_NOT_ENOUGH_MEMORY;
      goto Finish;
    }

    // Read version stamping info
    if ( GetFileVersionInfoW(wfilename.c_str(), dwHandle, dwVerInfoSize, prgbVersionInfo) )
    {
      // get the value for VS_FIXEDFILEINFO
      if ( VerQueryValueW(prgbVersionInfo, L"\\", (LPVOID*)&lpVSFixedFileInfo, &uiSize) && (uiSize != 0) )
      {
        dwMSVer = lpVSFixedFileInfo->dwFileVersionMS;
        dwLSVer = lpVSFixedFileInfo->dwFileVersionLS;
      }
    }
    else
    {
      dwResult = GetLastError();
      goto Finish;
    }
  }
  else
  {
    dwResult = GetLastError();
  }

  out_verb("%s is version %d.%d.%d.%d\n", filename, HIWORD(dwMSVer), LOWORD(dwMSVer), HIWORD(dwLSVer), LOWORD(dwLSVer));

Finish:
  if ( prgbVersionInfo != nullptr )
    qfree(prgbVersionInfo);
  if ( pdwMSVer != nullptr )
    *pdwMSVer = dwMSVer;
  if ( pdwLSVer != nullptr )
    *pdwLSVer = dwLSVer;

  return dwResult;
}
//-------------------------------------------------------------------------
static bool extract_version_from_dll(
        pylib_version_t *out,
        const char *fname)
{
  DWORD dwMSVer, dwLSVer;
  DWORD res = GetFileVersionNumber(fname, &dwMSVer, &dwLSVer);
  if ( res == NOERROR )
  {
    out->raw.sprnt("%d.%d.%d.%d", HIWORD(dwMSVer), LOWORD(dwMSVer), HIWORD(dwLSVer), LOWORD(dwLSVer));
    // Python DLL versions look like 3.7.4150.1013 -> 3.7.4
    out->major = HIWORD(dwMSVer);
    out->minor = LOWORD(dwMSVer);
    out->revision = HIWORD(dwLSVer) / 1000;
    return true;
  }
  else
  {
    out_verb("error getting  version of \"%s\": %s\n", fname, winerr(res));
    return false;
  }
}

//-------------------------------------------------------------------------
static bool is_python3Y_dll_file_name(const char *fname)
{
  return fname != nullptr
      && strnieq(fname, "python3", 7)
      && qisdigit(fname[7])
      && strieq(get_file_ext(fname), "dll");
}

#include <exehdr.h>

#include "../../ldr/pe/pe.h"
#include "../../ldr/pe/common.cpp"

//-------------------------------------------------------------------------
static bool check_dll_x86_64(const char *path)
{
  linput_t *linput = open_linput(path, /*remote=*/ false);
  if ( linput == nullptr )
    return false;
  linput_janitor_t lj(linput);
  pe_loader_t pl;
  return pl.read_header(linput, /*silent=*/ true)
      && pl.pe.machine == PECPU_AMD64;
}

//-------------------------------------------------------------------------
static bool probe_python_install_dir_from_dll_path(
        qstrvec_t *out_paths,
        pylib_version_t *out_version,
        const char *path,
        qstring *errbuf)
{
  char dir[QMAXPATH];
  if ( !qdirname(dir, sizeof(dir), path) )
  {
    errbuf->sprnt("Couldn't retrieve directory name from \"%s\"", path);
    return false;
  }
  extract_version_from_path(out_version, dir); // not fatal if this fails

  qstring found_python3_dll;
  qstring found_python3Y_dll;
  static const char dll_pattern[] = "python3*.dll";

  char pattern_path[QMAXPATH];
  qmakepath(pattern_path, sizeof(pattern_path), dir, dll_pattern, nullptr);
  qstring verbuf;
  qffblk64_t fb;
  for ( int code = qfindfirst(pattern_path, &fb, 0);
        code == 0;
        code = qfindnext(&fb) )
  {
    const char *ext = get_file_ext(fb.ff_name);
    if ( ext != nullptr && strieq(ext, "dll") )
    {
      char dll_path[QMAXPATH];
      qmakepath(dll_path, sizeof(dll_path), dir, fb.ff_name, nullptr);
      pylib_version_t tmp;
      if ( extract_version_from_dll(&tmp, path) )
        *out_version = tmp;

      if ( out_version->major < args.major_version || out_version->minor < args.minor_version )
      {
        out_verb("Unsupported python version %s (%d.%d+ is required)",
                      out_version->str(&verbuf), args.major_version, args.minor_version);
        continue;
      }

      out_verb("Found: \"%s\" (version: %s)\n", dll_path, out_version->str(&verbuf));
      out_paths->push_back(dll_path);

      if ( found_python3_dll.empty() && strieq(fb.ff_name, PYTHON3_DLL) )
        found_python3_dll = dll_path;
      if ( found_python3Y_dll.empty() && is_python3Y_dll_file_name(fb.ff_name) )
        found_python3Y_dll = dll_path;
    }
  }

  if ( found_python3_dll.empty() )
  {
    errbuf->sprnt("No \"" PYTHON3_DLL "\" >= %d.%d found in directory \"%s\"",
      args.major_version, args.minor_version, dir);
    return false;
  }
  if ( found_python3Y_dll.empty() )
  {
    errbuf->sprnt("No \"python3[0-9].dll\" >= %d.%d found in directory \"%s\"",
      args.major_version, args.minor_version, dir);
    return false;
  }

  if ( !check_dll_x86_64(found_python3_dll.c_str())
    || !check_dll_x86_64(found_python3Y_dll.c_str()) )
  {
    errbuf->sprnt("DLLs in directory \"%s\" do not have the x86_64 architecture", dir);
    return false;
  }

  return true;
}

static bool has_appx_path(qstrvec_t paths)
{
  static qstring appx_path;
  if ( appx_path.empty() )
  {
    HKEY hkey;
    bool ok = false;
    if ( RegOpenKeyExW(HKEY_LOCAL_MACHINE, L"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Appx", 0, KEY_READ, &hkey) == ERROR_SUCCESS )
    {
      ok = read_string(&appx_path, hkey, L"PackageRoot");
      RegCloseKey(hkey);
    }
    if ( !ok )
    {
      // no Appx support, set to dummy value which doesn't occur in paths
      appx_path = "<none>";
    }
  }
  for ( const qstring &path : paths )
    if ( path.find(appx_path) != qstring::npos )
      return true;

  return false;
}
//-------------------------------------------------------------------------
// ignore known bad Pythons:
// 3.8.0 release (https://bugs.python.org/issue37633)
// Anaconda 2019.10, 2020.02 and 2020.11 (https://github.com/ContinuumIO/anaconda-issues/issues/11374)
// AppStore Python on Windows 10 (dll can't be loaded from outside of Appx package)
static bool bad_entry(const pylib_entry_t &e)
{
  if ( e.version.major < args.major_version || e.version.minor < args.minor_version )
  {
    qstring verbuf;
    out("Unsupported python version %s (%d.%d+ is required)\n",
        e.version.str(&verbuf), args.major_version, args.minor_version);
        return true;
  }

  if ( e.version.major == 3
    && e.version.minor == 8
    && e.version.revision == 0 )
  {
    out("Ignoring unusable Python 3.8.0 \"%s\"\n", !e.paths.empty() ? e.paths[0].c_str() : "?");
    return true;
  }
  if ( e.display_name.starts_with("Anaconda 20" ) )
  {
    out("Ignoring unusable %s \"%s\"\n", e.display_name.c_str(), !e.paths.empty() ? e.paths[0].c_str() : "?");
    return true;
  }
  if ( has_appx_path(e.paths) )
  {
    out("Ignoring unusable AppStore Python \"%s\"\n", !e.paths.empty() ? e.paths[0].c_str() : "?");
    return true;
  }
  return false;
}

//-------------------------------------------------------------------------
static void remove_bad_entries(pylib_entries_t *result)
{
  auto p  = result->entries.begin();
  while ( p != result->entries.end() )
  {
    if ( bad_entry(*p) )
      p = result->entries.erase(p);
    else
      ++p;
  }
}

//-------------------------------------------------------------------------
// given a key to a path like HKLM\SOFTWARE\Python\PythonCore,
// - enumerate subkeys and their check InstallPath subkey, using the default value as the directory to the installation
// e.g.
// HKLM\SOFTWARE\Python\PythonCore\3.6\InstallPath -> (Default) = C:\Python36\
// - check for python3.dll and python3Y.dll in it and add them to 'result'
static void enum_python_key(pylib_entries_t *result, const HKEY hkey, qstring *_errbuf, qstring *_verbuf)
{
  qstring &errbuf = *_errbuf;
  qstring &verbuf = *_verbuf;
  int index = 0;
  WCHAR subkey[MAXSTR];
  qstring displayname;
  if ( read_string(&displayname, hkey, PYTHON_DISPLAY_NAME_SUBKEY) )
  {
    out("Checking installs from \"%s\"\n", displayname.c_str());
  }
  while ( true )
  {
    DWORD subkey_sz = qnumber(subkey);
    if ( RegEnumKeyExW(hkey, index++, subkey, &subkey_sz,
                       nullptr, nullptr, nullptr, nullptr) != ERROR_SUCCESS )
    {
      // no more installs
      break;
    }
    HKEY ihkey;
    if ( RegOpenKeyExW(hkey, subkey, 0, KEY_READ, &ihkey) == ERROR_SUCCESS )
    {
      out_verb("Opened \"%ls\"\n", subkey);
      // opened an install. get its version from SysVersion value
      qstring sysver;
      pylib_version_t version;
      bool ok = read_string(&sysver, ihkey, PYTHON_SYSVER_SUBKEY);
      ok = ok && parse_python_version_str(&version, sysver.c_str());
      qstring arch;
      if ( read_string(&arch, ihkey, PYTHON_SYSARCH_SUBKEY) && arch != "64bit" )
        ok = false;
      if ( ok )
      {
        if ( read_string(&displayname, ihkey, PYTHON_DISPLAY_NAME_SUBKEY) )
        {
          out("Checking \"%s\" (%s)\n", displayname.c_str(), sysver.c_str());
        }
        HKEY vhkey;
        if ( RegOpenKeyExW(ihkey, PYTHON_INSTALL_PATH_SUBKEY, 0, KEY_READ, &vhkey) == ERROR_SUCCESS )
        {
          qstring install_path;
          if ( read_string(&install_path, vhkey, PYTHON_INSTALL_PATH_DEFAULT_VALUE, &errbuf) )
          {
            char probe[QMAXPATH];
            qmakepath(probe, sizeof(probe), install_path.c_str(), PYTHON3_DLL, nullptr);
            pylib_version_t unused;
            qstrvec_t paths;
            if ( probe_python_install_dir_from_dll_path(
              &paths,
              &version,
              probe,
              &errbuf) )
            {
              out("Found: \"%s\" (version: %s)\n", install_path.c_str(), version.str(&verbuf));
              pylib_entry_t &e = result->add_entry(version, paths);
              e.display_name = displayname;
            }
            else
            {
              out_verb("Ignoring directory \"%s\": %s\n", install_path.c_str(), errbuf.c_str());
            }
          }
          else
          {
            out_verb("Couldn't query \"%s\"'s value \"%ls\": %s\n",
              subkey,
              PYTHON_INSTALL_PATH_DEFAULT_VALUE,
              errbuf.c_str());
          }
          RegCloseKey(vhkey);
        }
        else
        {
          out("Couldn't open \"%s\"\n", subkey);
        }
      }
      else
      {
        out_verb("Not a 64-bit Python 3.x or invalid version info, skipping\n");
      }

      RegCloseKey(ihkey);
    }
  }
}
//-------------------------------------------------------------------------
void pyver_tool_t::do_find_python_libs(pylib_entries_t *result) const
{
  qstring errbuf;
  qstring verbuf;

  //
  // Enumerate:
  //   * HKEY_LOCAL_MACHINE\Software\Python\PythonCore\* versions
  //   * HKEY_CURRENT_USER\Software\Python\PythonCore\* versions
  //
  static const HKEY top_keys[] =
    {
      HKEY_LOCAL_MACHINE,
      HKEY_CURRENT_USER,
    };

  static const char * knames[] =
    {
      "HLKM",
      "HKCU",
    };
  for ( size_t i = 0; i < qnumber(top_keys); ++i )
  {
    HKEY hkey_python;
    out_verb("Searching for subkeys of \"%s\\%ls\"\n", knames[i], PYTHON_INSTALLS_KEY);
    if ( RegOpenKeyExW(top_keys[i], PYTHON_INSTALLS_KEY, 0, KEY_READ, &hkey_python) == ERROR_SUCCESS )
    {
      WCHAR subkey[MAXSTR];
      int index = 0;
      while ( true )
      {
        DWORD subkey_sz = qnumber(subkey);
        if ( RegEnumKeyExW(hkey_python, index++, subkey, &subkey_sz,
          nullptr, nullptr, nullptr, nullptr) != ERROR_SUCCESS )
        {
          break;
        }
        HKEY hkey;
        if ( RegOpenKeyExW(hkey_python, subkey, 0, KEY_READ, &hkey) == ERROR_SUCCESS )
        {
          out_verb("Found \"%s\\%ls\\%ls\"\n", knames[i], PYTHON_INSTALLS_KEY, subkey);
          enum_python_key(result, hkey, &errbuf, &verbuf);
          RegCloseKey(hkey);
        }
      }
      RegCloseKey(hkey_python);
    }
  }

  remove_bad_entries(result);

  //
  // See if we already have one registered for IDA
  //
  qstring existing;
  if ( reg_read_string(&existing, IDA_ADDLIB_VALUE) )
  {
    out_verb("Previously-used DLL: \"%s\"\n", existing.c_str());
    pylib_version_t version;
    qstrvec_t paths;
    if ( probe_python_install_dir_from_dll_path(
                 &paths,
                 &version,
                 existing.c_str(),
                 &errbuf) )
    {
      out("IDA previously used: \"%s\" (guessed version: %s). "
          "Making this the preferred version.\n",
          existing.c_str(), version.str(&verbuf));
      // do we have it in the list?
      bool found = false;
      for ( pylib_entry_t &e : result->entries )
      {
        if ( e.paths.has(existing) )
        {
          found = true;
          e.preferred = true;
        }
      }
      if ( !found )
      {
        // add a new one
        pylib_entry_t e(version);
        e.paths.swap(paths);
        e.preferred = true;
        result->entries.push_back(e);
      }
    }
    else
    {
      out_verb("Ignoring directory \"%s\": %s\n",
               existing.c_str(), errbuf.c_str());
    }
  }
  else
  {
    out("no \"%ls\" value found\n", IDA_ADDLIB_VALUE);
  }
}

//-------------------------------------------------------------------------
bool pyver_tool_t::do_path_to_pylib_entry(
        pylib_entry_t *entry,
        const char *path,
        qstring *errbuf) const
{
  return probe_python_install_dir_from_dll_path(&entry->paths, &entry->version, path, errbuf);
}

//-------------------------------------------------------------------------
bool pyver_tool_t::do_apply_version(
        const pylib_entry_t &entry,
        qstring *errbuf) const
{
  qstring replacement;
  bool ok = false;
  for ( const auto &path : entry.paths )
  {
    const char *candidate = qbasename(path.c_str());
    if ( is_python3Y_dll_file_name(candidate) )
    {
      replacement = candidate;
      reg_write_string(IDA_ADDLIB_VALUE, path.c_str());
      return true;
    }
  }
  if ( !ok )
  {
    errbuf->sprnt("Couldn't find a suitable python3Y.dll file or errror writing to registry");
    return false;
  }
  return true;
}
