#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <libmcpu.h>
#include <mcpu-elf.h>

#include "mcpu-melf.h"


typedef struct fixture_shdr
{
  __mpu_uint32_t name;
  __mpu_uint32_t type;
  __mpu_uint64_t flags;
  __mpu_uint64_t addr;
  __mpu_uint64_t offset;
  __mpu_uint64_t size;
  __mpu_uint32_t link;
  __mpu_uint32_t info;
  __mpu_uint64_t align;
  __mpu_uint64_t entsize;
} fixture_shdr;


static void
get_shdr( const unsigned char *p, int data, fixture_shdr *sh )
{
  memset( sh, 0, sizeof( *sh ) );
  sh->name = cpu_melf_get32( p + MELF64_SHDR_SH_NAME_OFF, data );
  sh->type = cpu_melf_get32( p + MELF64_SHDR_SH_TYPE_OFF, data );
  sh->flags = cpu_melf_get64( p + MELF64_SHDR_SH_FLAGS_OFF, data );
  sh->addr = cpu_melf_get64( p + MELF64_SHDR_SH_ADDR_OFF, data );
  sh->offset = cpu_melf_get64( p + MELF64_SHDR_SH_OFFSET_OFF, data );
  sh->size = cpu_melf_get64( p + MELF64_SHDR_SH_SIZE_OFF, data );
  sh->link = cpu_melf_get32( p + MELF64_SHDR_SH_LINK_OFF, data );
  sh->info = cpu_melf_get32( p + MELF64_SHDR_SH_INFO_OFF, data );
  sh->align = cpu_melf_get64( p + MELF64_SHDR_SH_ADDRALIGN_OFF, data );
  sh->entsize = cpu_melf_get64( p + MELF64_SHDR_SH_ENTSIZE_OFF, data );

} /* End of get_shdr() */


static void
put_shdr( unsigned char *p, int data, const fixture_shdr *sh )
{
  memset( p, 0, MELF64_SHDR_SIZE );
  cpu_melf_put32( p + MELF64_SHDR_SH_NAME_OFF, sh->name, data );
  cpu_melf_put32( p + MELF64_SHDR_SH_TYPE_OFF, sh->type, data );
  cpu_melf_put64( p + MELF64_SHDR_SH_FLAGS_OFF, sh->flags, data );
  cpu_melf_put64( p + MELF64_SHDR_SH_ADDR_OFF, sh->addr, data );
  cpu_melf_put64( p + MELF64_SHDR_SH_OFFSET_OFF, sh->offset, data );
  cpu_melf_put64( p + MELF64_SHDR_SH_SIZE_OFF, sh->size, data );
  cpu_melf_put32( p + MELF64_SHDR_SH_LINK_OFF, sh->link, data );
  cpu_melf_put32( p + MELF64_SHDR_SH_INFO_OFF, sh->info, data );
  cpu_melf_put64( p + MELF64_SHDR_SH_ADDRALIGN_OFF, sh->align, data );
  cpu_melf_put64( p + MELF64_SHDR_SH_ENTSIZE_OFF, sh->entsize, data );

} /* End of put_shdr() */


static const char *
name_at( const unsigned char *table, size_t size, __mpu_uint32_t offset )
{
  const unsigned char *p;

  if( offset >= size )
    return( NULL );
  p = table + offset;
  if( !memchr( p, 0, size - offset ) )
    return( NULL );
  return( (const char *)p );

} /* End of name_at() */


