//						PEFILES.CPP
//
//

#include <windows.h>

#include "pefiles.h"
#include "textout.h"
#include "dasm.h"

pefile::pefile(void)
{ h1.sigbytes=0;
  securityatt.nLength=sizeof(SECURITY_ATTRIBUTES);
  securityatt.lpSecurityDescriptor=NULL;
  securityatt.bInheritHandle=FALSE;
  efile=INVALID_HANDLE_VALUE;
  peobj=NULL;
  code=NULL;
  importdata=NULL;
  dataimportdata=NULL;
  exportdata=NULL;
  dataexportdata=NULL;
  initdata=NULL;
  fixupdata=NULL;
}

pefile::~pefile(void)
{ struct pe_object *nxt;
  struct codedata *nxtcode;
  CloseHandle(efile);
  while(peobj!=NULL)
  { nxt=peobj->next;
	 delete peobj;
	 peobj=nxt;
  }
  while(code!=NULL)
  { nxtcode=code->next;
	 delete code->dataptr;
	 delete code;
	 code=nxtcode;
  }
  while(importdata!=NULL)
  { nxtcode=importdata->next;
	 delete importdata->dataptr;
	 delete importdata;
	 importdata=nxtcode;
  }
  if(dataimportdata!=NULL)delete dataimportdata;
  while(exportdata!=NULL)
  { nxtcode=exportdata->next;
	 delete exportdata->dataptr;
	 delete exportdata;
	 exportdata=nxtcode;
  }
  if(dataexportdata!=NULL)delete dataexportdata;
  while(initdata!=NULL)
  { nxtcode=initdata->next;
	 delete initdata->dataptr;
	 delete initdata;
	 initdata=nxtcode;
  }
  while(fixupdata!=NULL)
  { nxtcode=fixupdata->next;
	 delete fixupdata->dataptr;
	 delete fixupdata;
	 fixupdata=nxtcode;
  }
}

	 // offset 3c = dword ptr offset of PE header.
	 // need to check that file starts 'MZ'|'ZM' also.

BOOL pefile::open(char *fname)
{ char mzhead[2];
  DWORD pe_offset;
  DWORD num;
  if(efile!=INVALID_HANDLE_VALUE)return FALSE;
  efile=CreateFile(fname,GENERIC_READ,1,&securityatt,OPEN_EXISTING,0,NULL);
  if(efile==INVALID_HANDLE_VALUE)return FALSE;
  if(GetFileType(efile)!=FILE_TYPE_DISK)return FALSE;
  if(ReadFile(efile,mzhead,2,&num,NULL))
  { if((num==2)&&(((mzhead[0]=='M')&&(mzhead[1]=='Z'))||
		 ((mzhead[0]=='Z')&&(mzhead[1]=='M'))))
	 { SetFilePointer(efile,0x3c,NULL,FILE_BEGIN);
		if(ReadFile(efile,&pe_offset,4,&num,NULL))
		SetFilePointer(efile,pe_offset,NULL,FILE_BEGIN);
		if(ReadFile(efile,&h1,sizeof(struct header1),&num,NULL))
		{ if(h1.sigbytes==0x4550)
		  { // check the number of RVA/SIZE entries in the PE Header.
			 // correct it if we've gone too far.
			 if(h1.numintitems<0x0a)
			 { FillMemory(&h1+sizeof(struct header1)-(0x0a-h1.numintitems)*8,
				  (0x0a-h1.numintitems)*8,0);
				SetFilePointer(efile,-(0x0a-h1.numintitems)*8,NULL,FILE_CURRENT);
			 }
			 else if(h1.numintitems>0x0a)
			 { SetFilePointer(efile,(h1.numintitems-0x0a)*8,NULL,FILE_CURRENT);
			 }
			 start=h1.entrypoint_rva+h1.image_base;
			 readobjects();
			 readobjectsdata();
			 PrintBuff("FileName  :%s",fname);
			 return TRUE;
		  }
		  else if((h1.sigbytes&0xffff)==0x454e)newexe=TRUE;
		  else if((h1.sigbytes&0xffff)==0x454c)leexe=TRUE;
		  else if((h1.sigbytes&0xffff)==0x584c)os2exe=TRUE;
		  else if((h1.sigbytes&0xffff)==0x3357)w3exe=TRUE;
		  else if((h1.sigbytes&0xffff)==0x5a53)szexe=TRUE;
		  else dosexe=TRUE;
		}
	 }
	 else
	 { if(strlen(fname)>3)
		  if(!lstrcmpi(fname+strlen(fname)-3,"com"))
		  { SetFilePointer(efile,0,NULL,FILE_BEGIN);
			 PrintBuff("FileName  :%s",fname);
			 readcomfile();
			 start=0x100;
			 comfile=TRUE;
			 return TRUE;
		  }
	 }
  }
  CloseHandle(efile);
  efile=INVALID_HANDLE_VALUE;
  return FALSE;
}

