locknote.cpp

// Steganos LockNote - self-modifying encrypted notepad
// Copyright (C) 2006-2023 Steganos GmbH
//
// 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 of the License, 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, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

#include "stdafx.h"

#include <atlframe.h>
#include <atlctrls.h>
#include <atldlgs.h>

#include <process.h>
#include <algorithm>
#include <array>

#include <atlfile.h>

#include "resource.h"

#include "locknoteView.h"
#include "AboutDlg.h"
#include "PasswordDlg.h"
#include "MainFrm.h"

#include "aeslayer.cpp" // prevents having to include precompiled header in aeslayer.cpp

CAppModule _Module;

namespace
{
	constexpr int kFileRetryCount = 300;
	constexpr DWORD kFileRetrySleepMs = 100;

	bool RetryCopyFile(const TCHAR* source, const TCHAR* target)
	{
		for (int attempt = 0; attempt < kFileRetryCount; ++attempt)
		{
			if (::CopyFile(source, target, FALSE))
			{
				return true;
			}

			const DWORD lastError = ::GetLastError();
			if (lastError != ERROR_SHARING_VIOLATION && lastError != ERROR_ACCESS_DENIED)
			{
				return false;
			}
			::Sleep(kFileRetrySleepMs);
		}
		return false;
	}

	bool RetryDeleteFile(const TCHAR* target)
	{
		for (int attempt = 0; attempt < kFileRetryCount; ++attempt)
		{
			if (!::PathFileExists(target))
			{
				return true;
			}
			if (::DeleteFile(target))
			{
				return true;
			}

			const DWORD lastError = ::GetLastError();
			if (lastError != ERROR_SHARING_VIOLATION && lastError != ERROR_ACCESS_DENIED)
			{
				return false;
			}
			::Sleep(kFileRetrySleepMs);
		}
		return !::PathFileExists(target);
	}

	bool StageWritebackFromMainFrame(const CMainFrame& wndMain, std::string password)
	{
		std::string encryptedData;
		if (!wndMain.m_text.empty())
		{
			Utils::EncryptString(
				wndMain.m_text,
				password,
				encryptedData,
				Utils::ParseKdfModeValue(wndMain.GetKdfMode()));
		}
		else
		{
			Utils::MessageBox(nullptr, WSTR(IDS_TEXT_IS_ENCRYPTED), MB_OK | MB_ICONINFORMATION);
		}

		std::array<wchar_t, MAX_PATH> modulePath{};
		std::array<wchar_t, MAX_PATH> tempPath{};
		std::array<wchar_t, MAX_PATH> fileName{};

		const DWORD modulePathLength = ::GetModuleFileNameW(Utils::GetModuleHandle(), modulePath.data(), static_cast<DWORD>(modulePath.size()));
		if (modulePathLength == 0 || modulePathLength >= modulePath.size())
		{
			return false;
		}

		const DWORD tempPathLength = ::GetTempPathW(static_cast<DWORD>(tempPath.size()), tempPath.data());
		if (tempPathLength == 0 || tempPathLength >= tempPath.size())
		{
			return false;
		}

		if (::GetTempFileNameW(tempPath.data(), L"STG", 0, fileName.data()) == 0)
		{
			return false;
		}

		if (!::CopyFileW(modulePath.data(), fileName.data(), FALSE))
		{
			::DeleteFileW(fileName.data());
			return false;
		}

		const std::string fileNameUtf8 = wstring_to_utf8(fileName.data());
		if (fileNameUtf8.empty())
		{
			::DeleteFileW(fileName.data());
			return false;
		}

		bool writeResult = Utils::UpdateResource(fileNameUtf8, "CONTENT", "PAYLOAD", encryptedData);

		LOCKNOTEWINTRAITS traits{};
		traits.m_nWindowSizeX = wndMain.m_nWindowSizeX;
		traits.m_nWindowSizeY = wndMain.m_nWindowSizeY;
		traits.m_nFontSize = wndMain.m_nFontSize;
		traits.m_nLangId = wndMain.GetLanguage();
		traits.m_nKdfMode = wndMain.GetKdfMode();
		traits.m_nThemeMode = wndMain.GetThemeMode();
		traits.m_strFontName = wndMain.m_strFontName;
		writeResult &= Utils::WriteWinTraitsResources(fileNameUtf8, traits);

		if (!writeResult)
		{
			::DeleteFileW(fileName.data());
			return false;
		}

		const intptr_t spawnResult = _tspawnl(
			_P_NOWAIT,
			fileName.data(),
			fileName.data(),
			_T("-writeback"),
			modulePath.data(),
			nullptr);
		if (spawnResult == -1)
		{
			::DeleteFileW(fileName.data());
			return false;
		}

		return true;
	}
}

