#ifdef _MSC_VER
#define UNICODE
#define _CRT_SECURE_NO_WARNINGS
#endif

#pragma comment(lib, "Kernel32.lib")
#include <Windows.h>
#include <Winternl.h>
#include <Winbase.h>
#include <psapi.h>
#include <strsafe.h>

#include <cstdio>
#include <cstdlib>
#include <ctime>
#include <cassert>

#include "dtypes.h"


#define SHARE_EXCLUSIVE             0
#define STATUS_INFO_LENGTH_MISMATCH 0xc0000004
#define DUPLICATE_SAME_ACCESS       0x00000002
#define DUPLICATE_SAME_ATTRIBUTES   0x00000004


typedef struct _SYSTEM_HANDLE
{
    ULONG ProcessId;
    BYTE ObjectTypeNumber;
    BYTE Flags;
    USHORT Handle;
    PVOID Object;
    ACCESS_MASK GrantedAccess;
} SYSTEM_HANDLE, *PSYSTEM_HANDLE;

typedef struct _SYSTEM_HANDLE_INFORMATION
{
    ULONG HandleCount;
    SYSTEM_HANDLE Handles[1];
} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;

typedef NTSTATUS(WINAPI* _NtQuerySystemInformation)(
    ULONG SystemInformationClass,
    PVOID SystemInformation,
    ULONG SystemInformationLength,
    PULONG ReturnLength
    );

typedef NTSTATUS(WINAPI* _NtDuplicateObject)(
    HANDLE SourceProcessHandle,
    HANDLE SourceHandle,
    HANDLE TargetProcessHandle,
    PHANDLE TargetHandle,
    ACCESS_MASK DesiredAccess,
    ULONG Attributes,
    ULONG Options
    );

typedef NTSTATUS(WINAPI* _NtQueryObject)(
    HANDLE ObjectHandle,
    ULONG ObjectInformationClass,
    PVOID ObjectInformation,
    ULONG ObjectInformationLength,
    PULONG ReturnLength
    );

typedef enum _POOL_TYPE
{
    NonPagedPool,
    PagedPool,
    NonPagedPoolMustSucceed,
    DontUseThisType,
    NonPagedPoolCacheAligned,
    PagedPoolCacheAligned,
    NonPagedPoolCacheAlignedMustS
} POOL_TYPE, *PPOOL_TYPE;

typedef struct _OBJECT_TYPE_INFORMATION
{
    UNICODE_STRING Name;
    ULONG TotalNumberOfObjects;
    ULONG TotalNumberOfHandles;
    ULONG TotalPagedPoolUsage;
    ULONG TotalNonPagedPoolUsage;
    ULONG TotalNamePoolUsage;
    ULONG TotalHandleTableUsage;
    ULONG HighWaterNumberOfObjects;
    ULONG HighWaterNumberOfHandles;
    ULONG HighWaterPagedPoolUsage;
    ULONG HighWaterNonPagedPoolUsage;
    ULONG HighWaterNamePoolUsage;
    ULONG HighWaterHandleTableUsage;
    ULONG InvalidAttributes;
    GENERIC_MAPPING GenericMapping;
    ULONG ValidAccess;
    BOOLEAN SecurityRequired;
    BOOLEAN MaintainHandleCount;
    USHORT MaintainTypeList;
    POOL_TYPE PoolType;
    ULONG PagedPoolUsage;
    ULONG NonPagedPoolUsage;
} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;

HANDLE fh;

HANDLE close_list[20];
DWORD close_count;

bool IsFileOpen(LPCWSTR fpath);
HANDLE GetFileHandle(LPCWSTR fpath);
PVOID GetLibraryProcAddr(LPCWSTR LibraryName, LPCSTR ProcName);
PSYSTEM_HANDLE_INFORMATION GetAllHandleData(_NtQuerySystemInformation NtQuerySystemInfo);