int
main( int argc, char **argv )
{
  unsigned char *file = NULL;
  unsigned char *out = NULL;
  fixture_shdr *sh = NULL;
  const unsigned char *old_shstr;
  size_t file_size = 0;
  size_t old_shstr_size;
  size_t new_shstr_size;
  size_t i;
  size_t text_index = 0;
  int data;
  int overflow = 0;
  __mpu_uint64_t old_shoff;
  __mpu_uint64_t cur;
  __mpu_uint64_t new_shoff;
  __mpu_uint64_t total;
  __mpu_uint16_t shnum;
  __mpu_uint16_t shentsize;
  __mpu_uint16_t shstrndx;
  fixture_shdr strtab;
  fixture_shdr symtab;
  fixture_shdr shstr;
  unsigned char symbols[2 * MELF64_SYM_SIZE];
  const char fixture_strtab[] = "\0fixture_symbol\0";
  const char strtab_name[] = ".strtab";
  const char symtab_name[] = ".symtab";
  char err[256];
  int rc = EXIT_FAILURE;

  if( argc != 3 )
  {
    fprintf( stderr, "usage: %s INPUT-ET_EXEC-OR-ET_DYN OUTPUT\n", argv[0] );
    return( EXIT_FAILURE );
  }

  if( cpu_melf_readall( argv[1], &file, &file_size, err, sizeof( err ) ) )
  {
    fprintf( stderr, "%s\n", err );
    return( EXIT_FAILURE );
  }

  if( file_size < MELF64_EHDR_SIZE ||
      file[MELF_EI_MAG0] != MELFMAG0 || file[MELF_EI_MAG1] != MELFMAG1 ||
      file[MELF_EI_MAG2] != MELFMAG2 || file[MELF_EI_MAG3] != MELFMAG3 ||
      file[MELF_EI_CLASS] != MELF_CLASS_64 )
    goto out;

  data = file[MELF_EI_DATA];
  if( data != MELF_DATA_LSB && data != MELF_DATA_MSB )
    goto out;
  {
    __mpu_uint16_t type = cpu_melf_get16(
      file + MELF64_EHDR_E_TYPE_OFF, data );

    if( type != MELF_ET_EXEC && type != MELF_ET_DYN )
      goto out;
  }

  old_shoff = cpu_melf_get64( file + MELF64_EHDR_E_SHOFF_OFF, data );
  shnum = cpu_melf_get16( file + MELF64_EHDR_E_SHNUM_OFF, data );
  shentsize = cpu_melf_get16( file + MELF64_EHDR_E_SHENTSIZE_OFF, data );
  shstrndx = cpu_melf_get16( file + MELF64_EHDR_E_SHSTRNDX_OFF, data );
  if( !shnum || shnum > UINT16_MAX - 2 || shentsize != MELF64_SHDR_SIZE ||
      shstrndx >= shnum ||
      !cpu_melf_span( file_size, old_shoff,
                      (__mpu_uint64_t)shnum * MELF64_SHDR_SIZE ) )
    goto out;

  sh = (fixture_shdr *)calloc( (size_t)shnum + 2, sizeof( *sh ) );
  if( !sh )
    goto out;
  for( i = 0; i < shnum; ++i )
    get_shdr( file + (size_t)old_shoff + i * MELF64_SHDR_SIZE,
              data, &sh[i] );

  if( sh[shstrndx].type != MELF_SHT_STRTAB ||
      !cpu_melf_span( file_size, sh[shstrndx].offset,
                      sh[shstrndx].size ) ||
      sh[shstrndx].size > SIZE_MAX )
    goto out;

  old_shstr = file + (size_t)sh[shstrndx].offset;
  old_shstr_size = (size_t)sh[shstrndx].size;
  for( i = 1; i < shnum; ++i )
  {
    const char *name = name_at( old_shstr, old_shstr_size, sh[i].name );

    if( !name )
      goto out;
    if( !strcmp( name, MELF_SEC_TEXT ) )
      text_index = i;
  }
  if( !text_index )
    goto out;

  cur = old_shoff;
  memset( &strtab, 0, sizeof( strtab ) );
  strtab.type = MELF_SHT_STRTAB;
  strtab.offset = cur;
  strtab.size = sizeof( fixture_strtab );
  strtab.align = 1;
  cur += strtab.size;

  cur = cpu_melf_align_up( cur, 8, &overflow );
  if( overflow )
    goto out;
  memset( &symtab, 0, sizeof( symtab ) );
  symtab.type = MELF_SHT_SYMTAB;
  symtab.offset = cur;
  symtab.size = sizeof( symbols );
  symtab.link = shnum;
  symtab.info = 1;
  symtab.align = 8;
  symtab.entsize = MELF64_SYM_SIZE;
  cur += symtab.size;

  if( old_shstr_size > SIZE_MAX - sizeof( strtab_name ) ||
      old_shstr_size + sizeof( strtab_name ) > SIZE_MAX - sizeof( symtab_name ) )
    goto out;
  new_shstr_size = old_shstr_size + sizeof( strtab_name ) +
                   sizeof( symtab_name );
  shstr = sh[shstrndx];
  shstr.offset = cur;
  shstr.size = new_shstr_size;
  cur += new_shstr_size;

  new_shoff = cpu_melf_align_up( cur, 8, &overflow );
  if( overflow || new_shoff > UINT64_MAX -
      (__mpu_uint64_t)( shnum + 2 ) * MELF64_SHDR_SIZE )
    goto out;
  total = new_shoff + (__mpu_uint64_t)( shnum + 2 ) * MELF64_SHDR_SIZE;
  if( total > SIZE_MAX )
    goto out;

  out = (unsigned char *)calloc( (size_t)total, 1 );
  if( !out )
    goto out;
  memcpy( out, file, (size_t)old_shoff );
  memcpy( out + (size_t)strtab.offset, fixture_strtab,
          sizeof( fixture_strtab ) );

  memset( symbols, 0, sizeof( symbols ) );
  cpu_melf_put32( symbols + MELF64_SYM_SIZE + MELF64_SYM_ST_NAME_OFF,
                  1, data );
  symbols[MELF64_SYM_SIZE + MELF64_SYM_ST_INFO_OFF] =
    MELF64_ST_INFO( MELF_STB_GLOBAL, MELF_STT_FUNC );
  symbols[MELF64_SYM_SIZE + MELF64_SYM_ST_OTHER_OFF] = MELF_STV_DEFAULT;
  cpu_melf_put16( symbols + MELF64_SYM_SIZE + MELF64_SYM_ST_SHNDX_OFF,
                  (__mpu_uint16_t)text_index, data );
  cpu_melf_put64( symbols + MELF64_SYM_SIZE + MELF64_SYM_ST_VALUE_OFF,
                  sh[text_index].addr, data );
  memcpy( out + (size_t)symtab.offset, symbols, sizeof( symbols ) );

  memcpy( out + (size_t)shstr.offset, old_shstr, old_shstr_size );
  memcpy( out + (size_t)shstr.offset + old_shstr_size,
          strtab_name, sizeof( strtab_name ) );
  memcpy( out + (size_t)shstr.offset + old_shstr_size + sizeof( strtab_name ),
          symtab_name, sizeof( symtab_name ) );

  strtab.name = (__mpu_uint32_t)old_shstr_size;
  symtab.name = (__mpu_uint32_t)( old_shstr_size + sizeof( strtab_name ) );
  sh[shstrndx] = shstr;
  sh[shnum] = strtab;
  sh[shnum + 1] = symtab;

  for( i = 0; i < (size_t)shnum + 2; ++i )
    put_shdr( out + (size_t)new_shoff + i * MELF64_SHDR_SIZE,
              data, &sh[i] );

  cpu_melf_put64( out + MELF64_EHDR_E_SHOFF_OFF, new_shoff, data );
  cpu_melf_put16( out + MELF64_EHDR_E_SHNUM_OFF,
                  (__mpu_uint16_t)( shnum + 2 ), data );

  if( cpu_melf_writeall( argv[2], out, (size_t)total, err, sizeof( err ) ) )
  {
    fprintf( stderr, "%s\n", err );
    goto out;
  }
  rc = EXIT_SUCCESS;

out:
  free( out );
  free( sh );
  free( file );
  return( rc );

} /* End of main() */