/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "stdafx.h"
#include "InstallerLauncher.h"
#include <string>
#include <sstream>
#include <fstream>
#include <algorithm>
#include <vector>
#include <iostream>
#include <locale>
#include <codecvt>
#include <array>
#include <memory>
#include <Shlwapi.h>

// Typedefs

typedef std::basic_string<TCHAR> tstring;
typedef std::basic_ostringstream<TCHAR> tostringstream;

// Enumeration of the launchers process exit codes.

enum ExitCode {
  Exit_Success = 0,
  Exit_InitializationFailed = 1,
  Exit_InvalidArguments = 2,
  Exit_CantWriteToInstallDir = 3,
  Exit_OpenLauncherLogFailed = 4,
  Exit_ElevationFailed = 5,
  Exit_ElevationCancelled = 6,
  Exit_UnarchivingFailed = 7,
  Exit_AccessingLogFailed = 8,
  Exit_InstallationFailed = 9,
};

// Forward declarations.

ExitCode RunInstaller(const tstring& installerPath,
                      const tstring& iniPath,
                      const tstring& targetPath,
                      std::wofstream& logFile);

// Utilities.

namespace util {

class AutoHandle  {
  HANDLE mHandle;
  AutoHandle(const AutoHandle&);
public:
  AutoHandle(HANDLE h = NULL) : mHandle(h) {}
  ~AutoHandle() { if (mHandle) ::CloseHandle(mHandle); }
  HANDLE get() const { return mHandle; }
  HANDLE* address() { return &mHandle; }
  AutoHandle& operator=(HANDLE h)  {
    mHandle = h;
    return *this;
  }
};

bool AppendPath(const tstring& first, const tstring& second, tstring& out)
{
  std::vector<TCHAR> buffer(MAX_PATH);
  const size_t length = std::min(first.length(), size_t(MAX_PATH)-1);
  first.copy(&buffer[0], length);
  buffer[length] = '\0';

  BOOL result = ::PathAppend(&buffer[0], second.c_str());
  if (!result) {
    return false;
  }

  out = &buffer[0];
  return true;
}

void Trim(tstring& str)
{
  tstring totrim = L" \n\r\t\f";

  size_t first = str.find_first_not_of(totrim);
  size_t last = str.find_last_not_of(totrim);

  if (first == tstring::npos) {
    str.clear();
    return;
  }

  if (last == tstring::npos)
    last = str.length() - 1;

  str = str.substr(first, last - first + 1);
}

bool CanUserElevate()
{
  AutoHandle token;
  if (!::OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, token.address())) {
     return false;
  }

  TOKEN_ELEVATION_TYPE elevationType;
  DWORD len;
  bool canElevate = ::GetTokenInformation(token.get(), TokenElevationType,
                                        &elevationType,
                                        sizeof(elevationType), &len) &&
                    (elevationType == TokenElevationTypeLimited);
  return canElevate;
}

bool IsVistaOrHigher()
{
  OSVERSIONINFO versionInfo;
  versionInfo.dwOSVersionInfoSize = sizeof(versionInfo);
  BOOL res = ::GetVersionEx(&versionInfo);
  if (!res) {
    return false;
  }

  return versionInfo.dwMajorVersion >= 6;
}

bool TempFileName(const tstring& dir, const tstring& prefix, tstring& str)
{
  std::vector<TCHAR> buffer(MAX_PATH);
  UINT result = ::GetTempFileName(dir.c_str(), prefix.c_str(), 0, &buffer[0]);
  if (result == 0) {
    return false;
  }

  str = &buffer[0];
  return true;
}

bool CanOpenFileForWriting(const tstring& path)
{
  {
    AutoHandle file = ::CreateFile(path.c_str(), FILE_GENERIC_WRITE, 0, NULL,
                                   OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (file.get() == INVALID_HANDLE_VALUE) {
      const DWORD err = ::GetLastError();
      tostringstream oss;

      if (err == ERROR_FILE_NOT_FOUND) {
        oss << L"CanOpenFileForWriting() - File not found: " << path << "\n";
      }
      else {
        oss << L"CanOpenFileForWriting() - Error: 0x" << std::hex << err << ", Path: " << path << "\n";
      }

      OutputDebugString(oss.str().c_str());
      return false;
    }
  }

  if (!::DeleteFile(path.c_str()))  {
    tostringstream oss;
    oss << "Failed to delete with error 0x" << std::hex << ::GetLastError() << ", path: " << path << "\n";
    OutputDebugString(oss.str().c_str());
  }

  return true;
}

bool CanWriteToInstallDir(const tstring& dir)
{
  tstring path;
  if (!TempFileName(dir, L"write-check", path))  {
    OutputDebugString(L"TempFileName() failed\n");
    return false;
  }

  return CanOpenFileForWriting(path);
}

bool GetFileLastWriteTime(const tstring& path, FILETIME& lastWriteTime)
{
  AutoHandle file = ::CreateFile(path.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL,
                                 OPEN_EXISTING, 0, NULL);
  if (file.get() == INVALID_HANDLE_VALUE)  {
    tostringstream oss;
    oss << L"Failed to open file " << path << L", error is 0x" << ::GetLastError() << L"\n";
    OutputDebugString(oss.str().c_str());
    return false;
  }

  FILETIME written;
  if (!::GetFileTime(file.get(), NULL, NULL, &written))  {
    tostringstream oss;
    oss << L"GetFileTime() failed for file " << path << L", error is 0x" << ::GetLastError() << L"\n";
    OutputDebugString(oss.str().c_str());
    return false;
  }

  lastWriteTime = written;
  return true;
}

} // namespace util

