/*
    libntldd - builds a dependency tree of a module, with symbols

    Copyright (C) 2010 LRN

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

/*
Code is mostly written after
"An In-Depth Look into the Win32 Portable Executable File Format"
MSDN Magazine articles
*/

#include <windows.h>

#include <imagehlp.h>

#include <winnt.h>

#include <string.h>
#include <stdio.h>
#include <stdint.h>

#include "libntldd.h"

typedef struct _soff_entry soff_entry;

struct _soff_entry
{
  DWORD start;
  DWORD end;
  void *off;
};

void *MapPointer (soff_entry *soffs, int soffs_len, DWORD in_ptr, int *section)
{
  int i;
  for (i = 0; i < soffs_len; i++)
    if (soffs[i].start <= in_ptr && soffs[i].end >= in_ptr)
    {
      if (section != NULL)
        *section = i;
      if (soffs[i].off)
        return soffs[i].off + in_ptr;
    }
  return NULL;
}

/*
int FindSectionID (IMAGE_OPTIONAL_HEADER *oh, DWORD address, DWORD size)
{
  int i;
  for (i = 0; i < oh->NumberOfRvaAndSizes; i++)
  {
    if (oh->DataDirectory[i].VirtualAddress == address &&
        oh->DataDirectory[i].Size == size)
      return i;
  }
  return -1;
}
*/

int FindSectionByRawData (LOADED_IMAGE *img, DWORD address)
{
  ULONG i;
  for (i = 0; i < img->NumberOfSections; i++)
  {
    DWORD start = img->Sections[i].VirtualAddress;
    DWORD end = start + img->Sections[i].SizeOfRawData;
    if (address >= start && address < end)
      return i;
  }
  return -1;
}

void ResizeArray (void **data, uint64_t *data_size, size_t sizeof_data)
{
  uint64_t new_size = (*data_size) > 0 ? (*data_size) * 2 : 64;
  void *new_data;
  new_data = realloc (*data, new_size * sizeof_data);
  memset (((unsigned char *) new_data) + (*data_size * sizeof_data), 0, (new_size - (*data_size)) * sizeof_data);
  *data = new_data;
  *data_size = new_size;
}

#define ResizeDepList(ptr_deptree, ptr_deptree_size) ResizeArray ((void **) ptr_deptree, ptr_deptree_size, sizeof (struct DepTreeElement *))
#define ResizeImportList(ptr_import_list, ptr_import_list_size) ResizeArray ((void **) ptr_import_list, ptr_import_list_size, sizeof (struct ImportTableItem))
#define ResizeStack(ptr_stack, ptr_stack_size) ResizeArray ((void **) ptr_stack, ptr_stack_size, sizeof (char *))

void AddDep (struct DepTreeElement *parent, struct DepTreeElement *child)
{
  if (parent->childs_len >= parent->childs_size)
  {
    ResizeDepList (&parent->childs, &parent->childs_size);
  }
  parent->childs[parent->childs_len] = child;
  parent->childs_len += 1;
}

struct ImportTableItem *AddImport (struct DepTreeElement *self)
{
  if (self->imports_len >= self->imports_size)
  {
    ResizeImportList (&self->imports, &self->imports_size);
  }
  self->imports_len += 1;
  return &self->imports[self->imports_len - 1];
}

int FindDep (struct DepTreeElement *root, char *name, struct DepTreeElement **result)
{
  int ret = -1;
  uint64_t i;
  if (root->flags & DEPTREE_VISITED)
  {
    return -2;
  }
  root->flags |= DEPTREE_VISITED;
  for (i = 0; i < root->childs_len; i++)
  {
    if (stricmp (root->childs[i]->module, name) == 0)
    {
      if (result != NULL)
        *result = root->childs[i];
      root->flags &= ~DEPTREE_VISITED;
      return (root->childs[i]->flags & DEPTREE_UNRESOLVED) ? 1 : 0;
    }
  }
  for (i = 0; i < root->childs_len && ret < 0; i++)
  {
    ret = FindDep (root->childs[i], name, result);
  }
  root->flags &= ~DEPTREE_VISITED;
  return ret;
}

int BuildDepTree (BuildTreeConfig* cfg, char *name, struct DepTreeElement *root, struct DepTreeElement *self);

