#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <fcntl.h>
#include <unistd.h>
#include <libelf.h>

#define	die() (__die (__FILE__, __LINE__))
extern /address@hidden@*/ void __die (const char *, int) __attribute__ ((noreturn));

void __die (const char *file, int line)
{
  fprintf (stderr, "FATAL ERROR: at line %d of file %s\n", line, file);
  exit (EXIT_FAILURE);
}

const int debug_elf = 0;

#define MAX_SECS 20

static void fail (void) __attribute__ ((noreturn));

static unsigned char *raw_text = NULL;
int code_size = -1;

static Elf32_Sym *symtab = NULL;
static int symtab_len = -1;

static char *strtab = NULL;
static int text_sec_num = -1;

int data_seg_size = 0;
int bss_seg_size = 0;

unsigned short get_insn_short (int addr)
{
  unsigned short ret = 
    raw_text[addr] |
    raw_text[addr+1] << 8;
  return ret;
}

static const char * sh_type_name (int sht)
{
  switch (sht) {
  case 0: return "SHT_NULL";
  case 1: return "SHT_PROGBITS";
  case 2: return "SHT_SYMTAB";
  case 3: return "SHT_STRTAB";	
  case 4: return "SHT_RELA";
  case 5: return "SHT_HASH";	
  case 6: return "SHT_DYNAMIC";		
  case 7: return "SHT_NOTE";
  case 8: return "SHT_NOBITS";		
  case 9: return "SHT_REL";	
  case 10: return "SHT_SHLIB";		
  case 11: return "SHT_DYNSYM";	
  case 14: return "SHT_INIT_ARRAY";	
  case 15: return "SHT_FINI_ARRAY";	
  case 16: return "SHT_PREINIT_ARRAY";
  case 17: return "SHT_GROUP";
  case 18: return "SHT_SYMTAB_SHNDX";
  case 19: return "SHT_NUM";
  }
  assert (0);
}	

void open_obj_file (char *fn)
{
  Elf32_Ehdr *ehdr;
  Elf *elf;
  Elf_Scn *scn;
  Elf_Data *data;
  int fd;
  unsigned int cnt;

  /* Open the input file */
  fd = open (fn, O_RDONLY);
  if (fd == -1) {
    printf ("oops: couldn't open '%s'\n", fn);
    die ();
  }

  /* Obtain the ELF descriptor */
  (void) elf_version (EV_CURRENT);
  elf = elf_begin (fd, ELF_C_READ, NULL);
  if (!elf) {
    fail ();
  }

  /* Obtain the .shstrtab data buffer */
  ehdr = elf32_getehdr (elf);
  if (!ehdr) fail ();

  if (debug_elf) {
    printf ("ehdr->e_machine = %d\n", ehdr->e_machine);
  }

  // make sure we're talking about an AVR here
  // assert (ehdr->e_machine == EM_AVR);

  scn = elf_getscn (elf, ehdr->e_shstrndx);
  if (!scn) fail ();

  data = elf_getdata (scn, NULL);
  if (!data) fail ();

  if (debug_elf) {
    printf ("sections:\n");
  }

  cnt = 0;
  scn = NULL; 

  while ((scn = elf_nextscn (elf, scn))) {
    char *name;
    Elf_Data *scn_data;
    Elf32_Shdr *shdr;
    void *dbuf;
    Elf_Type dtype;
    size_t dsize;

    cnt++;

    shdr = elf32_getshdr (scn);
    if (!shdr) {
      fail();
    }

    name = (char *) data->d_buf + shdr->sh_name;
    scn_data = elf_getdata (scn, NULL);
    if (!scn_data) fail ();

    dbuf = scn_data->d_buf;
    dtype = scn_data->d_type;
    dsize = scn_data->d_size;

    if (debug_elf) {
      printf ("\n");
      printf ("  [%d]\n", cnt);    
      printf ("    sh_name:      %s\n", name);
      printf ("    sh_type:      %s\n", sh_type_name(shdr->sh_type));
      printf ("    sh_flags:     %x\n", shdr->sh_flags);
      printf ("    sh_addr:      %x\n", shdr->sh_addr);
      printf ("    sh_offset:    %x\n", shdr->sh_offset);
      printf ("    sh_size:      %x\n", shdr->sh_size);
      printf ("    sh_link:      %x\n", shdr->sh_link);
      printf ("    sh_info:      %x\n", shdr->sh_info);
      printf ("    sh_addralign: %x\n", shdr->sh_addralign);
      printf ("    sh_entsize:   %x\n", shdr->sh_entsize);
      printf ("\n");
      printf ("    d_buf:  %p\n", dbuf);
      printf ("    d_type: %x\n", dtype);
      printf ("    d_size: %x\n", dsize);
    }

    if (strcmp (name, ".data") == 0) {
      data_seg_size = dsize;
    }

    if (strcmp (name, ".bss") == 0) {
      bss_seg_size = dsize;
    }

    if (strcmp (name, ".text") == 0) {
      text_sec_num = cnt;
      raw_text = (unsigned char *) dbuf;
      code_size = dsize;
      if (debug_elf) {
	printf ("    code_size = %x\n", code_size);
      }
    }

    if (strcmp (name, ".symtab") == 0) {
      symtab = (Elf32_Sym *) dbuf;
      symtab_len = dsize / sizeof (Elf32_Sym);
      if (debug_elf) {
	printf ("    symtab_len = %x\n", symtab_len);
      }
    }

    if (strcmp (name, ".strtab") == 0) {
      strtab = (char *) dbuf;
    }
  }

  assert (raw_text);
  assert (strtab);
  assert (symtab);

  close (fd);
}

int lookup_text_symbol_by_name (char *lookup_name)
{
  int i;
  int ret = -1;
  
  assert (text_sec_num != -1);

  for (i=0; i<symtab_len; i++) {
    char *name = strtab + symtab[i].st_name;
    int section = symtab[i].st_shndx;

    assert (name != NULL);

    if ((section == text_sec_num) &&
	(strlen (name) > 0) &&
	(strcmp (name, lookup_name) == 0)) {
      ret = symtab[i].st_value;
    }
  }
  
  return ret;
}

char *lookup_text_symbol_by_addr (unsigned int addr)
{
  int i;
  char *ret = NULL;
  
  assert (text_sec_num != -1);

  for (i=0; i<symtab_len; i++) {
    char *name = strtab + symtab[i].st_name;
    int section = symtab[i].st_shndx;

    assert (name != NULL);

    if ((section == text_sec_num) &&
	(symtab[i].st_value == addr) &&
	(strlen (name) > 0)) {
      ret = name;
    }
  }
  
  return ret;
}

static void fail (void)
{
  fprintf (stderr, "%s\n", elf_errmsg (elf_errno()));
  die ();
}

void find (char *name)
{
  printf ("%s ", name);
  int addr = lookup_text_symbol_by_name (name);
  if (addr != -1) {
    printf ("%x", addr);
  } else {
    printf ("NOT FOUND");
  }
  printf ("\n");
}

int main (int argc, char *argv[])
{
  if (argc != 2) {
    printf ("usage: readelf filename\n");
    return -1;
  }

  open_obj_file (argv[1]);
  find ("__success");
  find ("__failure");
  find ("main");
  return 0;
}