int main(int argc, WCHAR** argv)
{
    PSYSTEM_HANDLE_INFORMATION handleInfo;
    HANDLE close_handle = NULL;
    HANDLE proc_handle  = NULL;
    _NtQuerySystemInformation NtQuerySystemInformation = NULL;
    _NtDuplicateObject NtDuplicateObject = NULL;
    _NtQueryObject NtQueryObject = NULL;
    FILE* hLog = fopen("handle_log.txt", "w+");

    LPSYSTEMTIME win_time = new SYSTEMTIME;



    NtQuerySystemInformation    = static_cast<_NtQuerySystemInformation>(GetLibraryProcAddr(L"ntdll.dll", "NtQuerySystemInformation"));
    NtDuplicateObject           = static_cast<_NtDuplicateObject>(GetLibraryProcAddr(L"ntdll.dll", "NtDuplicateObject"));
    NtQueryObject               = static_cast<_NtQueryObject>(GetLibraryProcAddr(L"ntdll.dll", "NtQueryObject"));

    GetSystemTime(win_time);

    fprintf(hLog, "%02d/%02d/%04d\n", win_time->wMonth, win_time->wDay, win_time->wYear);
    fprintf(hLog, "%02d:%02d:%02d\n\n", win_time->wHour, win_time->wMinute, win_time->wSecond);

    // fprintf(hLog, "Getting handle data...\n");
    handleInfo = GetAllHandleData(NtQuerySystemInformation);

    // fprintf(hLog, "Preparing to parse handles...\n\n");
    ZeroMemory(close_list, sizeof(close_list));
    close_count = 0;

    // this will run each time we find that a file is still open - don't change handleInfo -
    // this will be pulled out and become a function or three in the end
    for (ULONG i = 0; i < handleInfo->HandleCount; ++i)
    {
        SYSTEM_HANDLE sys_h = handleInfo->Handles[i];
        HANDLE objdupHandle = NULL;
        POBJECT_TYPE_INFORMATION objectTypeInfo;
        HANDLE hFileMap;
        WCHAR pszFilename[MAX_PATH + 1];
        PVOID pMem = NULL;

        // open each process as needed for each handle on the list
        if (sys_h.ProcessId != 4)
        {
            if ( !(proc_handle = OpenProcess(PROCESS_DUP_HANDLE, FALSE, sys_h.ProcessId)) )
            {
                // fprintf(hLog, "Could not open PID %d!\n", sys_h.ProcessId);
                continue;
            }
        }

        // duplicate the handle into my process so i can read it
        if ( !NT_SUCCESS( NtDuplicateObject(proc_handle, (HANDLE)sys_h.Handle, GetCurrentProcess(), &objdupHandle, 0, 0, DUPLICATE_SAME_ATTRIBUTES | DUPLICATE_SAME_ACCESS) ) )
        {
            fprintf(hLog, "Duplication Error! %X\n", sys_h.Handle);
            continue;
        }

        objectTypeInfo = static_cast<POBJECT_TYPE_INFORMATION>(malloc(0x1000));

        // get the handle's type info for parsing
        if ( !NT_SUCCESS( NtQueryObject(objdupHandle, ObjectTypeInformation, objectTypeInfo, 0x1000, NULL) ) )
            {
                fprintf(hLog, "Query Error!\n", sys_h.Handle);
                CloseHandle(objdupHandle);
                continue;
            }

        if (sys_h.GrantedAccess == 0x0012019f)
        {
            /* skip these */
            fprintf(hLog, "Err: 0x12019F\n");
            free(objectTypeInfo);
            CloseHandle(objdupHandle);
            continue;
        }

        if ( wcscmp(L"File", objectTypeInfo->Name.Buffer) == 0 )
        {
            fprintf(hLog, "File!\n");

            hFileMap = CreateFileMapping(objdupHandle, NULL, PAGE_READONLY, 0, 1, NULL);

            if (!hFileMap)
            {
                fprintf(hLog, "Error: %X\n", GetLastError());
            }

            if (hFileMap)
            {
                // fprintf(hLog, "Mapped!\n");
                // create a file mapping to get the file name.
                pMem = MapViewOfFile(hFileMap, FILE_MAP_READ, 0, 0, 1);

                if (pMem)
                {
                    // fprintf(hLog, "View mapped!\n");
                    if (GetMappedFileNameW(GetCurrentProcess(), pMem, pszFilename, MAX_PATH))
                    {
                        // fprintf(hLog, "Name gotten!\n\n");
                        fwprintf(hLog, L"%ls\n", pszFilename);
                    }

                }
            }

        }
    }

    fclose(hLog);









    // printf("Current handle count: %d\n", handleInfo->HandleCount);

    // if (IsFileOpen(L"D:\\SBT\\PRO30\\APDATA\\apsysd10.dbf"))
    if ( IsFileOpen(L"D:\\Mics50\\sbtrindex.PJT") )
    {

    }

    printf("Success!");

    while (1);
    return 0;
}


bool IsFileOpen(LPCWSTR fpath)
{
    fh = CreateFile(fpath, GENERIC_READ, SHARE_EXCLUSIVE /* ask for exclusive access */, NULL, OPEN_EXISTING, 0, NULL);

    if ( fh == INVALID_HANDLE_VALUE )   // check if a file is open by trying to open it again with exclusive access
    {
        return true;
    }
    else
    {
        return false;
    }
}


PSYSTEM_HANDLE_INFORMATION GetAllHandleData(_NtQuerySystemInformation NtQuerySystemInfo)
{
    NTSTATUS query_state;
    PSYSTEM_HANDLE_INFORMATION hData;
    ULONG hDataSize = 0x80000;     // 512KB

    hData = static_cast<PSYSTEM_HANDLE_INFORMATION>(malloc(hDataSize)); // allocate for sys handle info

    while ( (query_state = NtQuerySystemInfo(static_cast<SYSTEM_INFORMATION_CLASS>(0x10), hData, hDataSize, NULL)) == STATUS_INFO_LENGTH_MISMATCH )
    {
        hData = static_cast<PSYSTEM_HANDLE_INFORMATION>(realloc(hData, hDataSize <<= 1));   // keep doubling the size of the allocation until we have enough
        printf("Bigger! %X\n", hDataSize);
    }

    if ( !NT_SUCCESS(query_state) )
    {
        printf("NtQuerySystemInformation for handle data failed!\n");
        while (1);
    }
    else
    {
        printf("Got all handles!\n");
        return hData;
    }
}


PVOID GetLibraryProcAddr(LPCWSTR LibraryName, LPCSTR ProcName)
{
    return GetProcAddress(GetModuleHandle(LibraryName), ProcName);
}