#if defined(_AIX) && !defined(_ALL_SOURCE)
#define _ALL_SOURCE
#endif
#include <stdio.h>
#include <stdlib.h>
#include <filehdr.h>
#include <aouthdr.h>
#include <scnhdr.h>
#ifdef __ultrix
typedef void *pCHDRR;
#endif
#include <ldfcn.h>
#include <sys/types.h>
#include <sys/stat.h>

#ifdef __ultrix
#include <alloca.h>
#define DBASE 0x10000000
#define OS_OK
#define	off_decrypt_code_length	0
#define	off_decrypt_start_addr	4
#define	off_decrypt_data_length	8
#define	off_entry		16
#endif

#ifdef _AIX
#define OS_OK
#define GLUE_IN_DATA
#define GLUE_SECTION ".data"
#endif

#ifndef OS_OK
 #error OS not supported
#endif

static int find_glue_section(struct ldfile *file, SCNHDR *scnhdr)
{
#ifdef DBASE
  int i;

  /* search through all sections for one that starts at DBASE */
  for (i = 1; i <= HEADER(file).f_nscns; i++)
    {
      if (ldshread(file, i, &scnhdr) == 0)
	{
	  fputs("unable to read program section headers\n", stderr);
	  return 1;
	}
      if (scnhdr.s_vaddr == DBASE)
	break;
    }
  if (i > HEADER(file).f_nscns ||
      fseek(IOPTR(file), scnhdr.s_scnptr, 0) == -1)
    {
      fprintf(stderr, "unable to find section with origin at %x\n", DBASE);
      return 1;
    }
#else
  if (ldnshread(file, ".data", &scnhdr) == 0)
    {
      fputs("unable to read data section header\n", stderr);
      return 1;
    }
#endif
}

static void encrypt_data(register unsigned char *data,
			 unsigned int data_length, int xor)
{
  register int i;
  for (i = 0; i < data_length; i++)
    *data++ ^= xor;
}

static struct ldfile *ldopen_rw(const char *filename)
{
  register struct ldfile *file;
  struct stat st;
  ino_t inum;

  file = ldopen((char *)filename, (struct ldfile *)0);
  if (file == 0)
    return 0;
  if (fstat(fileno(IOPTR(file)), &st) == -1)
    return 0;
  inum = st.st_ino;
  if (freopen(filename, "r+", IOPTR(file)) == 0)
    return 0;
  if (fstat(fileno(IOPTR(file)), &st) == -1)
    return 0;
  if (inum != st.st_ino)
    return 0;
  return file;
}

