1 files changed, 30 insertions, 20 deletions

diff --git a/disasm/extractfunction.cpp b/disasm/extractfunction.cpp
index 6d6f7bc..4fb675c 100644
--- a/disasm/extractfunction.cpp
+++ b/disasm/extractfunction.cpp
@@ -1,39 +1,49 @@
 #include "extractfunction.hpp"
 
-#include "parsefileandrewind.hpp"
+#include "leparseutil.hpp"
 
 #include <distorm.h>
 
 #include <iomanip>
+#include <fstream>
 
-void extract_function(std::istream& is, binparse::Value32 object_id, binparse::Offset32 function_offset) {
-	
-	auto file = parse_file_and_rewind(is);
+void extract_function(std::string file_path, binparse::Value32 object_id, binparse::Offset32 function_offset) {
+
+	std::ifstream ifs(file_path, std::ios::binary);
+	ifs.unsetf(std::ios::skipws);
 	
-	is.ignore(file.le_header.data_page_offset);
+	auto file = parse_file(ifs);
 	
-	_DInst decodedInstructions[1];
+	_DInst decinst;
 	_DecodeType dt = Decode32Bits;
 	unsigned int decodedInstructionsCount = 0;
 	
-	auto object = file.object_table.entries[object_id];
-	unsigned int object_size = object.nr_page_table_entries * file.le_header.page_size;
+	std::vector<uint8_t> code_buf = read_object(ifs, file, object_id);
 	
-	std::vector<uint8_t> code_buf;
-	code_buf.reserve(object_size);
-	std::copy_n(std::istream_iterator<uint8_t>(is), object_size, std::back_inserter(code_buf));
+	binparse::Offset32 base_reloc_offset = file.object_table.entries[object_id].reloc_base_address;
+	binparse::Offset32 buffer_offset = binparse::Offset32(function_offset - base_reloc_offset);
 	
 	_CodeInfo ci;
-	ci.code = code_buf.data() + function_offset;
-	ci.codeLen = code_buf.size() - function_offset;
-	ci.codeOffset = object.reloc_base_address + function_offset;
+	ci.code = code_buf.data() + buffer_offset;
+	ci.codeLen = code_buf.size() - buffer_offset;
+	ci.codeOffset = base_reloc_offset + buffer_offset;
 	ci.dt = dt;
 	ci.features = DF_NONE;
 	
-	distorm_decompose64(&ci, decodedInstructions, 1, &decodedInstructionsCount);
-	
-	_DecodedInst inst;
-	distorm_format64(&ci, &decodedInstructions[0], &inst);
-	
-	std::cout << std::hex << std::setw(8) << std::setfill('0') << inst.offset << ":\t" << inst.mnemonic.p << " " << inst.operands.p << std::endl;
+	while(true) {
+		distorm_decompose64(&ci, &decinst, 1, &decodedInstructionsCount);
+		
+		if(decinst.flags == FLAG_NOT_DECODABLE) {
+			break;
+		}
+	
+		_DecodedInst inst;
+		distorm_format64(&ci, &decinst, &inst);
+	
+		std::cout << std::hex << std::setw(8) << std::setfill('0') << inst.offset << ":\t" << inst.mnemonic.p << " " << inst.operands.p << std::endl;
+		
+		ci.code += decinst.size;
+		ci.codeLen -= decinst.size;
+		ci.codeOffset += decinst.size;
+	}
 }