// Main implementation.

int APIENTRY _tWinMain(_In_ HINSTANCE hInstance,
                     _In_opt_ HINSTANCE hPrevInstance,
                     _In_ LPTSTR    lpCmdLine,
                     _In_ int       nCmdShow)
{
  UNREFERENCED_PARAMETER(hPrevInstance);
  UNREFERENCED_PARAMETER(lpCmdLine);

  int argc;
  std::shared_ptr<LPWSTR> argv(::CommandLineToArgvW(lpCmdLine, &argc), &::LocalFree);
  if (!argv) {
    OutputDebugString(L"Failed to parse arguments.\n");
    return Exit_InvalidArguments;
  }

  if (argc < 4)  {
    OutputDebugString(L"Not enough arguments.\n");
    return Exit_InvalidArguments;
  }

  const tstring installerPath = argv.get()[0];
  const tstring iniPath = argv.get()[1];
  const tstring targetPath = argv.get()[2];
  std::wofstream logFile(argv.get()[3]);

  if (!logFile.good()) {
    OutputDebugString(L"Failed to open launcher log file.\n");;
    return Exit_OpenLauncherLogFailed;
  }

  return RunInstaller(installerPath, iniPath, targetPath, logFile);
}

ExitCode RunElevated(const tstring& file, const tstring& parameters,
                     std::wofstream& logFile, util::AutoHandle& handle)
{
  std::vector<TCHAR> fileCopy(file.begin(), file.end());
  fileCopy.push_back('\0');
  std::vector<TCHAR> parametersCopy(parameters.begin(), parameters.end());
  parametersCopy.push_back('\0');

  SHELLEXECUTEINFO sinfo;
  memset(&sinfo, 0, sizeof(SHELLEXECUTEINFO));
  sinfo.cbSize = sizeof(SHELLEXECUTEINFO);
  sinfo.fMask = SEE_MASK_FLAG_NO_UI |
    SEE_MASK_FLAG_DDEWAIT |
    SEE_MASK_NOCLOSEPROCESS;
  sinfo.hwnd = nullptr;
  sinfo.lpFile = &fileCopy[0];
  sinfo.lpParameters = &parametersCopy[0];
  sinfo.lpVerb = L"runas";
  sinfo.nShow = SW_SHOWNORMAL;

  OutputDebugString(L"Requesting Elevation\n");
  BOOL result = ::ShellExecuteEx(&sinfo);
  DWORD dwErr = ::GetLastError();

  if (!result) {
    OutputDebugString(L"Elevation failed\n");
    if (dwErr == ERROR_CANCELLED) {
      logFile << Exit_ElevationCancelled << "\n";
      logFile << L"Elevation was cancelled.\n";
      return Exit_ElevationCancelled;
    }

    logFile << Exit_ElevationFailed << "\n";
    logFile << "Elevation failed - GetLastError() returned " << dwErr << "\n";
    return Exit_ElevationFailed;
  }

  handle = sinfo.hProcess;
  return Exit_Success;
}

ExitCode RunNonElevated(const tstring& file, const tstring& parameters,
                        std::wofstream& logFile, util::AutoHandle& handle)
{
  tostringstream oss;
  oss << L"\"" << file << "\" " << parameters;
  tstring temp = oss.str();
  std::vector<TCHAR> cmdLine(temp.begin(), temp.end());
  cmdLine.push_back('\0');

  STARTUPINFO startupInfo;
  PROCESS_INFORMATION processInformation;

  ::ZeroMemory(&startupInfo, sizeof(startupInfo));

  BOOL bret = ::CreateProcess(NULL, &cmdLine[0],
                              NULL, NULL, FALSE, 0, NULL, NULL,
                              &startupInfo, &processInformation);
  if (!bret) {
    logFile << Exit_InstallationFailed << "\n";
    logFile << L"GetLastError() returns " << ::GetLastError() << "\n";
    return Exit_InstallationFailed;
  }

  ::CloseHandle(processInformation.hThread);
  handle = processInformation.hProcess;
  return Exit_Success;
}

