unit PEHeaderUnit;

interface   
         
uses
  Windows, SysUtils, 
  JwaWinnt, UtilUnit;

type
  TPEHeader = class
  private
    FHdrPtr: Pointer;
    FSections: Array of PImageSectionHeader;
    FDataDirectorys: Array of PImageDataDirectory;
    FDosHeader: PImageDosHeader;
    FFileHeader: PImageFileHeader;
    FOptionalHeader32: PImageOptionalHeader32;
    procedure DumpHeader;
    procedure DumpDataDirectorys;
    procedure DumpSectionHeaders;
    function GetDataDirectory(Index: Cardinal): PImageDataDirectory;
    function GetDirectoryCount: Cardinal;
    function GetSection(Index: Cardinal): PImageSectionHeader;
    function GetSectionCount: Cardinal;
    function GetHeaderSize: Cardinal;
    procedure SetHeaderSize(const Value: Cardinal);
    function GetFileAlignment: Cardinal;
    procedure SetFileAlignment(const Value: Cardinal);
    function GetSectionAlignment: Cardinal;
    procedure SetSectionAlignment(const Value: Cardinal);
  public
    constructor Create;
    destructor Destroy; override;
    procedure Clear;
    function ValidatePE(HdrPtr: Pointer): Boolean;
    function Assign(HdrPtr: Pointer): Bool;
    function OffsetInSection(Offset: Cardinal): Cardinal;
    function OffsettoRVA(Offset: Cardinal): Cardinal;
    function RVAInSection(RVA: Cardinal): Cardinal;
    function RVAtoOffset(RVA: Cardinal): Cardinal;
    function RVAtoPtr(RVA: Cardinal): Pointer;
    function AddSection(var NewSectionHeader: PImageSectionHeader): Bool;
    function ExtendHeader(NewSize: Cardinal; var RealSize: Cardinal): Bool;
    function WipeSectionHeader(Index: Cardinal): Bool;
    function CalcFileAilgnment(RawLength: Cardinal): Cardinal;
    function CalcSectionAilgnment(RawLength: Cardinal): Cardinal;
    procedure UpdateImageSize;
    function GetRealHeaderSize: Cardinal;
    property HeaderPtr: Pointer read FHdrPtr;
    property SectionCount: Cardinal read GetSectionCount;
    property DirectoryCount: Cardinal read GetDirectoryCount;
    property Sections[Index: Cardinal]: PImageSectionHeader read GetSection;
    property DataDirectorys[Index: Cardinal]: PImageDataDirectory read GetDataDirectory;
    property DosHeader: PImageDosHeader read FDosHeader;
    property FileHeader: PImageFileHeader read FFileHeader;
    property OptionalHeader32: PImageOptionalHeader32 read FOptionalHeader32;
    property HeaderSize: Cardinal read GetHeaderSize Write SetHeaderSize;
    property SA: Cardinal read GetSectionAlignment write SetSectionAlignment;
    property FA: Cardinal read GetFileAlignment write SetFileAlignment;
  end;

implementation

{ TPEHeader }

function TPEHeader.AddSection(var NewSectionHeader: PImageSectionHeader): Bool;
var
  NeedHdrSize: Cardinal;
  i: integer;
  pLastSection: PImageSectionHeader;
  lsOffset: Cardinal;
  l: Cardinal;
  NeedMovPtr: Pointer;
  NeedMovSize: Cardinal;
  NeedMovOffset: Cardinal;
