#include "pch.h"
#pragma hdrstop
#include <delayimp.h>
#include "peimage.h"
#include "persist.h"

//-------------------------------------------------------------------------
// Purpose:
//    Returns true if the current state of the passed in
//    WIN32_FIND_DATA structure represents a child directory.
//
BOOL
IsChildDir (
    IN const WIN32_FIND_DATAA* pFindData)
{
   return (pFindData->dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) &&
          (*pFindData->cFileName != '.');
}

BOOL
FFindNextChildDir (
    IN HANDLE hFind,
    IN OUT WIN32_FIND_DATAA* pFindData)
{
    BOOL fFound;

    do
    {
        fFound = !!FindNextFileA (hFind, pFindData);
    }
    while (fFound && !IsChildDir (pFindData));

    return fFound;
}

HANDLE
FFindFirstChildDir (
    IN PCSTR pszFileSpec,
    OUT WIN32_FIND_DATAA* pFindData)
{
    HANDLE hFind;

    hFind = FindFirstFileA (pszFileSpec, pFindData);

    if (INVALID_HANDLE_VALUE != hFind)
    {
        BOOL fFound = IsChildDir (pFindData);

        if (!fFound)
        {
            fFound = FFindNextChildDir (hFind, pFindData);
        }

        if (!fFound)
        {
            FindClose (hFind);
            hFind = INVALID_HANDLE_VALUE;
        }
    }

    return hFind;
}


// pszFilePath must be at least MAX_PATH characters
//
VOID
GetIndexFilePath (
    OUT PWSTR pszFilePath,
    IN UINT cchFilePath)
{
    GetSystemWindowsDirectory (pszFilePath, cchFilePath);
    wcscat (pszFilePath, L"dllquery.dat");
}

HRESULT
HrLoadModuleTree (
    OUT CModuleTree* pTree)
{
    HRESULT hr;
    WCHAR szIndexFile [MAX_PATH];
    HANDLE hFile;

    hr = E_UNEXPECTED;

    pTree->Modules.m_Granularity = 2048;

    // Open the index file if it exists.
    //
    GetIndexFilePath (szIndexFile, celems(szIndexFile));

    hFile = CreateFile (szIndexFile, GENERIC_READ, FILE_SHARE_READ,
                NULL, OPEN_EXISTING,
                FILE_ATTRIBUTE_NORMAL | FILE_FLAG_SEQUENTIAL_SCAN,
                NULL);

    if (INVALID_HANDLE_VALUE != hFile)
    {
        HANDLE hMapping;

        hMapping = CreateFileMapping (hFile, NULL, PAGE_READONLY, 0, 0, NULL);

        if (hMapping)
        {
            BYTE* pbBuf;
            UINT  cbBuf;
            DWORD cbSizeHi;

            pbBuf = (BYTE*)MapViewOfFile (hMapping, FILE_MAP_READ, 0, 0, 0);

            if (pbBuf)
            {
                cbBuf = GetFileSize (hFile, &cbSizeHi);
                hr = HrLoadModuleTreeFromBuffer (pbBuf, cbBuf, pTree);

                UnmapViewOfFile (pbBuf);
            }

            CloseHandle (hMapping);
        }

        CloseHandle (hFile);
    }
    else
    {
        hr = HrLoadModuleTreeFromFileSystem (pTree);
    }

    return hr;
}

HRESULT
HrLoadModuleTreeFromBuffer (
    IN const BYTE* pbBuf,
    IN ULONG cbBuf,
    OUT CModuleTree* pTree)
{
    HRESULT hr = S_OK;
    return hr;
}



