//
// Script that parses TIB/PEB/PEB_LDR/LDR_MODULE structures
// Load it when the process is being debugged
// We don't apply structures, instead, we only name some interesting entries in those structures
//
// (c) Hex-Rays.com
//

#include <idc.idc>

static GetThreadTib(tid)
{
  auto ea;

  ea = SegByName(form("TIB[%08X]", tid));
  if (ea == BADADDR)
    ea = SegByName("TIB");
  return SegByBase(ea);
}

static MakeNameWithType(ea, type)
{
  auto old_type;
  old_type = GetLongPrm(INF_STRTYPE);
  SetLongPrm(INF_STRTYPE, type);
  MakeStr(ea, BADADDR);
  SetLongPrm(INF_STRTYPE, old_type);
}

static MakeWordName(ea, name)
{
  MakeWord(ea);
  MakeNameEx(ea, name, SN_NOWARN);
}

static MakeDwordName(ea, name)
{
  MakeDword(ea);
  MakeNameEx(ea, name, SN_NOWARN);
}

static MakeByteName(ea, name)
{
  MakeByte(ea);
  MakeNameEx(ea, name, SN_NOWARN);
}

static WalkExceptionRegRecords(tid, strt)
{
  auto cnt, t;

  MakeName(strt, form("seh_list_%08X", tid));
  MakeDword(strt);

  cnt = 1;
  strt = Dword(strt);
  while (strt != BADADDR)
  {
    MakeDword(strt+4);
    MakeDword(strt);

    t = Dword(strt + 4);

    Message("%08x: exception handler %d\n", t, cnt);

    MakeCode(t);

    // Go to next
    strt = Dword(strt);
    cnt++;
  }
}

static DescribePEB(peb_ea)
{
  auto t;
  MakeName(peb_ea, "PEB");
  Message("--PEB--\n%08X\n", peb_ea);

  MakeByteName(peb_ea+0x02, "PEB_ImageBase");
  MakeDwordName(peb_ea+0x08, "PEB_ImageBase");
  MakeDwordName(peb_ea+0x68, "PEB_NtGlobalFlag");

  t = peb_ea+0xC;
  MakeDwordName(t, "o_PEB_LDR");
  DescribePebLdr(Dword(t));
}

static DescribePebLdr(pebldr_ea)
{
  auto t, s;
  MakeDwordName(pebldr_ea, "PEB_LDR");
  MakeDword(pebldr_ea+4);
  MakeDword(pebldr_ea+8);
  Message("--PEB_LDR--\n%08X\n", pebldr_ea);

  t = Dword(pebldr_ea);
  if ((t != 0x28))
  {
    if (AskYN(1, "Unknown PEB_LDR_DATA structure size, continue?") == 0)
      return;
  }

  t = pebldr_ea + 0x0C;
  s = "InLoadOrderModuleList";
  Message("%s list ; head=%08X\n", s, t);
  DescribeModuleListLinks(t, s);
  WalkModuleList(t, 0x0, 1);

  t = pebldr_ea + 0x14;
  s = "InMemoryOrderModuleList";
  Message("%s list ; head=%08X\n", s, t);
  DescribeModuleListLinks(t, s);
  WalkModuleList(t, 0x08, 0);

  t = pebldr_ea + 0x1C;
  s = "InInitializationOrderModuleList";
  Message("%s list ; head=%08X\n", s, t);
  DescribeModuleListLinks(t, s);
  WalkModuleList(t, 0x10, 0);
}

static DescribeModuleListLinks(ofs, name)
{
  auto fmt;
  fmt = "o_%s_%s";
  MakeDwordName(ofs, form(fmt, name, "Flink"));
  MakeDwordName(ofs+4, form(fmt, name, "Blink"));
}

static WalkModuleList(head, ofs, give_names)
{
  auto m, cnt, e;

  m = Dword(head) - ofs;
  cnt = 1;

  do
  {
    if (give_names)
    {
      MakeDwordName(m+0x00, form("module_%04d", cnt));
      MakeDwordName(m+0x04, form("module_%04d_InLoadOrderModuleList_Blink", cnt));
      MakeDwordName(m+0x08, form("module_%04d_InMemoryOrderModuleList_Flink", cnt));
      MakeDwordName(m+0x0C, form("module_%04d_InMemoryOrderModuleList_Blink", cnt));
      MakeDwordName(m+0x10, form("module_%04d_InInitializationOrderModuleList_Flink", cnt));
      MakeDwordName(m+0x14, form("module_%04d_InInitializationOrderModuleList_Blink", cnt));
      MakeDwordName(m+0x18, form("module_%04d_base", cnt));
      MakeDwordName(m+0x1C, form("module_%04d_entrypoint", cnt));
      MakeDwordName(m+0x20, form("module_%04d_sizeofimage", cnt));
      MakeDwordName(m+0x24, form("module_%04d_unicode_len", cnt));
      MakeDwordName(m+0x28, form("module_%04d_fulldllname", cnt));
      MakeNameWithType(Dword(m+0x28), ASCSTR_UNICODE);
      MakeDwordName(m+0x2C, form("module_%04d_base_unicode_len", cnt));
      MakeDwordName(m+0x30, form("module_%04d_basedllname", cnt));
      MakeDwordName(m+0x34, form("module_%04d_flags", cnt));
      MakeWordName(m+0x38, form("module_%04d_loadcount", cnt));
      MakeWordName(m+0x3A, form("module_%04d_TlsIndex", cnt));
    }

    Message(" module(%04d): ldrmodule(%08X) base(%08x) entry(%08X) name(%s)\n", cnt, m, Dword(m+0x18), Dword(m+0x1C), GetString(Dword(m+0x24+4), BADADDR, ASCSTR_UNICODE));
    e = Dword(m + ofs);
    if (e == BADADDR || e == 0)
      break;
    m = e - ofs;
    ++cnt;
  } while (e != head);
}

static main()
{
  auto tib_ea, tid, peb_ea, t;

  tid = GetCurrentThreadId();
  if (tid == -1)
  {
    Message("Could not determine thread id\n");
    return BADADDR;
  }

  tib_ea = GetThreadTib(tid);
  if (tib_ea == BADADDR)
  {
    Message("Could not determine thread TIB\n");
    return -1;
  }

  // TIB
  Message("---TIB---\n%08X\n", tib_ea);

 // Jump(tib_ea);
  MakeName(tib_ea, form("TIB_%08X", tid));
  MakeDwordName(tib_ea+0x24, form("TID_%08X", tid));
  MakeDwordName(tib_ea+0x20, form("PID_%08X", tid));
  MakeDwordName(tib_ea+0x34, form("LastError_%08X", tid));

  Message("---SEH---\n%08X: head\n", tib_ea);
  WalkExceptionRegRecords(tid, tib_ea);

  // TIB.PEB
  t = tib_ea+0x30;
  peb_ea = Dword(t);
  MakeDwordName(t, form("PEB_%08X", tid));

  // PEB
  DescribePEB(peb_ea);
}