int Run(LPTSTR /*lpstrCmdLine*/ = NULL, int nCmdShow = SW_SHOWDEFAULT)
{
	TCHAR szModulePath[MAX_PATH] = {'\0'};
	const DWORD modulePathLength = ::GetModuleFileName(Utils::GetModuleHandle(), szModulePath, MAX_PATH);
	if (modulePathLength == 0 || modulePathLength >= MAX_PATH)
	{
		return -1;
	}

	CMessageLoop theLoop;
	_Module.AddMessageLoop(&theLoop);

	if (__argc == 3)
	{
#ifdef _UNICODE
		TCHAR* lpszCommand = __wargv[1];
		TCHAR* lpszPath = __wargv[2];
#else
		char* lpszCommand = __argv[1];
		char* lpszPath = __argv[2];
#endif
		if (!_tcscmp(lpszCommand, _T("-writeback")))
		{
			if (!RetryCopyFile(szModulePath, lpszPath))
			{
				return -1;
			}

			const intptr_t spawnResult = _tspawnl(_P_NOWAIT, lpszPath, lpszPath, _T("-erase"), szModulePath, nullptr);
			if (spawnResult == -1)
			{
				return -1;
			}

			return 0;
		}
		else if (!_tcscmp(lpszCommand, _T("-erase")))
		{
			if (!RetryDeleteFile(lpszPath))
			{
				return -1;
			}

			return 0;
		}
	}

	if (__argc > 1)
	{
		int nResult = Utils::MessageBox(NULL, WSTR(IDS_ASK_CONVERT_FILES), MB_YESNO | MB_ICONQUESTION);
		if (nResult == IDYES)
		{
			int nConverted = 0;
			std::string encryptPassword;
			for (int nIndex = 1; nIndex < __argc; nIndex++)
			{
#ifdef _UNICODE
				const std::string filename = wstring_to_utf8(__wargv[nIndex]);
#else
				const std::string filename = __argv[nIndex];
#endif
				if (HasExtension(filename, ".txt"))
				{
					std::string newfilename = filename.substr(0, filename.size() - 4);
					newfilename += ".exe";
					std::string text;
					std::string password;
					if (LoadTextFromFile(filename, text, password))
					{
						if (SaveTextToFile(newfilename, text, encryptPassword))
						{
							nConverted += 1;
						}
					}
				}
				else
				{
					std::array<wchar_t, 512> message{};
					const std::wstring wideFilename = utf8_to_wstring(filename);
					swprintf_s(message.data(), message.size(), WSTR(IDS_CANT_CONVERT).c_str(), wideFilename.c_str());
					Utils::MessageBox(NULL, message.data(), MB_OK | MB_ICONERROR);
				}
			}

			if (nConverted)
			{
				std::array<wchar_t, 256> message{};
				swprintf_s(message.data(), message.size(), WSTR(IDS_CONVERT_DONE).c_str(), nConverted);
				Utils::MessageBox(NULL, message.data(), MB_OK | MB_ICONINFORMATION);
			}
		}

		return 0;
	}

	std::string szFileMappingName = wstring_to_utf8(szModulePath);
	std::replace_if(
		szFileMappingName.begin(),
		szFileMappingName.end(),
		[](const char ch)
		{
			return ch == '\\' || ch == ':';
		},
		'_');

	CAtlFileMapping<HWND> fmSingleInstanceHWND;
	BOOL bAlreadyExisted = FALSE;
	fmSingleInstanceHWND.MapSharedMem(sizeof(HWND), utf8_to_wstring(szFileMappingName).c_str(), &bAlreadyExisted);
	if (bAlreadyExisted)
	{
		HWND hWndMain = *fmSingleInstanceHWND;
		::SetForegroundWindow(hWndMain);
		return 0;
	}

	std::string text;
	std::string data;
	std::string password;
	Utils::LoadResource("CONTENT", "PAYLOAD", data);
	if (!data.empty())
	{
		password = GetPasswordDlg();
		if (password.empty())
		{
			return -1;
		}
		if (!Utils::DecryptString(data, password, text))
		{
			MessageBox(NULL, WSTR(IDS_INVALID_PASSWORD), MB_OK | MB_ICONERROR);
			return -1;
		}		
	}
	else
	{
		text = STR(IDS_WELCOME);
	}

	CMainFrame wndMain;

	wndMain.m_password = password;
	wndMain.m_text = text;

	// get window sizes from resource
	std::string strSizeX;
	std::string strSizeY;
	Utils::LoadResource("SIZEX", "INFORMATION", strSizeX);
	Utils::LoadResource("SIZEY", "INFORMATION", strSizeY);
	int parsedValue = 0;
	if (Utils::TryParseInt(strSizeX, parsedValue))
	{
		wndMain.m_nWindowSizeX = parsedValue;
	}
	if (Utils::TryParseInt(strSizeY, parsedValue))
	{
		wndMain.m_nWindowSizeY = parsedValue;
	}

	// get font size
	std::string fontsize;
	Utils::LoadResource("FONTSIZE", "INFORMATION", fontsize);
	if (Utils::TryParseInt(fontsize, parsedValue))
	{
		wndMain.m_nFontSize = parsedValue;
	}
	else
	{
		wndMain.m_nFontSize = DEFAULT_FONT_SIZE;
	}

	// get font set
	std::string fontface;
	Utils::LoadResource("TYPEFACE", "INFORMATION", fontface);
	if (!fontface.empty())
	{
		// font typeface format: "Lucida Console" (example)
		wndMain.m_strFontName = fontface;
	}
	else
	{
		wndMain.m_strFontName = DEFAULT_FONT_NAME;
	}
	
	// get language id
	std::string language;
	Utils::LoadResource("LANGID", "INFORMATION", language);
	if (Utils::TryParseInt(language, parsedValue))
	{
		wndMain.SetLanguage(parsedValue);
	}
	else
	{
		wndMain.SetLanguage(0);
	}

	std::string kdfMode;
	Utils::LoadResource("KDFMODE", "INFORMATION", kdfMode);
	if (Utils::TryParseInt(kdfMode, parsedValue))
	{
		wndMain.SetKdfMode(parsedValue);
	}
	else
	{
		wndMain.SetKdfMode(static_cast<int>(AESLayer::KdfMode::Scrypt));
	}

	std::string themeMode;
	Utils::LoadResource("THEMEMODE", "INFORMATION", themeMode);
	if (Utils::TryParseInt(themeMode, parsedValue))
	{
		wndMain.SetThemeMode(parsedValue);
	}
	else
	{
		wndMain.SetThemeMode(static_cast<int>(ThemeMode::System));
	}

	wndMain.PrepareInitialWindowSizeForCreate();
	RECT initialWindowRect{ 0, 0, wndMain.m_nWindowSizeX, wndMain.m_nWindowSizeY };

	// create window
	if (wndMain.CreateEx(nullptr, &initialWindowRect) == nullptr)
	{
		ATLTRACE(_T("Main window creation failed!\n"));
		return 0;
	}
	wndMain.CheckFontSize();
	wndMain.CheckFontTypeFace();

	*fmSingleInstanceHWND = wndMain.m_hWnd;

	wndMain.ShowWindow(nCmdShow);

	int nRet = theLoop.Run();

	_Module.RemoveMessageLoop();

	if (((wndMain.m_text != text) || (wndMain.m_password != password) || wndMain.m_bTraitsChanged) && (wndMain.m_password.size()))
	{
		password = wndMain.m_password;
		if (!StageWritebackFromMainFrame(wndMain, password))
		{
			Utils::MessageBox(nullptr, L"Saving changes failed.", MB_OK | MB_ICONERROR);
		}
	}

	return nRet;
}

int WINAPI _tWinMain(HINSTANCE hInstance, HINSTANCE /*hPrevInstance*/, LPTSTR lpstrCmdLine, int nCmdShow)
{
	HRESULT hRes = ::CoInitialize(NULL);
// If you are running on NT 4.0 or higher you can use the following call instead to 
// make the EXE free threaded. This means that calls come in on a random RPC thread.
//	HRESULT hRes = ::CoInitializeEx(NULL, COINIT_MULTITHREADED);
	ATLASSERT(SUCCEEDED(hRes));

	// this resolves ATL window thunking problem when Microsoft Layer for Unicode (MSLU) is used
	::DefWindowProc(NULL, 0, 0, 0L);

	AtlInitCommonControls(ICC_BAR_CLASSES);	// add flags to support other controls

	hRes = _Module.Init(NULL, hInstance);
	ATLASSERT(SUCCEEDED(hRes));

	int nRet = Run(lpstrCmdLine, nCmdShow);

	_Module.Term();
	::CoUninitialize();

	return nRet;
}