ExitCode RunInstaller(const tstring& installerPath,
                      const tstring& iniPath,
                      const tstring& targetPath,
                      std::wofstream& logFile)
{
  // Get install.log timestamp. We parse that file to determine whether the installer succeeded.
  // If an old version exists and the installer fails to write one, we incorrectly declare success.

  tstring installerLogPath;
  if (!util::AppendPath(targetPath, L"install.log", installerLogPath))  {
    OutputDebugString(L"Failed to build path for install.log file.\n");
    logFile << Exit_AccessingLogFailed << "\n";
    logFile << L"Failed to build path for install.log file.\n";
    return Exit_AccessingLogFailed;
  }

  FILETIME installerLogWrittenPreInstall;
  ::ZeroMemory(&installerLogWrittenPreInstall, sizeof(installerLogWrittenPreInstall));
  util::GetFileLastWriteTime(installerLogPath, installerLogWrittenPreInstall);

  // Try running installer elevated

  const bool vistaOrHigher = util::IsVistaOrHigher();
  const bool canElevate = util::CanUserElevate();

  bool elevated = false;
  if (vistaOrHigher) {
    const bool canWriteToInstallDir = util::CanWriteToInstallDir(targetPath);
    elevated = canElevate || !canWriteToInstallDir;

    tostringstream oss;
    oss << "CanWriteToInstallDir=" << std::boolalpha << canWriteToInstallDir << "\n";
    OutputDebugString(oss.str().c_str());
  }

  {
    tostringstream oss;
    oss << L"elevated=" << std::boolalpha << elevated
        << L", IsVistaOrHigher=" << vistaOrHigher
        << L", CanUserElevate=" << canElevate << " \n";
    OutputDebugString(oss.str().c_str());
  }

  const tstring parameters = L"/INI=" + iniPath;
  util::AutoHandle process;
  ExitCode code;

  if (elevated) {
    code = RunElevated(installerPath, parameters, logFile, process);
  } else {
    code = RunNonElevated(installerPath, parameters, logFile, process);
  }
  if (code != Exit_Success) {
    return code;
  }

  OutputDebugString(L"Waiting for process\n");
  WaitForSingleObject(process.get(), INFINITE);
  OutputDebugString(L"Process completed\n");

  // Check unarchiving process exit code

  DWORD exitCode;
  BOOL result = ::GetExitCodeProcess(process.get(), &exitCode);

  if (result) {
    tostringstream oss;
    oss << "Process exit code: " << exitCode << "\n";
    OutputDebugString(oss.str().c_str());

    if (exitCode != 0) {
      OutputDebugString(L"Failed to unarchive.\n");
      logFile << Exit_UnarchivingFailed << "\n";
      logFile << L"Unarchiving failed.\n";
      return Exit_UnarchivingFailed;
    }
  }

  // Check if we actually got a new or modified install.log file.

  FILETIME installerLogWrittenPostInstall;
  ::ZeroMemory(&installerLogWrittenPostInstall, sizeof(installerLogWrittenPostInstall));
  util::GetFileLastWriteTime(installerLogPath, installerLogWrittenPostInstall);

  if (::CompareFileTime(&installerLogWrittenPreInstall, &installerLogWrittenPostInstall) != -1)  {
    OutputDebugString(L"install.log file was not modified or does not exist.\n");
    logFile << Exit_AccessingLogFailed << "\n";
    logFile << L"install.log file was not modified or does not exist.\n";
    return Exit_AccessingLogFailed;
  }

  // Check installer log file

  std::wifstream installerLogFile(installerLogPath);
  if (!installerLogFile.good()) {
    OutputDebugString(L"Failed to open install.log file.\n");
    logFile << Exit_AccessingLogFailed << "\n";
    logFile << L"Failed to open install.log file.\n";
    return Exit_AccessingLogFailed;
  }

  installerLogFile.imbue(std::locale(installerLogFile.getloc(),
    new std::codecvt_utf16<wchar_t, 0x10ffff, std::consume_header>));

  const std::array<tstring, 2> successLogLines = {
    L"Added Registry Key: HKLM | Software\\Microsoft\\MediaPlayer\\ShimInclusionList\\plugin-container.exe",
    L"Added Registry Key: HKCU | Software\\Microsoft\\MediaPlayer\\ShimInclusionList\\plugin-container.exe",
  };

  tstring line, nonemptyLine;

  while (std::getline(installerLogFile, line))
  {
    tostringstream oss;
    oss << L"line: " << line << L"\n";
    //OutputDebugString(oss.str().c_str());

    util::Trim(line);
    if (line.empty()) {
      continue;
    }

    if (std::find(successLogLines.begin(), successLogLines.end(), line) != successLogLines.end()) {
      OutputDebugString(L"Found success log line\n");
      logFile << Exit_Success << "\n";
      return Exit_Success;
    }

    std::swap(line, nonemptyLine);
  }

  tostringstream oss;
  oss << L"Installation failed. Details:" << nonemptyLine << "\n";
  OutputDebugString(oss.str().c_str());

  logFile << Exit_InstallationFailed << "\n";
  logFile << L"Details: " << nonemptyLine << "\n";
  return Exit_InstallationFailed;
}