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

#include "Common.h"



// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==
// Dynamic Array Implementation

PDYNAMIC_ARRAY DynamicArrayCreate(IN SIZE_T cbElementSize)
{
    PDYNAMIC_ARRAY pArray = NULL;

    if (!(pArray = (PDYNAMIC_ARRAY)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(DYNAMIC_ARRAY))))
    {
        printf("[!] HeapAlloc Failed With Error: %lu\n", GetLastError());
        return NULL;
    }

    pArray->cbElementSize   = cbElementSize;
    pArray->cbCapacity      = INITIAL_ARRAY_SIZE;
    pArray->cbElements      = 0x00;

    if (!(pArray->pData = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, pArray->cbCapacity * cbElementSize)))
    {
        printf("[!] HeapAlloc Failed With Error: %lu\n", GetLastError());
        HeapFree(GetProcessHeap(), 0x00, pArray);
        return NULL;
    }

    InitializeCriticalSection(&pArray->CriticalSection);
    return pArray;
}


BOOL DynamicArrayAdd(IN PDYNAMIC_ARRAY pArray, IN LPVOID pElement)
{
    LPVOID pNewData         = NULL;
    SIZE_T cNewCapacity     = 0x00;

    if (!pArray || !pElement)
        return FALSE;

    EnterCriticalSection(&pArray->CriticalSection);

    if (pArray->cbElements >= pArray->cbCapacity)
    {
        cNewCapacity = pArray->cbCapacity * ARRAY_GROWTH_FACTOR;
        
        if (!(pNewData = HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, pArray->pData, cNewCapacity * pArray->cbElementSize)))
        {
            printf("[!] HeapReAlloc Failed With Error: %lu\n", GetLastError());
            LeaveCriticalSection(&pArray->CriticalSection);
            return FALSE;
        }

        pArray->pData       = pNewData;
        pArray->cbCapacity  = cNewCapacity;
    }

    RtlCopyMemory((LPBYTE)pArray->pData + (pArray->cbElements * pArray->cbElementSize), pElement, pArray->cbElementSize);
    
    pArray->cbElements++;

    LeaveCriticalSection(&pArray->CriticalSection);

    return TRUE;
}

LPVOID DynamicArrayGet(IN PDYNAMIC_ARRAY pArray, IN SIZE_T cbElementIndex)
{
    LPVOID pElement = NULL;

    if (!pArray || cbElementIndex >= pArray->cbElements)
        return NULL;

    EnterCriticalSection(&pArray->CriticalSection);
    pElement = (LPBYTE)pArray->pData + (cbElementIndex * pArray->cbElementSize);
    LeaveCriticalSection(&pArray->CriticalSection);

    return pElement;
}

VOID DynamicArrayDestroy(IN PDYNAMIC_ARRAY pArray)
{
    if (!pArray) return;

    DeleteCriticalSection(&pArray->CriticalSection);

    if (pArray->pData)
    {
        HeapFree(GetProcessHeap(), 0x00, pArray->pData);
        pArray->pData = NULL;
    }

    HeapFree(GetProcessHeap(), 0x00, pArray);
}

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==
// Format the file size value

PCHAR FormatNumberUnit(IN OUT PCHAR pszBuffer, IN DWORD dwBufferSize, IN DWORD64 qwBytes)
{
    DWORD64 qwValue         = 0x00;
    DWORD   dwRemainder     = 0x00;
    PCSTR   pszUnit         = NULL;
    
    RtlSecureZeroMemory(pszBuffer, dwBufferSize);
    
    if (qwBytes < 1024) 
    {
        StringCbPrintfA(pszBuffer, dwBufferSize, "%llu B", qwBytes);
        return pszBuffer;
    }
    else if (qwBytes < (10 * 1024))  
    {
        StringCbPrintfA(pszBuffer, dwBufferSize, "%.2f KB", qwBytes / 1024.0);
        return pszBuffer;
    }
    else if (qwBytes < (1024 * 1024))  
    {
        StringCbPrintfA(pszBuffer, dwBufferSize, "%.2f KB", qwBytes / 1024.0);
        return pszBuffer;
    }
    else if (qwBytes < (1024LL * 1024 * 1024))  
    {
        qwValue = qwBytes / (1024 * 1024);
        dwRemainder = (DWORD)((qwBytes % (1024 * 1024)) * 100 / (1024 * 1024));
        pszUnit = "MB";
    }
    else  
    {
        qwValue = qwBytes / (1024LL * 1024 * 1024);
        dwRemainder = (DWORD)((qwBytes % (1024LL * 1024 * 1024)) * 100 / (1024LL * 1024 * 1024));
        pszUnit = "GB";
    }
    
    StringCbPrintfA(pszBuffer, dwBufferSize, "%llu.%02u %s", qwValue, dwRemainder, pszUnit);

    return pszBuffer;
}

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==
// Duplicate a given string using a heap buffer 

LPWSTR HeapAllocStrDuplicateW(IN LPWSTR szStrToDuplicate) 
{
    SIZE_T      cchLength   = 0x00;
    LPWSTR      szDupStr    = NULL;

    if (!szStrToDuplicate) return NULL;

    cchLength = lstrlenW(szStrToDuplicate) + 0x01;

    if (cchLength > (SIZE_MAX / sizeof(WCHAR)))
        return NULL;

    if (!(szDupStr = (LPWSTR)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, cchLength * sizeof(WCHAR))))
        return NULL;

    if (FAILED(StringCchCopyW(szDupStr, cchLength, szStrToDuplicate))) 
    {
        HeapFree(GetProcessHeap(), 0x00, szDupStr);
        return NULL;
    }

    return szDupStr;
}

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==
// Check if the specified file has the specified extension

BOOL FileMatchExtension(IN LPWSTR szFileName, IN LPCWSTR pwszExtension)
{
    LPWSTR  pwszFileExt = NULL;

    if (!(pwszFileExt = PathFindExtensionW(szFileName)) || !*pwszFileExt)
        return FALSE;

    if (lstrcmpiW(pwszFileExt, pwszExtension) == 0x00)
        return TRUE;

    return FALSE;
}

   