Refactored argument parsing, and added stub for extractfunction.

8 files changed, 178 insertions, 41 deletions

diff --git a/disasm/CMakeLists.txt b/disasm/CMakeLists.txt
index d65f522..7c21409 100644
--- a/disasm/CMakeLists.txt
+++ b/disasm/CMakeLists.txt
@@ -1,7 +1,9 @@
 
 add_executable(disasm
 	disasm.cpp
+	parsefileandrewind.hpp parsefileandrewind.cpp
 	dumpobject.hpp dumpobject.cpp
+	extractfunction.hpp extractfunction.cpp
 )
 
 find_package(Boost COMPONENTS filesystem program_options system REQUIRED)
diff --git a/disasm/disasm.cpp b/disasm/disasm.cpp
index 53c74eb..d5b1333 100644
--- a/disasm/disasm.cpp
+++ b/disasm/disasm.cpp
@@ -1,5 +1,6 @@
 
 #include "dumpobject.hpp"
+#include "extractfunction.hpp"
 
 #include <boost/filesystem.hpp>
 #include <boost/program_options.hpp>
@@ -8,12 +9,14 @@
 #include <fstream>
 
 int parse_dumpobject_options(std::vector<std::string> arguments);
+int parse_extractfunction_options(std::vector<std::string> arguments);
 
 void print_top_level_help(std::string program_name) {
 	std::cout << "Usage " << program_name << ":" << std::endl;
 	std::cout << std::endl;
 	std::cout << "\t" << program_name << " " << "help" << std::endl;
 	std::cout << "\t" << program_name << " " << "dumpobject" << std::endl;
+	std::cout << "\t" << program_name << " " << "extractfunction" << std::endl;
 }
 
 int main(int argc, char* argv[]) {
@@ -64,12 +67,61 @@ int main(int argc, char* argv[]) {
 	
 	if(command == "dumpobject") {
 		return parse_dumpobject_options(options);
+	} else if (command == "extractfunction") {
+		return parse_extractfunction_options(options);
 	} else {
 		print_top_level_help(program_name);
 		return 0;
 	}
 }
 
+int parse_exe_file_option(boost::program_options::variables_map const& vm, boost::filesystem::path& file_path) {
+	if(vm.count("exe")) {
+		std::string exe_file = vm["exe"].as<std::string>();
+		
+		if(boost::filesystem::exists(exe_file)) {
+			if(!boost::filesystem::is_directory(exe_file)) {
+				file_path = exe_file;
+			} else {
+				std::cerr << exe_file << " is a folder" << std::endl;
+				return -1;
+			}
+		} else {
+			std::cerr << "file: " << exe_file << " does not exist" << std::endl;
+			return -1;
+		}
+	} else {
+		std::cerr << "exe_file option is required" << std::endl;
+		return -1;
+	}
+	return 0;
+}
+
+int parse_object_id_option(boost::program_options::variables_map const& vm, uint32_t& object_id) {
+	if(vm.count("object_id")) {
+		object_id = vm["object_id"].as<uint32_t>();
+		if(object_id == 0) {
+			std::cerr << "object id cannot be 0" << std::endl;
+			return -1;
+		}
+	} else {
+		std::cerr << "object_id option is required" << std::endl;
+		return -1;
+	}
+	return 0;
+}
+
+int parse_function_offset_option(boost::program_options::variables_map const& vm, uint32_t& function_offset) {
+	if(vm.count("function_offset")) {
+		std::string str = vm["function_offset"].as<std::string>();
+		function_offset = std::stoul(str, nullptr, 16);
+	} else {
+		std::cerr << "function_offset option is required" << std::endl;
+		return -1;
+	}
+	return 0;
+}
+
 int parse_dumpobject_options(std::vector<std::string> arguments) {
 	boost::program_options::options_description description;
 	description.add_options()
@@ -94,36 +146,60 @@ int parse_dumpobject_options(std::vector<std::string> arguments) {
 	}
 	
 	boost::filesystem::path file_path;
-	if(vm.count("exe")) {
-		std::string exe_file = vm["exe"].as<std::string>();
-		
-		if(boost::filesystem::exists(exe_file)) {
-			if(!boost::filesystem::is_directory(exe_file)) {
-				file_path = exe_file;
-			} else {
-				std::cerr << exe_file << " is a folder" << std::endl;
-				std::cerr << std::endl;
-				std::cerr << description << std::endl;
-				return -1;
-			}
-		} else {
-			std::cerr << "file: " << exe_file << " does not exist" << std::endl;
-			std::cerr << std::endl;
-			std::cerr << description << std::endl;
-			return -1;
-		}
+	if(int ret = parse_exe_file_option(vm, file_path)) {
+		return ret;
 	}
 	
-	uint32_t object_id;
-	if(vm.count("object_id")) {
-		object_id = vm["object_id"].as<uint32_t>();
-		if(object_id == 0) {
-			std::cerr << "object id cannot be 0" << std::endl;
-			return -1;
-		}
+	binparse::Value32 object_id;
+	if(int ret =parse_object_id_option(vm, object_id)) {
+		return ret;
 	}
 	
 	std::ifstream file_stream(file_path.string());
 	
 	dump_object(file_stream, object_id);
 }
+
+int parse_extractfunction_options(std::vector<std::string> arguments) {
+	boost::program_options::options_description description;
+	description.add_options()
+		("help,h", "produces this help message")
+		("exe,e", boost::program_options::value<std::string>()->required(), "The LE executable to parse the header for.")
+		("object_id,o", boost::program_options::value<uint32_t>()->required(), "The object number to disassemble")
+		("function_offset,f", boost::program_options::value<std::string>()->required(), "The function start offset in hexadecimal notation")
+	;
+	
+	boost::program_options::variables_map vm;
+	boost::program_options::store(boost::program_options::command_line_parser(arguments).options(description).run(), vm);
+	
+	if(vm.count("help")) {
+		std::cout << description << std::endl;
+		return 0;
+	}
+	
+	try {
+		boost::program_options::notify(vm);
+	} catch(boost::program_options::required_option const& e) {
+		std::cout << e.what() << std::endl;
+		return -1;
+	}
+	
+	boost::filesystem::path file_path;
+	if(int ret = parse_exe_file_option(vm, file_path)) {
+		return ret;
+	}
+	
+	binparse::Value32 object_id;
+	if(int ret =parse_object_id_option(vm, object_id)) {
+		return ret;
+	}
+	
+	binparse::Offset32 function_offset;
+	if(int ret = parse_function_offset_option(vm, function_offset)) {
+		return ret;
+	}
+	
+	std::ifstream file_stream(file_path.string());
+	
+	extract_function(file_stream, object_id, function_offset);
+}
diff --git a/disasm/dumpobject.cpp b/disasm/dumpobject.cpp
index 85fbf43..4f1a5cc 100644
--- a/disasm/dumpobject.cpp
+++ b/disasm/dumpobject.cpp
@@ -1,4 +1,5 @@
 #include "dumpobject.hpp"
+#include "parsefileandrewind.hpp"
 
 #include "le_file.hpp"
 
@@ -6,28 +7,25 @@
 
 #include <iomanip>
 
-void dump_object(std::ifstream& ifs, uint32_t object_id)
-{
-	auto x = le::parse_file(ifs);
-	
-	ifs.clear();
-	ifs.seekg(0, std::ios::beg);
-	
+void dump_object(std::ifstream& ifs, binparse::Value32 object_id)
+{	
+	auto file = parse_file_and_rewind(ifs);
+
 	std::vector<uint8_t> code(std::istreambuf_iterator<char>(ifs), {});
 	
 	std::vector<_DecodedInst> instructions;
 	instructions.resize(100000);
 	unsigned int read_inst; 
 
-	auto object = x.object_table.entries[object_id];
+	auto object = file.object_table.entries[object_id];
 	auto index = object.page_table_index;
 	
-	binparse::Offset32 offset = x.le_header.data_page_offset;
+	binparse::Offset32 offset = file.le_header.data_page_offset;
 	
 	auto result = distorm_decode64(
 		object.reloc_base_address,
-		code.data() + offset + (index - 1) * x.le_header.page_size,
-		x.object_table.entries[1].nr_page_table_entries * x.le_header.page_size,
+		code.data() + offset + (index - 1) * file.le_header.page_size,
+		file.object_table.entries[1].nr_page_table_entries * file.le_header.page_size,
 		Decode32Bits,
 		instructions.data(),
 		instructions.size(),
@@ -37,10 +35,6 @@ void dump_object(std::ifstream& ifs, uint32_t object_id)
 	instructions.resize(read_inst);
 	instructions.shrink_to_fit();
 	
-	if(result) {
-		
-	}
-	
 	for(auto&& inst : instructions) {
 		std::cout << std::hex << std::setw(8) << std::setfill('0') << inst.offset << ":\t" << inst.mnemonic.p << " " << inst.operands.p << std::endl;
 	}
diff --git a/disasm/dumpobject.hpp b/disasm/dumpobject.hpp
index b8042b1..8685588 100644
--- a/disasm/dumpobject.hpp
+++ b/disasm/dumpobject.hpp
@@ -1,6 +1,8 @@
 #pragma once
 
+#include "types.hpp"
+
 #include <fstream>
 #include <cstdint>
 
-void dump_object(std::ifstream& ifs, uint32_t object_id);
\ No newline at end of file
+void dump_object(std::ifstream& ifs, binparse::Value32 object_id);
\ No newline at end of file
diff --git a/disasm/extractfunction.cpp b/disasm/extractfunction.cpp
new file mode 100644
index 0000000..6d6f7bc
--- /dev/null
+++ b/disasm/extractfunction.cpp
@@ -0,0 +1,39 @@
+#include "extractfunction.hpp"
+
+#include "parsefileandrewind.hpp"
+
+#include <distorm.h>
+
+#include <iomanip>
+
+void extract_function(std::istream& is, binparse::Value32 object_id, binparse::Offset32 function_offset) {
+	
+	auto file = parse_file_and_rewind(is);
+	
+	is.ignore(file.le_header.data_page_offset);
+	
+	_DInst decodedInstructions[1];
+	_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;
+	code_buf.reserve(object_size);
+	std::copy_n(std::istream_iterator<uint8_t>(is), object_size, std::back_inserter(code_buf));
+	
+	_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.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;
+}
diff --git a/disasm/extractfunction.hpp b/disasm/extractfunction.hpp
new file mode 100644
index 0000000..57e0597
--- /dev/null
+++ b/disasm/extractfunction.hpp
@@ -0,0 +1,6 @@
+#pragma once
+
+#include "le_file.hpp"
+#include "types.hpp"
+
+void extract_function(std::istream& is, binparse::Value32 object_id, binparse::Offset32 function_offset);
diff --git a/disasm/parsefileandrewind.cpp b/disasm/parsefileandrewind.cpp
new file mode 100644
index 0000000..704e9da
--- /dev/null
+++ b/disasm/parsefileandrewind.cpp
@@ -0,0 +1,11 @@
+#include "parsefileandrewind.hpp"
+
+le::File parse_file_and_rewind(std::istream& is)
+{
+	auto file = le::parse_file(is);
+	
+	is.clear();
+	is.seekg(0, std::ios::beg);
+	
+	return file;
+}
diff --git a/disasm/parsefileandrewind.hpp b/disasm/parsefileandrewind.hpp
new file mode 100644
index 0000000..4d564ac
--- /dev/null
+++ b/disasm/parsefileandrewind.hpp
@@ -0,0 +1,7 @@
+#pragma once
+
+#include "le_file.hpp"
+
+#include <istream>
+
+le::File parse_file_and_rewind(std::istream& is);