From 76ab7fb7f281beca96a40c115ec3d43021bc9f09 Mon Sep 17 00:00:00 2001 From: Elisha Riedlinger Date: Sat, 1 Apr 2023 22:15:34 -0700 Subject: [PATCH] Fix some issues and add error checking --- Encryption.cpp | 90 ++++++++++++++++++++++++++++++++++++-------------- NeuEncrypt.cpp | 12 +++++-- Resource.h | 2 +- 3 files changed, 76 insertions(+), 28 deletions(-) diff --git a/Encryption.cpp b/Encryption.cpp index a322b65..13b2178 100644 --- a/Encryption.cpp +++ b/Encryption.cpp @@ -136,60 +136,90 @@ void RenameFile(const std::wstring& path, const std::wstring& srcname, const std _wrename(srcFileName.c_str(), dstFileName.c_str()); } -void EncryptSingleFile(const std::wstring& path, const std::wstring& name) +bool EncryptSingleFile(const std::wstring& path, const std::wstring& name) { std::wstring fileName(path + L"\\" + name); // Open file - std::fstream myfile; - myfile.open(fileName.c_str(), std::ios::binary | std::ios::in | std::ios::out); + HANDLE hFile = CreateFile(name.c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); - // Check if file was opened - if (!myfile.is_open()) + if (hFile == INVALID_HANDLE_VALUE) { - return; + DWORD lastError = GetLastError(); + + LPWSTR buffer = NULL; + FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, + NULL, lastError, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPWSTR)&buffer, 0, NULL); + + std::wcerr << L"Failed to open file \"" << fileName << L"\". Error code: " << GetLastError() << " " << buffer << std::endl; + + LocalFree(buffer); + + return false; } // Get file size - myfile.seekg(0, std::ios::end); - long long fileSize = myfile.tellg(); + LARGE_INTEGER fileSize = { 0 }; + if (!GetFileSizeEx(hFile, &fileSize)) + { + std::wcerr << L"Failed to get file size for \"" << fileName << L"\". Error code: " << GetLastError() << std::endl; + CloseHandle(hFile); + return false; + } // Check file size value - if (fileSize < 1 || fileSize >= 1024 * 1024 * 1024) + if (fileSize.QuadPart < 1 || fileSize.QuadPart >= 1024LL * 1024LL * 1024LL) { - return; + std::wcerr << L"File size is invalid for \"" << fileName << L"\"." << std::endl; + CloseHandle(hFile); + return false; } - long long fileLoc = 0; + LARGE_INTEGER fileLoc = { 0 }; unsigned char memblock[4 * 1024]; // Encrypt file 4k at a time do { // Get update size - int updateSize = (int)min(fileSize - fileLoc, 4 * 1024); + DWORD updateSize = (DWORD)min(fileSize.QuadPart - fileLoc.QuadPart, 4 * 1024); // Read 4k of the file into memory block - myfile.seekg(fileLoc); - myfile.read((char*)memblock, updateSize); + DWORD bytesRead = 0; + OVERLAPPED overlapped = { 0 }; + overlapped.Offset = fileLoc.LowPart; + overlapped.OffsetHigh = fileLoc.HighPart; + if (!ReadFile(hFile, memblock, updateSize, &bytesRead, &overlapped) || bytesRead != updateSize) + { + std::wcerr << L"Failed to read file \"" << fileName << L"\". Error code: " << GetLastError() << std::endl; + CloseHandle(hFile); + return false; + } // Encrypt memory block - for (int x = 0; x < updateSize; x++) + for (DWORD x = 0; x < updateSize; x++) { memblock[x] = memblock[x] xor rndblock[x]; } // Write encrypted memory back to file - myfile.seekg(fileLoc); - myfile.write((char*)memblock, updateSize); + DWORD bytesWritten = 0; + if (!WriteFile(hFile, memblock, updateSize, &bytesWritten, &overlapped) || bytesWritten != updateSize) + { + std::wcerr << L"Failed to write file \"" << fileName << L"\". Error code: " << GetLastError() << std::endl; + CloseHandle(hFile); + return false; + } // Set new file location - fileLoc += updateSize; + fileLoc.QuadPart += updateSize; - } while (myfile.tellg() < fileSize); + } while (fileLoc.QuadPart < fileSize.QuadPart); // Close file - myfile.close(); + CloseHandle(hFile); + + return true; } void EncryptFiles() @@ -210,8 +240,14 @@ void EncryptFiles() // Rename and encrypt file RenameFile(LR"(\\?)", originalName.c_str(), tempName.c_str()); - EncryptSingleFile(LR"(\\?)", tempName.c_str()); - RenameFile(LR"(\\?)", tempName.c_str(), encryptName.c_str()); + if (EncryptSingleFile(LR"(\\?)", tempName.c_str())) + { + RenameFile(LR"(\\?)", tempName.c_str(), encryptName.c_str()); + } + else + { + RenameFile(LR"(\\?)", tempName.c_str(), originalName.c_str()); + } } } @@ -234,7 +270,13 @@ void DecryptFiles() // Rename and encrypt file RenameFile(LR"(\\?)", encryptName.c_str(), tempName.c_str()); - EncryptSingleFile(LR"(\\?)", tempName.c_str()); - RenameFile(LR"(\\?)", tempName.c_str(), originalName.c_str()); + if (EncryptSingleFile(LR"(\\?)", tempName.c_str())) + { + RenameFile(LR"(\\?)", tempName.c_str(), originalName.c_str()); + } + else + { + RenameFile(LR"(\\?)", tempName.c_str(), encryptName.c_str()); + } } } diff --git a/NeuEncrypt.cpp b/NeuEncrypt.cpp index 1a44288..81cab34 100644 --- a/NeuEncrypt.cpp +++ b/NeuEncrypt.cpp @@ -77,6 +77,12 @@ int APIENTRY wWinMain(_In_ HINSTANCE hInstance, void GetArgumentOptions(LPCWSTR lpCmdLine) { + // Error checking + if (!lpCmdLine || wcslen(lpCmdLine) == 0) + { + return; + } + // Retrieve command line arguments LPWSTR* szArgList; int argCount; @@ -177,7 +183,7 @@ bool ReadRegistryStruct(const std::wstring& lpzSection, const std::wstring& lpzK hr = RegQueryValueEx(hKey, lpzKey.c_str(), 0, 0, (BYTE*)lpStruct, &DataSize); - if (FAILED(hr) || !DataSize) + if (FAILED(hr) || DataSize == 0) { return false; } @@ -204,7 +210,7 @@ void LoadSettings() else { wchar_t path[MAX_PATH]; - if (ReadRegistryStruct(L"NeuEncrypt", L"LastEncryptFolder", (LPVOID)path, MAX_PATH)) + if (ReadRegistryStruct(L"NeuEncrypt", L"LastEncryptFolder", (LPVOID)path, MAX_PATH * sizeof(wchar_t))) { SetEncryptFolder(path); SetEncryptFolderText(); @@ -220,7 +226,7 @@ void SaveSettings() WriteRegistryStruct(L"NeuEncrypt", L"WindowPlacement", &wndpl, sizeof(WINDOWPLACEMENT)); - WriteRegistryStruct(L"NeuEncrypt", L"LastEncryptFolder", (LPVOID)GetEncryptFolder(), MAX_PATH); + WriteRegistryStruct(L"NeuEncrypt", L"LastEncryptFolder", (LPVOID)GetEncryptFolder(), (wcslen(GetEncryptFolder()) + 1) * sizeof(wchar_t)); } void SetEncryptFolderText() diff --git a/Resource.h b/Resource.h index a337966..d523fd4 100644 --- a/Resource.h +++ b/Resource.h @@ -23,7 +23,7 @@ // Main resource file details #define APP_NAME "NeuEncrypt" #define APP_MAJOR 1 -#define APP_MINOR 2 +#define APP_MINOR 3 #define APP_BUILDNUMBER 0 #define APP_REVISION 0 #define APP_COMPANYNAME "NeuShield Inc."