#include "ElfReader.hpp"
#include "sal/SALConfig.hpp"
#include <stdio.h>
#include <cstdlib>
#include <cstring> // memcpy, cstring
#include <algorithm>
#include "Demangler.hpp"
#include <limits.h>

namespace fail {

const std::string ELF::NOTFOUND = "[ELFReader] Function not found.";
static const ElfSymbol g_SymbolNotFound;

bool operator==(const std::string & str, const ElfSymbol & sym) {
	return sym.getName() == str;
}

bool operator==(guest_address_t address, const ElfSymbol & sym) {
	return sym.getAddress() == address;
}

std::ostream& operator<< (std::ostream &out, const ElfSymbol &symbol) {
	return (out << symbol.getName()
			<< " @ 0x" << std::hex << symbol.getAddress()
			<< " size " << std::dec << symbol.getSize());
}



ElfReader::ElfReader() : m_log("ElfReader", false) {
	// try to open elf file from environment variable
	char * elfpath = getenv("FAIL_ELF_PATH");
	if (elfpath == NULL) {
		m_log << "FAIL_ELF_PATH not set :(" << std::endl;
	} else {
		setup(elfpath);
	}
}

ElfReader::ElfReader(const char* path) : m_log("FAIL*Elfinfo", false) {
	setup(path);
}

void ElfReader::setup(const char* path) {
	// Try to open the ELF file
	FILE *fp = fopen(path, "rb");
	if (!fp) {
		m_log << "Error: Could not open " << path << std::endl;
		return;
	}

	m_filename = std::string(path);

	// Evaluate headers
	Elf64_Ehdr ehdr;
	Elf64_Shdr sec_hdr;
	int num_hdrs;
	if (!read_ELF_file_header(fp, &ehdr)) {
		m_log << "Error: " << path << " is not an ELF file" << std::endl;
		return;
	}
	num_hdrs = ehdr.e_shnum;
	m_log << "Evaluating "
		<< (m_elfclass == ELFCLASS32 ? "32-bit" : "64-bit")
		<< " ELF file: " << path << std::endl;

	// Parse ELF segments
	Elf64_Phdr seg_hdr;
	for(unsigned i = 0; i < ehdr.e_phnum; ++i) {
		if(!read_ELF_segment_header(fp, &ehdr, i, &seg_hdr)) {
			m_log << "Invalid segment header at index " << i << std::endl;
			continue;
		}
		process_segment(&seg_hdr);
	}

	// Parse symbol table and generate internal map
	for (int i = 0; i < num_hdrs; ++i) {
		if (!read_ELF_section_header(fp, &ehdr, i, &sec_hdr)) {
			m_log << "Wrong section to read" << std::endl;
			continue;
		} else if (sec_hdr.sh_type != SHT_SYMTAB && sec_hdr.sh_type != SHT_DYNSYM) {
			continue;
		}
		process_symboltable(fp, &ehdr, i);
	}
	// Parse section information
	if (!read_ELF_section_header(fp, &ehdr, ehdr.e_shstrndx, &sec_hdr)) {
		m_log << "Error: reading section string table sect_num = " << ehdr.e_shstrndx << std::endl;
		return;
	}

	char *buff = (char*)malloc(sec_hdr.sh_size);
	if (!buff) {
		m_log << "Error: malloc failed to allocate buffer for shstrtab" << std::endl;
		return;
	}
	// seek to the offset in the file,
	fseek(fp, (off_t)sec_hdr.sh_offset, SEEK_SET);
	if (fread(buff, sec_hdr.sh_size, 1, fp) != 1) {
		printf("Couldn't read complete shstrtab\n");
		exit(0);
	}

	m_log << "Total number of sections: " << num_hdrs << std::endl;
	for (int i = 0; i < num_hdrs; ++i) {
		if (!read_ELF_section_header(fp, &ehdr, i, &sec_hdr)) {
			printf("Error: wrong Section to read\n");
		} else {
			process_section(&sec_hdr, buff);
		}
	}

	free(buff);
	fclose(fp);
}


int ElfReader::process_section(Elf64_Shdr *sect_hdr, char *sect_name_buff) {
	// Add section name, start address and size to list
	int idx=sect_hdr->sh_name;
	// m_sections_map.push_back( sect_hdr->sh_addr, sect_hdr->sh_size, sect_name_buff+idx );
	m_sectiontable.push_back( ElfSymbol(sect_name_buff+idx, sect_hdr->sh_addr, sect_hdr->sh_size, ElfSymbol::SECTION)  );

	return 0;
}

void ElfReader::process_segment(Elf64_Phdr *seg_hdr) {
	// FIXME: this does not handle the PT_PHDR segment type where the
	//        program header is actually located somewhere else.
	if(seg_hdr->p_type == PT_LOAD) {
		m_segmenttable.emplace_back(seg_hdr);
	}
}

static void Elf32to64_Sym(Elf32_Sym const *src, Elf64_Sym *dest)
{
	dest->st_name = src->st_name;
	dest->st_value = src->st_value;
	dest->st_size = src->st_size;
	dest->st_info = src->st_info;
	dest->st_other = src->st_other;
	dest->st_shndx = src->st_shndx;
}

bool ElfReader::process_symboltable(FILE *fp, Elf64_Ehdr const *ehdr, int sect_num) {

	Elf64_Shdr sect_hdr;
	Elf32_Sym mysym32;
	Elf64_Sym mysym;
	char *name_buf = 0;
	int num_sym, link, idx;
	off_t sym_data_offset;
	int sym_data_size;

	if (!read_ELF_section_header(fp, ehdr, sect_num, &sect_hdr)) {
		return false;
	}
	// we have to check to which strtab it is linked
	link = sect_hdr.sh_link;
	sym_data_offset = sect_hdr.sh_offset;
	sym_data_size = sect_hdr.sh_size;
	num_sym = sym_data_size /
		(m_elfclass == ELFCLASS32 ? sizeof(Elf32_Sym) : sizeof(Elf64_Sym));

	// read the corresponding strtab
	if (!read_ELF_section_header(fp, ehdr, link, &sect_hdr)) {
		return false;
	}
	// get the size of strtab in file and allocate a buffer
	name_buf = (char *) malloc(sect_hdr.sh_size);
	if (!name_buf) {
		return false;
	}
	// get the offset of strtab in file and seek to it
	fseek(fp, sect_hdr.sh_offset, SEEK_SET);
	// read all data from the section to the buffer.
	if (fread(name_buf, sect_hdr.sh_size, 1, fp) != 1) {
		return false;
	}
	// so we have the namebuf now seek to symtab data
	fseek(fp, sym_data_offset, SEEK_SET);

	for (int i = 0; i < num_sym; ++i) {
		int type = ELFCLASSNONE;
		if (m_elfclass == ELFCLASS32) {
			if (fread(&mysym32, sizeof(mysym32), 1, fp) != 1) {
				return false;
			}
			Elf32to64_Sym(&mysym32, &mysym);
			type = ELF32_ST_TYPE(mysym32.st_info);
		} else if (m_elfclass == ELFCLASS64) {
			if (fread(&mysym, sizeof(mysym), 1, fp) != 1) {
				return false;
			}
			type = ELF64_ST_TYPE(mysym.st_info);
		}
		idx = mysym.st_name;

		if (type != STT_SECTION && type != STT_FILE) {
			m_symboltable.push_back(
				ElfSymbol(name_buf + idx, mysym.st_value, mysym.st_size, ElfSymbol::SYMBOL, type));
		}
	}
	free(name_buf);
	return true;
}

static void Elf32to64_Ehdr(Elf32_Ehdr const *src, Elf64_Ehdr *dest)
{
	memcpy(dest->e_ident, src->e_ident, sizeof(dest->e_ident));
	dest->e_type = src->e_type;
	dest->e_machine = src->e_machine;
	dest->e_version = src->e_version;
	dest->e_entry = src->e_entry;
	dest->e_phoff = src->e_phoff;
	dest->e_shoff = src->e_shoff;
	dest->e_flags = src->e_flags;
	dest->e_ehsize = src->e_ehsize;
	dest->e_phentsize = src->e_phentsize;
	dest->e_phnum = src->e_phnum;
	dest->e_shentsize = src->e_shentsize;
	dest->e_shnum = src->e_shnum;
	dest->e_shstrndx = src->e_shstrndx;
}

bool ElfReader::read_ELF_file_header(FILE *fp, Elf64_Ehdr *filehdr)
{
	Elf32_Ehdr filehdr32;
	size_t ret;

	rewind(fp);
	ret = fread(&filehdr32, sizeof(filehdr32), 1, fp);

	if (ret != 1 ||
		strncmp((const char *)filehdr32.e_ident, "\177ELF", 4) != 0) {
		return false;
	}

	m_machine = filehdr32.e_machine;
	m_elfclass = filehdr32.e_ident[EI_CLASS];
	if (m_elfclass == ELFCLASS32) {
		Elf32to64_Ehdr(&filehdr32, filehdr);
	} else if (m_elfclass == ELFCLASS64) {
		rewind(fp);
		ret = fread(filehdr, sizeof(*filehdr), 1, fp);

		if (ret != 1) {
			return false;
		}
	} else {
		return false;
	}
	return true;
}

static void Elf32to64_Shdr(Elf32_Shdr const *src, Elf64_Shdr *dest)
{
	dest->sh_name = src->sh_name;
	dest->sh_type = src->sh_type;
	dest->sh_flags = src->sh_flags;
	dest->sh_addr = src->sh_addr;
	dest->sh_offset = src->sh_offset;
	dest->sh_size = src->sh_size;
	dest->sh_link = src->sh_link;
	dest->sh_info = src->sh_info;
	dest->sh_addralign = src->sh_addralign;
	dest->sh_entsize = src->sh_entsize;
}

bool ElfReader::read_ELF_section_header(FILE *fp, Elf64_Ehdr const *filehdr, int section, Elf64_Shdr *sect_hdr)
{
	int numsect = filehdr->e_shnum;
	off_t sect_hdr_off;
	Elf32_Shdr sect_hdr32;

	if (numsect < section || section < 0) {
		return false;
	}
	sect_hdr_off = filehdr->e_shoff;
	sect_hdr_off += filehdr->e_shentsize * section;
	fseek(fp, (off_t) sect_hdr_off, SEEK_SET);
	if (m_elfclass == ELFCLASS32) {
		int ret = fread(&sect_hdr32, sizeof(sect_hdr32), 1, fp);
		if (ret != 1) {
			return false;
		}

		Elf32to64_Shdr(&sect_hdr32, sect_hdr);
	} else if (m_elfclass == ELFCLASS64) {
		int ret = fread(sect_hdr, sizeof(*sect_hdr), 1, fp);
		if (ret != 1) {
			return false;
		}
	} else {
		return false;
	}

	return true;
}

static void Elf32to64_Phdr(Elf32_Phdr const *src, Elf64_Phdr *dest) {
	dest->p_type   = src->p_type;
	dest->p_offset = src->p_offset;
	dest->p_vaddr  = src->p_vaddr;
	dest->p_paddr  = src->p_paddr;
	dest->p_filesz = src->p_filesz;
	dest->p_memsz  = src->p_memsz;
	dest->p_flags  = src->p_flags;
	dest->p_align  = src->p_align;
}

bool ElfReader::read_ELF_segment_header(FILE* fp, Elf64_Ehdr const * filehdr, unsigned segment, Elf64_Phdr * seg_hdr) {
	if(filehdr->e_phnum < segment) return false;
	off_t phdr_offset = filehdr->e_phoff;
	phdr_offset += filehdr->e_phentsize*segment;
	fseek(fp, phdr_offset, SEEK_SET);
	if(m_elfclass == ELFCLASS32) {
		Elf32_Phdr seg_hdr32;
		if(!fread(&seg_hdr32, sizeof(seg_hdr32), 1, fp)){
			return false;
		}
		Elf32to64_Phdr(&seg_hdr32,seg_hdr);
	} else {
		if(!fread(seg_hdr, sizeof(*seg_hdr), 1, fp)) {
			return false;
		}
	}

	return true;
}

const ElfSymbol& ElfReader::getSymbol(guest_address_t address) {
	for (container_t::const_iterator it = m_symboltable.begin(); it !=m_symboltable.end(); ++it) {
		if (it->contains(address)) {
			return *it;
		}
	}

	return g_SymbolNotFound;
}

// Symbol search
const ElfSymbol& ElfReader::getSymbol( const std::string& name ) {
	container_t::const_iterator it;
	// Fist, try to find as mangled symbol
	it = std::find(m_symboltable.begin(), m_symboltable.end(), name);
	if (it != m_symboltable.end()) {
		return *it;
	}

	// Then, try to find as demangled symbol
	std::string dname = Demangler::demangle(name);
	if (dname == Demangler::DEMANGLE_FAILED) {
		return g_SymbolNotFound;
	}

	it = std::find(m_symboltable.begin(), m_symboltable.end(), dname);
	if (it != m_symboltable.end()) {
		return *it;
	}

	return g_SymbolNotFound;
}

// Section search
const ElfSymbol& ElfReader::getSection(guest_address_t address) {
	for (container_t::const_iterator it = m_sectiontable.begin(); it != m_sectiontable.end(); ++it) {
		if (it->contains(address)) {
			return *it;
		}
	}
	return g_SymbolNotFound;
}

const ElfSymbol& ElfReader::getSection( const std::string& name ) {
	for (container_t::const_iterator it = m_sectiontable.begin(); it !=m_sectiontable.end(); ++it) {
		if (it->getName() == name) {
			return *it;
		}
	}
	return g_SymbolNotFound;
}

// "Pretty" Print
void ElfReader::printDemangled() {
	m_log << "Demangled: " << std::endl;
	for (container_t::const_iterator it = m_symboltable.begin(); it !=m_symboltable.end(); ++it) {
		std::string str = Demangler::demangle(it->getName());
		if (str == Demangler::DEMANGLE_FAILED) {
			str = it->getName();
		}
		m_log << "0x"  << std::hex << it->getAddress() << "\t" << str.c_str() << "\t" << it->getSize() << std::endl;
	}
}

void ElfReader::printMangled() {
	for (container_t::const_iterator it = m_symboltable.begin(); it !=m_symboltable.end(); ++it) {
		m_log  << "0x" << std::hex << it->getAddress() << "\t" << it->getName().c_str() << "\t" << it->getSize() << std::endl;
	}
}

void ElfReader::printSections() {
	for (container_t::const_iterator it = m_sectiontable.begin(); it !=m_sectiontable.end(); ++it) {
		m_log  << "0x" << it->getAddress() << "\t" << it->getName().c_str() << "\t" << it->getSize() << std::endl;
	}
}

void ElfReader::printSegments() {
	for(auto it = m_segmenttable.begin(); it != m_segmenttable.end(); ++it) {
		m_log << *it << std::endl;
	}
}

std::pair<guest_address_t,guest_address_t>
ElfReader::getValidAddressBounds() {
	guest_address_t min = std::numeric_limits<guest_address_t>::max();
	guest_address_t max = std::numeric_limits<guest_address_t>::min();

	for(const auto& seg: m_segmenttable) {
		min = std::min(min, seg.getStart());
		max = std::max(max, seg.getEnd());
	}
	if (min == std::numeric_limits<guest_address_t>::max())
		throw new std::invalid_argument("invalid minimum ELF address, either the segment header are invalid or the ELF will take no space in memory.");
	if (max == std::numeric_limits<guest_address_t>::min())
		throw new std::invalid_argument("invalid maximum ELF address, either the segment header are invalid or the ELF will take no space in memory.");

	return std::make_pair(min, max);
}

std::pair<guest_address_t,guest_address_t>
ElfReader::getTextSegmentBounds() {
	address_t minimal_ip;
	address_t maximal_ip;
	bool found_executable_segment = false;

	for (const auto &seg  : m_segmenttable) {
		if(seg.isExecutable()) {
			if(found_executable_segment)  {
				m_log << "warning: found multiple executable segments in ELF, we will only detect writes to the first executable segment!" << std::endl;
				continue;
			}
			found_executable_segment = true;
			minimal_ip = seg.getStart();
			maximal_ip = seg.getEnd();
		}
	}

	if (!found_executable_segment)
		throw new std::invalid_argument("ELF file has no executable segment. aborting!");

	return std::make_pair(minimal_ip, maximal_ip);
}



} // end-of-namespace fail