/* Test program for elf_newdata.
   Copyright (C) 2014 Red Hat, Inc.
   This file is part of elfutils.

   This file is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3 of the License, or
   (at your option) any later version.

   elfutils is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.  */

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

/* Adds a new "Hello World" string to the section.  */
static void
add_data (Elf_Scn *scn)
{
  printf ("Add Hello World\n");

  Elf_Data *data = elf_newdata (scn);
  if (data == NULL)
    error (1, 0, "elf_newdata: %s", elf_errmsg (-1));

  data->d_align = 1;
  data->d_buf = "Hello World";
  data->d_type = ELF_T_BYTE;
  data->d_size = sizeof ("Hello World");
  data->d_version = EV_CURRENT;
}

/* Creates a new minimal ELF file with a .data section containing
   "Hello World".  Returns the name of the new file.  */
static const char *
create_elf ()
{
  Elf *elf;
  int fd;

  static char name[] = "test.XXXXXX";
  fd = mkstemp (name);
  if (fd < 0)
    error (1, errno, "mkstemp");

  elf = elf_begin (fd, ELF_C_WRITE, NULL);
  if (elf == NULL)
    error (1, 0, "elfbegin: %s", elf_errmsg (-1));

  Elf32_Ehdr *ehdr = elf32_newehdr (elf);
  if (ehdr == NULL)
    error (1, 0, "elf32_newehdr: %s", elf_errmsg (-1));

  /* Our data.  */
  Elf_Scn *scn = elf_newscn (elf);
  if (scn == NULL)
    error (1, 0, "elf_newscn: %s", elf_errmsg (-1));

  add_data (scn);

  Elf32_Shdr *shdr = elf32_getshdr(scn);
  if (shdr == NULL)
    error (1, 0, "elf_getshdr: %s", elf_errmsg (-1));

  shdr->sh_name = 1; /* .data */
  shdr->sh_type = SHT_PROGBITS;
  shdr->sh_flags = SHF_STRINGS;
  shdr->sh_entsize = 0;

  /* The section string table.  */
  scn = elf_newscn (elf);
  if (scn == NULL)
    error (1, 0, "elf_newscn: %s", elf_errmsg (-1));

  Elf_Data *data = elf_newdata (scn);
  if (data == NULL)
    error (1, 0, "elf_newdata: %s", elf_errmsg (-1));

  // Quick and dirty string table for section zero, our new data section,
  // and this shstrtab section itself.
  char string_table[] = { '\0',
			  '.', 'd', 'a', 't', 'a', '\0',
			  '.', 's', 'h', 's', 't', 'r', 't', 'a', 'b', '\0' };

  data->d_align = 1;
  data->d_off = 0;
  data->d_buf = string_table;
  data->d_type = ELF_T_BYTE;
  data->d_size = sizeof (string_table);
  data->d_version = EV_CURRENT;

  shdr = elf32_getshdr(scn);
  if (shdr == NULL)
    error (1, 0, "elf_getshdr: %s", elf_errmsg (-1));

  shdr->sh_name = 7; /* .shstrtab */
  shdr->sh_type = SHT_STRTAB;
  shdr->sh_flags = SHF_STRINGS;
  shdr->sh_entsize = 0;

  size_t ndx = elf_ndxscn (scn);
  ehdr->e_shstrndx = ndx;
  ehdr->e_version = EV_CURRENT;

  /* Write it all out. */
  if (elf_update (elf, ELF_C_WRITE) < 0)
    error (1, 0, "elf_update: %s", elf_errmsg (-1));

  if (elf_end (elf) < 0)
    error (1, 0, "elf_end: %s", elf_errmsg (-1));

  if (close (fd) < 0)
    error (1, errno, "close");

  return name;
}

/* Opens an existing ELF file, checks that it has a .data section, if
   read is true checks that it contains one or more "Hello World"
   strings, and adds an extra "Hello World" to that section.  Gets the
   section data either through elf_getdata () or elf_rawdata () based
   on the last argument.  */
static void
modify_elf (const char *name, bool read, bool raw)
{
  int fd = open (name, O_RDWR);
  if (fd < 0)
    error (1, errno, "open '%s'", name);

  Elf *elf = elf_begin (fd, ELF_C_RDWR, NULL);
  if (elf == NULL)
    error (1, 0, "elf_begin: %s", elf_errmsg (-1));

  /* Get the section header string table index.  */
  size_t ndx;
  if (elf_getshdrstrndx (elf, &ndx) < 0)
    error (1, 0, "elf_getshdrstrndx: %s", elf_errmsg (-1));

  /* Get the .data section. */
  Elf_Scn *scn = NULL;
  Elf32_Shdr *shdr = NULL;
  while ((scn = elf_nextscn (elf, scn)) != NULL)
    {
      shdr = elf32_getshdr (scn);
      if (shdr != NULL && shdr->sh_type == SHT_PROGBITS)
	{
	  const char *sname = elf_strptr (elf, ndx, shdr->sh_name);
	  if (sname != NULL && strcmp (sname, ".data") == 0)
	    break;
	}
    }

  if (scn == NULL)
    error (1, 0, "No PROGBITS .data section.");

  if (read)
    {
      Elf_Data *data;
      if (raw)
	{
	  data  = elf_rawdata (scn, NULL);
	  if (data == NULL)
	    error (1, 0, "elf_rawdata: %s", elf_errmsg (-1));
	}
      else
	{
	  data  = elf_getdata (scn, NULL);
	  if (data == NULL)
	    error (1, 0, "elf_getdata: %s", elf_errmsg (-1));
	}

      /* Is it all Hello World?  */
      printf ("data size: %zd\n", data->d_size);
      int nr = data->d_size / sizeof ("Hello World");
      printf ("strings: %d\n", nr);
      char *str = data->d_buf;
      for (int i = 0; i < nr; i++)
	{
	  printf ("%s\n", (char *) str);
	  str += sizeof ("Hello World");
	}
    }

  /* And add one more string... */
  add_data (scn);

  /* Write it all out. */
  if (elf_update (elf, ELF_C_WRITE) < 0)
    error (1, 0, "elf_update: %s", elf_errmsg (-1));

  if (elf_end (elf) < 0)
    error (1, 0, "elf_end: %s", elf_errmsg (-1));

  if (close (fd) < 0)
    error (1, errno, "close");
}

int
main (int argc, char **argv)
{
  // Initialize libelf.
  if (elf_version (EV_CURRENT) == EV_NONE)
    error (1, 0, "elf_version: %s", elf_errmsg (-1));

  const char *name = create_elf ();
  printf ("name: %s\n", name);
  modify_elf (name, false, false); /* Don't read, just add.   */
  modify_elf (name, true, false);  /* Read with elf_getdata.  */
  modify_elf (name, true, true);   /* Read with elf_rawdata.  */
  modify_elf (name, true, false);  /* Read with elf_getdata again.  */

  return 0;
}