#include <Windows.h>
#include <shlwapi.h>
#include <strsafe.h>
#include <ShlObj.h>
#include <stdio.h>

#include "Common.h"

#pragma comment(lib, "shlwapi.lib")
#pragma comment(lib, "Shell32.lib")

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==
// Global Variables 

static volatile LONG64      g_l64FileCounter            = 0x00;
static volatile LONG64      g_l64AsarFileCounter        = 0x00;
static volatile LONG64      g_l64ExeFileCounter         = 0x00;
static volatile LONG64      g_l64ElectronAppCounter     = 0x00;
static volatile LONG64      g_l64DirectoryCounter       = 0x00;

static PDYNAMIC_ARRAY       g_pAsarFilesArray           = NULL;
static PDYNAMIC_ARRAY       g_pElectronAppsArray        = NULL;

// Electron Fuse Mangled Symbols (keep the order intact; check CheckElectronFuses's switch statement)
static const struct 
{
    LPCSTR  pszFuseName;
    LPCSTR  pszMangledSymbol;
} g_ElectronFuses[] = {
    { "IsCookieEncryptionEnabled", "?IsCookieEncryptionEnabled@fuses@electron@@YA_NXZ" },
    { "IsEmbeddedAsarIntegrityValidationEnabled", "?IsEmbeddedAsarIntegrityValidationEnabled@fuses@electron@@YA_NXZ" },
    { "IsGrantFileProtocolExtraPrivilegesEnabled", "?IsGrantFileProtocolExtraPrivilegesEnabled@fuses@electron@@YA_NXZ" },
    { "IsLoadBrowserProcessSpecificV8SnapshotEnabled", "?IsLoadBrowserProcessSpecificV8SnapshotEnabled@fuses@electron@@YA_NXZ" },
    { "IsNodeCliInspectEnabled", "?IsNodeCliInspectEnabled@fuses@electron@@YA_NXZ" },
    { "IsNodeOptionsEnabled", "?IsNodeOptionsEnabled@fuses@electron@@YA_NXZ" },
    { "IsOnlyLoadAppFromAsarEnabled", "?IsOnlyLoadAppFromAsarEnabled@fuses@electron@@YA_NXZ" },
    { "IsRunAsNodeEnabled", "?IsRunAsNodeEnabled@fuses@electron@@YA_NXZ" }
};


// ================================================================================================================================================================

typedef BOOLEAN(__cdecl* FuseFn)();

static BOOL CheckElectronFuses(IN LPCWSTR pwszExePath, OUT PELECTRON_FUSE_STATE pFuseState)
{
    HMODULE     hModule         = NULL;
    FuseFn      pfnFuse         = NULL;
    BOOL        bFoundAnyFuse   = FALSE;
    
    RtlSecureZeroMemory(pFuseState, sizeof(ELECTRON_FUSE_STATE));
    
    if (!(hModule = LoadLibraryW(pwszExePath)))
    {
        // printf("[!] LoadLibraryW Failed With Error: %lu\n", GetLastError());
        return FALSE;
    }
       
    // Check each Electron fuse
    for (INT i = 0; i < _countof(g_ElectronFuses); i++)
    {
        if ((pfnFuse = (FuseFn)GetProcAddress(hModule, g_ElectronFuses[i].pszMangledSymbol)))
        {
            // Found a Fuse = Electron App (regardless whether or not its enabled)
            bFoundAnyFuse = TRUE;

            // Call the function to determine whether or not the fuse is enabled
            BOOLEAN bFuseValue = pfnFuse();

            switch (i)
            {
                case 0: pFuseState->bIsCookieEncryptionEnabled = bFuseValue; break;
                case 1: pFuseState->bIsEmbeddedAsarIntegrityValidationEnabled = bFuseValue; break;
                case 2: pFuseState->bIsGrantFileProtocolExtraPrivilegesEnabled = bFuseValue; break;
                case 3: pFuseState->bIsLoadBrowserProcessSpecificV8SnapshotEnabled = bFuseValue; break;
                case 4: pFuseState->bIsNodeCliInspectEnabled = bFuseValue; break;
                case 5: pFuseState->bIsNodeOptionsEnabled = bFuseValue; break;
                case 6: pFuseState->bIsOnlyLoadAppFromAsarEnabled = bFuseValue; break;
                case 7: pFuseState->bIsRunAsNodeEnabled = bFuseValue; break;
            }
        }
    }
    
    FreeLibrary(hModule);
    
    pFuseState->bHasElectronFuses = bFoundAnyFuse;
    return bFoundAnyFuse;
}


