/*
 * Copyright 2007, Ingo Weinhold <bonefish@cs.tu-berlin.de>.
 * All rights reserved. Distributed under the terms of the MIT License.
 */

#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <algorithm>
#include <string>

#include <system_revision.h>

// We use htonl(), which is defined in <ByteOrder.h> on BeOS R5.
#ifdef HAIKU_HOST_PLATFORM_BEOS
	#include <ByteOrder.h>
#else
	#include <arpa/inet.h>
#endif


using std::string;
using std::max;
using std::min;


// #pragma mark - ELF definitions


// types
typedef uint32_t	Elf32_Addr;
typedef uint16_t	Elf32_Half;
typedef uint32_t	Elf32_Off;
typedef int32_t		Elf32_Sword;
typedef uint32_t	Elf32_Word;
typedef uint64_t	Elf64_Addr;
typedef uint64_t	Elf64_Off;
typedef uint16_t	Elf64_Half;
typedef uint32_t	Elf64_Word;
typedef int32_t		Elf64_Sword;
typedef uint64_t	Elf64_Xword;
typedef int64_t		Elf64_Sxword;

// e_ident indices
#define EI_MAG0		0
#define EI_MAG1		1
#define EI_MAG2		2
#define EI_MAG3		3
#define EI_CLASS	4
#define EI_DATA		5
#define EI_VERSION	6
#define EI_PAD		7
#define EI_NIDENT	16

// object file header
typedef struct {
	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;
} Elf32_Ehdr;

typedef struct {
	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;
} Elf64_Ehdr;

// e_ident EI_CLASS and EI_DATA values
#define ELFCLASSNONE	0
#define ELFCLASS32		1
#define ELFCLASS64		2
#define ELFDATANONE		0
#define ELFDATA2LSB		1
#define ELFDATA2MSB		2

// program header
typedef struct {
	Elf32_Word	p_type;
	Elf32_Off	p_offset;
	Elf32_Addr	p_vaddr;
	Elf32_Addr	p_paddr;
	Elf32_Word	p_filesz;
	Elf32_Word	p_memsz;
	Elf32_Word	p_flags;
	Elf32_Word	p_align;
} Elf32_Phdr;

typedef struct {
	Elf64_Word	p_type;
	Elf64_Word	p_flags;
	Elf64_Off	p_offset;
	Elf64_Addr	p_vaddr;
	Elf64_Addr	p_paddr;
	Elf64_Xword	p_filesz;
	Elf64_Xword	p_memsz;
	Elf64_Xword	p_align;
} Elf64_Phdr;

// p_type
#define PT_NULL		0
#define PT_LOAD		1
#define PT_DYNAMIC	2
#define PT_INTERP	3
#define PT_NOTE		4
#define PT_SHLIB	5
#define PT_PHDIR	6
#define PT_LOPROC	0x70000000
#define PT_HIPROC	0x7fffffff

// section header
typedef struct {
	Elf32_Word	sh_name;
	Elf32_Word	sh_type;
	Elf32_Word	sh_flags;
	Elf32_Addr	sh_addr;
	Elf32_Off	sh_offset;
	Elf32_Word	sh_size;
	Elf32_Word	sh_link;
	Elf32_Word	sh_info;
	Elf32_Word	sh_addralign;
	Elf32_Word	sh_entsize;
} Elf32_Shdr;

typedef struct {
	Elf64_Word	sh_name;
	Elf64_Word	sh_type;
	Elf64_Xword	sh_flags;
	Elf64_Addr	sh_addr;
	Elf64_Off	sh_offset;
	Elf64_Xword	sh_size;
	Elf64_Word	sh_link;
	Elf64_Word	sh_info;
	Elf64_Xword	sh_addralign;
	Elf64_Xword	sh_entsize;
} Elf64_Shdr;

// sh_type values
#define SHT_NULL		0
#define SHT_PROGBITS	1
#define SHT_SYMTAB		2
#define SHT_STRTAB		3
#define SHT_RELA		4
#define SHT_HASH		5
#define SHT_DYNAMIC		6
#define SHT_NOTE		7
#define SHT_NOBITS		8
#define SHT_REL			9
#define SHT_SHLIB		10
#define SHT_DYNSYM		11
#define SHT_LOPROC		0x70000000
#define SHT_HIPROC		0x7fffffff
#define SHT_LOUSER		0x80000000
#define SHT_HIUSER		0xffffffff

// special section indexes
#define SHN_UNDEF		0

static const char		kELFFileMagic[4]	= { 0x7f, 'E', 'L', 'F' };


// #pragma mark - Usage

