Fix support for PE32+ files.

src/file_analysis/analyzer/pe/pe-file-headers.pac

@@ -65,8 +65,14 @@ type Optional_Header(len: uint16, number_of_sections: uint16) = record {
 	size_of_uninit_data     : uint32;
 	addr_of_entry_point     : uint32;
 	base_of_code            : uint32;
-	base_of_data            : uint32;
-	image_base              : uint32;
+	have_base_of_data: case pe_format of {
+		PE32    -> base_of_data: uint32;
+		default -> not_present:  empty;
+	};
+	is_pe32: case pe_format of {
+		PE32_PLUS -> image_base_64: uint64;
+		default   -> image_base_32: uint32;
+	};
 	section_alignment       : uint32;
 	file_alignment          : uint32;
 	os_version_major        : uint16;
@@ -81,14 +87,17 @@ type Optional_Header(len: uint16, number_of_sections: uint16) = record {
 	checksum                : uint32;
 	subsystem               : uint16;
 	dll_characteristics     : uint16;
-	mem: case magic of {
-		267  -> i32           : Mem_Info32;
-		268  -> i64           : Mem_Info64;
+	mem: case pe_format of {
+		PE32      -> i32: Mem_Info32;
+		PE32_PLUS -> i64: Mem_Info64;
 		default -> InvalidPEFile : empty;
 	};
 	loader_flags            : uint32;
 	number_of_rva_and_sizes : uint32;
 	rvas			: RVAS(number_of_rva_and_sizes);
+} &let {
+	pe_format: uint8 = $context.connection.set_pe32_format(magic);
+	image_base: uint64 = pe_format == PE32_PLUS ? image_base_64 : image_base_32;
 } &length=len;
 
 type Section_Headers(num: uint16) = record {
@@ -112,3 +121,43 @@ type Section_Header = record {
 	proc: bool = $context.connection.proc_section(this);
 } &length=40;
 
+refine connection MockConnection += {
+	%member{
+		uint64 max_file_location_;
+		uint8  pe32_format_;
+	%}
+
+	%init{
+		max_file_location_ = 0;
+		pe32_format_ = UNKNOWN_VERSION;;
+	%}
+
+	function proc_section(h: Section_Header): bool
+		%{
+		if ( ${h.size_of_raw_data} + ${h.ptr_to_raw_data} > max_file_location_ )
+			max_file_location_ = ${h.size_of_raw_data} + ${h.ptr_to_raw_data};
+		
+		if ( ${h.virtual_addr} > 0 && ${h.virtual_addr} == import_table_rva_ )
+			import_table_va_ = ${h.ptr_to_raw_data};
+		return true;
+		%}
+
+	function set_pe32_format(magic: uint16): uint8
+		%{
+		if ( ${magic} == 0x10b )
+			pe32_format_ = PE32;
+		if ( ${magic} == 0x20b )
+			pe32_format_ = PE32_PLUS;
+		return pe32_format_;
+		%}
+
+	function get_max_file_location(): uint64
+		%{
+		return max_file_location_;
+		%}
+
+	function get_pe32_format(): uint8
+		%{
+		return pe32_format_;
+		%}
+};

src/file_analysis/analyzer/pe/pe-file-idata.pac

@@ -11,12 +11,17 @@ type import_directory = record {
 	proc: bool = $context.connection.proc_image_import_directory(this);
 } &length=20;
 
-type import_lookup_attrs = record {
-	attrs: uint32;
-} &length=4;
+type import_lookup_attrs(pe32_format: uint8) = record {
+	is_pe32_plus: case pe32_format of {
+		PE32_PLUS -> attrs_64: uint64;
+		default   -> attrs_32: uint32;
+	};
+} &let {
+	attrs: uint64 = (pe32_format == PE32_PLUS) ? attrs_64 : attrs_32;
+} &length=(pe32_format == PE32_PLUS ? 8 : 4);
 
 type import_lookup_table = record {
-	attrs: import_lookup_attrs[] &until($element.attrs == 0);
+	attrs: import_lookup_attrs($context.connection.get_pe32_format())[] &until($element.attrs == 0);
 } &let {
 	proc: bool = $context.connection.proc_import_lookup_table(this);
 };
@@ -28,6 +33,8 @@ type import_entry(is_module: bool, pad_align: uint8) = record {
 		false -> index: uint16;
 	};
 	name: null_terminated_string;
+} &let {
+	proc: bool = $context.connection.proc_import_hint(name);
 };
 
 type idata = record {
@@ -68,6 +75,8 @@ refine connection MockConnection += {
 		next_hint_align_ = 0;
 	%}
 
+	# When we read the section header, store the relative virtual address and
+	# size of the .idata section, so we know when we get there.
 	function proc_idata_rva(r: RVA): bool
 		%{
 		import_table_rva_ = ${r.virtual_address};
@@ -76,19 +85,50 @@ refine connection MockConnection += {
 		return true;
 		%}
 
-	function proc_section(h: Section_Header): bool
-		%{
-		if ( ${h.virtual_addr} > 0 && ${h.virtual_addr} == import_table_rva_ )
-			import_table_va_ = ${h.ptr_to_raw_data};
-		return true;
-		%}
-
+	# Each import directory means another module we're importing from.
 	function proc_image_import_directory(i: import_directory): bool
 		%{
 		num_imports_++;
 		return true;
 		%}
 
+	# Store the number of functions imported in each module lookup table.
+	function proc_import_lookup_table(t: import_lookup_table): bool
+		%{
+		--num_imports_;
+		imports_per_module_.push_back(${t.attrs}->size());
+		return true;
+		%}		
+
+	# We need to calculate the length of the next padding field
+	function proc_import_hint(hint_name: bytestring): bool
+		%{
+		next_hint_align_ = ${hint_name}.length() % 2;
+		printf("Imported %s\n", ${hint_name}.data());
+		return true;
+		%}
+
+	# Functions have an index field, modules don't. Which one is this?
+	function get_next_hint_type(): bool
+		%{
+		if ( next_hint_is_module_ )
+			{
+			next_hint_is_module_ = false;
+			return true;
+			}
+		if ( --imports_per_module_[next_hint_index_] == 0)
+			{
+			++next_hint_index_;
+			return true;
+			}
+		return false;
+		%}
+
+	function imports_done(): bool
+		%{
+		return next_hint_index_ == imports_per_module_.size();
+		%}
+
 	function get_import_table_addr(): uint32
 		%{
 		return import_table_va_ > 0 ? import_table_va_ : 0;
@@ -109,37 +149,4 @@ refine connection MockConnection += {
 		return next_hint_align_;
 		%}		
 
-	function proc_import_lookup_table(t: import_lookup_table): bool
-		%{
-		--num_imports_;
-		imports_per_module_.push_back(${t.attrs}->size());
-		return true;
-		%}		
-
-	function get_next_hint_type(): bool
-		%{
-		if ( next_hint_is_module_ )
-			{
-			next_hint_is_module_ = false;
-			return true;
-			}
-		if ( --imports_per_module_[next_hint_index_] == 0)
-			{
-			++next_hint_index_;
-			return true;
-			}
-		return false;
-		%}
-
-	function imports_done(): bool
-		%{
-		return next_hint_index_ == imports_per_module_.size();
-		%}
-
-	function proc_import_hint(hint_name: bytestring): bool
-		%{
-		next_hint_align_ = ${hint_name}.length() % 2;
-		printf("Import function: %s\n", ${hint_name}.data());
-		return true;
-		%}
 };

src/file_analysis/analyzer/pe/pe-file-types.pac

@@ -1,5 +1,11 @@
 # Basic PE types
 
+enum PE_File_Format {
+	UNKNOWN_VERSION = 0,
+	PE32            = 1,
+	PE32_PLUS       = 2,
+};
+
 type Mem_Info32 = record {
 	size_of_stack_reserve : uint32;
 	size_of_stack_commit  : uint32;
@@ -24,7 +30,7 @@ type RVA = record {
 } &length=8;
 
 # The BinPAC padding type doesn't work here.
-type Padding(length: uint32) = record {
+type Padding(length: uint64) = record {
 	pad: bytestring &length=length &transient;
 };
 

src/file_analysis/analyzer/pe/pe-file.pac

@@ -5,10 +5,12 @@
 # The base record for a Portable Executable file
 type PE_File = record {
 	headers : Headers;
-	pad		: Padding(iat_loc);
+	pad1	: Padding(iat_loc);
 	iat     : idata &length=$context.connection.get_import_table_len();
+	pad2    : Padding(restofdata);
 } &let {
 	unparsed_hdr_len: uint32 = headers.pe_header.optional_header.size_of_headers - headers.length;
-	iat_loc: uint32 = $context.connection.get_import_table_addr() - headers.pe_header.optional_header.size_of_headers + unparsed_hdr_len;
+	iat_loc: uint64 = $context.connection.get_import_table_addr() - headers.pe_header.optional_header.size_of_headers + unparsed_hdr_len;
+	restofdata: uint64 = $context.connection.get_max_file_location() - $context.connection.get_import_table_addr() - $context.connection.get_import_table_len();
 } &byteorder=littleendian;