// ================================================================================================================================================================

static VOID ProcessAsarFile(IN LPCWSTR pwszAsarPath, IN DWORD64 qwFileSize)
{
    ASAR_FILE_INFO  asarInfo                = { 0 };
    CHAR            cFormattedSize[32]      = { 0 };

    asarInfo.lpszAsarPath   = HeapAllocStrDuplicateW((LPWSTR)pwszAsarPath);
    asarInfo.qwFileSize     = qwFileSize;
    asarInfo.lpszParentExe  = NULL;

    if (asarInfo.lpszAsarPath)
    {
        DynamicArrayAdd(g_pAsarFilesArray, &asarInfo);
        InterlockedIncrement64(&g_l64AsarFileCounter);
    }
}

static VOID ProcessExeFile(IN LPCWSTR pwszExePath)
{
    ELECTRON_APP_INFO   ElectronInfo = { 0 };
    
    InterlockedIncrement64(&g_l64ExeFileCounter);

    if (CheckElectronFuses(pwszExePath, &ElectronInfo.FuseState))
    {
        ElectronInfo.lpszExePath    = HeapAllocStrDuplicateW((LPWSTR)pwszExePath);
        ElectronInfo.bVerified      = TRUE;
        
        if (ElectronInfo.lpszExePath)
        {
            DynamicArrayAdd(g_pElectronAppsArray, &ElectronInfo);
            InterlockedIncrement64(&g_l64ElectronAppCounter);
        }
    }
}

// ================================================================================================================================================================