struct DepTreeElement *ProcessDep (BuildTreeConfig* cfg, soff_entry *soffs, int soffs_len, DWORD name, struct DepTreeElement *root, struct DepTreeElement *self, int deep)
{
  struct DepTreeElement *child = NULL;
  int found;
  int64_t i;
  char *dllname = (char *) MapPointer (soffs, soffs_len, name, NULL);
  if (dllname == NULL)
    return NULL;
  if (strlen (dllname) > 10 && strnicmp ("api-ms-win", dllname, 10) == 0)
  {
    /* TODO: find a better way to identify api stubs. Versioninfo, maybe? */
    return NULL;
  }
  for (i = (int64_t)*cfg->stack_len - 1; i >= 0; i--)
  {
    if ((*cfg->stack)[i] && stricmp ((*cfg->stack)[i], dllname) == 0)
      return NULL;
    if (i == 0)
      break;
  }
  found = FindDep (root, dllname, &child);
  if (found < 0)
  {
    child = (struct DepTreeElement *) malloc (sizeof (struct DepTreeElement));
    memset (child, 0, sizeof (struct DepTreeElement));
    if (deep == 0)
    {
      child->module = strdup (dllname);
      AddDep (self, child);
    }
  }
  if (deep == 1)
  {
    BuildDepTree (cfg, dllname, root, child);
  }
  return child;
}


/*
struct ExportTableItem *FindExportForward (struct DepTreeElement *self, char *dllname, char *export_name, DWORD export_ordinal)
{
  return NULL;
}
*/

int ClearDepStatus (struct DepTreeElement *self, uint64_t flags)
{
  uint64_t i;
  for (i = 0; i < self->childs_len; i++)
    ClearDepStatus (self->childs[i], flags);
  self->flags &= ~flags;
  return 0;
}

void PushStack (char ***stack, uint64_t *stack_len, uint64_t *stack_size, char *name)
{
  if (*stack_len >= *stack_size)
  {
    ResizeStack (stack, stack_size);
  }
  (*stack)[*stack_len] = strdup (name);
  (*stack_len) += 1;
}

void PopStack (char ***stack, uint64_t *stack_len, uint64_t *stack_size, char *name)
{
  (*stack)[*stack_len] = NULL;
  (*stack_len) -= 1;
}

static uint64_t thunk_data_u1_function (void *thunk_array, DWORD index, BuildTreeConfig *cfg)
{
  if (cfg->machineType == IMAGE_FILE_MACHINE_I386)
    return ((IMAGE_THUNK_DATA32 *) thunk_array)[index].u1.Function;
  else
    return ((IMAGE_THUNK_DATA64 *) thunk_array)[index].u1.Function;
}

static void *opt_header_get_dd_entry (void *opt_header, DWORD entry_type, BuildTreeConfig *cfg)
{
  if (cfg->machineType == IMAGE_FILE_MACHINE_I386)
    return &(((PIMAGE_OPTIONAL_HEADER32) opt_header)->DataDirectory[entry_type]);
  else
    return &(((PIMAGE_OPTIONAL_HEADER64) opt_header)->DataDirectory[entry_type]);
}