// usage
static const char *kUsage =
"Usage: %s  <file> <revision>\n"
"\n"
"Finds the haiku revision section in ELF object file <file> and replaces the\n"
"writes the number given by <revision> into the first 32 bits of the\n"
"section.\n"
;

// command line args
static int sArgc;
static const char *const *sArgv;

// print_usage
void
print_usage(bool error)
{
	// get nice program name
	const char *programName = (sArgc > 0 ? sArgv[0] : "resattr");
	if (const char *lastSlash = strrchr(programName, '/'))
		programName = lastSlash + 1;

	// print usage
	fprintf((error ? stderr : stdout), kUsage, programName);
}

// print_usage_and_exit
static void
print_usage_and_exit(bool error)
{
	print_usage(error);
	exit(error ? 1 : 0);
}


// #pragma mark - Exception


class Exception {
public:
	// constructor
	Exception()
		: fError(errno),
		  fDescription()
	{
	}

	// constructor
	Exception(const char* format,...)
		: fError(errno),
		  fDescription()
	{
		va_list args;
		va_start(args, format);
		SetTo(errno, format, args);
		va_end(args);
	}

	// constructor
	Exception(int error)
		: fError(error),
		  fDescription()
	{
	}

	// constructor
	Exception(int error, const char* format,...)
		: fError(error),
		  fDescription()
	{
		va_list args;
		va_start(args, format);
		SetTo(error, format, args);
		va_end(args);
	}

	// copy constructor
	Exception(const Exception& exception)
		: fError(exception.fError),
		  fDescription(exception.fDescription)
	{
	}

	// destructor
	~Exception()
	{
	}

	// SetTo
	void SetTo(int error, const char* format, va_list arg)
	{
		char buffer[2048];
		vsprintf(buffer, format, arg);
		fError = error;
		fDescription = buffer;
	}

	// Error
	int Error() const
	{
		return fError;
	}

	// Description
	const string& Description() const
	{
		return fDescription;
	}

private:
	int		fError;
	string	fDescription;
};


// #pragma mark - ELFObject


struct SectionInfo {
	uint32_t	type;
	off_t		offset;
	size_t		size;
	const char*	name;
};


class ELFObject {
public:
	ELFObject()
		: fFD(-1),
		  fSectionHeaderStrings(NULL),
		  fSectionHeaderStringsLength(0)
	{
	}

	~ELFObject()
	{
		Unset();
	}

	void Unset()
	{
		if (fFD >= 0) {
			close(fFD);
			fFD = -1;
		}

		delete[] fSectionHeaderStrings;
		fSectionHeaderStrings = NULL;
	}

	void SetTo(const char* fileName)
	{
		Unset();

		// open the file
		fFD = open(fileName, O_RDWR);
		if (fFD < 0)
			throw Exception("Failed to open \"%s\"", fileName);

		// get the file size
		fFileSize = FileSize();
		if (fFileSize < 0)
			throw Exception("Failed to get the file size.");

		// Read identification information
		unsigned char ident[EI_NIDENT];
		Read(0, ident, sizeof(ident), "Failed to read ELF identification.");
		if (memcmp(ident, kELFFileMagic, sizeof(kELFFileMagic)) != 0)
			throw Exception("Not a valid ELF file.");
		fELFClass = ident[EI_CLASS];

		if (fELFClass == ELFCLASS64)
			_ParseELFHeader<Elf64_Ehdr, Elf64_Shdr>();
		else
			_ParseELFHeader<Elf32_Ehdr, Elf32_Shdr>();
	}

	bool FindSectionByName(const char* name, SectionInfo& foundInfo)
	{
		// can't find the section by name without section names
		if (!fSectionHeaderStrings)
			return false;

		// iterate through the section headers
		for (size_t i = 0; i < fSectionHeaderCount; i++) {
			SectionInfo info;

			bool result;
			if (fELFClass == ELFCLASS64)
				result = _ReadSectionHeader<Elf64_Shdr>(i, info);
			else
				result = _ReadSectionHeader<Elf32_Shdr>(i, info);

			if (result) {
//printf("section %3d: offset: %7d, size: %7d, name: %s\n", i, info.offset, info.size, info.name);
				if (strcmp(info.name, name) == 0) {
					foundInfo = info;
					return true;
				}
			}
		}

		return false;
	}

	void Read(off_t position, void* buffer, size_t size,
		const char *errorMessage = NULL)
	{
		if (lseek(fFD, position, SEEK_SET) < 0)
			throw Exception(errorMessage);

		ssize_t bytesRead = read(fFD, buffer, size);
		if (bytesRead < 0)
			throw Exception(errorMessage);

		if ((size_t)bytesRead != size) {
			if (errorMessage) {
				throw Exception("%s Read too few bytes (%zd/%zu).",
					errorMessage, bytesRead, size);
			} else {
				throw Exception("Read too few bytes (%zd/%zu).",
					bytesRead, size);
			}
		}
	}

