#include "ModuleScanner.h"
#include <stdlib.h>
#include <string.h>
#include <elf.h>
#include <dlfcn.h>

CModuleScanner::CModuleScanner(void *addr)
  : m_addr(addr)
{}

void *CModuleScanner::FindSignature(const char *sig, const char *pattern)
{
  DynLibInfo lib;
  bool found;
  char *ptr, *end;
  size_t len = strlen(pattern);

  memset(&lib, 0, sizeof(DynLibInfo));

  if (!GetLibraryInfo(m_addr, lib))
    return NULL;

  ptr = reinterpret_cast<char *>(lib.baseAddress);
  end = ptr + lib.memorySize;

  while (ptr < end)
  {
    found = true;
    for (register size_t i = 0; i < len; i++)
    {
      if (pattern[i] != '?' && sig[i] != ptr[i])
      {
        found = false;
        break;
      }
    }

    if (found)
      return ptr;

    ptr++;
  }

  return NULL;
}

bool CModuleScanner::GetLibraryInfo(const void *libPtr, DynLibInfo &lib)
{
  unsigned long baseAddr;
  Dl_info info;
  Elf32_Ehdr *file;
  Elf32_Phdr *phdr;
  uint16_t phdrCount;

  if (libPtr == NULL)
    return false;

  if (!dladdr(libPtr, &info))
    return false;

  if (!info.dli_fbase || !info.dli_fname)
    return false;

  /* This is for our insane sanity checks :o */
  baseAddr = reinterpret_cast<unsigned long>(info.dli_fbase);
  file = reinterpret_cast<Elf32_Ehdr *>(baseAddr);

  /* Check ELF magic */
  if (memcmp(ELFMAG, file->e_ident, SELFMAG) != 0)
    return false;

  /* Check ELF version */
  if (file->e_ident[EI_VERSION] != EV_CURRENT)
    return false;

  /* Check ELF architecture, which is 32-bit/x86 right now
   * Should change this for 64-bit if Valve gets their act together
   */
  if (file->e_ident[EI_CLASS] != ELFCLASS32 || file->e_machine != EM_386 || file->e_ident[EI_DATA] != ELFDATA2LSB)
    return false;

  /* For our purposes, this must be a dynamic library/shared object */
  if (file->e_type != ET_DYN)
    return false;

  phdrCount = file->e_phnum;
  phdr = reinterpret_cast<Elf32_Phdr *>(baseAddr + file->e_phoff);

  /* Add up the memory sizes of segments marked as PT_LOAD as those are the only ones that should be in memory */
  for (uint16_t i = 0; i < phdrCount; i++)
  {
    Elf32_Phdr &hdr = phdr[i];
    if (hdr.p_type == PT_LOAD)
      lib.memorySize += hdr.p_memsz;
  }

  lib.baseAddress = reinterpret_cast<void *>(baseAddr);
  return true;
}