HRESULT
HrLoadImportsForModule (
    IN CModule* pMod,
    IN CModuleTree* pTree,
    IN PVOID pvImage)
{
    HRESULT hr;
    PIMAGE_NT_HEADERS pNtHeaders;
    PIMAGE_SECTION_HEADER pNtSection = NULL;
    PIMAGE_IMPORT_DESCRIPTOR pImportDescriptor;
    PIMAGE_THUNK_DATA pThunk;
    ULONG cbImportDescriptor;
    PCSTR pszImportFileName;
    CModule* pImportMod;

    hr = S_OK;

    pNtHeaders = RtlImageNtHeader(pvImage);
    Assert (pNtHeaders);

    pImportDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)
        RtlImageDirectoryEntryToData (
            pvImage,
            FALSE,
            IMAGE_DIRECTORY_ENTRY_IMPORT,
            &cbImportDescriptor);

    if (pImportDescriptor)
    {
        // Characteristics is zero for the terminating descriptor.
        //
        while (pImportDescriptor->Characteristics &&
               pImportDescriptor->FirstThunk)
        {
            Assert ((ULONG_PTR)pImportDescriptor <
                    (ULONG_PTR)pImportDescriptor + cbImportDescriptor);

            Assert (pImportDescriptor->Name);

            pszImportFileName = (PCSTR)RtlImageRvaToVa (pNtHeaders,
                                        pvImage,
                                        pImportDescriptor->Name,
                                        &pNtSection);

            Assert(*pszImportFileName);

            // Add this import to the list if we don't already have it
            // present.
            //
            hr = pTree->Modules.HrInsertNewModule (
                            pszImportFileName,
                            0,
                            INS_IGNORE_IF_DUP | INS_SORTED,
                            &pImportMod);

            if (S_OK == hr)
            {
                PSZ Function = NULL;

                hr = pTree->HrAddEntry (pMod, pImportMod, MTE_DEFAULT);

                if (S_OK != hr)
                {
                    break;
                }

                pThunk = (PIMAGE_THUNK_DATA)RtlImageRvaToVa (pNtHeaders,
                                                pvImage,
                                                pImportDescriptor->OriginalFirstThunk,
                                                &pNtSection);
                if (!IMAGE_SNAP_BY_ORDINAL(pThunk->u1.Ordinal))
                {
                    PIMAGE_IMPORT_BY_NAME ImportName =
                        (PIMAGE_IMPORT_BY_NAME)RtlImageRvaToVa (pNtHeaders,
                                                pvImage,
                                                pThunk->u1.AddressOfData,
                                                &pNtSection);
                    Function = (PSZ)ImportName->Name;

                }

                pThunk++;
                if (!pThunk->u1.AddressOfData)
                {
                    printf("%-16s  only imports 1 function from  %s  (%s)\n",
                        pMod->m_pszFileName,
                        pImportMod->m_pszFileName,
                        Function);
                }
            }

            pImportDescriptor++;
        }
    }

    return hr;
}

HRESULT
HrLoadDelayImportsForModule (
    IN CModule* pMod,
    IN CModuleTree* pTree,
    IN PVOID pvImage)
{
    HRESULT hr;
    PIMAGE_NT_HEADERS pNtHeaders;
    PIMAGE_SECTION_HEADER pNtSection = NULL;
    PImgDelayDescr pDelayDescriptor;
    ULONG cbDelayDescriptor;
    PCSTR pszImportFileName;
    CModule* pImportMod;

    hr = S_OK;

    pNtHeaders = RtlImageNtHeader(pvImage);
    Assert (pNtHeaders);

    pDelayDescriptor = (PImgDelayDescr)
        RtlImageDirectoryEntryToData (
            pvImage,
            FALSE,
            IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT,
            &cbDelayDescriptor);

    if (pDelayDescriptor)
    {
        Assert (pNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT].Size);

        ULONG_PTR ImageBase = pNtHeaders->OptionalHeader.ImageBase;

        // Characteristics is zero for the terminating descriptor.
        //
        while (pDelayDescriptor->pIAT &&
               pDelayDescriptor->pINT &&
               pDelayDescriptor->phmod)
        {
            Assert ((ULONG_PTR)pDelayDescriptor <
                    (ULONG_PTR)pDelayDescriptor + cbDelayDescriptor);

            Assert (pDelayDescriptor->szName);

            pNtSection = IMAGE_FIRST_SECTION (pNtHeaders);
            ULONG_PTR Rva = (ULONG_PTR)pDelayDescriptor->szName;
            ULONG_PTR Base = 0;

            for (INT i = 0; i < pNtHeaders->FileHeader.NumberOfSections; i++)
            {
                if ((Rva >= (ImageBase + pNtSection->VirtualAddress)) &&
                    (Rva <  (ImageBase + pNtSection->VirtualAddress + pNtSection->SizeOfRawData)))
                {
                    Base = (ULONG_PTR)pvImage
                        + pNtSection->PointerToRawData
                        - pNtSection->VirtualAddress - ImageBase;
                    break;
                }
                pNtSection++;
            }

            if (0 == Base)
            {
                break;
            }

            pszImportFileName = (PCSTR)(Base + Rva);

            Assert(*pszImportFileName);

            // Add this import to the list if we don't already have it
            // present.
            //
            hr = pTree->Modules.HrInsertNewModule (
                            pszImportFileName,
                            0,
                            INS_IGNORE_IF_DUP | INS_SORTED,
                            &pImportMod);

            if (S_OK == hr)
            {
                hr = pTree->HrAddEntry (pMod, pImportMod, MTE_DELAY_LOADED);

                if (S_OK != hr)
                {
                    break;
                }
            }

            pDelayDescriptor++;
        }
    }

    return hr;
}