	void Write(off_t position, const void* buffer, size_t size,
		const char *errorMessage = NULL)
	{
		if (lseek(fFD, position, SEEK_SET) < 0)
			throw Exception(errorMessage);

		ssize_t bytesWritten = write(fFD, buffer, size);
		if (bytesWritten < 0)
			throw Exception(errorMessage);

		if ((size_t)bytesWritten != size) {
			if (errorMessage) {
				throw Exception("%s Wrote too few bytes (%zd/%zu).",
					errorMessage, bytesWritten, size);
			} else {
				throw Exception("Wrote too few bytes (%zd/%zu).",
					bytesWritten, size);
			}
		}
	}

	off_t FileSize()
	{
		off_t currentPos = lseek(fFD, 0, SEEK_END);
		if (currentPos < 0)
			return -1;

		return lseek(fFD, currentPos, SEEK_SET);
	}

	template<typename Type>
	Type GetValue(Type& value);

private:
	template<typename EhdrType, typename ShdrType>
	void _ParseELFHeader();

	template<typename ShdrType>
	bool _ReadSectionHeader(int index, SectionInfo& info);

	// _SwapUInt16
	static inline uint16_t _SwapUInt16(uint16_t value)
	{
		return ((value & 0xff) << 8) | (value >> 8);
	}

	// _SwapUInt32
	static inline uint32_t _SwapUInt32(uint32_t value)
	{
		return ((uint32_t)_SwapUInt16(value & 0xffff) << 16)
			| _SwapUInt16(uint16_t(value >> 16));
	}

	// _SwapUInt64
	static inline uint64_t _SwapUInt64(uint64_t value)
	{
		return ((uint64_t)_SwapUInt32(value & 0xffffffff) << 32)
			| _SwapUInt32(uint32_t(value >> 32));
	}

private:
	int			fFD;
	uint8_t		fELFClass;
	bool		fHostEndianess;
	off_t		fFileSize;
	size_t		fELFHeaderSize;
	off_t		fSectionHeaderTableOffset;
	size_t		fSectionHeaderSize;
	size_t		fSectionHeaderCount;
	char*		fSectionHeaderStrings;
	uint32_t	fSectionHeaderStringsLength;
};


template<>
int16_t ELFObject::GetValue(int16_t& value)
{
	return (fHostEndianess ? value : _SwapUInt16(value));
}


template<>
uint16_t ELFObject::GetValue(uint16_t& value)
{
	return (fHostEndianess ? value : _SwapUInt16(value));
}


template<>
int32_t ELFObject::GetValue(int32_t& value)
{
	return (fHostEndianess ? value : _SwapUInt32(value));
}


template<>
uint32_t ELFObject::GetValue(uint32_t& value)
{
	return (fHostEndianess ? value : _SwapUInt32(value));
}


template<>
int64_t ELFObject::GetValue(int64_t& value)
{
	return (fHostEndianess ? value : _SwapUInt64(value));
}


template<>
uint64_t ELFObject::GetValue(uint64_t& value)
{
	return (fHostEndianess ? value : _SwapUInt64(value));
}