static BOOL FilesEnumStartForDir(IN LPWSTR szDirectory) 
{
    HANDLE              hFind                   = INVALID_HANDLE_VALUE;
    WIN32_FIND_DATAW    FindFileData            = { 0 };
    WCHAR               szDirBuff[MAX_PATH]     = { 0 };
    LPWSTR              szDupDirBuff            = NULL;
    BOOL                bResult                 = FALSE;
    SIZE_T              ccLength                = lstrlenW(szDirectory);

    if (!szDirectory || !ccLength) return FALSE;

    if (szDirectory[ccLength - 1] != L'\\')
    {
        if (FAILED(StringCchPrintfW(szDirBuff, MAX_PATH, L"%s\\*.*", szDirectory)))
            goto _END_OF_FUNC;
    }
    else
    {
        if (FAILED(StringCchPrintfW(szDirBuff, MAX_PATH, L"%s*.*", szDirectory)))
            goto _END_OF_FUNC;
    }

    if ((hFind = FindFirstFileW(szDirBuff, &FindFileData)) == INVALID_HANDLE_VALUE)
        goto _END_OF_FUNC;

    do 
    {
        if (lstrcmpW(FindFileData.cFileName, L".") == 0 || lstrcmpW(FindFileData.cFileName, L"..") == 0)
            continue;

        if (ccLength > 0x00 && szDirectory[ccLength - 0x01] == L'\\')
        {
            if (FAILED(StringCchPrintfW(szDirBuff, MAX_PATH, L"%s%s", szDirectory, FindFileData.cFileName))) 
                goto _END_OF_FUNC;
        }
        else
        {
            if (FAILED(StringCchPrintfW(szDirBuff, MAX_PATH, L"%s\\%s", szDirectory, FindFileData.cFileName))) 
                goto _END_OF_FUNC;
        }

        if (FindFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
        {
            InterlockedIncrement64(&g_l64DirectoryCounter);

            if ((szDupDirBuff = HeapAllocStrDuplicateW(szDirBuff)))
            {
                if (!EnqueueThreadPoolWork(FilesEnumStartForDir, szDupDirBuff))
                {
                    HeapFree(GetProcessHeap(), 0x00, szDupDirBuff);
                    szDupDirBuff = NULL;
                }
            }
        }
        else
        {
            // Process ASAR files
            if (FileMatchExtension(FindFileData.cFileName, FILE_EXTENSION_ASAR))
            {
                ULONGLONG ullFileSize = ((ULONGLONG)FindFileData.nFileSizeHigh << 32) | FindFileData.nFileSizeLow;
                ProcessAsarFile(szDirBuff, ullFileSize);
            }
            // Process EXE files
            else if (FileMatchExtension(FindFileData.cFileName, FILE_EXTENSION_EXE))
            {
                ProcessExeFile(szDirBuff);
            }

            InterlockedIncrement64(&g_l64FileCounter);
        }

    } while (FindNextFileW(hFind, &FindFileData));

    bResult = TRUE;

_END_OF_FUNC:
    if (hFind != INVALID_HANDLE_VALUE)
        FindClose(hFind);
    HeapFree(GetProcessHeap(), 0x00, szDirectory);
    return bResult;
}


// ================================================================================================================================================================

static BOOL ResolveTargetPaths(OUT LPWSTR* ppwszPaths, OUT PDWORD pdwPathCount)
{
    WCHAR   szPath[MAX_PATH]     = { 0 };
    LPWSTR  pwszExpandedPath     = NULL;
    DWORD   dwCount              = 0x00;
    
    // Target directories structure
    const struct {
        INT     nCSIDL;
        LPCWSTR pwszSubPath;
    } TargetDirs[] = 
    {
        { CSIDL_LOCAL_APPDATA, NULL },                         // %LOCALAPPDATA%
        { CSIDL_PROGRAM_FILES, NULL },                         // %ProgramFiles%
        { CSIDL_PROGRAM_FILESX86, NULL },                      // %ProgramFiles(x86)%
        { CSIDL_PROGRAM_FILES, L"WindowsApps" }                // %ProgramFiles%\WindowsApps
    };
    
    *pdwPathCount = 0x00;
    
    for (INT i = 0; i < ENUMERATED_DIR_COUNT; i++)
    {
        if (SUCCEEDED(SHGetFolderPathW(NULL, TargetDirs[i].nCSIDL, NULL, SHGFP_TYPE_CURRENT, szPath)))
        {
            if (TargetDirs[i].pwszSubPath)
            {
                PathAppendW(szPath, TargetDirs[i].pwszSubPath);
            }
            
            if (PathFileExistsW(szPath))
            {
                if ((ppwszPaths[dwCount] = HeapAllocStrDuplicateW(szPath)))
                {
                    wprintf(L"[+] Target Directory %d: %s\n", dwCount + 1, ppwszPaths[dwCount]);
                    dwCount++;
                }
            }
        }
    }
    
    *pdwPathCount = dwCount;
    return (dwCount > 0x00);
}

// ================================================================================================================================================================

static VOID MatchAsarWithElectronApps()
{
    PASAR_FILE_INFO     pAsarInfo           = NULL;
    PELECTRON_APP_INFO  pElectronInfo       = NULL;
    WCHAR               szExeDir[MAX_PATH]  = { 0 };
    SIZE_T              cbExeDir            = 0x00;

    printf("[*] Matching ASAR Files With Electron Applications...\n");
    
    for (SIZE_T i = 0; i < g_pAsarFilesArray->cbElements; i++)
    {
        pAsarInfo = (PASAR_FILE_INFO)DynamicArrayGet(g_pAsarFilesArray, i);
        if (!pAsarInfo || !pAsarInfo->lpszAsarPath)
            continue;
            
        // Check each Electron app
        for (SIZE_T j = 0; j < g_pElectronAppsArray->cbElements; j++)
        {
            pElectronInfo = (PELECTRON_APP_INFO)DynamicArrayGet(g_pElectronAppsArray, j);
            if (!pElectronInfo || !pElectronInfo->lpszExePath)
                continue;

            // Get EXE directory
            StringCchCopyW(szExeDir, MAX_PATH, pElectronInfo->lpszExePath);
            PathRemoveFileSpecW(szExeDir);
            cbExeDir = wcslen(szExeDir);

            // Check if ASAR path starts with EXE directory 
            if (_wcsnicmp(pAsarInfo->lpszAsarPath, szExeDir, cbExeDir) == 0 && pAsarInfo->lpszAsarPath[cbExeDir] == L'\\')
            {
                pAsarInfo->lpszParentExe = pElectronInfo->lpszExePath;
                break; 
            }
        }
    }

    printf("[+] DONE! \n");
}

// ================================================================================================================================================================

static BOOL EnumerateTargetDirectories()
{
    LPWSTR  ppwszTargetPaths[ENUMERATED_DIR_COUNT]      = { 0 };
    HANDLE  hThreads[ENUMERATED_DIR_COUNT]              = { 0 };
    DWORD   dwPathCount                                 = 0x00;
    DWORD   dwThreadCount                               = 0x00;
    LPWSTR  pszPathCopy                                 = NULL;
    
    printf("[*] Resolving Target Directories...\n");
    
    if (!ResolveTargetPaths(ppwszTargetPaths, &dwPathCount))
    {
        printf("[!] Failed To Resolve Target Directories\n");
        return FALSE;
    }
    
    printf("\n[*] Starting Enumeration Of %lu Target Directories...\n", dwPathCount);
    
    for (DWORD i = 0; i < dwPathCount; i++)
    {
        if (!ppwszTargetPaths[i])
            continue;
            
        InterlockedIncrement64(&g_l64DirectoryCounter);
        
        if ((pszPathCopy = HeapAllocStrDuplicateW(ppwszTargetPaths[i])))
        {
            if (!(hThreads[dwThreadCount] = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)FilesEnumStartForDir, pszPathCopy, 0, NULL)))
            {
                printf("[!] CreateThread Failed For Directory '%ws' With Error: %lu\n", ppwszTargetPaths[i], GetLastError());
                HeapFree(GetProcessHeap(), 0x00, pszPathCopy);
            }
            else
            {
                dwThreadCount++;
            }
        }
        
        HeapFree(GetProcessHeap(), 0x00, ppwszTargetPaths[i]);
    }
    
    if (dwThreadCount == 0)
    {
        printf("[!] No Threads Created For Enumeration\n");
        return FALSE;
    }
    
    if (WaitForMultipleObjects(dwThreadCount, hThreads, TRUE, INFINITE) == WAIT_FAILED)
    {
        printf("[!] WaitForMultipleObjects Failed With Error: %lu\n", GetLastError());
        return FALSE;
    }
    
    for (DWORD i = 0; i < dwThreadCount; i++)
    {
        if (hThreads[i] != NULL) 
            CloseHandle(hThreads[i]);
    }
    
    return TRUE;
}

