#include <stdio.h>
#include <windows.h>
#include <wchar.h>
#include <process.h>
#include "SystemHandles.h"


HANDLE OpenProcessWrapper(DWORD processId)
{
    // Open the process for handle duplication
    return OpenProcess( PROCESS_DUP_HANDLE, TRUE, processId );
}


HANDLE DuplicateHandleWrapper(HANDLE hProcess, HANDLE hRemote)
{
    HANDLE hDup = NULL;
    // Duplicate the remote handle for our process
    DuplicateHandle( hProcess, hRemote, GetCurrentProcess(), &hDup, 0, FALSE, DUPLICATE_SAME_ACCESS );
    return hDup;
}


void GetFileNameThread( PVOID pParam )
{
    int len;
    GetFileNameThreadParam* p = (GetFileNameThreadParam*)pParam;
    p->rc = FALSE;

    UCHAR lpBuffer[MAX_PATH*2];
    PFILE_NAME_INFORMATION pFNI = (PFILE_NAME_INFORMATION)lpBuffer; 
    DWORD iob[2];

    DWORD ntRet = NtQueryInformationFile( p->hFile, iob, lpBuffer, MAX_PATH*2, 9 );
    if ( ntRet == STATUS_SUCCESS){
        len  = pFNI->FileNameLength/2;
        if (len > MAX_PATH)
        {
            len = MAX_PATH-1;
        }
        wcsncpy(p->pName,  pFNI->FileName, len );
        p->pName [len] = L'\0';
        p->rc  = TRUE;
    }
}


BOOL GetFileName(HANDLE h, PWCHAR pName, DWORD processId )
{
    ULONG size = 0x2000;
    UCHAR* lpBuffer = NULL;
    BOOL ret = FALSE;

    HANDLE handle;
    HANDLE hRemoteProcess = NULL;
    BOOL remote = processId != GetCurrentProcessId();

    if ( remote )
    {
        // Open the remote process
        hRemoteProcess = OpenProcessWrapper( processId );
        if ( hRemoteProcess == NULL )
        {
            return FALSE;
        }
        // Duplicate the handle
        handle = DuplicateHandleWrapper( hRemoteProcess, h );
    }
    else
    {
        handle = h;
    }


        // Query the info size
    NtQueryObject( handle, 2, NULL, 0, &size );


    lpBuffer = (UCHAR*)VirtualAlloc(NULL, sizeof(UCHAR)*size, MEM_COMMIT, PAGE_READWRITE );
    // Query the info size ( type = 2 = ObjectTypeInformation)
    if ( NtQueryObject( handle, 2, lpBuffer, size, NULL ) == 0 )
    {
        int i = _wcsicmp (L"File", (wchar_t*)(lpBuffer+0x60)); // * 0x60 to skip data from OBJECT_INFORMATION_CLASS
        if (i == 0)
        {
            GetFileNameThreadParam tp;
            tp.hFile = handle;
            tp.pName = pName;


            //Start the thread to get the file name
            HANDLE hThread = (HANDLE)_beginthread( GetFileNameThread, 0, &tp );
            if ( hThread != NULL )
            {
                if ( WaitForSingleObject(hThread, 50 ) == WAIT_TIMEOUT)
                {     
                    // Access denied, terminate the thread
                    TerminateThread(  hThread, 0 );
                    wcscpy(pName,  L"");
                    ret  = FALSE;
                }
                else
                {
                    ret = tp.rc;
                }
            } // else goto cleanup;
        } // else goto cleanup;
     }


//cleanup:
    if ( remote )
    {
        if ( hRemoteProcess != NULL )
            CloseHandle( hRemoteProcess );


        if ( handle != NULL )
            CloseHandle( handle );
    }
    if ( lpBuffer != NULL )
        VirtualFree(lpBuffer, 0, MEM_RELEASE);
    return ret;
}




JNIEXPORT jboolean JNICALL Java_SystemHandles_initNative (JNIEnv *env, jobject self)
{
    HANDLE ntdll = LoadLibrary("ntdll.dll");
    NtQuerySystemInformation = (PNtQuerySystemInformation)
                            GetProcAddress((HINSTANCE)ntdll,     "NtQuerySystemInformation" );
    NtQueryObject = (PNtQueryObject)
                            GetProcAddress((HINSTANCE)ntdll,     "NtQueryObject");
    NtQueryInformationFile = (PNtQueryInformationFile)
                            GetProcAddress((HINSTANCE)ntdll,     "NtQueryInformationFile" );
    HANDLE psapi = LoadLibrary("PSAPI.DLL");
    GetProcessImageFileName = (PGetProcessImageFileName)
                            GetProcAddress((HINSTANCE)psapi,    "GetProcessImageFileNameW");

    EnumProcessModules = (PEnumProcessModules)
                            GetProcAddress((HINSTANCE)psapi, "EnumProcessModules");


    GetModuleBaseName = (PGetModuleBaseName)
                            GetProcAddress((HINSTANCE)psapi, "GetModuleBaseNameW");

    GetMappedFileName = (PGetMappedFileName)                        
                            GetProcAddress((HINSTANCE)psapi, "GetMappedFileNameW");

    jboolean ret = NtQuerySystemInformation && NtQueryObject && NtQueryInformationFile && 
        GetProcessImageFileName && EnumProcessModules && GetModuleBaseName && GetMappedFileName;
    return ret;
}