void pefile::readobjects(void)
{ int i=h1.objects;
  struct pe_object *nxt;
  DWORD num;
  nxt=peobj;
  while(i)
  { if(nxt==NULL)
	 { peobj=new struct pe_object;
		nxt=peobj;
	 }
	 else
	 { nxt->next=new struct pe_object;
		nxt=nxt->next;
	 }
	 nxt->type=0;
	 ReadFile(efile,nxt,sizeof(struct peobjdata),&num,NULL);
	 if(nxt->d.obj_flags&0x40)nxt->type=INIT_DATA;
	 if(nxt->d.obj_flags&0x80)nxt->type=UNINIT_DATA;
	 if(nxt->d.obj_flags&0x20)nxt->type=CODE_DATA;
	 if(nxt->d.rva==h1.exporttable_rva)nxt->type=EXPORT_INFO;
	 if(nxt->d.rva==h1.importtable_rva)nxt->type=IMPORT_INFO;
	 if(nxt->d.rva==h1.tls_rva)nxt->type=TLS_INFO;
	 if(nxt->d.rva==h1.resourcetable_rva)nxt->type=RESOURCE_INFO;
	 if(nxt->d.rva==h1.fixuptable_rva)nxt->type=FIXUP_INFO;
	 if(nxt->d.rva==h1.debugtable_rva)nxt->type=DEBUG_INFO;
    if((!nxt->type)&&(nxt->d.obj_flags&0x20000000))nxt->type=CODE_DATA;
	 nxt->next=NULL;
	 i--;
  }
}

