// This is free and unencumbered software released into the public domain.
// 
// Anyone is free to copy, modify, publish, use, compile, sell, or
// distribute this software, either in source code form or as a compiled
// binary, for any purpose, commercial or non-commercial, and by any
// means.
// 
// In jurisdictions that recognize copyright laws, the author or authors
// of this software dedicate any and all copyright interest in the
// software to the public domain. We make this dedication for the benefit
// of the public at large and to the detriment of our heirs and
// successors. We intend this dedication to be an overt act of
// relinquishment in perpetuity of all present and future rights to this
// software under copyright law.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
// IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
// 
// For more information, please refer to <https://unlicense.org/>
// 
// @author xamidev <xamidev@riseup.net>
// @brief Recon tool for ELF32/64

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <stdint.h>
#include <string.h>
#include <stdbool.h>

// All data structures & constants we need are defined here
#include <elf.h>

bool verbose = false;
bool pie = false;
bool nx = false;
bool canary = false;
bool relro = false;
bool bind_now = false;
bool stripped = true;
bool interp = false;

int parse_elf_header32(Elf32_Ehdr* header, FILE* fp)
{
	if (!fread(header, sizeof(Elf32_Ehdr), 1, fp))
	{
		return -EINVAL;
	}
	return 0;
}

int parse_elf_header64(Elf64_Ehdr* header, FILE* fp)
{
	if (!fread(header, sizeof(Elf64_Ehdr), 1, fp))
	{
		return -EINVAL;
	}
	return 0;
}