begin
  result := true;
  NewSectionHeader := nil;

  NeedHdrSize := CalcFileAilgnment(Cardinal(PtrDiff(FDataDirectorys[DirectoryCount - 1],
    FHdrPtr)) + SizeOf(TImageDataDirectory) + (SectionCount + 1) *
    SizeOf(TImageSectionHeader));
  if (HeaderSize < NeedHdrSize)  then
    if not ExtendHeader(NeedHdrSize, l) then
      result := false;
  
  NeedMovSize := 0;
  if SectionCount > 0 then
    NeedMovOffset := PtrDiff(PtrAdd(Sections[SectionCount - 1],
      SizeOf(TImageSectionHeader) * 3), FHdrPtr)
  else
    NeedMovOffset := PtrDiff(PtrAdd(FDataDirectorys[DirectoryCount - 1],
      SizeOf(TImageDataDirectory)), FHdrPtr);
  for i := 0 to DirectoryCount - 1 do
  begin
    if DataDirectorys[i]^.VirtualAddress = 0 then Continue;
    if (DataDirectorys[i]^.VirtualAddress <= HeaderSize) and
       (HeaderSize - NeedMovOffset < DataDirectorys[i]^.Size) then
       inc(NeedMovSize, DataDirectorys[i]^.Size);
  end;
  if (NeedHdrSize < NeedMovOffset + NeedMovSize) then
    if not ExtendHeader(NeedMovOffset + NeedMovSize, l) then
      result := false;
  if result then
  begin
    if SectionCount > 0 then
    begin
      pLastSection := Sections[SectionCount - 1];
      lsOffset := PtrDiff(pLastSection, FHdrPtr) + SizeOf(TImageSectionHeader);
      NeedMovPtr := PtrAdd(pLastSection, SizeOf(TImageSectionHeader) * 2);
      inc(pLastSection);
    end
    else
    begin
      pLastSection := PtrAdd(FDataDirectorys[DirectoryCount - 1], SizeOf(TImageDataDirectory));
      lsOffset := PtrDiff(pLastSection, FHdrPtr);
      NeedMovPtr := PtrAdd(pLastSection, SizeOf(TImageSectionHeader));
    end;
    for i := 0 to DirectoryCount - 1 do
    begin
      if DataDirectorys[i]^.VirtualAddress = 0 then Continue;
      if DataDirectorys[i]^.VirtualAddress <= lsOffset then
      begin
        CopyMemory(NeedMovPtr, RVAtoPtr(DataDirectorys[i]^.VirtualAddress),
          DataDirectorys[i]^.Size);
        DataDirectorys[i]^.VirtualAddress := PtrDiff(NeedMovPtr, FHdrPtr);
        PtrAdd(NeedMovPtr, DataDirectorys[i]^.Size);
      end;
    end;

    ZeroMemory(pLastSection, SizeOf(TImageSectionHeader) * 2);
    inc(FFileHeader^.NumberOfSections);
    NewSectionHeader := pLastSection;
    if SectionCount > 1 then
      NewSectionHeader^.VirtualAddress :=
        CalcSectionAilgnment(Sections[SectionCount - 2]^.VirtualAddress +
        Sections[SectionCount - 2]^.Misc.VirtualSize)
    else       
      NewSectionHeader^.VirtualAddress := 0;

    NewSectionHeader^.Misc.VirtualSize := SA;
    NewSectionHeader^.SizeOfRawData := 0;
    NewSectionHeader^.PointerToRawData := 0;
    NewSectionHeader^.Characteristics := $E00000E0;
    DumpSectionHeaders;
  end;
end;

function TPEHeader.Assign(HdrPtr: Pointer): Bool;
begin
  result := ValidatePE(HdrPtr);
  if not result then exit;
  Clear;
  FHdrPtr := HdrPtr;
  DumpHeader;
  GetMem(FHdrPtr, HeaderSize);
  CopyMemory(FHdrPtr, HdrPtr, HeaderSize);
  DumpHeader;
  DumpDataDirectorys;
  DumpSectionHeaders;
  HeaderSize := GetRealHeaderSize;
end;

function TPEHeader.CalcFileAilgnment(RawLength: Cardinal): Cardinal;
begin
  result := RawLength;
  if result mod FA <> 0 then
    result := (result div FA + 1) * FA;
end;

function TPEHeader.CalcSectionAilgnment(RawLength: Cardinal): Cardinal;
begin
  result := RawLength;
  if result mod SA <> 0 then
    result := (result div SA + 1) * SA;
end;