static void BuildDepTree32or64 (LOADED_IMAGE *img, BuildTreeConfig* cfg, struct DepTreeElement *root, struct DepTreeElement *self, soff_entry *soffs, int soffs_len)
{
  IMAGE_DATA_DIRECTORY *idata;
  IMAGE_IMPORT_DESCRIPTOR *iid;
  IMAGE_EXPORT_DIRECTORY *ied;
  IMAGE_DELAYLOAD_DESCRIPTOR *idd;
  void *ith, *oith;
  void *opt_header = &img->FileHeader->OptionalHeader;
  DWORD i, j;

  idata = opt_header_get_dd_entry (opt_header, IMAGE_DIRECTORY_ENTRY_EXPORT, cfg);
  if (idata->Size > 0 && idata->VirtualAddress != 0)
  {
    int export_section = -2;
    ied = (IMAGE_EXPORT_DIRECTORY *) MapPointer (soffs, soffs_len, idata->VirtualAddress, &export_section);
    if (ied && ied->Name != 0)
    {
      char *export_module = MapPointer (soffs, soffs_len, ied->Name, NULL);
      if (export_module != NULL)
      {
        if (self->export_module == NULL)
          self->export_module = strdup (export_module);
      }
    }
    if (ied && ied->NumberOfFunctions > 0)
    {
      DWORD *addrs, *names;
      WORD *ords;
      int section = -1;
      self->exports_len = ied->NumberOfFunctions;
      self->exports = (struct ExportTableItem *) malloc (sizeof (struct ExportTableItem) * self->exports_len);
      memset (self->exports, 0, sizeof (struct ExportTableItem) * self->exports_len);
      addrs = (DWORD *) MapPointer (soffs, soffs_len, ied->AddressOfFunctions, NULL);
      ords = (WORD *) MapPointer (soffs, soffs_len, ied->AddressOfNameOrdinals, NULL);
      names = (DWORD *) MapPointer (soffs, soffs_len, ied->AddressOfNames, NULL);
      for (i = 0; i < ied->NumberOfNames; i++)
      {
        self->exports[ords[i]].ordinal = ords[i] + ied->Base;
        if (names[i] != 0)
        {
          char *s_name = (char *) MapPointer (soffs, soffs_len, names[i], NULL);
          if (s_name != NULL)
            self->exports[ords[i]].name = strdup (s_name);
        }
      }
      for (i = 0; i < ied->NumberOfFunctions; i++)
      {
        if (addrs[i] != 0)
        {
          int section_index = FindSectionByRawData (img, addrs[i]);
          if ((idata->VirtualAddress <= addrs[i]) && (idata->VirtualAddress + idata->Size > addrs[i]))
          {
            self->exports[i].address = NULL;
            self->exports[i].forward_str = strdup ((char *) MapPointer (soffs, soffs_len, addrs[i], NULL));
          }
          else
            self->exports[i].address = MapPointer (soffs, soffs_len, addrs[i], &section);
          self->exports[i].ordinal = i + ied->Base;
          self->exports[i].section_index = section_index;
          self->exports[i].address_offset = addrs[i];
        }
      }
    }
  }

  idata = opt_header_get_dd_entry (opt_header, IMAGE_DIRECTORY_ENTRY_IMPORT, cfg);
  if (idata->Size > 0 && idata->VirtualAddress != 0)
  {
    iid = (IMAGE_IMPORT_DESCRIPTOR *) MapPointer (soffs, soffs_len,
        idata->VirtualAddress, NULL);
    if (iid)
      for (i = 0; iid[i].Characteristics || iid[i].TimeDateStamp ||
          iid[i].ForwarderChain || iid[i].Name || iid[i].FirstThunk; i++)
      {
        struct DepTreeElement *dll;
        uint64_t impaddress;
        dll = ProcessDep (cfg, soffs, soffs_len, iid[i].Name, root, self, 0);
        if (dll == NULL)
          continue;
        ith = (void *) MapPointer (soffs, soffs_len, iid[i].FirstThunk, NULL);
        oith = (void *) MapPointer (soffs, soffs_len, iid[i].OriginalFirstThunk, NULL);
        for (j = 0; (impaddress = thunk_data_u1_function (ith, j, cfg)) != 0; j++)
        {
          struct ImportTableItem *imp = AddImport (self);
          imp->dll = dll;
          imp->ordinal = -1;
          if (oith)
            imp->orig_address = thunk_data_u1_function (oith, j, cfg);
          if (cfg->on_self)
          {
            imp->address = impaddress;
          }
          if (oith && imp->orig_address & (1 << (sizeof (DWORD) * 8 - 1)))
          {
            imp->ordinal = imp->orig_address & ~(1 << (sizeof (DWORD) * 8 - 1));
          }
          else if (oith)
          {
            IMAGE_IMPORT_BY_NAME *byname = (IMAGE_IMPORT_BY_NAME *) MapPointer (soffs, soffs_len, imp->orig_address, NULL);
            if (byname != NULL)
              imp->name = strdup ((char *) byname->Name);
          }
        }
      }
  }

  idata = opt_header_get_dd_entry (opt_header, IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT, cfg);
  if (idata->Size > 0 && idata->VirtualAddress != 0)
  {
    idd = (IMAGE_DELAYLOAD_DESCRIPTOR *) MapPointer (soffs, soffs_len, idata->VirtualAddress, NULL);
    if (idd)
      for (i = 0; idd[i].Attributes.AllAttributes || idd[i].DllNameRVA ||
          idd[i].ModuleHandleRVA || idd[i].ImportAddressTableRVA || idd[i].ImportNameTableRVA ||
          idd[i].BoundImportAddressTableRVA || idd[i].UnloadInformationTableRVA ||
          idd[i].TimeDateStamp; i++)
      {
        struct DepTreeElement *dll;
        uint64_t impaddress;
        dll = ProcessDep (cfg, soffs, soffs_len, idd[i].DllNameRVA, root, self, 0);
        if (dll == NULL)
          continue;
        if (idd[i].Attributes.AllAttributes & 0x00000001)
        {
          ith = (void *) MapPointer (soffs, soffs_len, idd[i].ImportAddressTableRVA, NULL);
          oith = (void *) MapPointer (soffs, soffs_len, idd[i].ImportNameTableRVA, NULL);
        }
        else
        {
          ith = (void *) idd[i].ImportAddressTableRVA;
          oith = (void *) idd[i].ImportNameTableRVA;
        }
        for (j = 0; (impaddress = thunk_data_u1_function (ith, j, cfg)) != 0; j++)
        {
          struct ImportTableItem *imp = AddImport (self);
          imp->dll = dll;
          imp->ordinal = -1;
          if (oith)
            imp->orig_address = thunk_data_u1_function (oith, j, cfg);
          if (cfg->on_self)
          {
            imp->address = impaddress;
          }
          if (oith && imp->orig_address & (1 << (sizeof (DWORD) * 8 - 1)))
          {
            imp->ordinal = imp->orig_address & ~(1 << (sizeof (DWORD) * 8 - 1));
          }
          else if (oith)
          {
            IMAGE_IMPORT_BY_NAME *byname = (IMAGE_IMPORT_BY_NAME *) MapPointer (soffs, soffs_len, imp->orig_address, NULL);
            if (byname != NULL)
              imp->name = strdup ((char *) byname->Name);
          }
        }
      }
  }

  idata = opt_header_get_dd_entry (opt_header, IMAGE_DIRECTORY_ENTRY_IMPORT, cfg);
  if (idata->Size > 0 && idata->VirtualAddress != 0)
  {
    iid = (IMAGE_IMPORT_DESCRIPTOR *) MapPointer (soffs, soffs_len,
        idata->VirtualAddress, NULL);
    if (iid)
      for (i = 0; iid[i].Characteristics || iid[i].TimeDateStamp ||
          iid[i].ForwarderChain || iid[i].Name || iid[i].FirstThunk; i++)
        ProcessDep (cfg, soffs, soffs_len, iid[i].Name, root, self, 1);
  }

  idata = opt_header_get_dd_entry (opt_header, IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT, cfg);
  if (idata->Size > 0 && idata->VirtualAddress != 0)
  {
    idd = (IMAGE_DELAYLOAD_DESCRIPTOR *) MapPointer (soffs, soffs_len, idata->VirtualAddress, NULL);
    if (idd)
      for (i = 0; idd[i].Attributes.AllAttributes || idd[i].DllNameRVA ||
          idd[i].ModuleHandleRVA || idd[i].ImportAddressTableRVA || idd[i].ImportNameTableRVA ||
          idd[i].BoundImportAddressTableRVA || idd[i].UnloadInformationTableRVA ||
          idd[i].TimeDateStamp; i++)
        ProcessDep (cfg, soffs, soffs_len, idd[i].DllNameRVA, root, self, 1);
  }
}