int read_elf_magic(FILE* fp)
{
	char buf[5] = {0};
	if (!fread(buf, 1, 4, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}
	if (memcmp(buf, "\177ELF", 4) == 0) return 1;
	return 0;
}

void display_elf_common(unsigned char e_ident[EI_NIDENT])
{
	// EI_DATA
	switch(e_ident[5])
	{
		case 0x01: printf("(little-endian) "); break;
		case 0x02: printf("(big-endian) "); break;
		default: printf("\nInvalid data encoding!\n"); exit(-EINVAL);
	}

	if (verbose)
	{
		// EI_VERSION
		printf("version %d ", e_ident[6]);
	}
}

void check_dynamic32(Elf32_Ehdr* header, FILE* fp)
{
	// Is dynamically linked? (has PT_INTERP program header?)
	// read program headers	
	fseek(fp, header->e_phoff, SEEK_SET);
	Elf32_Phdr *p_headers = malloc(header->e_phnum * sizeof(Elf32_Phdr));
	if (!p_headers)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	if (!fread(p_headers, sizeof(Elf32_Phdr), header->e_phnum, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	const char* interpstr = NULL;

	for (int i=0; i<header->e_phnum; i++)
	{
		if (p_headers[i].p_type == PT_INTERP)
		{
			interp = true;

			// Get interpreter path
			interpstr = malloc(p_headers[i].p_filesz);
			if (!interpstr)
			{
				puts("Memory allocation error.");
				exit(-ENOMEM);
			}
			fseek(fp, p_headers[i].p_offset, SEEK_SET);
			if (!fread((char*)interpstr, 1, p_headers[i].p_filesz, fp))
			{
				puts("Read error.");
				exit(-EIO);
			}
			break;
		}
	}

	if (interp)
	{
		printf("dynamically linked (interpreter: %s), ", interpstr);
	}
	else
	{
		printf("statically linked, ");
	}

	free(p_headers);
	free((char*)interpstr);
}

void check_dynamic64(Elf64_Ehdr* header, FILE* fp)
{
	// Is dynamically linked? (has PT_INTERP program header?)
	// read program headers	
	fseek(fp, header->e_phoff, SEEK_SET);
	Elf64_Phdr *p_headers = malloc(header->e_phnum * sizeof(Elf64_Phdr));
	if (!p_headers)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	if (!fread(p_headers, sizeof(Elf64_Phdr), header->e_phnum, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	const char* interpstr = NULL;

	for (int i=0; i<header->e_phnum; i++)
	{
		if (p_headers[i].p_type == PT_INTERP)
		{
			interp = true;

			// Get interpreter path
			interpstr = malloc(p_headers[i].p_filesz);
			if (!interpstr)
			{
				puts("Memory allocation error.");
				exit(-ENOMEM);
			}
			fseek(fp, p_headers[i].p_offset, SEEK_SET);
			if (!fread((char*)interpstr, 1, p_headers[i].p_filesz, fp))
			{
				puts("Read error.");
				exit(-EIO);
			}
			break;
		}
	}

	if (interp)
	{
		printf("dynamically linked (interpreter: %s), ", interpstr);
	}
	else
	{
		printf("statically linked, ");
	}

	free(p_headers);
	free((char*)interpstr);
}

void display_elf32(Elf32_Ehdr* header, FILE* fp)
{
	printf("32-bit ELF ");
	display_elf_common(header->e_ident);

	if (verbose)
	{
		switch (header->e_type)
		{
			case 0x00: printf("(no file type),"); break;
			case 0x01: printf("(relocatable file),"); break;
			case 0x02: printf("(executable file),"); break;
			// (ET_DYN, prob means PIE)
			case 0x03: printf("(shared object file),"); pie = true; break;
			case 0x04: printf("(core file),"); break;
			default: printf("(unknown),"); break;
		}
		printf(" ");
	} else if (header->e_type == 0x03)
	{
		pie = true;
	}

	switch (header->e_machine)
	{
		case 0x00: printf("No architecture, "); break;
		case 0x01: printf("AT&T WE 32100, "); break;
		case 0x02: printf("SPARC, "); break;
		case 0x03: printf("i386, "); break;
		case 0x04: printf("Motorola 68000, "); break;
		case 0x05: printf("Motorola 88000, "); break;
		case 0x07: printf("Intel 80860, "); break;
		case 0x08: printf("MIPS RS3000, "); break;
		case 0x09: printf("MIPS RS4000, "); break;
		default: printf("Unknown architecture, "); break;
	}
	
	if (verbose)
	{
		check_dynamic32(header, fp);
	}
}

// Duplicate code.. gotta find a solution to this
// (properly differentiate 32/64b)
void display_elf64(Elf64_Ehdr* header, FILE* fp)
{
	printf("64-bit ELF ");
	display_elf_common(header->e_ident);

	if (verbose)
	{
		switch (header->e_type)
		{
			case 0x00: printf("(no file type),"); break;
			case 0x01: printf("(relocatable file),"); break;
			case 0x02: printf("(executable file),"); break;
			// (ET_DYN, prob means PIE)
			case 0x03: printf("(shared object file),"); pie = true; break;
			case 0x04: printf("(core file),"); break;
			default: printf("(unknown),"); break;
		}
		printf(" ");
	} else if (header->e_type == 0x03)
	{
		pie = true;
	}
	
	switch (header->e_machine)
	{
		// A few only.. full list here (many exotic stuff, might implement it one day)
		// https://gist.github.com/x0nu11byt3/bcb35c3de461e5fb66173071a2379779
		case 40: printf("ARM, "); break;
		case 21: printf("PowerPC, "); break;
		case 62: printf("amd64, "); break;
		default: printf("Unknown architecture, "); break;
	}
	
	if (verbose)
	{
		check_dynamic64(header, fp);
	}
}

void check_sec32(Elf32_Ehdr* header, FILE* fp)
{
	/* NX:
	 * Browse thru program header table (e_phnum items) from e_phoff (program header offset)
	 * if its stack and has executable bit then we know NX isnt there
	*/
	fseek(fp, header->e_phoff, SEEK_SET);
	for (size_t i=0; i<header->e_phnum; i++)
	{
		Elf32_Phdr p_header;
		if (!fread(&p_header, 1, sizeof(Elf32_Phdr), fp))
		{
			puts("Read error.");
			exit(-EIO);
		}
		if (p_header.p_type == PT_GNU_STACK)
		{
			if (p_header.p_flags & PF_X)
			{
				nx = false;
			} 
			else
			{
				nx = true;
			}
		}
	}

	/* Stack canary:
	 * Same for section header table (e_shnum items) from e_shoff (section header offset)
	 * find sections .dynsym or .symtab:
	 * and look for __stack_chk_fail (call for stack smashing)
	*/
	fseek(fp, header->e_shoff, SEEK_SET);

	// Read section headers
	Elf32_Shdr *s_headers = malloc(header->e_shnum * sizeof(Elf32_Shdr));
	if (!s_headers)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	if (!fread(s_headers, sizeof(Elf32_Shdr), header->e_shnum, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	// Find section name string table (at e_shstrndx) and read it
	Elf32_Shdr s_header_str = s_headers[header->e_shstrndx];
	char* s_header_strtab = malloc(s_header_str.sh_size);
	if (!s_header_strtab)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	fseek(fp, s_header_str.sh_offset, SEEK_SET);
	if (!fread(s_header_strtab, 1, s_header_str.sh_size, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	for (size_t i=0; i<header->e_shnum; i++)
	{
		// additional small check for stripping as we're already browsing section headers
		if (s_headers[i].sh_type == SHT_SYMTAB && s_headers[i].sh_size > 0) stripped = false;

		if (s_headers[i].sh_type == SHT_DYNSYM || s_headers[i].sh_type == SHT_SYMTAB)
		{
			// load symbol table
			unsigned int nsyms = s_headers[i].sh_size / s_headers[i].sh_entsize;
			Elf32_Sym *syms = malloc(s_headers[i].sh_size);
			if (!syms)
			{
				puts("Memory allocation error.");
				exit(-EIO);
			}
			fseek(fp, s_headers[i].sh_offset, SEEK_SET);
			if (!fread(syms, s_headers[i].sh_size, 1, fp))
			{
				puts("Read error.");
				exit(-EIO);
			}

			// load string table
			Elf32_Shdr str_header = s_headers[s_headers[i].sh_link];
			char* strtab = malloc(str_header.sh_size);
			if (!strtab)
			{
				puts("Memory allocation error.");
				exit(-ENOMEM);
			}
			fseek(fp, str_header.sh_offset, SEEK_SET);
			if (!fread(strtab, 1, str_header.sh_size, fp))
			{
				puts("Read error.");
				exit(-EIO);
			}

			// iterate through symbols, check for __stack_chk_fail
			for (size_t j=0; j<nsyms; j++)
			{
				const char* name = strtab + syms[j].st_name;
				if (strcmp(name, "__stack_chk_fail") == 0)
				{
					canary = true;	
				}
			}
			
			free(syms);
			free(strtab);
		}
	}

	/* RELRO:
	 * Segment type PT_GNU_RELRO should be present
	*/

	// Read program headers
	fseek(fp, header->e_phoff, SEEK_SET);
	Elf32_Phdr *p_headers = malloc(header->e_phnum * sizeof(Elf32_Phdr));
	if (!p_headers)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	if (!fread (p_headers, sizeof(Elf32_Phdr), header->e_phnum, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	Elf32_Off dyn_off = 0;
	Elf32_Word dyn_size = 0;

	for (int i=0; i<header->e_phnum; i++)
	{
		if (p_headers[i].p_type == PT_GNU_RELRO)
		{
			relro = true;
		}
		if (p_headers[i].p_type == PT_DYNAMIC)
		{
			dyn_off = p_headers[i].p_offset;
			dyn_size = p_headers[i].p_filesz;
		}
	}

	if (dyn_off && dyn_size)
	{
		int n = dyn_size / sizeof(Elf32_Dyn);
		Elf32_Dyn *dyns = malloc(dyn_size);
		if (!dyns)
		{
			puts("Memory allocation error.");
			exit(-ENOMEM);
		}
		fseek(fp, dyn_off, SEEK_SET);
		if (!fread(dyns, sizeof(Elf32_Dyn), n, fp))
		{
			puts("Read error.");
			exit(-EIO);
		}

		for (int i=0; i<n; i++)
		{
			// Has full RELRO?
			if (dyns[i].d_tag == DT_BIND_NOW)
			{
				bind_now = true;
				break;
			}
		}
		free(dyns);
	}

	free(s_headers);
	free(s_header_strtab);
	free(p_headers);
}

// Duplicate code. Should find a workaround
void check_sec64(Elf64_Ehdr* header, FILE* fp)
{
	/* NX:
	 * Browse thru program header table (e_phnum items) from e_phoff (program header offset)
	 * if its stack and has executable bit then we know NX isnt there
	*/
	fseek(fp, header->e_phoff, SEEK_SET);
	for (size_t i=0; i<header->e_phnum; i++)
	{
		Elf64_Phdr p_header;
		if (!fread(&p_header, 1, sizeof(Elf64_Phdr), fp))
		{
			puts("Read error.");
			exit(-EIO);
		}
		if (p_header.p_type == PT_GNU_STACK)
		{
			if (p_header.p_flags & PF_X)
			{
				nx = false;
			} 
			else
			{
				nx = true;
			}
		}
	}

	/* Stack canary:
	 * Same for section header table (e_shnum items) from e_shoff (section header offset)
	 * find sections .dynsym or .symtab:
	 * and look for __stack_chk_fail (call for stack smashing)
	*/
	fseek(fp, header->e_shoff, SEEK_SET);

	// Read section headers
	Elf64_Shdr *s_headers = malloc(header->e_shnum * sizeof(Elf64_Shdr));
	if (!s_headers)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	if (!fread(s_headers, sizeof(Elf64_Shdr), header->e_shnum, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	// Find section name string table (at e_shstrndx) and read it
	Elf64_Shdr s_header_str = s_headers[header->e_shstrndx];
	char* s_header_strtab = malloc(s_header_str.sh_size);
	if (!s_header_strtab)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	fseek(fp, s_header_str.sh_offset, SEEK_SET);
	if (!fread(s_header_strtab, 1, s_header_str.sh_size, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	for (size_t i=0; i<header->e_shnum; i++)
	{
		// additional small check for stripping as we're already browsing section headers
		if (s_headers[i].sh_type == SHT_SYMTAB && s_headers[i].sh_size > 0) stripped = false;

		if (s_headers[i].sh_type == SHT_DYNSYM || s_headers[i].sh_type == SHT_SYMTAB)
		{
			// load symbol table
			unsigned int nsyms = s_headers[i].sh_size / s_headers[i].sh_entsize;
			Elf64_Sym *syms = malloc(s_headers[i].sh_size);
			if (!syms)
			{
				puts("Memory allocation error.");
				exit(-EIO);
			}
			fseek(fp, s_headers[i].sh_offset, SEEK_SET);
			if (!fread(syms, s_headers[i].sh_size, 1, fp))
			{
				puts("Read error.");
				exit(-EIO);
			}

			// load string table
			Elf64_Shdr str_header = s_headers[s_headers[i].sh_link];
			char* strtab = malloc(str_header.sh_size);
			if (!strtab)
			{
				puts("Memory allocation error.");
				exit(-ENOMEM);
			}
			fseek(fp, str_header.sh_offset, SEEK_SET);
			if (!fread(strtab, 1, str_header.sh_size, fp))
			{
				puts("Read error.");
				exit(-EIO);
			}

			// iterate through symbols, check for __stack_chk_fail
			for (size_t j=0; j<nsyms; j++)
			{
				const char* name = strtab + syms[j].st_name;
				if (strcmp(name, "__stack_chk_fail") == 0)
				{
					canary = true;	
				}
			}
			
			free(syms);
			free(strtab);
		}
	}

	/* RELRO:
	 * Segment type PT_GNU_RELRO should be present
	*/

	// Read program headers
	fseek(fp, header->e_phoff, SEEK_SET);
	Elf64_Phdr *p_headers = malloc(header->e_phnum * sizeof(Elf64_Phdr));
	if (!p_headers)
	{
		puts("Memory allocation error.");
		exit(-ENOMEM);
	}
	if (!fread (p_headers, sizeof(Elf64_Phdr), header->e_phnum, fp))
	{
		puts("Read error.");
		exit(-EIO);
	}

	Elf64_Off dyn_off = 0;
	Elf64_Word dyn_size = 0;

	for (int i=0; i<header->e_phnum; i++)
	{
		if (p_headers[i].p_type == PT_GNU_RELRO)
		{
			relro = true;
		}
		if (p_headers[i].p_type == PT_DYNAMIC)
		{
			dyn_off = p_headers[i].p_offset;
			dyn_size = p_headers[i].p_filesz;
		}
	}

	if (dyn_off && dyn_size)
	{
		int n = dyn_size / sizeof(Elf64_Dyn);
		Elf64_Dyn *dyns = malloc(dyn_size);
		if (!dyns)
		{
			puts("Memory allocation error.");
			exit(-ENOMEM);
		}
		fseek(fp, dyn_off, SEEK_SET);
		if (!fread(dyns, sizeof(Elf64_Dyn), n, fp))
		{
			puts("Read error.");
			exit(-EIO);
		}

		for (int i=0; i<n; i++)
		{
			// Has full RELRO?
			if (dyns[i].d_tag == DT_BIND_NOW)
			{
				bind_now = true;
				break;
			}
		}
		free(dyns);
	}

	free(s_headers);
	free(s_header_strtab);
	free(p_headers);
}

unsigned char read_elf_obj_size(FILE* fp)
{
	unsigned char local_elf_obj_size = 0;
	fseek(fp, 4, SEEK_SET);
	if (!fread(&local_elf_obj_size, 1, 1, fp))
	{
		puts("Failed to read ELF object size");
		exit(-EIO);
	}
	rewind(fp);
	return local_elf_obj_size;
}

void display_sec_common()
{
	printf("%s, ", nx ? "\x1b[32mNX enabled\x1b[0m" : "\x1b[31mNX disabled\x1b[0m");
	printf("%s, ", pie ? "\x1b[32mPIE enabled\x1b[0m" : "\x1b[31mPIE disabled\x1b[0m");
	printf("%s, ", canary ? "\x1b[32mStack canary enabled\x1b[0m" : "\x1b[31mStack canary disabled\x1b[0m");
	printf("%s, ", relro && bind_now ? "\x1b[32mFull RELRO\x1b[0m" : relro ? "\x1b[33mPartial RELRO\x1b[0m" : "\x1b[31mNo RELRO\x1b[0m");
	printf("%s\n", stripped ? "\x1b[32mStripped\x1b[0m" : "\x1b[31mNot stripped\x1b[0m");
}

int main(int argc, char* argv[])
{
	// Argument processing	
	if (argc < 2)
	{
		printf("Usage: %s <file> [options]\n -v --verbose             displays more information\n", argv[0]);
		return -EINVAL;
	}

	for (int i=1; i<argc; i++)
	{
		if (strncmp(argv[i], "-v", 2) == 0) verbose = true;
	}

	char* filename = argv[1];
	if (!filename)
	{
		puts("Error parsing filename.");
		return -EINVAL;
	}
	
	FILE* elf_file = fopen(filename, "rb");
	if (!elf_file)
	{
		printf("Couldn't open file '%s'\n", filename);
		return -ENOENT;
	}

	int is_elf = read_elf_magic(elf_file);
	if (!is_elf)
	{
		puts("Not a valid ELF file.");
		return -EINVAL;
	}

	printf("%s: ", filename);

	// ELF32 or ELF64?
	unsigned char elf_obj_size = read_elf_obj_size(elf_file);

	switch(elf_obj_size)
	{
		case 0x01:
			Elf32_Ehdr elf32_header = {0};
			if (parse_elf_header32(&elf32_header, elf_file))
			{
				puts("Failed to parse 32-bit ELF header");
				return -EINVAL;
			}
			display_elf32(&elf32_header, elf_file);
			check_sec32(&elf32_header, elf_file);
			break;
		case 0x02:
			Elf64_Ehdr elf64_header = {0};
			if (parse_elf_header64(&elf64_header, elf_file))
			{
				puts("Failed to parse 64-bit ELF header");
				return -EINVAL;
			}
			display_elf64(&elf64_header, elf_file);
			check_sec64(&elf64_header, elf_file);
			break;
		default:
			puts("Invalid ELF object size");
			return -EINVAL;
	}

	display_sec_common();

	return EXIT_SUCCESS;
}