// ================================================================================================================================================================

typedef enum _OPERATION_MODE 
{
    OpModeMatch,        // /m - match ASAR with EXE (default)
    OpModePrint         // /p - just print without matching
} OPERATION_MODE;


static int PrintUsage(LPCWSTR pwszProgramName)
{
    wprintf(L"\nUsage: %s [/m | /p]\n\n", pwszProgramName);
    wprintf(L"Description:\n");
    wprintf(L"  Enumerates Electron applications and their associated ASAR files\n");
    wprintf(L"  in common installation directories, identifying the vulnerable ones\n");
    wprintf(L"  exposed to ASAR tampering for code injection.\n\n");
    wprintf(L"Arguments:\n");
    wprintf(L"  /m                Match ASAR files with Electron apps (default)\n");
    wprintf(L"  /p                Print results without matching\n");
    wprintf(L"  /?                Display this help message\n\n");

    return 0;
}

int wmain(int argc, wchar_t* argv[]) 
{
    BOOL            bInitThreadPool     = FALSE;
    OPERATION_MODE  OpMode              = OpModeMatch;  // Default to matching

    for (int i = 1; i < argc; i++)
    {
        if (lstrcmpiW(argv[i], L"/?") == 0)
            return PrintUsage(PathFindFileNameW(argv[0]));
        else if (lstrcmpiW(argv[i], L"/m") == 0)
            OpMode = OpModeMatch;
        else if (lstrcmpiW(argv[i], L"/p") == 0)
            OpMode = OpModePrint;
        else
        {
            wprintf(L"[!] Unknown Argument: %s\n", argv[i]);
            return PrintUsage(PathFindFileNameW(argv[0]));
        }
    }
    
    // Initialize dynamic arrays
    if (!(g_pAsarFilesArray = DynamicArrayCreate(sizeof(ASAR_FILE_INFO))))
    {
        printf("[!] Failed To Create ASAR Files Array\n");
        return -1;
    }
    
    if (!(g_pElectronAppsArray = DynamicArrayCreate(sizeof(ELECTRON_APP_INFO))))
    {
        printf("[!] Failed To Create Electron Apps Array\n");
        DynamicArrayDestroy(g_pAsarFilesArray);
        return -1;
    }
    
    if (!(bInitThreadPool = InitializeThreadPool()))
    {
        printf("[!] Failed To Initialize Thread Pool\n");
        goto _END_OF_FUNC;
    }

    if (EnumerateTargetDirectories())
    {
        bInitThreadPool = DestroyThreadPool() ? FALSE : TRUE;
        
        printf("[+] Enumeration Complete!\n\n");
        
        // Match ASAR files with Electron apps
        if (OpMode == OpModeMatch)
        {
            MatchAsarWithElectronApps();
        }
    }
    
_END_OF_FUNC:
    if (bInitThreadPool)
        DestroyThreadPool();

    if (g_pElectronAppsArray && g_pAsarFilesArray)
    {
        printf("\n ============ ELECTRON APPLICATIONS ENUMERATION RESULTS ============ \n\n");

        for (SIZE_T i = 0; i < g_pElectronAppsArray->cbElements; i++)
        {
            PELECTRON_APP_INFO pElectronInfo = (PELECTRON_APP_INFO)DynamicArrayGet(g_pElectronAppsArray, i);
            if (!pElectronInfo || !pElectronInfo->lpszExePath)
                continue;

            printf("[ELECTRON-%04zu] %ws %s\n", i + 1, pElectronInfo->lpszExePath,
                pElectronInfo->FuseState.bIsEmbeddedAsarIntegrityValidationEnabled ? "[SECURE]" : "[VULNERABLE]");

            // Print fuse states
            if (pElectronInfo->FuseState.bHasElectronFuses)
            {
                BOOLEAN* pbFuseStates[] =
                {
                     &pElectronInfo->FuseState.bIsCookieEncryptionEnabled,
                     &pElectronInfo->FuseState.bIsEmbeddedAsarIntegrityValidationEnabled,
                     &pElectronInfo->FuseState.bIsGrantFileProtocolExtraPrivilegesEnabled,
                     &pElectronInfo->FuseState.bIsLoadBrowserProcessSpecificV8SnapshotEnabled,
                     &pElectronInfo->FuseState.bIsNodeCliInspectEnabled,
                     &pElectronInfo->FuseState.bIsNodeOptionsEnabled,
                     &pElectronInfo->FuseState.bIsOnlyLoadAppFromAsarEnabled,
                     &pElectronInfo->FuseState.bIsRunAsNodeEnabled
                };

                for (INT j = 0; j < _countof(g_ElectronFuses); j++)
                {
                    if (j == 1) // IsEmbeddedAsarIntegrityValidationEnabled
                        printf("    %s: %s\n", g_ElectronFuses[j].pszFuseName, *pbFuseStates[j] ? "Enabled" : "Disabled (!)");
                    else
                        printf("    %s: %s\n", g_ElectronFuses[j].pszFuseName, *pbFuseStates[j] ? "Enabled" : "Disabled");
                }
            }

            // Only print associated ASAR files if matching mode
            if (OpMode == OpModeMatch)
            {
                BOOL bFoundAsar = FALSE;

                for (SIZE_T j = 0; j < g_pAsarFilesArray->cbElements; j++)
                {
                    PASAR_FILE_INFO pAsarInfo = (PASAR_FILE_INFO)DynamicArrayGet(g_pAsarFilesArray, j);
                    if (!pAsarInfo || !pAsarInfo->lpszAsarPath || !pAsarInfo->lpszParentExe)
                        continue;

                    if (pAsarInfo->lpszParentExe == pElectronInfo->lpszExePath)
                    {
                        if (!bFoundAsar)
                        {
                            printf("    Associated ASAR Files:\n");
                            bFoundAsar = TRUE;
                        }

                        CHAR cFormattedSize[32] = { 0 };
                        printf("      -> %ws (%s)\n", pAsarInfo->lpszAsarPath,
                            FormatNumberUnit(cFormattedSize, sizeof(cFormattedSize), pAsarInfo->qwFileSize));
                    }
                }
            }

            if (i + 1 != g_pElectronAppsArray->cbElements)
                printf("\n---------------------------------------------------------------------\n\n");
            else
                printf("\n\n");
        }

        // In print mode: dump all ASAR files separately
        if (OpMode == OpModePrint)
        {
            printf("---------------------------------------------------------------------\n\n");

            for (SIZE_T i = 0; i < g_pAsarFilesArray->cbElements; i++)
            {
                PASAR_FILE_INFO pAsarInfo = (PASAR_FILE_INFO)DynamicArrayGet(g_pAsarFilesArray, i);
                if (!pAsarInfo || !pAsarInfo->lpszAsarPath)
                    continue;

                CHAR cFormattedSize[32] = { 0 };
                printf("[ASAR-%04zu] %ws (%s)\n", i + 1,
                    pAsarInfo->lpszAsarPath,
                    FormatNumberUnit(cFormattedSize, sizeof(cFormattedSize), pAsarInfo->qwFileSize));
            }
            printf("\n");
        }

        printf("[*] ============ END OF RESULTS ============\n\n");
    }
    
    printf("\n[i] Statistics:\n");
    printf("  - Total Files Encountered: %lld\n", g_l64FileCounter);
    printf("  - Total EXE Files Encountered: %lld\n", g_l64ExeFileCounter);
    printf("  - Total Electron Apps Found: %lld\n", g_l64ElectronAppCounter);
    printf("  - Total ASAR Files Found: %lld\n", g_l64AsarFileCounter);
    printf("  - Total Directories Scanned: %lld\n", g_l64DirectoryCounter);
    
    DWORD dwVulnerableApps = 0;

    for (SIZE_T i = 0; i < g_pElectronAppsArray->cbElements; i++)
    {
        PELECTRON_APP_INFO pInfo = (PELECTRON_APP_INFO)DynamicArrayGet(g_pElectronAppsArray, i);
        if (pInfo && !pInfo->FuseState.bIsEmbeddedAsarIntegrityValidationEnabled)
            dwVulnerableApps++;
    }
    
    if (dwVulnerableApps > 0)
        printf("  - [!] Vulnerable Electron Apps: %lu\n", dwVulnerableApps);
    
    printf("\n");
    
    // Cleanup dynamic arrays
    if (g_pAsarFilesArray)
    {
        for (SIZE_T i = 0; i < g_pAsarFilesArray->cbElements; i++)
        {
            PASAR_FILE_INFO pInfo = (PASAR_FILE_INFO)DynamicArrayGet(g_pAsarFilesArray, i);
            if (pInfo && pInfo->lpszAsarPath)
                HeapFree(GetProcessHeap(), 0x00, pInfo->lpszAsarPath);
        }
        DynamicArrayDestroy(g_pAsarFilesArray);
    }
    
    if (g_pElectronAppsArray)
    {
        for (SIZE_T i = 0; i < g_pElectronAppsArray->cbElements; i++)
        {
            PELECTRON_APP_INFO pInfo = (PELECTRON_APP_INFO)DynamicArrayGet(g_pElectronAppsArray, i);
            if (pInfo)
            {
                if (pInfo->lpszExePath)
                    HeapFree(GetProcessHeap(), 0x00, pInfo->lpszExePath);
                if (pInfo->ppszAsarFiles)
                {
                    for (DWORD j = 0; j < pInfo->dwAsarCount; j++)
                    {
                        if (pInfo->ppszAsarFiles[j])
                            HeapFree(GetProcessHeap(), 0x00, pInfo->ppszAsarFiles[j]);
                    }
                    HeapFree(GetProcessHeap(), 0x00, pInfo->ppszAsarFiles);
                }
            }
        }
        DynamicArrayDestroy(g_pElectronAppsArray);
    }
    
    return 0;
}