void pefile::peinfo(void)
{ struct pe_object *nxt;
  if(comfile)
  { PrintBuff("Format    :COM File");
	 PrintBuff("");
  }
  else
  { PrintBuff("Format    :PE File");
	 PrintBuff("");
	 PrintBuff("File Data:");
	 PrintBuff("==========");
	 PrintBuff(" Objects\t\t:%#lx",(unsigned long)h1.objects);
	 PrintBuff(" Time Date Stamp\t:%#lx",h1.timedatestamp);
	 PrintBuff(" Flags\t\t\t:%#lx",(unsigned long)h1.flags);
	 PrintBuff(" Entrypoint RVA\t\t:%#lx",h1.entrypoint_rva);
	 PrintBuff(" Image Base\t\t:%#lx",h1.image_base);
	 PrintBuff(" Object Align\t\t:%#lx",h1.objectalign);
	 PrintBuff(" File Align\t\t:%#lx",h1.filealign);
	 PrintBuff(" Image Size\t\t:%#lx",h1.imagesize);
	 PrintBuff(" Header Size\t\t:%#lx",h1.headersize);
	 PrintBuff(" Export RVA\t\t:%#lx",h1.exporttable_rva);
	 PrintBuff(" Export Data Size\t:%#lx",h1.export_datasize);
	 PrintBuff(" Import RVA\t\t:%#lx",h1.importtable_rva);
	 PrintBuff(" Import Data Size\t:%#lx",h1.import_datasize);
	 PrintBuff(" Resource RVA\t\t:%#lx",h1.resourcetable_rva);
	 PrintBuff(" Resource Data Size\t:%#lx",h1.resource_datasize);
	 PrintBuff(" Relocation RVA\t\t:%#lx",h1.fixuptable_rva);
    PrintBuff(" Relocation Data Size\t:%#lx",h1.fixup_datasize);
	 PrintBuff("");
  }
  nxt=peobj;
  PrintBuff("Objects Summary");
  PrintBuff("===============");
  while(nxt!=NULL)
  { PrintBuff("");
	 PrintBuff(" Object : %.8s",nxt->d.name);
	 switch(nxt->type)
	 { case EXPORT_INFO:
		  LastPrintBuff("  Export Information");
		  break;
		case IMPORT_INFO:
		  LastPrintBuff("  Import Information");
		  break;
		case TLS_INFO:
		  LastPrintBuff("  Thread Local Storage Information");
		  break;
		case RESOURCE_INFO:
		  LastPrintBuff("  Resource Information");
		  break;
		case FIXUP_INFO:
		  LastPrintBuff("  Fixup Information");
		  break;
		case DEBUG_INFO:
		  LastPrintBuff("  Debug Information");
		  break;
		case CODE_DATA:
		  LastPrintBuff("  Code");
		  break;
		case INIT_DATA:
		  LastPrintBuff("  Initialised Data");
		  break;
		case UNINIT_DATA:
		  LastPrintBuff("  Uninitialised Data");
		  break;
		default:
		  break;
	 }
	 PrintBuff("   Virtual Size    :%#lx",nxt->d.virt_size);
	 PrintBuff("   RVA             :%#lx",nxt->d.rva);
	 PrintBuff("   Physical Size   :%#lx",nxt->d.phys_size);
	 PrintBuff("   Physical Offset :%#lx",nxt->d.phys_offset);
	 PrintBuff("   Object Flags    :%#lx",nxt->d.obj_flags);
	 if(nxt->d.obj_flags&0x20)PrintBuff("     Code Object");
	 if(nxt->d.obj_flags&0x40)PrintBuff("     Initialised Data Object");
	 if(nxt->d.obj_flags&0x80)PrintBuff("     Uninitialised Data Object");
	 if(nxt->d.obj_flags&0x4000000)PrintBuff("     Object Must Be Cached");
	 if(nxt->d.obj_flags&0x8000000)PrintBuff("     Object Is Not Pageable");
	 if(nxt->d.obj_flags&0x10000000)PrintBuff("     Object Is Shared");
	 if(nxt->d.obj_flags&0x20000000)PrintBuff("     Executable Object");
	 if(nxt->d.obj_flags&0x40000000)PrintBuff("     Readable Object");
	 if(nxt->d.obj_flags&0x80000000)PrintBuff("     Writeable Object");
	 nxt=nxt->next;
  }
  PrintBuff("");
}              

void pefile::readobjectsdata(void)
{ struct pe_object *nxt=peobj;
  while(nxt!=NULL)
  { switch(nxt->type)
	 { case CODE_DATA:
		  readcode(&nxt->d);
		  break;
		case INIT_DATA:
		  if(loaddata) readinitdata(&nxt->d);
		  break;
		case UNINIT_DATA:
		  break;
		case EXPORT_INFO:
		  readexportdata(&nxt->d);
		  break;
		case IMPORT_INFO:
		  readimportdata(&nxt->d);
		  break;
		case TLS_INFO:
		  break;
		case RESOURCE_INFO:
		  readresources(&nxt->d);
		  break;
		case FIXUP_INFO:
		  readfixupdata(&nxt->d);
		  break;
		case DEBUG_INFO:
		  break;
		default:
		  break;
	 }
	 nxt=nxt->next;
  }
}

void pefile::readcomfile(void)
{ h1.objects=1;
  h1.entrypoint_rva=0;
  h1.image_base=0x100;
  peobj=new struct pe_object;
  peobj->next=NULL;
  peobj->type=CODE_DATA;
  strcpy(peobj->d.name,"CODESEG");
  peobj->d.virt_size=GetFileSize(efile,NULL);
  peobj->d.phys_size=peobj->d.virt_size;
  peobj->d.rva=0;
  peobj->d.phys_offset=0;
  peobj->d.obj_flags=0;
  readcode(&peobj->d);
}

