#include <windows.h>
#include <cstdlib>
#include <iostream>
#include <tlhelp32.h>

using namespace std;

typedef HANDLE (__stdcall *PfOpenT)(DWORD, BOOL, DWORD);
DWORD GetPIDByProcess(char argv[]);
void SetDebugPrivilege();

void SetDebugPrivilege(){
TOKEN_PRIVILEGES privilege;
LUID Luid;
HANDLE handle1;
HANDLE handle2;
handle1 = OpenProcess(PROCESS_ALL_ACCESS, FALSE, GetCurrentProcessId());
OpenProcessToken(handle1, TOKEN_ALL_ACCESS, &handle2);
LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &Luid);
privilege.PrivilegeCount = 1;
privilege.Privileges[0].Luid = Luid;
privilege.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
AdjustTokenPrivileges(handle2, FALSE, &privilege, sizeof(privilege), NULL, NULL);
CloseHandle(handle2);
CloseHandle(handle1);
} 

DWORD GetPIDByProcess(char argv[])
{
      printf("\n[*] Getting PID...");
      HANDLE hSnapShot;
      PROCESSENTRY32 uProcess;
      
      hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
      uProcess.dwSize = (DWORD) sizeof(PROCESSENTRY32);
      
      int p = 0;
      DWORD PID;
      
      p = Process32First(hSnapShot, &uProcess);
      
      while(p)
      {
              if(strstr(uProcess.szExeFile, argv))
              {
                                           printf("\n   [+]Process found : %s", uProcess.szExeFile);
                                           PID = uProcess.th32ProcessID;
                                           printf("\n   [+]Process ID : 0x%x", PID);
                                           return PID;
              }
              p = Process32Next(hSnapShot, &uProcess);
      }
      CloseHandle(hSnapShot);
      return 0;
}

