/*============================================================================
**
**    ERG TRANSIT SYSTEMS      Licensed software
**    (C) 2004                 All rights reserved
**
**============================================================================
**
**  Project/Product : TOOLS
**  Filename        : ShowDLLs.cpp
**  Author(s)       : DDP
**
**  Description     : Display DLLs required by a DLL or EXE
**
**
**  Information     :
**   Compiler       : ANSI C
**   Target         :
**
**
**  History         :
**   Date           Author  Description
**   -------------------------------------------------------------------------
**   10-Nov-04      DDP     Created
**==========================================================================*/

#include <windows.h>
#include <winnt.h>
#include <stdio.h>
#include <math.h>
#include <ctype.h>
#include  <io.h>
#include  <stdlib.h>


//----------------------------------------------------------------------------
//  A structure to hold information on each DLL discovered
//
typedef struct dlist
{
    struct dlist * next;
    char * name;
    char * path;
    bool   processed;
} dlist_t;

dlist_t *dlist = NULL;

//----------------------------------------------------------------------------
//  A structure to hold information on each PATH in which to locate DLLs
//
typedef struct plist
{
    struct plist * next;
    char * name;
} plist_t;

plist_t *plist = NULL;
plist_t *plist_last = NULL;

//----------------------------------------------------------------------------
//  Prototypes
//
void main(int argn,char ** argv);
void ShowHelp(void);
void ParseCmdLine(int argn,char ** argv);
void ReportAndExit(void);
void DisplayImageFile(void);
DWORD ConvertRVAtoFileOffset(DWORD rva);
DWORD MapRVA(DWORD rva);
void display_IMAGE_IMPORT_DIRECTORY(_IMAGE_DATA_DIRECTORY * pDDImport);
void AddUserDll( char * name );
void AddFoundDll( char * name );
void AddEnvPath ( void );
dlist_t *ProcessNextDll( void );
void ProcessFile ( char * name );
void PrintSummary( void );
void AddPath ( char * name );


//----------------------------------------------------------------------------
// Macros
//
#define CHECK(x) if(x) ReportAndExit();

//----------------------------------------------------------------------------
//  Globals
//
char * pfile;           // the address of the PE file in memory and
char * szfname;         // the files name
DWORD  cb;              // its size in bytes
int    verbose = 0;     // Debug support


//----------------------------------------------------------------------------
// FUNCTION           : main
//
// DESCRIPTION        : Main entry point
//
// INPUTS             : argn
//                      argv
//
// RETURNS            :
//
void main(int argn,char ** argv)
{
    dlist_t *dptr;

    /*
    **  check is help needs to be shown
    */
    if(argn == 1 || (argn ==2 && strcmp(argv[1],"/?")==0))
        ShowHelp();

    /*
    **  set all our global flags here
    */
    ParseCmdLine(argn-1,argv+1);

    /*
    **  Prime the paths to be searched
    */
    AddEnvPath();

    /*
    **  Let the magic will then happen.
    */
    while ( (dptr = ProcessNextDll()) )
    {
        ProcessFile ( dptr->path );
    }

    /*
    **  Display useful data to the user
    */
    PrintSummary();

}

//----------------------------------------------------------------------------
// FUNCTION           : ShowHelp
//
// DESCRIPTION        : Display basic help and exit
//
//
// INPUTS             : None
//
// RETURNS            : Does not return
//
void ShowHelp(void)
{
    printf("Locate DLLs required by a DLL or EXE\n");
    printf("Last update: " __DATE__ " "  __TIME__ "\n\n");
        
    printf("usage: ShowDlls  [options] <filename>\n\n");
    printf("\t-h        Display help\n");
    printf("\t-p path   Add Path to search list\n");
    printf("\t-v        Verbose operation\n");

    exit(1);
}

//----------------------------------------------------------------------------
// FUNCTION           : ParseCmdLine
//
// DESCRIPTION        : Extract user options from the command line
//
// INPUTS             : argn        - Number of arguments
//                      argv        - Array of arg pointers
//
// RETURNS            : Nothing
//
void ParseCmdLine(int argn,char ** argv)
{
    for(int i=0;i<argn;i++)
    {
        if(strcmp("-v",argv[i])==0)
        {
            verbose++;
        }
        else if ( strcmp("-p",argv[i])==0 )
        {
            if ( ++i >= argn )
            {
                printf( "Error: -p without a path\n" );
                exit(1) ;
            }
            AddPath( argv[i] );
        }
        else if ( strcmp("-h",argv[i])==0 )
        {
            ShowHelp();
        }
        else if ( strcmp("-",argv[i])==0 )
        {
            printf("Unknown Switch %s. Ignored.\n",argv[i]);
        }
        else
        {
            AddUserDll( argv[i] );
        }
    }
}