template<typename EhdrType, typename ShdrType>
void ELFObject::_ParseELFHeader()
{
	// read ELF header
	EhdrType fileHeader;
	Read(0, &fileHeader, sizeof(EhdrType), "Failed to read ELF header.");

	// check data encoding (endianess)
	switch (fileHeader.e_ident[EI_DATA]) {
		case ELFDATA2LSB:
			fHostEndianess = (htonl(1) != 1);
			break;
		case ELFDATA2MSB:
			fHostEndianess = (htonl(1) == 1);
			break;
		default:
		case ELFDATANONE:
			throw Exception(EIO, "Unsupported ELF data encoding.");
			break;
	}

	// get the header values
	fELFHeaderSize	= GetValue(fileHeader.e_ehsize);
	fSectionHeaderTableOffset = GetValue(fileHeader.e_shoff);
	fSectionHeaderSize	= GetValue(fileHeader.e_shentsize);
	fSectionHeaderCount = GetValue(fileHeader.e_shnum);
	bool hasSectionHeaderTable = (fSectionHeaderTableOffset != 0);

	// check the sanity of the header values
	// ELF header size
	if (fELFHeaderSize < sizeof(EhdrType)) {
		throw Exception(EIO,
			"Invalid ELF header: invalid ELF header size: %lu.",
			fELFHeaderSize);
	}

	// section header table offset and entry count/size
	if (hasSectionHeaderTable) {
		if (fSectionHeaderTableOffset < (off_t)fELFHeaderSize
			|| fSectionHeaderTableOffset > fFileSize) {
			throw Exception(EIO, "Invalid ELF header: invalid section "
							"header table offset: %llu.",
							fSectionHeaderTableOffset);
		}
		size_t sectionHeaderTableSize
			= fSectionHeaderSize * fSectionHeaderCount;
		if (fSectionHeaderSize < (off_t)sizeof(ShdrType)
			|| fSectionHeaderTableOffset + (off_t)sectionHeaderTableSize
				> fFileSize) {
			throw Exception(EIO, "Invalid ELF header: section header "
							"table exceeds file: %llu.",
							fSectionHeaderTableOffset
								+ sectionHeaderTableSize);
		}


		// load section header string section
		uint16_t sectionHeaderStringSectionIndex
			= GetValue(fileHeader.e_shstrndx);
		if (sectionHeaderStringSectionIndex != SHN_UNDEF) {
			if (sectionHeaderStringSectionIndex >= fSectionHeaderCount) {
				throw Exception(EIO, "Invalid ELF header: invalid section "
								"header string section index: %u.",
								sectionHeaderStringSectionIndex);
			}

			// get the section info
			SectionInfo info;
			if (_ReadSectionHeader<ShdrType>(sectionHeaderStringSectionIndex,
					info)) {
				fSectionHeaderStrings = new char[info.size + 1];
				Read(info.offset, fSectionHeaderStrings, info.size,
					"Failed to read section header string section.");
				fSectionHeaderStringsLength = info.size;
				// null-terminate to be on the safe side
				fSectionHeaderStrings[info.size] = '\0';
			}
		}
	}
}


template<typename ShdrType>
bool ELFObject::_ReadSectionHeader(int index, SectionInfo& info)
{
	off_t shOffset = fSectionHeaderTableOffset
		+ index * fSectionHeaderSize;
	ShdrType sectionHeader;
	Read(shOffset, &sectionHeader, sizeof(ShdrType),
		"Failed to read ELF section header.");

	// get the header values
	uint32_t type		= GetValue(sectionHeader.sh_type);
	off_t offset		= GetValue(sectionHeader.sh_offset);
	size_t size			= GetValue(sectionHeader.sh_size);
	uint32_t nameIndex	= GetValue(sectionHeader.sh_name);

	// check the values
	// SHT_NULL marks the header unused,
	if (type == SHT_NULL)
		return false;

	// SHT_NOBITS sections take no space in the file
	if (type != SHT_NOBITS) {
		if (offset < (off_t)fELFHeaderSize || offset > fFileSize) {
			throw Exception(EIO, "Invalid ELF section header: "
							"invalid section offset: %llu.", offset);
		}
		off_t sectionEnd = offset + size;
		if (sectionEnd > fFileSize) {
			throw Exception(EIO, "Invalid ELF section header: "
							"section exceeds file: %llu.", sectionEnd);
		}
	}

	// get name, if we have a string section
	if (fSectionHeaderStrings) {
		if (nameIndex >= (uint32_t)fSectionHeaderStringsLength) {
			throw Exception(EIO, "Invalid ELF section header: "
							"invalid name index: %" PRIu32 ".", nameIndex);
		}
		info.name = fSectionHeaderStrings + nameIndex;
	} else {
		info.name = "";
	}

	info.type = type;
	info.offset = offset;
	info.size = size;


	return true;
}


// main
int
main(int argc, const char* const* argv)
{
	sArgc = argc;
	sArgv = argv;

	if (argc < 3)
		print_usage_and_exit(true);

	// parameters
	const char* fileName = argv[1];
	const char* revisionString = argv[2];

	try {
		ELFObject elfObject;
		elfObject.SetTo(fileName);

		// find haiku revision section
		SectionInfo info;
		if (!elfObject.FindSectionByName("_haiku_revision", info)) {
			fprintf(stderr, "haiku revision section not found\n");
			exit(1);
		}

		// write revision string to section
		elfObject.Write(info.offset, revisionString,
			min((size_t)SYSTEM_REVISION_LENGTH, strlen(revisionString) + 1),
			"Failed to write revision.");

	} catch (Exception& exception) {
		if (exception.Description() == "") {
			fprintf(stderr, "%s\n", strerror(exception.Error()));
		} else {
			fprintf(stderr, "%s: %s\n", exception.Description().c_str(),
				strerror(exception.Error()));
		}
		exit(1);
	}

	return 0;
}