#define STRICT
//#define UNICODE
//#define _UNICODE

#include <windows.h>
#include <conio.h>
#include <stdio.h>

#include "..\..\He4HookDriverHide\He4HookDriverHide.hpp"
#include "..\..\..\NtDllTest\NtProcessList.h"
//#include "..\He4HookControlDll\ShareMemory.h"
//#include "..\He4HookControlDll\He4Message.h"

#define HE4_SERVICE_NAME           "He4HookInv"
#define HE4_SERVICE_FILE_NAME      "He4HookInv.sys" 

void ConsoleErrorMessage(void);
void ShowDeviceCurrentVersion(void);
void ShowHelp(void);
void ShowProtectedFiles(He4HookDriverHide* pNtDriverControl);
void AddProcessToUnLockList(char* pszProcessFileName, He4HookDriverHide* pNtDriverControl);
void DeleteProcessFromUnLockList(char* pszProcessFileName, He4HookDriverHide* pNtDriverControl);
void ShowUnlockThreads(He4HookDriverHide* pNtDriverControl);
BOOL InstallNewDriver(He4HookDriverHide* pNtDriverControl);
BOOL EnableDebugPriv(VOID);
void ShowStatistic(He4HookDriverHide* pNtDriverControl);

DWORD GetVersion(DWORD *pdwVer);
DWORD GetLocalBase(DWORD *pdwBase);


He4HookDriverHide* pNtDriverControl = NULL;
BOOL               bForceLoadDriver = FALSE;
BOOL               bShowProtectedFiles = FALSE;
int                bHookFileSystem = -1;
BOOL               bLockSaveObjectsForAllThreads = FALSE;
BOOL               bClearSaveList = FALSE;
DWORD              dwFileAccessType = 0;
DWORD              dwProcessAccessType = HE4_UNLOCK_READ | HE4_UNLOCK_WRITE | HE4_UNLOCK_DELETE | HE4_UNLOCK_VISIBLE;
BOOL               bShowUnlockThreadInfo = FALSE;
BOOL               bShowStatistic = FALSE;

//
// -h              - Show help
// -i:n            - Install driver type: 
//                   n = 0  - open exist driver (default)
//                   n != 0 - load new image driver force.
// -s              - Show protected files
// -u:process_name - Add process to unlock list
// -l:process_name - Remove process from unlock list
// -cp:_string_    - Access type for process:
//                   R - read
//                   W - write
//                   D - delete
//                   V - visible
//                   (Example: -c:RV)
// -la             - Remove all processes from unlock list
// -hk:n           - Hook/Unhook file system:
//                   n = 1 - hook Zw*/Nt* func.\n\
//                   n = 2 - hook DRIVER_OBJECT\n\
// -a:             - Add file to save list
// -d:             - Delete file from save list
// -da             - Delete all files from save list
// -c:             - Access type
// -t              - Show unlock thread info
// -q              - Show statistic
//

//#pragma pack(push)
//#pragma pack(16)
//typedef struct _TEST_ALIGN
//{
//  char    m_Symbol;
//  __int64    m_Long;
//} TEST_ALIGN, *PTEST_ALIGN;
//#pragma pack(pop)