PSYSTEM_HANDLE_INFORMATION SystemSnapshot()
{
    DWORD size = 0x6000;
    DWORD needed = 0;
    PSYSTEM_HANDLE_INFORMATION pHI;

    pHI = (PSYSTEM_HANDLE_INFORMATION)
             VirtualAlloc( NULL, size, MEM_COMMIT, PAGE_READWRITE );
    if (! pHI)
    {
        return NULL;
    }

    // SYSTEM_INFORMATION_CLASS SystemHandleInformation <-- 16
    if ( NtQuerySystemInformation( 16, pHI, size, &needed ) != 0 )
    {
           if ( needed == 0 )
           {
                VirtualFree(  pHI, 0, MEM_RELEASE );
                return NULL;
           }

           // The size was not enough
           VirtualFree( pHI, 0, MEM_RELEASE );

           pHI = (PSYSTEM_HANDLE_INFORMATION)
                          VirtualAlloc( NULL, size = needed + 256, MEM_COMMIT, PAGE_READWRITE );
    }
    if ( pHI == NULL )
    {
        return NULL;
    }

    // Query the objects ( system wide )
    if ( NtQuerySystemInformation( 16, pHI, size, NULL ) != 0 ) {
           VirtualFree( pHI, 0, MEM_RELEASE );
           return NULL;
    }

    return pHI; 
}


JNIEXPORT jobject JNICALL Java_SystemHandles_getOpenFilesNative (JNIEnv *env, jobject self)
{
    jclass cls;
    jmethodID constructor;
    jmethodID addFkt;
    jvalue args[3];
    jobject list;

    // reference of the class SystemHandleList
    cls = (*env)->FindClass(env, "SystemFileHandleList");
    // get a reference to the constructor; the name is <init>
    constructor = (*env)->GetMethodID(env, cls, "<init>", "()V");
    // get a reference to the add(String, int, int) method
    addFkt = (*env)->GetMethodID(env, cls, "add", "(Ljava/lang/String;II)V");
    // create a new instance of the class
    list = (*env)->NewObject(env, cls, constructor);


    

    PSYSTEM_HANDLE_INFORMATION pHI = SystemSnapshot();
    if (pHI != NULL)
    {
        int i = 0;
        for (i = 0; i < pHI->NumberOfHandles; i++ )
        {    
            SYSTEM_HANDLE_TABLE_ENTRY_INFO  h = pHI->Handles[i];

            WCHAR fileName[MAX_PATH];
            if (!GetFileName((HANDLE)h.HandleValue, fileName, h.UniqueProcessId))
            {
                continue;
            }

//            WCHAR fn[MAX_PATH];
//            GetFileNameFromHandle(h.HandleValue, h.UniqueProcessId, fn);

            args[0].l = (*env)->NewString(env, fileName, wcslen(fileName));
            args[1].i = h.HandleValue;
            args[2].i = h.UniqueProcessId;

            (*env)->CallVoidMethodA(env, list, addFkt, args);
        }
    }
    return list;
}




JNIEXPORT jstring JNICALL Java_SystemHandles_getProcessFullNameNative (JNIEnv *env, jobject self, jint pid)
{
    BOOL res;
    int k = 0;
    WCHAR name[MAX_PATH];
    HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid );
    int count = GetProcessImageFileName(hProcess, name, MAX_PATH);


    res = count  != 0;
    if (! res)
    {
        k = GetLastError();
        return (*env)->NewStringUTF(env, "<unknown>");
    }

    return (*env)->NewString(env, name, wcslen(name));
}


JNIEXPORT jstring JNICALL Java_SystemHandles_getProcessNameNative (JNIEnv *env, jobject self, jint processID)
{
    WCHAR szProcessName[MAX_PATH] = L"<unknown>";

    // Get a handle to the process.

    HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
                                   PROCESS_VM_READ,
                                   FALSE, processID );

    // Get the process name.

    if (hProcess != NULL)
    {
        HMODULE hMod;
        DWORD cbNeeded;

        if ( EnumProcessModules( hProcess, &hMod, sizeof(hMod), &cbNeeded) )
        {
            GetModuleBaseName( hProcess, hMod, szProcessName, sizeof(szProcessName)/sizeof(TCHAR) );
        }
        CloseHandle( hProcess );
    }
    return (*env)->NewString(env, szProcessName, wcslen(szProcessName));
}