procedure TPEHeader.Clear;
begin
  if FHdrPtr <> nil then
    FreeMem(FHdrPtr);
  FHdrPtr := nil;
  FDosHeader := nil;
  FFileHeader := nil;
  FOptionalHeader32 := nil;
end;

constructor TPEHeader.Create;
begin
  inherited;
end;

destructor TPEHeader.Destroy;
begin
  Clear;
  inherited;
end;

procedure TPEHeader.DumpDataDirectorys;  
var
  i: integer;
  DDPtr: PImageDataDirectory;
begin
  FDataDirectorys := nil;
  SetLength(FDataDirectorys, DirectoryCount);
  DDPtr := @FOptionalHeader32^.DataDirectory[0];
  for i := 0 to DirectoryCount - 1 do
  begin
    FDataDirectorys[i] := DDPtr;
    inc(DDPtr);
  end;
end;

procedure TPEHeader.DumpHeader;
begin
  FDosHeader := FHdrPtr;
  FFileHeader := PtrAdd(FDosHeader , FDosHeader^.e_lfanew + 4);
  FOptionalHeader32 := PtrAdd(FFileHeader, IMAGE_SIZEOF_FILE_HEADER);
end;

procedure TPEHeader.DumpSectionHeaders;
var
  i: integer;
  SHPtr: PImageSectionHeader;
begin
  FSections := nil;
  SetLength(FSections, SectionCount);
  SHPtr := PtrAdd(FOptionalHeader32, FFileHeader^.SizeOfOptionalHeader);
  for i := 0 to SectionCount - 1 do
  begin
    FSections[i] := SHPtr;
    inc(SHPtr);
  end;
end;

function TPEHeader.ExtendHeader(NewSize: Cardinal; var RealSize: Cardinal): Bool;
var
  i: integer;
  AddSize: Cardinal;
  TmpHdrPtr: Pointer;
begin
  NewSize := CalcFileAilgnment(NewSize);
  if (NewSize > HeaderSize) and (NewSize <= $1000) then
  begin
    AddSize := NewSize - HeaderSize;
    TmpHdrPtr := AllocMem(NewSize);
    CopyMemory(TmpHdrPtr, DosHeader, HeaderSize);
    Clear;
    FHdrPtr := TmpHdrPtr;
    DumpHeader;
    DumpDataDirectorys;
    DumpSectionHeaders;
    for i := 0 to SectionCount - 1 do
    begin
      inc(FSections[i]^.PointerToRawData, AddSize);
    end;
    inc(FOptionalHeader32^.SizeOfImage, AddSize);
    HeaderSize := NewSize;
    result := true;
  end
  else
    result := false;
  RealSize := HeaderSize;
end;

function TPEHeader.GetDataDirectory(Index: Cardinal): PImageDataDirectory;
begin
  if Index >= DirectoryCount then
    result := nil
  else
    result := FDataDirectorys[index];
end;

function TPEHeader.GetDirectoryCount: Cardinal;
begin
  result := FOptionalHeader32^.NumberOfRvaAndSizes;
end;

function TPEHeader.GetFileAlignment: Cardinal;
begin
  result := FOptionalHeader32^.FileAlignment;
end;

function TPEHeader.GetHeaderSize: Cardinal;
begin
  result := FOptionalHeader32^.SizeOfHeaders;
end;

function TPEHeader.GetRealHeaderSize: Cardinal;
var
  i: Cardinal;
begin
  result := $FFFFFFFF;
  for i := 0 to SectionCount - 1 do
  begin
    if Sections[i].PointerToRawData = 0 then Continue;
    result := Min(result, Sections[i].PointerToRawData);
  end;
end;

function TPEHeader.GetSection(Index: Cardinal): PImageSectionHeader;
begin
  if Index >= SectionCount then
    result := nil
  else
    result := FSections[index];
end;

function TPEHeader.GetSectionAlignment: Cardinal;
begin
  result := FOptionalHeader32^.SectionAlignment;
end;

function TPEHeader.GetSectionCount: Cardinal;
begin
  result := FFileHeader^.NumberOfSections;
end;

function TPEHeader.OffsetInSection(Offset: Cardinal): Cardinal;
var
  i: Cardinal;
