/*
  Name: ProcessInfo Class
  Author: Rik
  Description: Used to store process information
*/



#ifndef _PINFO_H_
#define _PINFO_H_


#include <string>
#include <vector>
#include <windows.h>
#include <Tlhelp32.h>
#include <cstdio>


class ProcessInformation;

typedef std::vector<ProcessInformation> ProcessList;


class ProcessInformation
{
public:
        ProcessInformation();
        ~ProcessInformation();

        std::string GetProcessName();
        unsigned int GetProcessId();

        friend bool GetProcessInformation(ProcessList&);
private:
        std::string ProcessName;
        unsigned int uiProcessId;
};

#endif



/*
  Name: ProcessInfo Class
  Author: Rik
  Description: Used to store process information
*/

#include "ProcessInfo.h"

#include <iostream>
using namespace std;


ProcessInformation::ProcessInformation()
 {
   #ifdef DEBUG
   printf("ProcessInformation Constructor called!\n");
   #endif
   this->uiProcessId = 0;
 }


ProcessInformation::~ProcessInformation()
 {
   #ifdef DEBUG
   printf("ProcessInformation Destructor called!\n");
   #endif
 }


unsigned int ProcessInformation::GetProcessId()
 {
   return this->uiProcessId;
 }


std::string ProcessInformation::GetProcessName()
 {
   return this->ProcessName;
 }







bool GetProcessInformation(ProcessList &List)
 {
    ProcessInformation ProcInfo;
    HANDLE hProcSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS,0);

    if( hProcSnapshot == INVALID_HANDLE_VALUE )
     {
        return false;
     }

    PROCESSENTRY32 ProcessInfo;
    memset((LPPROCESSENTRY32)&ProcessInfo,0,sizeof(PROCESSENTRY32));
    ProcessInfo.dwSize = sizeof(PROCESSENTRY32);

   if(!Process32First(hProcSnapshot,&ProcessInfo))
     {
       return false;
     }

    do
    {
      ProcInfo.ProcessName = ProcessInfo.szExeFile;
      ProcInfo.uiProcessId = ProcessInfo.th32ProcessID;
      List.push_back(ProcInfo);
    } while( Process32Next(hProcSnapshot,&ProcessInfo) ) ;

  return true;
 }


/*
  Name: FuncMap functions
  Author: Rik
  Description: Used to map function names to addresses
*/
#ifndef _FMAP_H_
#define _FMAP_H_

#include <windows.h>
#include <winnt.h>
#include <map>
#include <string>


typedef std::map<std::string,unsigned int> FuncMap;

unsigned int MapExportedFunctionsFromLibrary(std::string LibraryName,FuncMap& FuncToAddr);



#endif


/*
  Name: FuncMap functions
  Author: Rik
  Description: Used to map function names to addresses
*/

#include "FuncMap.h"