HRESULT
HrLoadModuleTreeFromCurrentDirectory (
    OUT CModuleTree* pTree)
{
    HRESULT hr;
    HANDLE hFind;
    WIN32_FIND_DATAA FindData;

    hr = S_OK;

    // Enumerate all of the files in the current directory.
    //
    hFind = FindFirstFileA ("*", &FindData);

    if (INVALID_HANDLE_VALUE != hFind)
    {
        CPeImage PeImage;
        ZeroMemory (&PeImage, sizeof(PeImage));

        do
        {
            CModule* pMod;

            // Don't process directories.
            //
            if (FindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
            {
                continue;
            }

            hr = PeImage.HrOpenFile (FindData.cFileName);

            if (S_OK == hr)
            {
                // Add this module to the list if we don't already have it
                // present.  (We can already have it if, while procesing the
                // imports of an earlier module, this module was one of those
                // imports.)
                //
                hr = pTree->Modules.HrInsertNewModule (
                                FindData.cFileName,
                                FindData.nFileSizeLow,
                                INS_IGNORE_IF_DUP | INS_SORTED,
                                &pMod);
                if (S_OK == hr)
                {
                    (VOID) HrLoadImportsForModule (
                                pMod, pTree, PeImage.m_pvImage);

                    (VOID) HrLoadDelayImportsForModule (
                                pMod, pTree, PeImage.m_pvImage);

                }

                PeImage.CloseFile ();
            }
        }
        while (FindNextFileA (hFind, &FindData));

        FindClose (hFind);
    }

    // Now recurse into subdirectories.
    //
    hFind = FFindFirstChildDir("*", &FindData);

    if (INVALID_HANDLE_VALUE != hFind)
    {
        do
        {
            fprintf(stderr, "Processing %s...\n", FindData.cFileName);

            if (SetCurrentDirectoryA (FindData.cFileName))
            {
                (VOID) HrLoadModuleTreeFromCurrentDirectory (pTree);

                SetCurrentDirectoryA ("..");
            }
        }
        while (FFindNextChildDir (hFind, &FindData));

        FindClose (hFind);
    }

    return hr;
}

HRESULT
HrLoadModuleTreeFromFileSystem (
    OUT CModuleTree* pTree)
{
    HRESULT hr;
    WCHAR szCurrentDir [MAX_PATH];
    WCHAR szSystemDir [MAX_PATH];

    // Save the current directory because we will be changing it.
    //
    GetCurrentDirectory (celems(szCurrentDir), szCurrentDir);
    GetSystemWindowsDirectory (szSystemDir, celems(szSystemDir));
    //wcscat (szSystemDir, L"\\system32");
    SetCurrentDirectory (szSystemDir);

    hr = HrLoadModuleTreeFromCurrentDirectory (pTree);

    SetCurrentDirectory (szCurrentDir);

    fprintf(stderr, "Processed %d PE modules.  (%d total imports)\n",
        pTree->Modules.size(),
        pTree->size());

    return hr;
}