void pefile::readcode(struct peobjdata *data)
{ struct codedata *inf=code;
  DWORD num;
  if(inf==NULL)
  { code=new struct codedata;
	 inf=code;
  }
  else
  { while(inf->next!=NULL) inf=inf->next;
	 inf->next=new struct codedata;
	 inf=inf->next;
  }
  inf->next=NULL;
  inf->size=data->phys_size;
  if(data->virt_size<data->phys_size)data->virt_size=data->phys_size;
  inf->virt_size=data->virt_size;
  SetFilePointer(efile,data->phys_offset,NULL,FILE_BEGIN);
  inf->addr=data->rva+h1.image_base;
  inf->dataptr=new unsigned char[inf->size+10];
  FillMemory(inf->dataptr+inf->size,10,0);
  ReadFile(efile,inf->dataptr,inf->size,&num,NULL);
}

void pefile::readimportdata(struct peobjdata *data)
{ struct codedata *inf=importdata;
  DWORD num;
  if(inf==NULL)
  { importdata=new struct codedata;
	 inf=importdata;
  }
  else
  { while(inf->next!=NULL) inf=inf->next;
	 inf->next=new struct codedata;
	 inf=inf->next;
  }
  inf->next=NULL;
  inf->size=data->phys_size;
  if(data->virt_size<data->phys_size)data->virt_size=data->phys_size;
  inf->virt_size=data->virt_size;
  SetFilePointer(efile,data->phys_offset,NULL,FILE_BEGIN);
  inf->addr=data->rva+h1.image_base;
  inf->dataptr=new unsigned char[inf->size];
  ReadFile(efile,inf->dataptr,inf->size,&num,NULL);
}

void pefile::readexportdata(struct peobjdata *data)
{ struct codedata *inf=exportdata;
  DWORD num;
  if(inf==NULL)
  { exportdata=new struct codedata;
	 inf=exportdata;
  }
  else
  { while(inf->next!=NULL) inf=inf->next;
	 inf->next=new struct codedata;
	 inf=inf->next;
  }
  inf->next=NULL;
  inf->size=data->phys_size;
  if(data->virt_size<data->phys_size)data->virt_size=data->phys_size;
  inf->virt_size=data->virt_size;
  SetFilePointer(efile,data->phys_offset,NULL,FILE_BEGIN);
  inf->addr=data->rva+h1.image_base;
  inf->dataptr=new unsigned char[inf->size];
  ReadFile(efile,inf->dataptr,inf->size,&num,NULL);
}

void pefile::readresources(struct peobjdata *data)
{ struct codedata *inf=resourcedata;
  DWORD num;
  if(inf==NULL)
  { resourcedata=new struct codedata;
	 inf=resourcedata;
  }
  else
  { while(inf->next!=NULL) inf=inf->next;
	 inf->next=new struct codedata;
	 inf=inf->next;
  }
  inf->next=NULL;
  inf->size=data->phys_size;
  if(data->virt_size<data->phys_size)data->virt_size=data->phys_size;
  inf->virt_size=data->virt_size;
  SetFilePointer(efile,data->phys_offset,NULL,FILE_BEGIN);
  inf->addr=data->rva+h1.image_base;
  inf->dataptr=new unsigned char[inf->size];
  ReadFile(efile,inf->dataptr,inf->size,&num,NULL);
}

void pefile::readinitdata(struct peobjdata *data)
{ struct codedata *inf=initdata;
  DWORD num;
  if(inf==NULL)
  { initdata=new struct codedata;
	 inf=initdata;
  }
  else
  { while(inf->next!=NULL) inf=inf->next;
	 inf->next=new struct codedata;
	 inf=inf->next;
  }
  inf->next=NULL;
  inf->size=data->phys_size;
  if(data->virt_size<data->phys_size)data->virt_size=data->phys_size;
  inf->virt_size=data->virt_size;
  SetFilePointer(efile,data->phys_offset,NULL,FILE_BEGIN);
  inf->addr=data->rva+h1.image_base;
  inf->dataptr=new unsigned char[inf->size];
  ReadFile(efile,inf->dataptr,inf->size,&num,NULL);
}

