/*
===========================================================================
Copyright (C) 2014 Essential Nature (Thilo Schulz)

Fix the command parse functions so as to replace ret
with an rjmp to the appropriate location
===========================================================================
*/

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

#include <libelf.h>
#include <gelf.h>

int isbe = 0;
int found_resp, found_noresp;
Elf32_Sym fresp, fnoresp;

void do_help(char *progname)
{
  printf("Usage: %s [options] <filename>\n"
         "Modify the elf files for the command parser so as to use rjmp instead\n"
         "of ret instruction to return to main interrupt service routine.\n\n"
         "Options:\n"
         "  -h                          Print this help\n"
         "\n",
         progname);
}

Elf_Scn *get_section(Elf *obj, Elf32_Word sh_type, const char *secname)
{
  Elf_Scn *cursec;
  Elf32_Shdr *curhdr;
  size_t shstrndex;
  char *name;
  
  if(elf_getshdrstrndx(obj, &shstrndex))
  {
    fprintf(stderr, "Couldn't get list of section names: %s\n", elf_errmsg(-1));
    exit(EX_DATAERR);
  }

  cursec = NULL;
  
  while((cursec = elf_nextscn(obj, cursec)))
  {
    if(!(curhdr = elf32_getshdr(cursec)))
    {
      fprintf(stderr, "Couldn't get section header: %s\n", elf_errmsg(-1));
      exit(EX_DATAERR);
    }
    
    if(name = elf_strptr(obj, shstrndex, curhdr->sh_name))
    {
      if(curhdr->sh_type == sh_type && (!secname || !strcmp(name, secname)))
        return cursec;
    }
  }
  
  return NULL;
}

size_t add_stringtotable(Elf *obj, size_t tableind, char *str)
{
  Elf_Scn *strsec;
  Elf32_Shdr *strhdr;
  Elf_Data *data;
  size_t lastofs;

  if(!(strsec = elf_getscn(obj, tableind)))
  {
    fprintf(stderr, "Couldn't get section to string table %zd: %s\n", elf_errmsg(-1));
    exit(EX_DATAERR);
  }

  if(!(strhdr = elf32_getshdr(strsec)))
  {
    fprintf(stderr, "Couldn't get header of string table\n");
    exit(EX_DATAERR);
  }
  
  // Workaround for bug in libelf
//  data = elf_getdata(strsec, NULL);
  
  if(!(data = elf_newdata(strsec)))
  {
    fprintf(stderr, "Couldn't add data to string table: %s\n", elf_errmsg(-1));
    exit(EX_DATAERR);
  }
  
  data->d_align = 1;
  data->d_buf = str;
  data->d_size = strlen(str) + 1;
  data->d_type = ELF_T_BYTE;
  
  lastofs = strhdr->sh_size;
  strhdr->sh_size += data->d_size;
  
  return lastofs;
}

void add_undefinedsym(Elf *obj, Elf_Scn *scn, Elf32_Sym *dest, size_t tableind, char *str)
{
  Elf_Data *data;
  GElf_Sym undef;

  if(!(data = elf_newdata(scn)))
  {
    fprintf(stderr, "Couldn't add symbol %s to symbol table: %s\n", str, elf_errmsg(-1));
    exit(EX_DATAERR);
  }

  undef.st_name = add_stringtotable(obj, tableind, str);
  undef.st_value = 0;
  undef.st_size = 0;
  undef.st_info = GELF_ST_INFO(STB_GLOBAL, STT_NOTYPE);
  undef.st_other = STV_DEFAULT;
  undef.st_shndx = 0;

  data->d_align = 4;
  data->d_buf = dest;
  data->d_size = sizeof(*dest);
  data->d_type = ELF_T_SYM;

  if(!gelf_update_sym(data, 0, &undef))
  {
    fprintf(stderr, "Couldn't add symbol %s to symbol table: %s\n", str, elf_errmsg(-1));
    exit(EX_DATAERR);
  }
}

int modify_elf(const char *fname)
{
  int fd;
  Elf *obj;
  Elf32_Ehdr *ehdr;
  Elf_Scn *symsec, *progsec;
  Elf32_Shdr *symsechdr;
  Elf_Data *data, *progdata;
  GElf_Sym sym;
  char *name;
  int numsym, index;
  
  if(elf_version(EV_CURRENT) == EV_NONE)
  {
    fprintf(stderr, "Could not initialize ELF library: %s\n", elf_errmsg(-1));
    return EX_SOFTWARE;
  }
  
  if((fd = open(fname, O_RDWR, 0)) < 0)
  {
    fprintf(stderr, "Couldn't open elf file: %s\n", strerror(errno));
    return EX_NOINPUT;
  }
  
  if(!(obj = elf_begin(fd, ELF_C_RDWR, NULL)))
  {
    fprintf(stderr, "Couldn't parse elf file: %s\n", elf_errmsg(-1));
    return EX_DATAERR;
  }

  if(!(ehdr = elf32_getehdr(obj)))
  {
    fprintf(stderr, "Failed to retrieve elf header: %s\n", elf_errmsg(-1));
    return EX_DATAERR;
  }

  if((ehdr->e_ident[EI_DATA] & ELFDATA2MSB))
    isbe = 1;

  symsec = get_section(obj, SHT_SYMTAB, ".symtab");
  
  if(!symsec)
  {
    fprintf(stderr, "No symbol section found\n");
    return EX_DATAERR;
  }
  
  data = NULL;
  if(!(data = elf_getdata(symsec, data)) || !data->d_size)
  {
    fprintf(stderr, "No data associated with symbol section\n");
    return EX_DATAERR;
  }

  if(!(symsechdr = elf32_getshdr(symsec)))
  {
    fprintf(stderr, "Couldn't get header of symbol section\n");
    return EX_DATAERR;
  }
  
  numsym = 0;
  found_resp = 0;
  found_noresp = 0;
  
  while(gelf_getsym(data, numsym, &sym))
  {
    name = elf_strptr(obj, symsechdr->sh_link, sym.st_name);

    if(!strcmp(name, "cmd_response"))
      found_resp = numsym;
    else if(!strcmp(name, "cmd_noresponse"))    
      found_noresp = numsym;

    numsym++;
  }
  
  /*
   * Make sure symbol table contains undefined symbols cmd_response and cmd_noresponse
   */

  if(!found_noresp)
  {
    add_undefinedsym(obj, symsec, &fnoresp, symsechdr->sh_link, "cmd_noresponse");
    found_noresp = numsym++;
  }

  if(!found_resp)
  {
    add_undefinedsym(obj, symsec, &fresp, symsechdr->sh_link, "cmd_response");
    found_resp = numsym++;
  }

  elf_update(obj, ELF_C_WRITE);
  
  elf_end(obj);
  
  if(close(fd) < 0)
  {
    fprintf(stderr, "Couldn't close elf file: %s\n", strerror(errno));
    return EX_NOINPUT;
  }
}

int main(int argc, char **argv)
{
  int opt;
  
   if(argc < 2)
  {
    if(argc < 1)
      do_help("elf-fixcmd");
    else
      do_help(argv[0]);

    return EX_OK;
  }
  
  while((opt = getopt(argc, argv, "h")) != -1)
  {
    switch(opt)
    {
      case 'h':
        do_help(argv[0]);
        return EX_OK;
      break;
      default:
        return EX_USAGE;
    }
  }
  
  return modify_elf(argv[optind]);
}