int main(int args, char *argv[])
{
    SetDebugPrivilege();
    
    //Recup de la fonction OpenThread dans kernel32.dll
    HMODULE hModDll = NULL;

    hModDll = LoadLibrary("kernel32.dll");
    PfOpenT OpenThread = (PfOpenT)GetProcAddress(hModDll, "OpenThread");
    
    ////////////////////////////////////

    char Process[100];
    DWORD PID, TID;
    HANDLE hSnapShot, hThread;
    THREADENTRY32 th32;
    CONTEXT Context;

	Context.ContextFlags = CONTEXT_FULL;
    printf("            ********* Stack Watcher *********");
    
    cout<<"\n\nProcess you want to analyse : ";
    cin>>Process;
    
    PID = GetPIDByProcess(Process);
    
    //////////////////////////////////////
    
    //On établit la liste de tous les threads
    hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if(hSnapShot == INVALID_HANDLE_VALUE)
    {
        printf("\n[!]Cannot create process list.\n---- Error with CreateToolhelp32Snapshot() : 0x%x\n\n", GetLastError());
        system("pause");
        return 0;
    }
    
    th32.dwSize = (DWORD) sizeof(THREADENTRY32);
    
    //On se place sur le premier thread
    Thread32First(hSnapShot, &th32);
    
    do
    {
        if(th32.th32OwnerProcessID == PID)
        {
            //On a trouvé un thread qui appartient au processus qu'on veut analyser
            
            TID = th32.th32ThreadID;
            
            printf("\n--------------------\\ Thread ID : 0x%x /--------------------", TID);
            
            /*
            HANDLE WINAPI OpenThread(
                  __in  DWORD dwDesiredAccess,
                  __in  BOOL bInheritHandle,
                  __in  DWORD dwThreadId
            );
            */

            //Droits d'accès à un thread : http://msdn2.microsoft.com/en-us/library/ms686769(VS.85).aspx
           
           hThread=OpenThread(THREAD_QUERY_INFORMATION|THREAD_SUSPEND_RESUME|THREAD_GET_CONTEXT, false, TID);
           if(hThread != NULL)
            {
                //On suspend le thread pour récupérer les infos dessus
                SuspendThread(hThread);
                
                //C'est avec cette fonction qu'on va pouvoir récupérer toutes les infos qui nous interressent
                if(!GetThreadContext(hThread, &Context))
                {
                    printf("\n[!]Error with GetThreadContext : 0x%x\n\n", GetLastError());
                    system("pause");
                    return 0;
                }
                
                /*
                typedef struct _CONTEXT {
                    	DWORD	ContextFlags;
                    	DWORD	Dr0;
                    	DWORD	Dr1;
                    	DWORD	Dr2;
                    	DWORD	Dr3;
                    	DWORD	Dr6;
                    	DWORD	Dr7;
                    	FLOATING_SAVE_AREA FloatSave;
                    	DWORD	SegGs;
                    	DWORD	SegFs;
                    	DWORD	SegEs;
                    	DWORD	SegDs;
                    	DWORD	Edi;
                    	DWORD	Esi;
                    	DWORD	Ebx;
                    	DWORD	Edx;
                    	DWORD	Ecx;
                    	DWORD	Eax;
                    	DWORD	Ebp;
                    	DWORD	Eip;
                    	DWORD	SegCs;
                    	DWORD	EFlags;
                    	DWORD	Esp;
                    	DWORD	SegSs;
                    	BYTE	ExtendedRegisters[MAXIMUM_SUPPORTED_EXTENSION];
                } CONTEXT;
                */
                
                printf("\n\n        ***** REGISTERS *****\n");
                printf("\n  EAX : 0x%x ---> %d", Context.Eax, Context.Eax);
                printf("\n  EBX : 0x%x ---> %d", Context.Ebx, Context.Ebx);
                printf("\n  ECX : 0x%x ---> %d", Context.Ecx, Context.Ecx);
                printf("\n  EDX : 0x%x ---> %d", Context.Edx, Context.Edx);
                printf("\n  ESI : 0x%x ---> %d", Context.Esi, Context.Esi);
                printf("\n  EDI : 0x%x ---> %d", Context.Edi, Context.Edi);
                printf("\n  EBP : 0x%x ---> %d", Context.Ebp, Context.Ebp);
                printf("\n  ESP : 0x%x ---> %d", Context.Esp, Context.Esp);
                
                printf("\n\n        ***** DEBUG REGISTERS *****\n");
                printf("\n  DR0 : 0x%x ---> %d", Context.Dr0, Context.Dr0);
                printf("\n  DR1 : 0x%x ---> %d", Context.Dr1, Context.Dr1);
                printf("\n  DR2 : 0x%x ---> %d", Context.Dr2, Context.Dr2);
                printf("\n  DR3 : 0x%x ---> %d", Context.Dr3, Context.Dr3);
                /*
                printf("\n  DR4 : 0x%x ---> %d", Context.Dr4, Context.Dr4);
                printf("\n  DR5 : 0x%x ---> %d", Context.Dr5, Context.Dr5);
                */
                printf("\n  DR6 : 0x%x ---> %d", Context.Dr6, Context.Dr6);
                printf("\n  DR7 : 0x%x ---> %d", Context.Dr7, Context.Dr7);
                
                printf("\n\n          ***** SEGMENTS *****\n");
                printf("\n  CS : 0x%x ---> %d", Context.SegCs, Context.SegCs);
                printf("\n  DS : 0x%x ---> %d", Context.SegDs, Context.SegDs);
                printf("\n  SS : 0x%x ---> %d", Context.SegSs, Context.SegSs);
                printf("\n  ES : 0x%x ---> %d", Context.SegEs, Context.SegEs);
                printf("\n  FS : 0x%x ---> %d", Context.SegFs, Context.SegFs);
                printf("\n  GS : 0x%x ---> %d", Context.SegGs, Context.SegGs);
                
                printf("\n\nEFLAG : 0x%x", Context.EFlags);
                
                //On relance le thread
                ResumeThread(hThread);
                
                CloseHandle(hThread);
            }
            else
            {
                printf("\n[!]Cannot open thread.\n---- Error with OpenThread() : 0x%x\n\n", GetLastError());
                system("pause");
                return 0;
            }
        }
    }while(Thread32Next(hSnapShot, &th32)); // On parcourt tous les threads
    
    cout<<"\n\n";
    CloseHandle(hSnapShot);
    system("pause");
    return 0;
}