static unsigned char *read_glue(const char *glue_file,
				unsigned int *length_return,
				unsigned int entry,
				unsigned int data_start,
				unsigned int data_length)
{
  register struct ldfile *file;
  int i;
  unsigned char *data, *data0;
#ifdef __ultrix
  unsigned char *rdata;
#endif
  unsigned int vaddr;
  struct aouthdr aouthdr;
  unsigned int len, total_len;

  file = ldopen((char *)glue_file, (struct ldfile *)0);

  if (file == 0)
    {
      fputs("glue file unreadable or not in executable format\n", stderr);
      return 0;
    }

  if (HEADER(file).f_opthdr < 28)
    {
      fputs("glue file lacks a.out header\n", stderr);
      return 0;
    }

  if (ldohseek(file) == 0 ||
      fread(&aouthdr, 28, 1, IOPTR(file)) != 1)
    {
      fputs("unable to read glue file a.out header\n", stderr);
      return 0;
    }

  if (aouthdr.bsize)
    {
      fputs("glue file has BSS\n", stderr);
      return 0;
    }

#ifdef DBASE
  if (aouthdr.entry != DBASE)
    {
      fprintf(stderr, "glue file entry point is %x, should be %x\n",
	      aouthdr.entry, DBASE);
      return 0;
    }
#endif

  total_len = aouthdr.tsize + aouthdr.dsize;

  data0 = malloc(total_len);
  if (data0 == 0)
    {
      fputs("unable to allocate memory for glue file\n", stderr);
      return 0;
    }

  len = 0;
#ifdef DBASE
  vaddr = DBASE;
#else
  vaddr = 0;
#endif
  data = data0;
  /* read all sections */
  for (i = 1; i <= HEADER(file).f_nscns; i++)
    {
      SCNHDR scnhdr;
      int flags;
      if (ldshread(file, i, &scnhdr) == 0)
	{
	  fprintf(stderr, "unable to read section %d from glue file\n", i);
	  return 0;
	}
#ifdef DEBUG
      printf("section \"%s\", addr %x, %d bytes\n", scnhdr.s_name,
	     scnhdr.s_vaddr, scnhdr.s_size);
#endif
      if (scnhdr.s_vaddr != vaddr)
	{
	  fprintf(stderr, "vaddr mismatch %x != %x\n", vaddr, scnhdr.s_vaddr);
	  return 0;
	}
      if (len + scnhdr.s_size > total_len)
	{
	  fputs("glue file a.out header has incorrect size\n", stderr);
	  return 0;
	}
      if (fseek(IOPTR(file), scnhdr.s_scnptr, 0) == -1
	  || fread(data, 1, scnhdr.s_size, IOPTR(file)) != scnhdr.s_size)
	{
	  fputs("unable to read glue file\n", stderr);
	  return 0;
	}
#ifdef __ultrix
      if (!strcmp(scnhdr.s_name, ".rdata"))
	rdata = data;
#endif
      data += scnhdr.s_size;
      len += scnhdr.s_size;
      vaddr += scnhdr.s_size;
    }
#ifdef __ultrix
  *(int *)(rdata + off_decrypt_code_length) = (len + 7) & ~7;
  *(int *)(rdata + off_decrypt_start_addr) = data_start;
  *(int *)(rdata + off_decrypt_data_length) = data_length;
  *(int *)(rdata + off_entry) = entry;
#endif
  *length_return = (len + 7) & ~7;
  return data0;
}


