#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include <unistd.h>
#include <fcntl.h>
#include <elf.h>
#include <libelf/libelf.h>

#define ELF_STRTAB	".strtab"
#define ELF_SYMTAB	".symtab"

#define NOP		0
#define LOCALIZE	1
#define DELETE		2
#define NULLIFY		3
#define RENAME		4

int
main(int argc, char* argv[]) {
    int ifd, ofd = -1;
    Elf *elf_i, *elf_o;
    Elf32_Ehdr *ehdr_i, *ehdr_o;
    char* sname, *iname, *oname, *strtab;
    int c, i, j, strtlen, strdex = -1, symdex = -1, action = NOP, move = 0;
    int verbose = 0;
    Elf_Scn *scn_i, *scn_o;
    Elf32_Shdr *shdr_i, *shdr_o;
    Elf_Data* data_i, *data_o, *str_data_o, *sym_data_o;
    Elf32_Sym* sym, *symtab;
    char* target = NULL;
    char* newname = "";
    char __bind[] = "lgwn ";
    char __type[] = "uofsfn ";

    while ((c = getopt(argc, argv, "ldnr:s:v")) != EOF) {
	switch(c) {
	 case 'l':
	    action = LOCALIZE;
	    break;
	case 'd':
	    action = DELETE;
	    break;
	case 'n':
	    action = NULLIFY;
	    break;
	case 'r':
	    action = RENAME;
	    newname = optarg;
	    break;
	case 's':
	    target = optarg;
	    break;
	case 'v':
	    verbose = 1;
	    break;
	default:
	    goto usage;
	}
    }

    if (/*action == NOP ||*/ !target || (optind != argc-1 && optind != argc-2)) {
usage:
	fprintf(stderr, "Usage: symedit [-ldn|-r name] -s symbol infile [outfile]\n");
	return 1;
    }
    iname = argv[optind];
    if (optind == argc-2) {
	oname = argv[optind+1];
    } else {
	static char tmp[] = "/tmp/syXXXXXX";
	ofd = mkstemp(oname = tmp);
	if (ofd < 0) {
	    fprintf(stderr, "symedit: mkstemp failed\n");
	    return 1;
	}
	move = 1;
    }
    if (elf_version(EV_CURRENT) == EV_NONE) {
	fprintf(stderr, "elf library out of date\n");
	return 1;
    }

    if ((ifd = open(iname, O_RDONLY)) < 0) {
	perror(iname);
	return 1;
    }
//printf("ofd = %d %s\n", ofd, oname);
    if (ofd < 0 && (ofd = open(oname, O_WRONLY|O_CREAT|O_TRUNC, 0666)) < 0) {
	perror(oname);
	return 1;
    }


    elf_i = elf_begin(ifd, ELF_C_READ, (Elf*)NULL);
    if ((elf_kind(elf_i) != ELF_K_ELF) ||
	((ehdr_i = elf32_getehdr(elf_i)) == NULL)) {
	fprintf(stderr, "%s: not a valid ELF object file\n", iname);
	return 1;
    }

    /* copy elf header */
    elf_o = elf_begin(ofd, ELF_C_WRITE, (Elf*)NULL);
    ehdr_o = elf32_newehdr(elf_o);
    *ehdr_o = *ehdr_i;
    if (ehdr_i->e_phnum > 0) {
	Elf32_Phdr* phdr_i, *phdr_o;
	phdr_i = elf32_getphdr(elf_i);
	phdr_o = elf32_newphdr(elf_o, ehdr_i->e_phnum);
	memcpy(phdr_o, phdr_i, ehdr_i->e_phnum* sizeof *phdr_i);
    }


    /* copy sections */
    for(i = 1; i < ehdr_i->e_shnum; i++) {
	if (((scn_i = elf_getscn(elf_i, i)) == NULL) ||
	    ((shdr_i = elf32_getshdr(scn_i)) == NULL)) {
	    fprintf(stderr, "%s: can't get section number %d\n", iname, i);
	    return 1;
	}
	sname = elf_strptr(elf_i, ehdr_i->e_shstrndx, (size_t)shdr_i->sh_name);
	if (verbose) printf("%d %s\n", i, sname);
	/* handle symbol table and string table specially */
	if (strcmp(sname, ELF_STRTAB) == 0) {
	    strdex = i; continue;
	} else if(strcmp(sname, ELF_SYMTAB) == 0) {
	    symdex = i; continue;
	}
	scn_o = elf_newscn(elf_o);
	shdr_o = elf32_getshdr(scn_o);
	*shdr_o = *shdr_i;
	data_i = elf_getdata(scn_i, NULL);
	data_o = elf_newdata(scn_o);
	*data_o = *data_i;
    }

    /* do symbol table and string table */
    if (strdex == -1 || symdex == -1) {
	fprintf(stderr, "%s: can't find symbol table or string table\n", iname);
	return 1;
    }

    /* symbol table */
    scn_i = elf_getscn(elf_i, symdex);
    shdr_i =  elf32_getshdr(scn_i);
    data_i = elf_getdata(scn_i, NULL);
    scn_o = elf_newscn(elf_o);
    shdr_o = elf32_getshdr(scn_o);
    *shdr_o = *shdr_i;
    sym_data_o = data_o = elf_newdata(scn_o);
    *data_o = *data_i;
    symtab = sym = (Elf32_Sym*)data_o->d_buf;

    /* string table */
    scn_i = elf_getscn(elf_i, strdex);
    shdr_i =  elf32_getshdr(scn_i);
    data_i = elf_getdata(scn_i, NULL);
    strtab = data_i->d_buf; strtlen = data_i->d_size;
    /* copy to out for now */
    scn_o = elf_newscn(elf_o);
    shdr_o = elf32_getshdr(scn_o);
    *shdr_o = *shdr_i;
    str_data_o = data_o = elf_newdata(scn_o);
    *data_o = *data_i;

    for(j = 0; j < sym_data_o->d_size/sizeof *sym; j++) {
	if (verbose) printf("%08x  %c %c %3d  %s\n", sym->st_value,
	    ELF32_ST_BIND(sym->st_info) <= STB_NUM ? __bind[ELF32_ST_BIND(sym->st_info)] : '?',
	    ELF32_ST_TYPE(sym->st_info) <= STT_NUM ?__type[ELF32_ST_TYPE(sym->st_info)] : '?',
	    sym->st_size, strtab+sym->st_name);
	if (strcmp(strtab+sym->st_name, target) == 0) goto found;
	sym++;
    }
    fprintf(stderr, "%s: symbol %s not found\n", iname, target);
    unlink(oname);
    return 1;
found:
    switch(action) {
    case NOP:
	break;
    case DELETE: {
	/* delete strtab entry */
	int k, l = strlen(target)+1;
	str_data_o->d_size -= l;
	memcpy(strtab+sym->st_name, strtab+sym->st_name+l, strtlen-sym->st_name-l);
	/* update offsets */
	for(k = 0; k < sym_data_o->d_size/sizeof *sym; k++)
	    if (symtab[k].st_name > sym->st_name) symtab[k].st_name -= l;
	/* delete symtab entry */
	memcpy(sym, sym+1, sym_data_o->d_size - (j+1)*sizeof *sym);
	sym_data_o->d_size -= sizeof *sym;
	}
	break;
    case RENAME: {
	int diff, l = strlen(target)+1;
	if ((diff = l-1 - strlen(newname)) == 0 && 0) {
	    strcpy(strtab+sym->st_name, newname);
	} else {
	    /* allocate new table and copy */
	    char* nstrtab = malloc(strtlen-diff);
	    int k;
	    /* update offsets */
	    for(k = 0; k < sym_data_o->d_size/sizeof *sym; k++)
		if (symtab[k].st_name > sym->st_name) symtab[k].st_name -= diff;
	    memcpy(nstrtab, strtab, sym->st_name);
	    strcpy(nstrtab+sym->st_name, newname);
	    memcpy(nstrtab+sym->st_name+l-diff,
		   strtab+sym->st_name+l, strtlen-sym->st_name-l);
	    str_data_o->d_size -= diff;
	    str_data_o->d_buf = nstrtab;
	}
	}
	break;
    case NULLIFY:
	memset(strtab+sym->st_name, '\0', strlen(strtab+sym->st_name));
	break;
    case LOCALIZE:
	if (ELF32_ST_BIND(sym->st_info) != STB_GLOBAL)
	    printf("%s: warning symbol %s not global\n", iname, target);
	sym->st_info = ELF32_ST_INFO(STB_LOCAL, ELF32_ST_TYPE(sym->st_info));
	break;
    }
    elf_update(elf_o, ELF_C_WRITE);
    elf_end(elf_o);
    elf_end(elf_i);
    if (move)
	rename(oname, iname);
    return 0;
}