BOOL TryMapAndLoad (PCSTR name, PCSTR path, PLOADED_IMAGE loadedImage, int requiredMachineType)
{
    BOOL success = MapAndLoad (name, path, loadedImage, FALSE, TRUE);
    if (!success && GetLastError () == ERROR_FILE_NOT_FOUND)
        success = MapAndLoad (name, path, loadedImage, TRUE, TRUE);
    if (success && requiredMachineType != -1 && (int)loadedImage->FileHeader->FileHeader.Machine != requiredMachineType)
    {
        UnMapAndLoad (loadedImage);
        return FALSE;
    }
    return success;
}

int BuildDepTree (BuildTreeConfig* cfg, char *name, struct DepTreeElement *root, struct DepTreeElement *self)
{
  LOADED_IMAGE loaded_image;
  LOADED_IMAGE *img;
  IMAGE_DOS_HEADER *dos;
  HMODULE hmod;
  BOOL success;

  DWORD i, j;
  int soffs_len;
  soff_entry *soffs;

  if (self->flags & DEPTREE_PROCESSED)
  {
    return 0;
  }

  if (cfg->on_self)
  {
    char modpath[MAX_PATH];
    success = GetModuleHandleExA (GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, name, &hmod);
    if (!success)
      return 1;
    if (GetModuleFileNameA (hmod, modpath, MAX_PATH) == 0)
      return 1;
    if (self->resolved_module == NULL)
      self->resolved_module = strdup (modpath);

    dos = (IMAGE_DOS_HEADER *) hmod;
    loaded_image.FileHeader = (IMAGE_NT_HEADERS *) ((char *) hmod + dos->e_lfanew);
    loaded_image.Sections = (IMAGE_SECTION_HEADER *) ((char *) hmod + dos->e_lfanew + sizeof (IMAGE_NT_HEADERS));
    loaded_image.NumberOfSections = loaded_image.FileHeader->FileHeader.NumberOfSections;
    loaded_image.MappedAddress = (void *) hmod;
    if (cfg->machineType != -1 && (int)loaded_image.FileHeader->FileHeader.Machine != cfg->machineType)
        return 1;
  }
  else
  {
    success = FALSE;
    for (i = 0; i < cfg->searchPaths->count && !success; ++i)
    {
      success = TryMapAndLoad (name, cfg->searchPaths->path[i], &loaded_image, cfg->machineType);
    }
    if (!success)
        success = TryMapAndLoad (name, NULL, &loaded_image, cfg->machineType);
    if (!success)
    {
      self->flags |= DEPTREE_UNRESOLVED;
      return 1;
    }
    if (self->resolved_module == NULL)
      self->resolved_module = strdup (loaded_image.ModuleName);
  }
  if (cfg->machineType == -1)
    cfg->machineType = (int)loaded_image.FileHeader->FileHeader.Machine;
  img = &loaded_image;

  PushStack (cfg->stack, cfg->stack_len, cfg->stack_size, name);

  self->mapped_address = loaded_image.MappedAddress;

  self->flags |= DEPTREE_PROCESSED;

  soffs_len = img->NumberOfSections;
  soffs = (soff_entry *) malloc (sizeof(soff_entry) * (soffs_len + 1));
  for (i = 0; i < img->NumberOfSections; i++)
  {
    soffs[i].start = img->Sections[i].VirtualAddress;
    soffs[i].end = soffs[i].start + img->Sections[i].Misc.VirtualSize;
    if (cfg->on_self)
      soffs[i].off = img->MappedAddress/* + img->Sections[i].VirtualAddress*/;
    else if (img->Sections[i].PointerToRawData != 0)
      soffs[i].off = img->MappedAddress + img->Sections[i].PointerToRawData - 
          img->Sections[i].VirtualAddress;
    else
      soffs[i].off = NULL;
  }
  soffs[img->NumberOfSections].start = 0;
  soffs[img->NumberOfSections].end = 0;
  soffs[img->NumberOfSections].off = 0;

  BuildDepTree32or64 (img, cfg, root, self, soffs, soffs_len);
  free (soffs);

  if (!cfg->on_self)
    UnMapAndLoad (&loaded_image);

  /* Not sure if a forwarded export warrants an import. If it doesn't, then the dll to which the export is forwarded will NOT
   * be among the dependencies of this dll and it will be necessary to do yet another ProcessDep...
  for (i = 0; i < self->exports_len; i++)
  {
    if (self->exports[i]->forward_str != NULL && self-.exports[i]->forward == NULL)
    {
      char *forward_str_copy = NULL, *export_name = NULL, *rdot = NULL;
      DWORD export_ordinal = 0;
      forward_str_copy = strdup (self->exports[i]->forward_str);
      rdot = strrchr (forward_str_copy, '.');
      if (rdot != NULL && rdot[1] != 0)
      {
        rdot[0] = 0;
        export_name = &rdot[1];
        if (export_name[0] == '#' && export_name[1] >= '0' && export_name[1] <= '9')
        {
          export_ordinal = strtol (&export_name[1], NULL, 10);
          export_name = NULL;
        }
        self->exports[i]->forward = FindExportForward (forward_str_copy, export_name, export_ordinal);
      }
      free (forward_str_copy);
    }
  }
  */
  for (i = 0; i < self->imports_len; i++)
  {
    if (self->imports[i].mapped == NULL && self->imports[i].dll != NULL && (self->imports[i].name != NULL || self->imports[i].ordinal > 0))
    {
      struct DepTreeElement *dll = self->imports[i].dll;
      for (j = 0; j < dll->exports_len; j++)
      {
        if ((self->imports[i].name != NULL && dll->exports[j].name != NULL && strcmp (self->imports[i].name, dll->exports[j].name) == 0) ||
            (self->imports[i].ordinal > 0 && dll->exports[j].ordinal > 0 && self->imports[i].ordinal == dll->exports[j].ordinal))
        {
          self->imports[i].mapped = &dll->exports[j];
          break;
        }
      }
/*
      if (self->imports[i].mapped == NULL)
        printf ("Could not match %s (%d) in %s to %s\n", self->imports[i].name, self->imports[i].ordinal, self->module, dll->module);
*/
    }
  }
  /* By keeping items in the stack we turn it into a list of all
   * processed modules, this should be more effective at preventing
   * us from processing modules multiple times
   */
  /*PopStack (stack, stack_len, stack_size, name);*/
  return 0;
}