//----------------------------------------------------------------------------
// FUNCTION           : ReportAndExit
//
// DESCRIPTION        : Error Check and reporting
//
// INPUTS             : None
//
// RETURNS            : Does not return
//
void ReportAndExit(void)
{
    printf("Error:");
    //shamelessly copied from MSDN
    LPVOID lpMsgBuf;
    FormatMessage( 
        FORMAT_MESSAGE_ALLOCATE_BUFFER | 
        FORMAT_MESSAGE_FROM_SYSTEM | 
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        GetLastError(),
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), // Default language
        (LPTSTR) &lpMsgBuf,
        0,
        NULL 
    );
    // Display the string.
    printf((char *)lpMsgBuf);
    // Free the buffer.
    LocalFree( lpMsgBuf );  
    
    exit(1);
}

//----------------------------------------------------------------------------
// FUNCTION           : ConvertRVAtoFileOffset
//
// DESCRIPTION        : Given the RVA this function will calculate the file offset
//
// INPUTS             : rva     - An RVA
//
// RETURNS            : File offset
//
DWORD ConvertRVAtoFileOffset(DWORD rva)
{
    if(!rva)
    {
        return 0;
    }
        
    _IMAGE_DOS_HEADER * pDosH;
    _IMAGE_FILE_HEADER *    pCoffH;
    _IMAGE_OPTIONAL_HEADER * pOptH;
    _IMAGE_SECTION_HEADER *  pSecH;
    
    pDosH = (_IMAGE_DOS_HEADER *)pfile;
    pCoffH = (_IMAGE_FILE_HEADER *)((char *)pfile + pDosH->e_lfanew + 4);   //4 byte PE signature
    pOptH = (_IMAGE_OPTIONAL_HEADER *)((char *)pCoffH + sizeof(_IMAGE_FILE_HEADER)); //these lie one after the other
    pSecH = (_IMAGE_SECTION_HEADER *)((char *)pOptH + sizeof(_IMAGE_OPTIONAL_HEADER)); // ^
        
    for(int i = 0; i < pCoffH->NumberOfSections; i++ )
    {
        DWORD rvadiff = rva - pSecH[i].VirtualAddress;
        if( rvadiff >= 0 && rvadiff <= pSecH[i].SizeOfRawData)
        {
            return rvadiff + pSecH[i].PointerToRawData;
        }
    }
    return 0;
}
//----------------------------------------------------------------------------
// FUNCTION           : MapRVA
//
// DESCRIPTION        : Given an RVA this function will give the actually
//                      current memory location which is:
//                          file base address + file offset
//
// INPUTS             : rva     - RVA to convert
//
// RETURNS            : Mapped address
//
DWORD MapRVA(DWORD rva)
{
    if(!rva)
    {
        return 0;
    }

    return ((DWORD)pfile)+ConvertRVAtoFileOffset(rva);
}

//----------------------------------------------------------------------------
// FUNCTION           : DisplayImageFile
//
// DESCRIPTION        : Process an file
//                      The file has been mapped into memory
//                      Calculate file offsets for later processing
//
// INPUTS             : None
//
// RETURNS            : Nothing
//
void DisplayImageFile(void)
{
    _IMAGE_DOS_HEADER * pDosH;
    char *          pPESign;
    _IMAGE_FILE_HEADER *    pCoffH;
    _IMAGE_OPTIONAL_HEADER * pOptH;
    _IMAGE_DATA_DIRECTORY * pDDImport;

    /*
    **  Map various structures to locations inside the PE file
    **  The PE file is already memory mapped
    */
    pDosH = (_IMAGE_DOS_HEADER *)pfile;
    pPESign = pfile + pDosH->e_lfanew;
    
    if(!(   pPESign[0] == 'P' &&
        pPESign[1] == 'E' &&
        pPESign[2] == 0 &&
        pPESign[3] == 0 ))
    {
        printf("This is not a PE file according to the file signature\n");
        exit(1);
    }
    
    pCoffH = (_IMAGE_FILE_HEADER *)(pPESign + 4) ;  //4 byte PE signature
    pOptH = (_IMAGE_OPTIONAL_HEADER *)((char *)pCoffH + sizeof(_IMAGE_FILE_HEADER)); //these lie one after the other
    pDDImport = &(pOptH->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]);

    /*
    **  Extact data from the IMPORT DIRECTORY
    */
    display_IMAGE_IMPORT_DIRECTORY(pDDImport);
}