void pefile::readfixupdata(struct peobjdata *data)
{ struct codedata *inf=fixupdata;
  DWORD num;
  if(inf==NULL)
  { fixupdata=new struct codedata;
	 inf=fixupdata;
  }
  else
  { while(inf->next!=NULL) inf=inf->next;
	 inf->next=new struct codedata;
	 inf=inf->next;
  }
  inf->next=NULL;
  inf->size=data->phys_size;
  if(data->virt_size<data->phys_size)data->virt_size=data->phys_size;
  inf->virt_size=data->virt_size;
  SetFilePointer(efile,data->phys_offset,NULL,FILE_BEGIN);
  inf->addr=data->rva+h1.image_base;
  inf->dataptr=new unsigned char[inf->size];
  ReadFile(efile,inf->dataptr,inf->size,&num,NULL);
}

struct codedata *pefile::getcode(void)
{ return code;
}

struct codedata *pefile::getimports(void)
{ struct codedata *dta;
  if(importdata!=NULL)return importdata;
  if(dataimportdata!=NULL)return dataimportdata;
  if(h1.importtable_rva)
  { dta=getinitdata();
	 while(dta!=NULL)
	 { if((h1.importtable_rva>=dta->addr-h1.image_base)&&(h1.importtable_rva<dta->addr+dta->size-h1.image_base))
		{ dataimportdata=new struct codedata;
		  dataimportdata->next=NULL;
		  dataimportdata->dataptr=dta->dataptr+h1.importtable_rva-dta->addr+h1.image_base;
		  dataimportdata->size=h1.import_datasize;
		  dataimportdata->virt_size=dataimportdata->size;
		  dataimportdata->addr=h1.importtable_rva+h1.image_base;
		  PrintBuff("NB: Import Data is part of initialised data object");
		  return dataimportdata;
		}
		dta=dta->next;
	 }
	 dta=getcode();
	 while(dta!=NULL)
	 { if((h1.importtable_rva>=dta->addr-h1.image_base)&&(h1.importtable_rva<dta->addr+dta->size-h1.image_base))
		{ dataimportdata=new struct codedata;
		  dataimportdata->next=NULL;
		  dataimportdata->dataptr=dta->dataptr+h1.importtable_rva-dta->addr+h1.image_base;
		  dataimportdata->size=h1.import_datasize;
		  dataimportdata->virt_size=dataimportdata->size;
		  dataimportdata->addr=h1.importtable_rva+h1.image_base;
		  PrintBuff("Warning: Import Data is part of code object");
		  return dataimportdata;
		}
		dta=dta->next;
	 }
  }
  return NULL;
}

struct codedata *pefile::getexports(void)
{ struct codedata *dta;
  if(exportdata!=NULL)return exportdata;
  if(dataexportdata!=NULL)return dataexportdata;
  if(h1.exporttable_rva)
  { dta=getinitdata();
	 while(dta!=NULL)
	 { if((h1.exporttable_rva>=dta->addr-h1.image_base)&&(h1.exporttable_rva<dta->addr+dta->size-h1.image_base))
		{ dataexportdata=new struct codedata;
		  dataexportdata->next=NULL;
		  dataexportdata->dataptr=dta->dataptr+h1.exporttable_rva-dta->addr+h1.image_base;
		  dataexportdata->size=h1.export_datasize;
		  dataexportdata->virt_size=dataexportdata->size;
		  dataexportdata->addr=h1.exporttable_rva+h1.image_base;
		  PrintBuff("NB: Export Data is part of initialised data object");
		  return dataexportdata;
		}
		dta=dta->next;
	 }
	 dta=getcode();
	 while(dta!=NULL)
	 { if((h1.exporttable_rva>=dta->addr-h1.image_base)&&(h1.exporttable_rva<dta->addr+dta->size-h1.image_base))
		{ dataexportdata=new struct codedata;
		  dataexportdata->next=NULL;
		  dataexportdata->dataptr=dta->dataptr+h1.exporttable_rva-dta->addr+h1.image_base;
		  dataexportdata->size=h1.export_datasize;
		  dataexportdata->virt_size=dataexportdata->size;
		  dataexportdata->addr=h1.exporttable_rva+h1.image_base;
		  PrintBuff("Warning: Export Data is part of code object");
		  return dataexportdata;
		}
		dta=dta->next;
	 }                
  }
  return NULL;
}