int main(int nCountArg, char* lpszArg[], char* lpszEnv[])
{
//  EnableDebugPriv();
//  CSharedMemory cSharedMemory;
//
//  if (cSharedMemory.Open(HE4_CONTROL_DLL_SHARED_MEMORY, 4096, 100))
//  {
//    char Msg[2048];
//    PHE4_CONTROL_MSG pMsg = (PHE4_CONTROL_MSG) Msg;
//    
//    pMsg->m_dwMessageId = 2;
//    strcpy((char*)pMsg->m_MessageBody, "D:\\tmp\\guard\\He4HookInv.sys");
//    pMsg->m_dwSizeByBytes = SIZE_OF_HE4_CONTROL_MSG_REAL + strlen((char*) pMsg->m_MessageBody) + sizeof(char);
//
//    cSharedMemory.Write(pMsg, pMsg->m_dwSizeByBytes, 500);
//
//    memset(Msg, 0, sizeof(Msg));
//
//    cSharedMemory.Write(pMsg, pMsg->m_dwSizeByBytes, 500);
//
//
//
//    cSharedMemory.Close();
//    printf("He4Dev@hotmal.com\n\n");
//  }
  
//  TEST_ALIGN  TestAlign;
//  PTEST_ALIGN pTestAlign = &TestAlign;
//  PCHAR pSymbol = (PCHAR)&(pTestAlign->m_Symbol);
//  PCHAR pLong = (PCHAR)&(pTestAlign->m_Long);
//  int nRest = pLong-pSymbol;

  printf("\nHe4HookControl v2.03 - control utility for He4HookInv\n");
  printf("Copyright (C) 2000 He4 developers team\n");
  printf("He4Dev@hotmail.com\n\n");

  ShowDeviceCurrentVersion();

  if (nCountArg <= 1)
  {
    ShowHelp();
    return 0;
  }

  for (int i=1; i<nCountArg; i++)
  {
    if (!stricmp(lpszArg[i], "-h"))
    {
      ShowHelp();
    }
    if (!strnicmp(lpszArg[i], "-i:", 3))
    {
      if (strlen(lpszArg[i]) > 3)
        bForceLoadDriver = (BOOL)atoi(lpszArg[i]+3);
    }
    if (!stricmp(lpszArg[i], "-s"))
    {
      bShowProtectedFiles = TRUE;
    }
    if (!strnicmp(lpszArg[i], "-hk:", 4))
    {
      if (strlen(lpszArg[i]) > 4)
        bHookFileSystem = atoi(lpszArg[i]+4);
    }
    if (!strnicmp(lpszArg[i], "-c:", 3))
    {
      if (strlen(lpszArg[i]) > 3)  
      {
        dwFileAccessType = 0;
        char* pszAcc = lpszArg[i]+3;
        for (; *pszAcc; pszAcc++)
        {
          if (*pszAcc == 'R' || *pszAcc == 'r')
            dwFileAccessType |= ACC_TYPE_READ;
          if (*pszAcc == 'W' || *pszAcc == 'w')
            dwFileAccessType |= ACC_TYPE_WRITE;
          if (*pszAcc == 'D' || *pszAcc == 'd')
            dwFileAccessType |= ACC_TYPE_DELETE;
          if (*pszAcc == 'V' || *pszAcc == 'v')
            dwFileAccessType |= ACC_TYPE_VISIBLE;
          if (*pszAcc == 'E' || *pszAcc == 'e')
            dwFileAccessType |= FILE_ACC_TYPE_EXCHANGE;
        }
      }
    }
    if (!strnicmp(lpszArg[i], "-cp:", 4))
    {
      if (strlen(lpszArg[i]) > 4)  
      {
//        dwFileAccessType = 0;
        char* pszAcc = lpszArg[i]+4;
        dwProcessAccessType = 0;
        for (; *pszAcc; pszAcc++)
        {
          if (*pszAcc == 'R' || *pszAcc == 'r')
            dwProcessAccessType |= HE4_UNLOCK_READ;
          if (*pszAcc == 'W' || *pszAcc == 'w')
            dwProcessAccessType |= HE4_UNLOCK_WRITE;
          if (*pszAcc == 'D' || *pszAcc == 'd')
            dwProcessAccessType |= HE4_UNLOCK_DELETE;
          if (*pszAcc == 'V' || *pszAcc == 'v')
            dwProcessAccessType |= HE4_UNLOCK_VISIBLE;
        }
      }
    }
    if (!stricmp(lpszArg[i], "-la"))
      bLockSaveObjectsForAllThreads = TRUE;
    if (!stricmp(lpszArg[i], "-da"))
      bClearSaveList = TRUE;
    if (!stricmp(lpszArg[i], "-t"))
      bShowUnlockThreadInfo = TRUE;
    if (!stricmp(lpszArg[i], "-q"))
      bShowStatistic = TRUE;
  }


  TCHAR szExeFileName[2048];
  TCHAR szDeviceFileName[2048];

  GetModuleFileName(NULL, szExeFileName, 2048);

  TCHAR drive[_MAX_DRIVE];   
  TCHAR dir[_MAX_DIR];
  TCHAR fname[_MAX_FNAME];   
  TCHAR ext[_MAX_EXT];

  _splitpath(szExeFileName, drive, dir, fname, ext);
  lstrcpy(szDeviceFileName, _T("\\??\\"));
  _makepath(szDeviceFileName+sizeof(_T("\\??\\"))-sizeof(TCHAR), drive, dir, HE4_SERVICE_FILE_NAME, _T(""));

  pNtDriverControl = new He4HookDriverHide(szDeviceFileName);
  if (pNtDriverControl == NULL)
  {
    printf("\nNo memory for create class He4HookDriverHide!!!\n");
    return -1;
  }

  if (pNtDriverControl->Result == FALSE)
  {
    delete pNtDriverControl;
    printf("\nCreate class He4HookDriverHide - ERROR!!!\n");
    return -1;
  }

  if (bForceLoadDriver)
  {
    if (InstallNewDriver(pNtDriverControl) == FALSE)
    {
      delete pNtDriverControl;
      printf("\nLoad new driver - ERROR!!!\n");
      return -1;
    }

    bShowProtectedFiles = TRUE;
    
    printf("\nNew version driver:\n");
    ShowDeviceCurrentVersion();
  }

  if (bHookFileSystem != -1)
  {
    DWORD dwDrivesMaskReal;
    if ((dwDrivesMaskReal = pNtDriverControl->HookFileSystem((DWORD)bHookFileSystem)) == (DWORD)-1)
    {
      if (bHookFileSystem)
        printf("Hook file system - ERROR!!!\n");
      else
        printf("Unhook file system - ERROR!!!\n");
    }
    else
    {
      if (bHookFileSystem)
        printf("File system - hooked\n");
      else
        printf("File system - unhooked\n");
    }
  }

  if (bLockSaveObjectsForAllThreads == TRUE)
    pNtDriverControl->LockSaveObjectsForAllThreads();

  if (bClearSaveList == TRUE)
    pNtDriverControl->ClearSaveList();
  
  for (i=1; i<nCountArg; i++)
  {
    if (!strnicmp(lpszArg[i], "-u:", 3))
    {
      if (strlen(lpszArg[i]) > 3)
      {
        AddProcessToUnLockList(lpszArg[i]+3, pNtDriverControl);
        bShowUnlockThreadInfo = TRUE;
      }
    }
    if (!strnicmp(lpszArg[i], "-l:", 3))
    {
      if (strlen(lpszArg[i]) > 3)
      {
        DeleteProcessFromUnLockList(lpszArg[i]+3, pNtDriverControl);
        bShowUnlockThreadInfo = TRUE;
      }
    }
    
    if (!strnicmp(lpszArg[i], "-a:", 3))
    {
      if (strlen(lpszArg[i]) > 3)
      {
        if (dwFileAccessType & FILE_ACC_TYPE_EXCHANGE)
        {
          char* pszFirstFile = strtok(lpszArg[i]+3, "=");
          if (pszFirstFile != NULL)
            pNtDriverControl->AddToSaveList(pszFirstFile, dwFileAccessType, strtok(NULL, "="));
        }
        else
          pNtDriverControl->AddToSaveList(lpszArg[i]+3, dwFileAccessType);
        bShowProtectedFiles = TRUE;
      }
    }
    if (!strnicmp(lpszArg[i], "-d:", 3))
    {
      if (strlen(lpszArg[i]) > 3)
      {
        pNtDriverControl->DelFromSaveList(lpszArg[i]+3);
        bShowProtectedFiles = TRUE;
      }
    }
  }

  if (bShowProtectedFiles)
  {
    printf("\nProtected files list:\n");
    ShowProtectedFiles(pNtDriverControl);
  }

  if (bShowUnlockThreadInfo)
  {
    ShowUnlockThreads(pNtDriverControl);
  }
  
  if (bShowStatistic)
  {
    ShowStatistic(pNtDriverControl);
  }


  delete pNtDriverControl;

  return 0;
}

