/*
* Copyright (c) 2007, Insomniac Games
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*     * Redistributions of source code must retain the above copyright
*       notice, this list of conditions and the following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright
*       notice, this list of conditions and the following disclaimer in the
*       documentation and/or other materials provided with the distribution.
*     * Neither the name of the <organization> nor the
*       names of its contributors may be used to endorse or promote products
*       derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY INSOMNIAC GAMES ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL INSOMNIAC GAMES BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <windows.h>
#include <stdio.h>
#include "psapi.h"
#include <map>
#include <string>

// PAGE_WRITECOMBINE missing in MinGW version of winnt.h
#ifndef PAGE_WRITECOMBINE
#define PAGE_WRITECOMBINE 0x4000
#endif

typedef std::map<DWORD,std::string> MM;
MM module_map;

//-----------------------------------------------------------------------------
// Get all the DLLs loaded in the specified process
void GetModules( DWORD process_id )
{
  HMODULE mods[1024];
  HANDLE process;
  DWORD needed;
  unsigned int i;

  // Print the process identifier.
  printf( "\nProcess ID: %u\n", process_id );

  // Get a list of all the modules in this process.
  process = OpenProcess(  PROCESS_QUERY_INFORMATION |  PROCESS_VM_READ, FALSE, process_id );
  if (process==0)
    return;

  if( EnumProcessModules(process, mods, sizeof(mods), &needed))
  {
    for ( i = 0; i < (needed / sizeof(HMODULE)); i++ )
    {
      char module_name[MAX_PATH];

      // Get the full path to the module's file.

      if ( GetModuleFileNameEx( process, mods[i], module_name,sizeof(module_name)))
      {
        // Print the module name and handle value.
        DWORD idx = (DWORD)mods[i];       // the module handle is the address in process space
        std::string name = module_name;   // make a map from address to module name
        module_map[idx] = name;
      }
    }
  }

  CloseHandle( process );
}

//-----------------------------------------------------------------------------
// Print the name of a process from a process id
void PrintProcessNameAndID( DWORD process_id )
{
  TCHAR process_name[MAX_PATH] = TEXT("<unknown>");

  // Get a handle to the process
  HANDLE process = OpenProcess( PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,false, process_id );

  // Get the process name.
  if (process!=0)
  {
      HMODULE mod;
      DWORD needed;

      if ( EnumProcessModules( process, &mod, sizeof(mod), &needed) )
      {
        GetModuleBaseName( process, mod, process_name, sizeof(process_name)/sizeof(TCHAR) );
      }
  }

  // Print the process name and identifier.
  printf( TEXT("%s  (PID: %u)\n"), process_name, process_id );
  CloseHandle( process );
}


//-----------------------------------------------------------------------------
void DumpAddressMap( DWORD process_id )
{
  DWORD total_free=0;
  DWORD total_reserved=0;
  DWORD total_commit =0;
  DWORD largest_free = 0;

  TCHAR process_name[MAX_PATH];

  // Get a handle to the process.
  HANDLE process = OpenProcess( PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,false, process_id );

  // Get the process name.
  if (process!=0 )
  {
    HMODULE mod;
    DWORD needed;

    if ( EnumProcessModules( process, &mod, sizeof(mod), &needed) )
    {
      GetModuleBaseName( process, mod, process_name,sizeof(process_name)/sizeof(TCHAR) );
    }
    printf( TEXT("%s  (PID: %u)\n"), process_name, process_id );

    DWORD base = 0;
    DWORD res = 1;
    while (res)
    {
      MEMORY_BASIC_INFORMATION mem;
      res = VirtualQueryEx(process,(void*)base,&mem,sizeof(mem));
      if (res)
      {
        printf("Address 0x%08x: Size 0x%08x",mem.BaseAddress,mem.RegionSize);
        MM::iterator i = module_map.find((DWORD)mem.AllocationBase);
        if (i!=module_map.end())
        {
          printf(" [%s]",(*i).second.c_str());
        }
        printf("\n");

        if (mem.State & MEM_FREE)
        {
          printf("  FREE\n");
          total_free+=mem.RegionSize;

          // we should not have to deal with different free blocks being consecutive
          if (mem.RegionSize>largest_free)
            largest_free = mem.RegionSize;
        }
        else
        {
          if (mem.State & MEM_RESERVE)
          {
            printf("  RESERVED: ");
            total_reserved+=mem.RegionSize;
          }
          if (mem.State & MEM_COMMIT)
          {
            printf("  COMMIT: ");
            total_commit+=mem.RegionSize;
          }
          if (mem.AllocationProtect)
          {
            if (mem.AllocationProtect & PAGE_NOACCESS)
              printf("NO ACCESS, ");
            if (mem.AllocationProtect & PAGE_READONLY)
              printf("READ ONLY, ");
            if (mem.AllocationProtect & PAGE_READWRITE)
              printf("READ WRITE, ");
            if (mem.AllocationProtect & PAGE_WRITECOPY)
              printf("WRITE COPY, ");
            if (mem.AllocationProtect & PAGE_EXECUTE)
              printf("EXECUTE - NO ACCESS");
            if (mem.AllocationProtect & PAGE_EXECUTE_READ)
              printf("EXECUTE - READ, ");
            if (mem.AllocationProtect & PAGE_EXECUTE_READWRITE)
              printf("EXECUTE - READ WRITE, ");
            if (mem.AllocationProtect & PAGE_EXECUTE_WRITECOPY)
              printf("EXECUTE - WRITE COPY, ");
            if (mem.AllocationProtect & PAGE_GUARD)
              printf("GUARD, ");
            if (mem.AllocationProtect & PAGE_NOCACHE)
              printf("NO CACHE, ");
            if (mem.AllocationProtect & PAGE_WRITECOMBINE)
              printf("WRITE COMBINE, ");
          }
          printf("\n");
          if (mem.Type)
          {
            printf("  ");
            if (mem.Type& MEM_IMAGE)
            {
              printf("MAPPED TO A CODE IMAGE, ");
            }
            if (mem.Type& MEM_MAPPED)
            {
              printf("MAPPED TO A NAMED SECTION, ");
            }
            if (mem.Type& MEM_PRIVATE)
            {
              printf("PROCESS PRIVATE, ");
            }
            printf("\n");
          }
        }

        base+=mem.RegionSize;
      }
    }

    printf("\n\n");
    printf("Total Commit   = %uK bytes\n",total_commit>>10);
    printf("Total Reserved = %uK bytes\n",total_reserved>>10);
    printf("Total Free     = %uK bytes\n",total_free>>10);
    printf("Largest Free   = %uK bytes\n",largest_free>>10);

    CloseHandle(process);
  }
}


//-----------------------------------------------------------------------------
int 
main(int argc,char* argv[])
{
  // get the process id from the first argument on the command line
  DWORD pid;
  if (argc<2)
  {
    pid=0xffffffff;
  }
  else
  {
    pid=atoi(argv[1]);
  }

  if (pid==0xffffffff)
  {
    // if no process ID was specified on the command line,
    // dump all the processes

    // Get the list of process identifiers.
    DWORD aProcesses[1024], cbNeeded, cProcesses;
    unsigned int i;  

    if ( !EnumProcesses( aProcesses, sizeof(aProcesses), &cbNeeded ) )
      return (-1);

    // Calculate how many process identifiers were returned.
    cProcesses = cbNeeded / sizeof(DWORD);

    for ( i = 0; i < cProcesses; i++ )
      PrintProcessNameAndID( aProcesses[i] );
  }
  else
  {
    // get and dump info about the specified process
    GetModules(pid);
    DumpAddressMap(pid);
  }

  return (0);
}