struct codedata *pefile::getresources(void)
{ struct codedata *dta;
  if(resourcedata!=NULL)return resourcedata;
  if(dataresourcedata!=NULL)return dataresourcedata;
  if(h1.resourcetable_rva)
  { dta=getinitdata();
	 while(dta!=NULL)
	 { if((h1.resourcetable_rva>=dta->addr-h1.image_base)&&(h1.resourcetable_rva<dta->addr+dta->size-h1.image_base))
		{ dataresourcedata=new struct codedata;
		  dataresourcedata->next=NULL;
		  dataresourcedata->dataptr=dta->dataptr+h1.resourcetable_rva-dta->addr+h1.image_base;
		  dataresourcedata->size=h1.resource_datasize;
		  dataresourcedata->virt_size=dataresourcedata->size;
		  dataresourcedata->addr=h1.resourcetable_rva+h1.image_base;
		  PrintBuff("NB: Resource Data is part of initialised data object");
		  return dataresourcedata;
		}
		dta=dta->next;
	 }
	 dta=getcode();
	 while(dta!=NULL)
	 { if((h1.resourcetable_rva>=dta->addr-h1.image_base)&&(h1.resourcetable_rva<dta->addr+dta->size-h1.image_base))
		{ dataresourcedata=new struct codedata;
		  dataresourcedata->next=NULL;
		  dataresourcedata->dataptr=dta->dataptr+h1.resourcetable_rva-dta->addr+h1.image_base;
		  dataresourcedata->size=h1.resource_datasize;
		  dataresourcedata->virt_size=dataresourcedata->size;
		  dataresourcedata->addr=h1.resourcetable_rva+h1.image_base;
		  PrintBuff("Warning: Resource Data is part of code object");
		  return dataresourcedata;
		}
		dta=dta->next;
	 }
  }
  return NULL;
}

struct codedata *pefile::getinitdata(void)
{ return initdata;
}

struct codedata *pefile::getrelocdata(void)
{ return fixupdata;
}

unsigned long pefile::getstart(void)
{ return start;
}

unsigned long pefile::getimagebase(void)
{ return h1.image_base;
}

unsigned char *pefile::findrva(unsigned long rva)
{ struct codedata *idata;
  unsigned long datarva;
  idata=getcode();
  while(idata!=NULL)
  { datarva=idata->addr-h1.image_base;
	 if((rva>=datarva)&&(rva<datarva+idata->size))break;
	 idata=idata->next;
  }
  if(idata!=NULL)
  { return rva-datarva+idata->dataptr;
  }
  idata=getinitdata();
  while(idata!=NULL)
  { datarva=idata->addr-h1.image_base;
	 if((rva>=datarva)&&(rva<datarva+idata->size))break;
	 idata=idata->next;
  }
  if(idata!=NULL)
  { return rva-datarva+idata->dataptr;
  }
  idata=exportdata;
  while(idata!=NULL)
  { datarva=idata->addr-h1.image_base;
	 if((rva>=datarva)&&(rva<datarva+idata->size))break;
	 idata=idata->next;
  }
  if(idata!=NULL)
  { return rva-datarva+idata->dataptr;
  }
  idata=importdata;
  while(idata!=NULL)
  { datarva=idata->addr-h1.image_base;
	 if((rva>=datarva)&&(rva<datarva+idata->size))break;
	 idata=idata->next;
  }
  if(idata!=NULL)
  { return rva-datarva+idata->dataptr;
  }
  return NULL;
}

unsigned char *pefile::finddata(unsigned long dat)
{ struct codedata *idata;
  unsigned long datar;
  idata=getcode();
  while(idata!=NULL)
  { datar=idata->addr;
	 if((dat>=datar)&&(dat<datar+idata->size))break;
	 idata=idata->next;
  }
  if(idata!=NULL)
  { return dat-datar+idata->dataptr;
  }
  idata=getinitdata();
  while(idata!=NULL)
  { datar=idata->addr;
	 if((dat>=datar)&&(dat<datar+idata->size))break;
	 idata=idata->next;
  }
  if(idata!=NULL)
  { return dat-datar+idata->dataptr;
  }
  return NULL;
}