void ShowHelp(void)
{
  printf("\n\
   -h                - Show help\n\
   -i:n              - Install driver type: \n\
                       n = 0  - open exist driver (default)\n\
                       n != 0 - load new image driver force.\n\
   -s                - Show protected files\n\
   -u:process_name   - Add process to unlock list\n\
   -l:process_name   - Remove process from unlock list\n\
   -cp:_string_      - Access type for process:\n\
                       R - read\n\
                       W - write\n\
                       D - delete\n\
                       V - visible\n\
                       (Example: -c:RV)\n\
   -la               - Remove all processes from unlock list\n\
   -hk:n             - Hook/Unhook file system:\n\
                       n = 0  - unhook\n\
                       n = 1 - hook Zw*/Nt* func.\n\
                       n = 2 - hook DRIVER_OBJECT\n\
   -a:full_file_name - Add file to save list\n\
   -d:full_file_name - Delete file from save list\n\
   -da               - Delete all files from save list\n\
   -c:_string_       - Access type for file:\n\
                       R - read\n\
                       W - write\n\
                       D - delete\n\
                       V - visible\n\
                       E - exchange (hook method = 1 (-hk:1))\n\
                       (Example: -c:RV)\n\
   -t                - Show unlock thread info\n\
   -q                - Show statistic\n\
   Examples:\n\
   He4HookControl.exe -a:c:\\MyFile -c:RV\n\
   He4HookControl.exe -a:c:\\MyFile=c:\\MyFileNew -c:ERV\n"
        );
}

void AddProcessToUnLockList(char* pszProcessFileName, He4HookDriverHide* pNtDriverControl)
{
  if (pszProcessFileName == NULL || pNtDriverControl == NULL)
    return;

  NTSTATUS           NtStatus;
  PROCESS_INFO       ProcInfo[1024];
  PPROCESS_INFO      lpProcInfo;
  ULONG              LehgthReturned = 0;
//  ULONG              i;

  NtStatus = NtGetProcessList(ProcInfo, sizeof(ProcInfo), &LehgthReturned);
  if (NtStatus == STATUS_SUCCESS)
  {
    lpProcInfo = ProcInfo;
    EnableDebugPriv();
    char szProcName[1024];
    while (1)
    {
      if (lpProcInfo->ProcessName.Buffer != NULL)