//----------------------------------------------------------------------------
// FUNCTION           : display_IMAGE_IMPORT_DIRECTORY
//
// DESCRIPTION        : Sacn the IMPORT directory and add any DLLs located
//                      to the list of DLLs used by the base EXE
//
// INPUTS             : pDDImport   Address of the Import Data
//
// RETURNS            :
//
// PRECONDITION       : The file has been memory mapped
//
void display_IMAGE_IMPORT_DIRECTORY(_IMAGE_DATA_DIRECTORY * pDDImport)
{
    if (verbose) printf("IMPORT DIRECTORY ENTRY: %s\n", szfname);

    if(pDDImport->VirtualAddress == 0)
    {
        if (verbose) printf("\tNo Import Table Present\n");
        return;
    }
        
    _IMAGE_IMPORT_DESCRIPTOR * pImpDesc = (_IMAGE_IMPORT_DESCRIPTOR *)MapRVA(pDDImport->VirtualAddress);
    while(pImpDesc->Characteristics)
    {
        if (verbose) printf("\tName = %#x --> \"%s\"\n",pImpDesc->Name,MapRVA(pImpDesc->Name));
        AddFoundDll( (char *) MapRVA(pImpDesc->Name) );
        pImpDesc++;
    }
}

//----------------------------------------------------------------------------
// FUNCTION           : ProcessFile
//
// DESCRIPTION        : Process one EXE or DLL file
//                      Examine the file for DLLs and add them to the list
//                      of discovered DLLs as we go
//
//
// INPUTS             : name    - Name of the file
//
// RETURNS            : Nothing
//
void ProcessFile ( char * name )
{

    /*
    ** open and memory map the named file
    */
    HANDLE hPE = CreateFile(name,GENERIC_READ,FILE_SHARE_READ,0,OPEN_EXISTING,0,0);
    CHECK(hPE == INVALID_HANDLE_VALUE);
    HANDLE hmPE = CreateFileMapping(hPE,0,PAGE_READONLY,0,0,"PEfile");
    CHECK(hmPE == 0);

    /*
    **  file address and size : there are assigned to globals
    */
    pfile = (char*)MapViewOfFile(hmPE,FILE_MAP_READ,0,0,0);
    CHECK(pfile == 0);
    cb = GetFileSize(hPE,0);
    szfname = name;

    /*
    **  at this point the file is opened and mapped to a memory location pointed
    **  to by 'pfile' and is valid for 'cb' bytes which is the file size          
    **  read the file and do what we want to                                      
    **
    */
    DisplayImageFile();

    /*
    **  close the open handles
    */
    CloseHandle(hPE);
    CloseHandle(hmPE);
    
}


//----------------------------------------------------------------------------
// FUNCTION           : AddPath
//
// DESCRIPTION        : Add a path to a linked list of paths
//                      New entries are added to the end of the list
//
//
// INPUTS             : name    - The path name to add
//
// RETURNS            : Nothing
//
void AddPath ( char * name )
{
    plist_t *ptr;

    /*
    **  Create a new entry and insert information
    */
    ptr = (plist_t *)malloc( sizeof(plist_t) );
    memset( (void *)ptr, 0, sizeof( *ptr));

    ptr->name = (char *)malloc(strlen(name)+1);
    strcpy( ptr->name, name );

    /*
    **  Add entry to the END of the list to preserve search order
    */
    if ( plist_last )
    {
        plist_last->next = ptr;
    }
    plist_last = ptr;
    if ( ! plist )
    {
        plist = ptr;
    }

    if (verbose > 1) printf ("Adding Path: %s\n", name);
}

//----------------------------------------------------------------------------
// FUNCTION           : AddEnvPath
//
// DESCRIPTION        : Add the users PATH to the list of seach paths
//                      This is done to prime the search path
//
//
// INPUTS             : None
//
// RETURNS            : Nothing
//
void AddEnvPath ( void )
{
    char *path;
    char *next;

    /*
    **  Extract the path from the environment
    */
    if ( (path = getenv( "PATH" )) == NULL )
    {
        return;
    }

    /*
    **  Extract ';' substrings from the path and add them to the
    **  list of paths for later searching
    */
    if (verbose > 2) printf ("Raw Path: %s\n", path);
    while ( *path )
    {
        next = strchr( path, ';' );
        if ( next )
        {
            *next = 0;
        }
        AddPath( path );

        if ( !next )
        {
            break;
        }
        path = next + 1;
    }
}

