// goal:
// replace "file" + "checksec"

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

#define MAX_FILENAME   256
#define EI_NIDENT      16
#define MAX_HEADER_LEN 2048

#define Elf32_Half uint16_t
#define Elf32_Word uint32_t
#define Elf32_Addr uint32_t
#define Elf32_Off  uint32_t

#define Elf64_Half uint32_t
#define Elf64_Word uint64_t
#define Elf64_Addr uint64_t
#define Elf64_Off  uint64_t

struct Elf32_Ehdr
{
	unsigned char e_ident[EI_NIDENT];
	Elf32_Half e_type;
	Elf32_Half e_machine;
	Elf32_Word e_version;
	Elf32_Addr e_entry; 
	Elf32_Off e_phoff;
	Elf32_Off e_shoff;
	Elf32_Word e_flags;
	Elf32_Half e_ehsize;
	Elf32_Half e_phentsize;
	Elf32_Half e_phnum;
	Elf32_Half e_shentsize;
	Elf32_Half e_shnum;
	Elf32_Half e_shstrndx;
};

struct Elf64_Ehdr
{
	unsigned char e_ident[EI_NIDENT];
	Elf64_Half e_type;
	Elf64_Half e_machine;
	Elf64_Word e_version;
	Elf64_Addr e_entry; 
	Elf64_Off e_phoff;
	Elf64_Off e_shoff;
	Elf64_Word e_flags;
	Elf64_Half e_ehsize;
	Elf64_Half e_phentsize;
	Elf64_Half e_phnum;
	Elf64_Half e_shentsize;
	Elf64_Half e_shnum;
	Elf64_Half e_shstrndx;
};

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

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

int read_elf_magic(FILE* fp)
{
	char buf[5] = {0};
	fread(buf, 1, 4, fp);
	if (memcmp(buf, "\177ELF", 4) == 0) return 1;
	return 0;
}

void display_elf32(struct Elf32_Ehdr* header)
{
	printf("32-bit ELF\n");
}

void display_elf64(struct Elf64_Ehdr* header)
{
	printf("64-bit ELF\n");
}

unsigned char read_elf_obj_size(FILE* fp)
{
	unsigned char local_elf_obj_size = 0;
	fseek(fp, 4, SEEK_SET);
	fread(&local_elf_obj_size, 1, 1, fp);
	rewind(fp);
	return local_elf_obj_size;
}

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

	char* filename = argv[1];
	if (!filename)
	{
		printf("Error parsing filename.\n");
		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)
	{
		printf("Not a valid ELF file.\n");
		return -EINVAL;
	}

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

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

	return EXIT_SUCCESS;
}