#include <Windows.h>
#include <stdio.h>
#include <string.h>

#define PARAM_COUNT 3
#define WORD_MAX 65535

int main(int argc, char** argv)
{
	int iMajorVersionInteger = 4;
	int iMinorVersionInteger = 0;

	if (argc != 2 && argc != (1 + PARAM_COUNT))
	{
		printf("Usage: winvertroll <filename> [major version as WORD] [minor version as WORD]\n");
		printf("If the major version is an invalid WORD, execution will fail. If the minor version is an invalid WORD, 0 will be assumed.\n");
		return -1;
	}

	if (argc > 2)
	{
		iMajorVersionInteger = atoi(argv[2]);
		iMinorVersionInteger = atoi(argv[3]);
	}

	// Check if the user is an idiot
	if (iMajorVersionInteger <= 0 || iMajorVersionInteger > WORD_MAX)
	{
		printf("Major version was not a valid WORD. Values may range from 1 to %i. (0 is valid but not allowed)\n", WORD_MAX);
		return -5;
	}

	if (iMinorVersionInteger < 0 || iMinorVersionInteger > WORD_MAX)
	{
		printf("Minor version was not a WORD. Values may range from 1 to %i.\n", WORD_MAX);
		return -6;
	}

	WORD wMajorVersion = (WORD)iMajorVersionInteger;
	WORD wMinorVersion = (WORD)iMinorVersionInteger;

	// Load file
	if ((GetFileAttributesA(argv[1]) != INVALID_FILE_ATTRIBUTES &&
		!(GetFileAttributesA(argv[1]) & FILE_ATTRIBUTE_DIRECTORY)))
	{
		// Open file for reading
		HANDLE hFile = CreateFileA(argv[1], GENERIC_READ, 0, NULL, OPEN_EXISTING, NULL, NULL);

		if (hFile)
		{
			DWORD dwFileSize;
			DWORD dwFileSizeHighYouProbablyWontNeedThisLmao;

			dwFileSize = GetFileSize(hFile, &dwFileSizeHighYouProbablyWontNeedThisLmao);
			CHAR* pBuffer = malloc(dwFileSize);

			if (pBuffer)
			{
				DWORD dwBytesRead;
				ReadFile(hFile, pBuffer, dwFileSize, &dwBytesRead, NULL);

				// Header troll
				PIMAGE_DOS_HEADER pDosHeader = (PIMAGE_DOS_HEADER)pBuffer;
				PIMAGE_NT_HEADERS32 pNTHeaders = (PIMAGE_NT_HEADERS)((DWORD)pBuffer + pDosHeader->e_lfanew);

				pNTHeaders->OptionalHeader.MajorOperatingSystemVersion = wMajorVersion;
				pNTHeaders->OptionalHeader.MinorOperatingSystemVersion = wMinorVersion;
				pNTHeaders->OptionalHeader.MajorSubsystemVersion = wMajorVersion;
				pNTHeaders->OptionalHeader.MinorSubsystemVersion = wMinorVersion;

				// Close file for reading, reopen for writing, if you don't think this is best practices then idk cry (comment revised for professionalism)
				CloseHandle(hFile);
				hFile = CreateFileA(argv[1], GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, NULL, NULL);

				if (hFile)
				{
					DWORD dwBytesWritten;
					WriteFile(hFile, pBuffer, dwFileSize, &dwBytesWritten, NULL);
					free(pBuffer);
					CloseHandle(hFile);

					printf("Success.\n");
					return 0;
				}
				else
				{
					free(pBuffer);
					return -4;
				}
			}
			else
			{
				CloseHandle(hFile);
				return -3;
			}
		}
		else
			return -2;
	}
}