int main(int argc, char *argv[])
{
  register struct ldfile *file;
  register unsigned char xor;
  FILE *output;		/* data needed to decrypt program */
  unsigned char *glue, *data;
  struct aouthdr aouthdr;
  SCNHDR scnhdr, data_scnhdr;
  int i;
  unsigned char key[256];
  unsigned int data_start, data_length, glue_length;

  if (argc != 5)
    {
      fputs("usage: enc program glue key output\n", stderr);
      return 1;
    }

  file = ldopen_rw(argv[1]);
  if (file == 0)
    {
      fprintf(stderr, "Unable to open program \"%s\" for write.\n", argv[1]);
      return 2;
    }

  if (HEADER(file).f_opthdr != 28 &&
      HEADER(file).f_opthdr != sizeof aouthdr)
    {
      fputs("a.out header size mismatch\n", stderr);
      return 3;
    }

  if (ldohseek(file) == 0 ||
      fread(&aouthdr, sizeof aouthdr, 1, IOPTR(file)) != 1)
    {
      fputs("a.out header read failed\n", stderr);
      return 3;
    }

#ifdef DEBUG
  printf("magic = %x, version %d, %d+%d+%d, entry %x\n", aouthdr.magic,
	 aouthdr.vstamp, aouthdr.tsize, aouthdr.dsize, aouthdr.bsize,
	 aouthdr.entry);
  if (HEADER(file).f_opthdr == sizeof aouthdr)
    {
#ifdef _AIX
      printf("o_toc = %x, o_snentry = %d, o_sntext = %d, o_sndata = %d\n",
	     aouthdr.o_toc, aouthdr.o_snentry, aouthdr.o_sntext,
	     aouthdr.o_sndata);
      printf("modtype = %.2s, maxstack = %x\n", aouthdr.o_modtype,
	     aouthdr.o_maxstack);
#endif
#ifdef __ultrix
      printf("text start %x, data start %x, bss start %x, GP %x\n",
	     aouthdr.text_start, aouthdr.data_start, aouthdr.bss_start,
	     aouthdr.gp_value);
      printf("GPR mask %08x, CPR mask %08x %08x %08x %08x\n",
	     aouthdr.gprmask, aouthdr.cprmask[0], aouthdr.cprmask[1],
	     aouthdr.cprmask[2], aouthdr.cprmask[3]);
#endif
    }
#endif

  /* find start address and length of .data section  */
  if (ldnshread(file, ".data", &data_scnhdr) == 0)
    {
      fputs("unable to read program .data section header\n", stderr);
      return 3;
    }
  data_start = data_scnhdr.s_vaddr;
  data_length = data_scnhdr.s_size;
  /* read glue code, substituting variables */
  glue = read_glue(argv[2], &glue_length, aouthdr.entry,
		   data_start, data_length);
  if (glue == 0)
    return 4;
#ifdef DEBUG
  printf("glue code: %8x %8x %8x %8x\n",
	 ((int *)glue)[0], ((int *)glue)[1],
	 ((int *)glue)[2], ((int *)glue)[3]);
#endif
  /* replace start of data with glue code */
  if (find_glue_section(file, &scnhdr))
    return 3;

  /* If glue does not go into the .data section, save the data which
     is being replaced.  */
#ifndef GLUE_IN_DATA
  data = malloc(scnhdr.s_size);
  if (data == 0)
    {
      fputs("unable to allocate memory for program data\n", stderr);
      return 5;
    }
  if (glue_length > scnhdr.s_size)
    {
#ifdef DEBUG
      fprintf(stderr, "%u > %u\n", glue_length, scnhdr.s_size);
#endif
      fputs("glue file too large for program\n", stderr);
      return 4;
    }

  if (fread(data, 1, glue_length, IOPTR(file)) != glue_length)
    {
      fputs("unable to read program data\n", stderr);
      return 3;
    }
  output = fopen(argv[4], "w");
  if (output == 0)
    {
      fputs("unable to open output file for write\n", stderr);
      return 6;
    }
  if (fwrite(data, 1, glue_length, output) != glue_length)
    {
      fputs("unable to write output file\n", stderr);
      return 6;
    }
  free(data);

  if (fseek(IOPTR(file), scnhdr.s_scnptr, 0) == -1)
    {
      fprintf(stderr, "unable to find section with origin at %x\n", DBASE);
      return 3;
    }
  if (fwrite(glue, 1, glue_length, IOPTR(file)) != glue_length)
    {
      fputs("unable to write glue code to program\n", stderr);
      return 7;
    }
#endif /* ! GLUE_IN_DATA */

#ifdef DBASE
  aouthdr.entry = DBASE;
  if (ldohseek(file) == 0 ||
      fwrite(&aouthdr, sizeof aouthdr, 1, IOPTR(file)) != 1)
    {
      fputs("a.out header write failed\n", stderr);
      return 7;
    }
#endif

  data = malloc(data_length);
  if (data == 0)
    {
      fputs("unable to allocate memory for program data\n", stderr);
      return 5;
    }
  if (fseek(IOPTR(file), data_scnhdr.s_scnptr, 0) == -1 ||
      fread(data, 1, data_length, IOPTR(file)) != data_length)
    {
      fputs("unable to read program data\n", stderr);
      return 2;
    }

#ifdef GLUE_IN_DATA
  output = fopen(argv[4], "w");
  if (output == 0)
    {
      fputs("unable to open output file for write\n", stderr);
      return 6;
    }
  if (fwrite(data, 1, glue_length, output) != glue_length)
    {
      fputs("unable to write output file\n", stderr);
      return 6;
    }
  encrypt_data(data + glue_length, data_length - glue_length, atoi(argv[3]));
#else
  encrypt_data(data, data_length, atoi(argv[3]));
#endif

  if (fseek(IOPTR(file), data_scnhdr.s_scnptr, 0) == -1 ||
      fwrite(data, 1, data_length, IOPTR(file)) != data_length)
    {
      fputs("unable to write program data\n", stderr);
      return 2;
    }
  memset(key, xor, sizeof key);
  if (fwrite(key, 1, sizeof key, output) != sizeof key)
    {
      fputs("unable to write key to output file\n", stderr);
      return 7;
    }
  return 0;
}