//----------------------------------------------------------------------------
// FUNCTION           : AddUserDll
//
// DESCRIPTION        : Add the initial user DLL to the list of EXE/DLLs to
//                      process. If the user has specified an absolute path
//                      then use the path as a part of the search path.
//
//
// INPUTS             : namep       - Pointer to the DLL's name
//
// RETURNS            :
//
void AddUserDll( char * name )
{
    dlist_t *dptr;
    char *ptr;
    char *path = NULL;
    char *name_ptr = NULL;

    if ( verbose )
        printf("AddUserDll: %s\n", name );

    /*
    **  Allow for '/' and '\' and locate the last '\' in the string
    */
    for ( ptr = name; *ptr; ptr++)
    {
        if ( *ptr == '/' )
            *ptr = '\\';

        if ( *ptr == '\\' )
            name_ptr = ptr;
    }

    /*
    **  Split into path and name
    */
    if ( name_ptr )
    {
        if (verbose > 1) printf( "Absolute path detected\n");

        path = (char *)malloc( name_ptr - name + 1 );
        path[0] = 0;

        strncat(path, name, name_ptr - name );
        name = name_ptr + 1;
    }
    else
    {
        path = ".";
    }

    /*
    **  Scan the list of existing source dirs
    */
    dptr = dlist;
    while ( dptr )
    {
        if ( strcmp ( dptr->name, name ) == 0 )
            return;

        dptr = dptr->next;
    }

    /*
    **  Entry not found
    **  Create a new one and insert information
    */
    dptr = (dlist_t *)malloc( sizeof(dlist_t) );
    memset( (void *)dptr, 0, sizeof( *dptr));

    dptr->next = dlist;
    dlist = dptr;

    dptr->name = (char *)malloc(strlen(name)+1);
    strcpy( dptr->name, name );

    if (verbose > 1) printf ("Adding: %s\n", name);

    /*
    **  Add the source path to the DLL to the search list
    */
    AddPath ( path );
}

//----------------------------------------------------------------------------
// FUNCTION           : AddFoundDll
//
// DESCRIPTION        : Add a new DLL to a list of DLL's
//                      If the DLL is already in the list just skip it.
//
//
// INPUTS             : namep       - Pointer to the DLL's name
//
// RETURNS            :
//
void AddFoundDll( char * name )
{
    dlist_t *dptr;

    if ( verbose )
        printf("AddFoundDll: %s\n", name );

    /*
    **  Scan the list of existing source dirs
    */
    dptr = dlist;
    while ( dptr )
    {
        if ( strcmp ( dptr->name, name ) == 0 )
            return;

        dptr = dptr->next;
    }

    /*
    **  Entry not found
    **  Create a new one and insert information
    */
    dptr = (dlist_t *)malloc( sizeof(dlist_t) );
    memset( (void *)dptr, 0, sizeof( *dptr));

    dptr->next = dlist;
    dlist = dptr;

    dptr->name = (char *)malloc(strlen(name)+1);
    strcpy( dptr->name, name );

    if (verbose > 1) printf ("Adding: %s\n", name);
}


//----------------------------------------------------------------------------
// FUNCTION           : ProcessNextDll
//
// DESCRIPTION        : Locate the next DLL to explore
//
//
// INPUTS             : None
//
// RETURNS            : A pointer to a DLL entry to process
//                      The entry will have a full path attached
//

dlist_t *ProcessNextDll( void )
{
    dlist_t *dptr;
    plist_t *pptr;

    for ( ; ; )
    {
        /*
        **  Scan the list of existing source dirs
        */
        for ( dptr = dlist; dptr ; dptr = dptr->next)
        {
            if ( ! dptr->processed )
            {
                break;
            }
        }

        /*
        **  Exit if we do not have any more DLLs to process
        */
        if ( ! dptr )
        {
            return NULL;
        }

        /*
        **  Mark the entry as processed
        **  If we don't find the path to the DLL then the path will be left as
        **  NULL, but the entry will still have been processed
        */
        dptr->processed = TRUE;

        /*
        **  Locate the DLL in the users PATH and other specified directories
        */
        for ( pptr = plist; pptr; pptr = pptr->next )
        {
            int plen = strlen( dptr->name ) + 1 + strlen( pptr->name ) + 1;
            char *buf = (char *)malloc( plen );

            /*
            **  Create the name of the file
            */
            _snprintf( buf, plen, "%s\\%s", pptr->name, dptr->name );


            /*
            **  Attempt to open it
            */
            if (verbose > 2) printf( "Looking for: %s\n", buf );
            if ( _access ( buf, 0 ) != -1 )
            {
                /*
                **  File does exist
                **  Add the path information to the entry
                */
                dptr->path = buf;
                if (verbose > 1) printf( "   FOUND: %s\n", buf );
                return dptr;
            }
        }

        /*
        **  File not found
        **  Try another
        */
    }
}

//----------------------------------------------------------------------------
// FUNCTION           : PrintSummary
//
// DESCRIPTION        : Display summary information on DLLs discovered
//
//
// INPUTS             : None
//
// RETURNS            : Nothing
//
void PrintSummary( void )
{
    dlist_t *dptr;

    for ( dptr = dlist; dptr ; dptr = dptr->next)
    {
        printf( "%30s", dptr->name );
        if ( dptr->path )
        {
            printf( " : %s\n", dptr->path );
        }
        else
        {
            printf( " : NOT FOUND\n");
        }
    }
}