begin
  result := NULL_VALUE;
  if Offset = 0 then exit;
  for i := 0 to SectionCount - 1 do
    if (Offset >= Sections[i]^.PointerToRawData) and
       (Offset - Sections[i]^.PointerToRawData < Sections[i]^.SizeOfRawData) then
    begin
      result := i;
      break;
    end;
end;

function TPEHeader.OffsettoRVA(Offset: Cardinal): Cardinal;
begin
  result := OffsetInSection(Offset);
  if result <> NULL_VALUE then
    result := Offset - Sections[result]^.SizeOfRawData +
              Sections[result]^.VirtualAddress;
end;

function TPEHeader.RVAInSection(RVA: Cardinal): Cardinal;
var
  i: Cardinal;
begin
  result := NULL_VALUE;
  if RVA = 0 then exit;
  for i := 0 to SectionCount - 1 do
    if (RVA >= Sections[i]^.VirtualAddress) and
       (RVA - Sections[i]^.VirtualAddress < Sections[i]^.Misc.VirtualSize) then
    begin
      result := i;
      break;
    end;
end;

function TPEHeader.RVAtoOffset(RVA: Cardinal): Cardinal;
var
  i: Cardinal;
begin
  if integer(RVA - Sections[SectionCount - 1]^.VirtualAddress) >
     integer(Sections[SectionCount - 1]^.Misc.VirtualSize) then
    result := NULL_VALUE
  else
  begin
    i := RVAInSection(RVA);
    if i = NULL_VALUE then
      result := RVA
    else
      result := RVA - Sections[i]^.VirtualAddress +
                Sections[i]^.PointerToRawData;
  end;
end;

function TPEHeader.RVAtoPtr(RVA: Cardinal): Pointer;
begin
  result := PtrAdd(FHdrPtr, RVAtoOffset(RVA));
end;

procedure TPEHeader.SetFileAlignment(const Value: Cardinal);
begin
  FOptionalHeader32^.FileAlignment := Value;
end;

procedure TPEHeader.SetHeaderSize(const Value: Cardinal);
begin
  FOptionalHeader32^.SizeOfHeaders := Value;
end;

procedure TPEHeader.SetSectionAlignment(const Value: Cardinal);
begin
  FOptionalHeader32^.SectionAlignment := Value;
end;

procedure TPEHeader.UpdateImageSize;
begin
  FOptionalHeader32^.SizeOfImage := Sections[SectionCount - 1]^.VirtualAddress +
    Sections[SectionCount - 1]^.Misc.VirtualSize;
end;

function TPEHeader.ValidatePE(HdrPtr: Pointer): Boolean;
begin
  result := true;
  try
    result := not IsBadReadPtr(HdrPtr, 2) and result;
    if result then
    begin
      result := (PWORD(HdrPtr)^ = $5A4D) and result;
      HdrPtr := PtrAdd(HdrPtr, PImageDosHeader(HdrPtr)^.e_lfanew);
      result := not IsBadReadPtr(HdrPtr, 2) and result;
      if result then
        result := (PWORD(HdrPtr)^ = $4550) and result;
    end;
  except
    result := false;
  end;
end;

function TPEHeader.WipeSectionHeader(Index: Cardinal): Bool;
var
  i: Cardinal;
  ps, ps2: PImageSectionHeader;
begin
  result := false;
  if Index >= SectionCount then exit;
  ps := Sections[Index];
  ps2 := PtrAdd(ps, SizeOf(TImageSectionHeader));
  for i := Index to SectionCount - 1 do
  begin                                              
    ps2^.VirtualAddress := ps^.VirtualAddress;
    CopyMemory(ps, ps2, SizeOf(TImageSectionHeader));
    ps2^.VirtualAddress := ps^.VirtualAddress + ps^.Misc.VirtualSize;
    inc(ps);
    inc(ps2);
  end;
  ZeroMemory(ps2, SizeOf(TImageSectionHeader));
  dec(FFileHeader^.NumberOfSections);
  result := true;
end;

end.