From 6b69dda15d1d628e073e688605fc8f3f5dcecf91 Mon Sep 17 00:00:00 2001
From: C-Sto <7466346+C-Sto@users.noreply.github.com>
Date: Wed, 21 Dec 2022 15:36:56 +0800
Subject: [PATCH] add delay imports functions, refactor section retrieval

---
 pe/imports.go | 175 ++++++++++++++++++++++++++++++++++----------------
 1 file changed, 119 insertions(+), 56 deletions(-)

diff --git a/pe/imports.go b/pe/imports.go
index d1bc9c5..ebcf881 100644
--- a/pe/imports.go
+++ b/pe/imports.go
@@ -16,70 +16,30 @@ type ImportDirectory struct {
 	DllName string
 }
 
+// ImgDelayDescr entry for delayloaded libraries
+type ImgDelayDescr struct {
+	GrAttrs,
+	RVADLLName,
+	RVAHmod,
+	RVAIAT,
+	RVAINT,
+	RVABoundIAT,
+	RVAUnloadIAT,
+	DwTimeStamp uint32
+
+	DllName string
+}
+
 // IAT returns the DataDirectory for the IAT
 func (f *File) IAT() *DataDirectory {
-	pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64
-
-	// grab the number of data directory entries
-	var ddLength uint32
-	if pe64 {
-		ddLength = f.OptionalHeader.(*OptionalHeader64).NumberOfRvaAndSizes
-	} else {
-		ddLength = f.OptionalHeader.(*OptionalHeader32).NumberOfRvaAndSizes
-	}
-
-	// check that the length of data directory entries is large
-	// enough to include the imports directory.
-	if ddLength < IMAGE_DIRECTORY_ENTRY_IAT+1 {
-		return nil
-	}
-
-	// grab the IAT entry
-	var idd DataDirectory
-	if pe64 {
-		idd = f.OptionalHeader.(*OptionalHeader64).DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT]
-	} else {
-		idd = f.OptionalHeader.(*OptionalHeader32).DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT]
-	}
+	_, idd := f.sectionFromDirectoryEntry(IMAGE_DIRECTORY_ENTRY_IAT)
 	return &idd
 }
 
 // ImportDirectoryTable - returns the Import Directory Table, a pointer to the section, and the section raw data
 func (f *File) ImportDirectoryTable() ([]ImportDirectory, *Section, *[]byte, error) {
 
-	pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64
-
-	// grab the number of data directory entries
-	var ddLength uint32
-	if pe64 {
-		ddLength = f.OptionalHeader.(*OptionalHeader64).NumberOfRvaAndSizes
-	} else {
-		ddLength = f.OptionalHeader.(*OptionalHeader32).NumberOfRvaAndSizes
-	}
-
-	// check that the length of data directory entries is large
-	// enough to include the imports directory.
-	if ddLength < IMAGE_DIRECTORY_ENTRY_IMPORT+1 {
-		return nil, nil, nil, nil
-	}
-
-	// grab the import data directory entry
-	var idd DataDirectory
-	if pe64 {
-		idd = f.OptionalHeader.(*OptionalHeader64).DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]
-	} else {
-		idd = f.OptionalHeader.(*OptionalHeader32).DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT]
-	}
-
-	// figure out which section contains the import directory table
-	var ds *Section
-	ds = nil
-	for _, s := range f.Sections {
-		if s.VirtualAddress <= idd.VirtualAddress && idd.VirtualAddress < s.VirtualAddress+s.VirtualSize {
-			ds = s
-			break
-		}
-	}
+	ds, idd := f.sectionFromDirectoryEntry(IMAGE_DIRECTORY_ENTRY_IMPORT)
 
 	// didn't find a section, so no import libraries were found
 	if ds == nil {
@@ -180,3 +140,106 @@ func (f *File) ImportedLibraries() ([]string, error) {
 	}
 	return all, nil
 }
+
+func (f File) sectionFromDirectoryEntry(directory uint32) (*Section, DataDirectory) {
+	pe64 := f.Machine == IMAGE_FILE_MACHINE_AMD64
+
+	// grab the number of data directory entries
+	var ddLength uint32
+	if pe64 {
+		ddLength = f.OptionalHeader.(*OptionalHeader64).NumberOfRvaAndSizes
+	} else {
+		ddLength = f.OptionalHeader.(*OptionalHeader32).NumberOfRvaAndSizes
+	}
+
+	// check that the length of data directory entries is large
+	// enough to include the directory.
+	if ddLength < directory+1 {
+		return nil, DataDirectory{}
+	}
+
+	// grab the directory entry
+	var idd DataDirectory
+	if pe64 {
+		idd = f.OptionalHeader.(*OptionalHeader64).DataDirectory[directory]
+	} else {
+		idd = f.OptionalHeader.(*OptionalHeader32).DataDirectory[directory]
+	}
+
+	// figure out which section contains the directory table
+	var ds *Section
+	for _, s := range f.Sections {
+		if s.VirtualAddress <= idd.VirtualAddress && idd.VirtualAddress < s.VirtualAddress+s.VirtualSize {
+			ds = s
+			break
+		}
+	}
+	return ds, idd
+}
+
+// ImportDelayDirectoryTable - returns the Import Directory Table, a pointer to the section, and the section raw data
+func (f *File) ImportDelayDirectoryTable() ([]ImgDelayDescr, *Section, *[]byte, error) {
+
+	ds, idd := f.sectionFromDirectoryEntry(IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT)
+
+	// didn't find a section, so no import libraries were found
+	if ds == nil {
+		return nil, nil, nil, nil
+	}
+
+	sectionData, err := ds.Data()
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	// seek to the virtual address specified in the import data directory
+	d := sectionData[idd.VirtualAddress-ds.VirtualAddress:]
+	var dida []ImgDelayDescr
+	for len(d) > 0 {
+		var dt ImgDelayDescr
+		idx := 0
+		dt.GrAttrs = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.RVADLLName = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.RVAHmod = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.RVAIAT = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.RVAINT = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.RVABoundIAT = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.RVAUnloadIAT = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+		dt.DwTimeStamp = binary.LittleEndian.Uint32(d[idx*4 : (idx*4)+4])
+		idx++
+
+		//check for nulls (termination entry) https://github.com/VirusTotal/yara/blob/master/libyara/modules/pe/pe.c#L1163
+		if dt.DwTimeStamp|dt.GrAttrs|dt.RVADLLName|dt.RVAHmod|dt.RVAIAT|dt.RVAINT|dt.RVABoundIAT|dt.RVAUnloadIAT|dt.DwTimeStamp == 0 {
+			break
+		}
+		if s, ok := getString(sectionData, int(dt.RVADLLName-ds.VirtualAddress)); ok {
+			dt.DllName = s
+		}
+		d = d[32:]
+		dida = append(dida, dt)
+	}
+
+	return dida, ds, &sectionData, nil
+}
+
+// ImportedDelayLibraries returns the names of all libraries referred to by the binary f
+// that are added to the delay imports directory. These libraries are not loaded at initialisation,
+// but may be loaded during runtime.
+func (f *File) ImportedDelayLibraries() ([]string, error) {
+	ida, _, _, err := f.ImportDelayDirectoryTable()
+	if err != nil {
+		return nil, err
+	}
+	var all []string
+	for _, dt := range ida {
+		all = append(all, dt.DllName)
+	}
+	return all, nil
+}