#include <sys/types.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <elf.h>
#include <errno.h>
#include <err.h>

/*
 * mcs { -c | -d | -p | -V | -a string | -n name } ... file ...
 *
 * -c compress. remove duplicates. Do not change order.
 * -d delete all.
 * -p print.
 * -a add comment.
 * -n use different section name.
 */

static int
sec_cmp(const void *aa, const void *bb)
{
	Elf_Shdr *a = *(Elf_Shdr * const *)aa;
	Elf_Shdr *b = *(Elf_Shdr * const *)bb;

	return a->sh_offset - b->sh_offset;
}

void
mcs(char *fn, char *section, int aflag, int cflag, int dflag, int pflag, char *string)
{
	int fd, newfd;
	char *newfn;
	int i, j;
	caddr_t ma;
	struct stat st;
	Elf_Ehdr *eh, *neh;
	Elf_Shdr *sh, *nsh, **order, *csect, *sectnames;
	char *secnames, *nsecnames;
	int comment_sect;
	char *c_start, *c_end, *c;
	int ncomments;
	char *newcomments, *comp;
	int newc;
	int newsize;
	int clen;
	int skipit;
	FILE *newfile;
	int delta;
	int align;
	Elf_Off oldoff;
	int once;
	int added_section;

	fd = open(fn, O_RDONLY);
	if (fd < 0)
		err(1, "open");
	if (fstat(fd, &st) < 0)
		err(1, "fstat");
	ma = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE|MAP_NOCORE, fd, 0);
	if (ma == MAP_FAILED)
		err(1, "mmap");
	eh = (Elf_Ehdr *)ma;
	if (!IS_ELF(*eh))
		err(1, "not an elf file");
	if (eh->e_version != EV_CURRENT)
		err(1, "incorrect version");
	if (!ELF_MACHINE_OK(eh->e_machine))
		err(1, "not native machine, sorry.");
	if (eh->e_shentsize != sizeof(Elf_Shdr))
		err(1, "incorrect shentsize");
	sh = (Elf_Shdr *)(ma + eh->e_shoff);
	secnames = ma + sh[eh->e_shstrndx].sh_offset;
	comment_sect = -1;
	added_section = -1;
	for (i = 0; i < eh->e_shnum; i++)
		if (strcmp(section, secnames + sh[i].sh_name) == 0)
			comment_sect = i;
	ncomments = 0;
	newsize = 0;
	newcomments = NULL;
	newc = 0;
	if (comment_sect != -1) {
		c_start = ma + sh[comment_sect].sh_offset;
		c_end = c_start + sh[comment_sect].sh_size;
		for (c = c_start; comment_sect != -1 && c < c_end; ) {
			if (*c == '\0') {
				c++;
				continue;
			}
			ncomments++;
			c += strlen(c) + 1;
		}
		for (c = c_start; c < c_end; ) {
			if (*c == '\0') {
				c++;
				continue;
			}
			skipit = 0;
			if (cflag) {
				for (comp = newcomments; comp < newcomments + newsize; ) {
					if (*comp == '\0') {
						comp++;
						continue;
					}
					if (strcmp(comp, c) == 0) {
						skipit = 1;
						break;
					}
					comp += strlen(comp) + 1;
				}
			}
			if (skipit) {
				c += strlen(c) + 1;
				continue;
			}
			clen = strlen(c) + 1;
			newsize += clen + 1;
			newcomments = newcomments ? realloc(newcomments, newsize) : malloc(newsize);
			newcomments[newc] = '\0';
			strcpy(newcomments + newc + 1, c);
			newc += clen + 1;
			c += clen;
		}
	}
	if (aflag) {
		clen = strlen(string) + 1;
		newsize += clen + 1;
		newcomments = newcomments ? realloc(newcomments, newsize) : malloc(newsize);
		newcomments[newc] = '\0';
		strcpy(newcomments + newc + 1, string);
	}
	if (dflag) {
		if (newcomments)
			free(newcomments);
		newcomments = NULL;
		newsize = 0;
	}
	if (pflag) {
		for (comp = newcomments; comp < newcomments + newsize; ) {
			if (*comp == '\0') {
				comp++;
				continue;
			}
			printf("%s\n", comp);
			comp += strlen(comp) + 1;
		}
	}
	neh = malloc(sizeof(Elf_Ehdr));
	memcpy(neh, eh, sizeof(Elf_Ehdr));
	added_section = -1;
	sectnames = NULL;
	delta = 0;
	if (comment_sect == -1) {
		nsh = malloc(sizeof(Elf_Shdr) * (neh->e_shnum + 1));
		memset(nsh, 0, sizeof(Elf_Shdr) * (neh->e_shnum + 1));
		for (i = 0, j = 0; i < eh->e_shnum; i++, j++) {
			if (j == eh->e_shstrndx)
				j++;
			nsh[j] = sh[i];
		}
		oldoff = sh[eh->e_shstrndx].sh_size;
		nsecnames = malloc(oldoff + strlen(section) + 1);
		memcpy(nsecnames, secnames, oldoff);
		strcpy(nsecnames + oldoff, section);
		comment_sect = eh->e_shstrndx;
		added_section = comment_sect;
		neh->e_shnum++;
		neh->e_shstrndx++;
		nsh[neh->e_shstrndx].sh_size += strlen(section) + 1;
		sectnames = &nsh[neh->e_shstrndx];
		nsh[comment_sect].sh_name = oldoff;
		nsh[comment_sect].sh_type = SHT_PROGBITS;
		nsh[comment_sect].sh_flags = 0;
		nsh[comment_sect].sh_addr = 0;
		nsh[comment_sect].sh_offset = nsh[neh->e_shstrndx].sh_offset;
		nsh[comment_sect].sh_size = newsize;
		nsh[comment_sect].sh_link = 0;
		nsh[comment_sect].sh_info = 0;
		nsh[comment_sect].sh_addralign = 1;
		nsh[comment_sect].sh_entsize = 0;
		delta = newsize;
	} else {
		nsh = malloc(sizeof(Elf_Shdr) * eh->e_shnum);
		memcpy(nsh, sh, sizeof(Elf_Shdr) * eh->e_shnum);
		nsecnames = malloc(sh[eh->e_shstrndx].sh_size);
		memcpy(nsecnames, secnames, sh[eh->e_shstrndx].sh_size);
		delta = newsize - nsh[comment_sect].sh_size;
		nsh[comment_sect].sh_size = newsize;
	}
	order = malloc(neh->e_shnum * sizeof(Elf_Shdr *));
	for (i = 0; i < neh->e_shnum; i++)
		order[i] = &nsh[i];
	mergesort(order, neh->e_shnum, sizeof(*order), sec_cmp);
	csect = &nsh[comment_sect];
	once = 1;
	for (i = 0; i < neh->e_shnum; i++) {
		if (order[i] == csect)
			continue;
		if (order[i]->sh_offset >= csect->sh_offset) {
			oldoff = order[i]->sh_offset;
			order[i]->sh_offset += delta;
			align = order[i]->sh_offset % order[i]->sh_addralign;
			order[i]->sh_offset += align;
			delta += align;
		}
		if (sectnames && order[i] == sectnames) {
			align = strlen(section) + 1;
			delta += align;
		}
		if (once && order[i]->sh_offset + order[i]->sh_size > neh->e_shoff) {
			once = 0;
			neh->e_shoff += delta;
			align = neh->e_shoff % sizeof(long);
			neh->e_shoff += align;
			delta += align;
			align = (neh->e_shnum - eh->e_shnum) * sizeof(Elf_Shdr);
			delta += align;
		}
	}
	if (added_section > -1)
		for (i = 0; i < neh->e_shnum; i++)
			if (nsh[i].sh_link != 0 && nsh[i].sh_link > added_section)
				nsh[i].sh_link++;
	if (!aflag && !cflag && !dflag) {
		if (nsh)
			free(nsh);
		if (neh)
			free(neh);
		if (nsecnames)
			free(nsecnames);
		if (newcomments)
			free(newcomments);
		return;
	}
	asprintf(&newfn, "%s.XXXXXX", fn);
	newfd = mkstemp(newfn);
	if (newfd < 0)
		err(1, "mkstemp");
	newfile = fdopen(newfd, "w+");
	if (newfile == NULL)
		err(1, "fopen");
	fseek(newfile, 0, SEEK_SET);
	fwrite(neh, sizeof(*neh), 1, newfile);
	if (neh->e_phoff > 0 && neh->e_phnum > 0) {
		fseek(newfile, neh->e_phoff, SEEK_SET);
		fwrite(ma + eh->e_phoff, eh->e_phentsize, eh->e_phnum, newfile);
	}
	fseek(newfile, neh->e_shoff, SEEK_SET);
	fwrite(nsh, sizeof(Elf_Shdr), neh->e_shnum, newfile);
	for (i = 0; i < neh->e_shnum; i++) {
		fseek(newfile, nsh[i].sh_offset, SEEK_SET);
		if (i == comment_sect) {
			fwrite(newcomments, newsize, 1, newfile);
		} else if (i == neh->e_shstrndx) {
			fwrite(nsecnames, nsh[neh->e_shstrndx].sh_size, 1, newfile);
		} else {
			if (added_section != -1 && i > added_section)
				fwrite(ma + sh[i - 1].sh_offset, sh[i - 1].sh_size, 1, newfile);
			else
				fwrite(ma + sh[i].sh_offset, sh[i].sh_size, 1, newfile);
		}
	}
	fclose(newfile);
	if (rename(newfn, fn) < 0) {
		warn("rename");
		unlink(newfn);
		exit(1);
	}
	if (nsh)
		free(nsh);
	if (neh)
		free(neh);
	if (nsecnames)
		free(nsecnames);
	if (newcomments)
		free(newcomments);
}

int
main(int ac, char **av)
{
	int c;
	char *string;
	char *section;
	int aflag;
	int cflag;
	int dflag;
	int pflag;

	aflag = 0;
	cflag = 0;
	dflag = 0;
	pflag = 0;
	string = NULL;
	section = ".comment";
	while ((c = getopt(ac, av, "a:cdn:pV")) != -1) {
		switch (c) {
		case 'a':
			aflag++;
			string = optarg;	/* XXX support multiple per pass */
			break;
		case 'c':
			cflag++;
			break;
		case 'd':
			dflag++;
			break;
		case 'n':
			section = optarg;
			break;
		case 'p':
			pflag++;
			break;
		case 'V':
			printf("mcs version $FreeBSD$\n");
			exit(0);
		default:
			err(1, "usage: mcs { -c | -d | -p | -V | -a string | -n name } ... file ...");
			break;
		}
	}
	ac -= optind;
	av += optind;
	for (; *av != NULL; av++)
		mcs(*av, section, aflag, cflag, dflag, pflag, string);
	exit(0);
}