unsigned int MapExportedFunctionsFromLibrary(std::string LibraryName,FuncMap& FuncToAddr)
 {
   IMAGE_DOS_HEADER* pDosHeader;
   IMAGE_NT_HEADERS* pNTHeader;
   IMAGE_EXPORT_DIRECTORY* pExportDir;

   HMODULE hMod = GetModuleHandle(LibraryName.c_str());
   if( !hMod ) return 0;

   pDosHeader = (IMAGE_DOS_HEADER*)hMod;
   if( pDosHeader->e_magic != IMAGE_DOS_SIGNATURE ) return 0;
   pNTHeader = (IMAGE_NT_HEADERS*)((DWORD)pDosHeader + (DWORD)pDosHeader->e_lfanew);

   pExportDir = (IMAGE_EXPORT_DIRECTORY*)((DWORD)pDosHeader+(DWORD) pNTHeader->OptionalHeader.DataDirectory
       [IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress);


  unsigned int* FuncNameAddr = (unsigned int*)((DWORD)hMod + pExportDir->AddressOfNames);
  unsigned int* FuncAddress = (unsigned int*) ((DWORD)hMod + pExportDir->AddressOfFunctions);

  for(unsigned int x = 0; x < pExportDir->NumberOfFunctions;x++)
  {
    char* pFuncName =(char*)((DWORD)hMod + FuncNameAddr[x]);
    FuncToAddr[pFuncName] = ((DWORD)hMod + FuncAddress[x]);
  }
  return pExportDir->NumberOfFunctions;

 }



/*
  Name: Serenity Detect
  Author: Rik
  Description: Finds and corrects Hooks in the Import Address Tables of executables
*/

#include <cstdlib>
#include <iostream>
#include <windows.h>
#include <winsock2.h>
#include <assert.h>
#include "ProcessInfo.h"
#include "FuncMap.h"

#define _DEBUG_

using namespace std;

// TYPEDEFS

typedef std::vector<std::string> DllList;

// GLOBALS

ProcessList ProcList; // Holds Process Names and Id's

FuncMap AddrOf; // Holds function names and the addresses of the functions

const char Library[] = "SerenityInject.dll";
const unsigned int uiPort = 1337;

unsigned int uiLibSock = 0;

// FUNCTIONS

unsigned int __stdcall explode(string str,string delimiter,vector<string> &container);
unsigned int __stdcall RecvModuleNames(int iSock,DllList& List);
DWORD __stdcall AcceptThread(void* listening_socket);

int main(int argc, char *argv[])
{

     // Set Up Listening Socket

    WSADATA Data;
    WSAStartup(MAKEWORD(2,0),&Data);

    sockaddr_in Config;
    memset((sockaddr*)&Config,'\0',sizeof(Config));

    Config.sin_family = AF_INET;
    Config.sin_port = htons(uiPort);
    Config.sin_addr.s_addr = inet_addr("127.0.0.1");

    unsigned int uiListenSocket = socket(AF_INET,SOCK_STREAM,0);

    if( bind(uiListenSocket,(sockaddr*)&Config,sizeof(Config)) != 0  )
     {
       cout<<"Bind failed with error: "<<WSAGetLastError()<<"!\n\n";
       system("pause");
       exit(1);
     }

    if( listen( uiListenSocket,5 ) != 0 )
     {
       cout<<"Listen failed with error: "<<WSAGetLastError()<<"!\n\n";
       system("pause");
       exit(1);
     }

    // Accept Thread

    HANDLE AcceptThreadHandle = CreateThread(0,0,AcceptThread,(void*)&uiListenSocket,0,0);

    //Get Process List
    bool ProcListRet = GetProcessInformation(ProcList);
    if( !ProcListRet )
      {
       cout<<"Failed to retrieve Process List!\n";
       system("pause");
       exit(0);
      }
    #ifdef _DEBUG_
    cout<<"GetProcessInformation called successfully!\n\n";
    #endif

    // Print List to User
    int iIndex = 0;
    ProcessList::iterator pIter;
    for( pIter = ProcList.begin(); pIter != ProcList.end(); pIter++,iIndex++ )
    cout<<"["<<iIndex<<"] "<<pIter->GetProcessName()<<"\n\n";


    // Get Desired Process
      while( true )
      {
        cout<<"Process Index: ";
        cin>>iIndex;
        if( iIndex >= ProcList.size() )
         {
           cout<<"Index not valid!\n\n";
           continue;
         }
        break;
     }

    // Open Handle to Desired Process

    HANDLE hProcHandle = OpenProcess(PROCESS_ALL_ACCESS,false,ProcList[iIndex].GetProcessId());

    if( !hProcHandle )
    {
    cout<<"Failed to open handle to "+ProcList[iIndex].GetProcessName()+" with error: ";
    cout<<GetLastError()<<"!\n\n";
    system("pause");
    exit(0);
    }


    // Write Name Of Library to Remote Process

   void* pMemoryAddress = VirtualAllocEx(hProcHandle,0,strlen(Library)+1,MEM_COMMIT,PAGE_READWRITE);
   if( !pMemoryAddress )
     {
       cout<<"Failed to allocate memory in "+ProcList[iIndex].GetProcessName()+"\n\n";
       system("pause");
       exit(1);
     }

   if( !WriteProcessMemory(hProcHandle,pMemoryAddress,Library,strlen(Library),NULL) )
     {
      cout<<"Failed to write to memory in "+ProcList[iIndex].GetProcessName()+"\n\n";
      system("pause");
      exit(1);
     }



    // LoadLibrary in Remote Process

   DWORD dwThreadId = 0;
   HANDLE RemoteThread = CreateRemoteThread(hProcHandle,
                       0,0,
                       (LPTHREAD_START_ROUTINE)GetProcAddress(LoadLibrary("Kernel32.dll"),"LoadLibraryA"),
                       pMemoryAddress,
                       0,&dwThreadId
                       );

   if( !RemoteThread )
       {
         cout<<"Failed to create remote process in "+ProcList[iIndex].GetProcessName()+"\n\n";
         system("pause");
         exit(1);
       }
   WaitForSingleObject(RemoteThread,INFINITE);

   // Wait for connection from Serenity Inject DLL
   WaitForSingleObject(AcceptThreadHandle,INFINITE);

   #ifdef _DEBUG_
   cout<<"Connection Established!\n\n";
   #endif
   // Connection established!

   /*
    TODO:  1)Receive DLL Names from Serenity Inject DLL
           2)MapExportedFunctionFromLib
           3)Recv Function Addresses from DLL and compare
   */

   DllList ModuleNames;
   RecvModuleNames(uiLibSock,ModuleNames);

   // Output Module Names

   DllList::iterator iter;
   for( iter = ModuleNames.begin(); iter != ModuleNames.end(); iter++ )
    {
      if( iter->length() > 1 ) cout<<*iter<<endl;
    }


   // Close all open handles and exit

   #ifdef _DEBUG_
   assert(uiListenSocket != INVALID_SOCKET);
   assert(uiLibSock != INVALID_SOCKET);
   assert(hProcHandle != NULL);
   #endif
   CloseHandle( hProcHandle );
   closesocket(uiListenSocket);
   closesocket(uiLibSock);

   system("pause");
   getchar();
   exit(0);
}






unsigned int __stdcall RecvModuleNames(int iSock,DllList& List)
{
  int bytes_received = 0;
  char Buffer[256] = {0x00};
  string DllNames;
  while( bytes_received = recv(iSock,Buffer,256,0) )
   {
    if( bytes_received == -1 ) break;
    DllNames += Buffer;
    ZeroMemory(Buffer,256);
   }
 return explode(DllNames,"\n",List);
}




unsigned int explode(string str,string delimiter,vector<string> &container)
 {
   int iPosition = 0;
   int iLastPos = 0;
   while( (iPosition = str.find(delimiter,iLastPos)) != string::npos )
    {
          container.push_back(str.substr(iLastPos,iPosition-(iLastPos)));
          iLastPos = iPosition+delimiter.length();
    }
   container.push_back(str.substr(iLastPos));
   return container.size();
  }



DWORD __stdcall AcceptThread(void* listening_socket)
{
  unsigned int* uiSock = (unsigned int*)listening_socket;
  uiLibSock = accept((*uiSock),0,0);
}