/*
 * PaX control
 * Copyright 2004 PaX Team <pageexec@freemail.hu>
 * Licensed under the GNU GPL version 2
 */

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

#include "paxctl.h"

static void report_flags(const Elf64_Word flags, const struct pax_state * const state)
{
  char buffer[13];

  /* the logic is: lower case: explicitly disabled, upper case: explicitly enabled, - : default */
  buffer[0] = (flags & PF_PAGEEXEC    ? 'P' : '-');
  buffer[1] = (flags & PF_NOPAGEEXEC  ? 'p' : '-');
  buffer[2] = (flags & PF_SEGMEXEC    ? 'S' : '-');
  buffer[3] = (flags & PF_NOSEGMEXEC  ? 's' : '-');
  buffer[4] = (flags & PF_MPROTECT    ? 'M' : '-');
  buffer[5] = (flags & PF_NOMPROTECT  ? 'm' : '-');
  buffer[6] = (flags & PF_RANDEXEC    ? 'X' : '-');
  buffer[7] = (flags & PF_NORANDEXEC  ? 'x' : '-');
  buffer[8] = (flags & PF_EMUTRAMP    ? 'E' : '-');
  buffer[9] = (flags & PF_NOEMUTRAMP  ? 'e' : '-');
  buffer[10] = (flags & PF_RANDMMAP   ? 'R' : '-');
  buffer[11] = (flags & PF_NORANDMMAP ? 'r' : '-');
  buffer[12] = 0;

  fprintf(stdout, "- PaX flags: %s [%s]\n", buffer, state->argv[state->files]);

  if (state->shortonly)
    return;

  if (flags & PF_PAGEEXEC)   fprintf(stdout, "\tPAGEEXEC is enabled\n");
  if (flags & PF_NOPAGEEXEC) fprintf(stdout, "\tPAGEEXEC is disabled\n");
  if (flags & PF_SEGMEXEC)   fprintf(stdout, "\tSEGMEXEC is enabled\n");
  if (flags & PF_NOSEGMEXEC) fprintf(stdout, "\tSEGMEXEC is disabled\n");
  if (flags & PF_MPROTECT)   fprintf(stdout, "\tMPROTECT is enabled\n");
  if (flags & PF_NOMPROTECT) fprintf(stdout, "\tMPROTECT is disabled\n");
  if (flags & PF_RANDEXEC)   fprintf(stdout, "\tRANDEXEC is enabled\n");
  if (flags & PF_NORANDEXEC) fprintf(stdout, "\tRANDEXEC is disabled\n");
  if (flags & PF_EMUTRAMP)   fprintf(stdout, "\tEMUTRAMP is enabled\n");
  if (flags & PF_NOEMUTRAMP) fprintf(stdout, "\tEMUTRAMP is disabled\n");
  if (flags & PF_RANDMMAP)   fprintf(stdout, "\tRANDMMAP is enabled\n");
  if (flags & PF_NORANDMMAP) fprintf(stdout, "\tRANDMMAP is disabled\n");
}

static void * elf32_map_phdr(const int fd, const struct pax_state * const state)
{
  return mmap(NULL, state->ops->phoff._32 + state->ops->phentsize._32 * state->ops->phnum._32,
               PROT_READ | ((state->flags_on | state->flags_off) ? PROT_WRITE : 0), MAP_SHARED, fd, (off_t)0);
}

static void elf32_modify_phdr(unsigned char * const map, const struct pax_state * const state)
{
  unsigned int i;
  Elf32_Phdr * const phdr = (Elf32_Phdr *)(map + state->ops->phoff._32);

  for (i = 0U; i < state->ops->phnum._32; i++)
    if (phdr[i].p_type == PT_PAX_FLAGS) {
      if (state->view)
        report_flags(phdr[i].p_flags, state);
      if (state->flags_on | state->flags_off) {
        const Elf32_Ehdr * const ehdr = (const Elf32_Ehdr *)map;

        if (ehdr->e_type == ET_DYN) {
          phdr[i].p_flags &= ~((state->flags_off | PF_RANDEXEC) & ~PF_NORANDEXEC);
          phdr[i].p_flags |= (state->flags_on | PF_NORANDEXEC) & ~PF_RANDEXEC;
        } else {
          phdr[i].p_flags &= ~state->flags_off;
          phdr[i].p_flags |= state->flags_on;
        }
      }
      return;
    }
}

static void * elf64_map_phdr(const int fd, const struct pax_state * const state)
{
  return mmap(NULL, state->ops->phoff._64 + state->ops->phentsize._64 * state->ops->phnum._64,
               PROT_READ | ((state->flags_on | state->flags_off) ? PROT_WRITE : 0), MAP_SHARED, fd, (off_t)0);
}

static void elf64_modify_phdr(unsigned char * const map, const struct pax_state * const state)
{
  unsigned int i;
  Elf64_Phdr * const phdr = (Elf64_Phdr *)(map + state->ops->phoff._64);

  for (i = 0U; i < state->ops->phnum._64; i++)
    if (phdr[i].p_type == PT_PAX_FLAGS) {
      if (state->view)
        report_flags(phdr[i].p_flags, state);
      if (state->flags_on | state->flags_off) {
        const Elf64_Ehdr * const ehdr = (const Elf64_Ehdr *)map;

        if (ehdr->e_type == ET_DYN) {
          phdr[i].p_flags &= ~((state->flags_off | PF_RANDEXEC) & ~PF_NORANDEXEC);
          phdr[i].p_flags |= (state->flags_on | PF_NORANDEXEC) & ~PF_RANDEXEC;
        } else {
          phdr[i].p_flags &= ~state->flags_off;
          phdr[i].p_flags |= state->flags_on;
        }
      }
      return;
    }
}

static struct elf_ops elf32 = {
  .map_phdr = elf32_map_phdr,
  .modify_phdr = elf32_modify_phdr,
};

static struct elf_ops elf64 = {
  .map_phdr = elf64_map_phdr,
  .modify_phdr = elf64_modify_phdr,
};

static void banner(void)
{
  fprintf(stderr,
    "PaX control v" PAXCTL_VERSION "\n"
    "Copyright 2004 PaX Team <pageexec@freemail.hu>\n\n");
}

static void usage(void)
{
  banner();
  fprintf(stderr,
    "usage: paxctl <options> <files>\n\n"
    "options:\n"
    "\t-p: disable PAGEEXEC\t\t-P: enable PAGEEXEC\n"
    "\t-e: disable EMUTRMAP\t\t-E: enable EMUTRMAP\n"
    "\t-m: disable MPROTECT\t\t-M: enable MPROTECT\n"
    "\t-r: disable RANDMMAP\t\t-R: enable RANDMMAP\n"
    "\t-x: disable RANDEXEC\t\t-X: enable RANDEXEC\n"
    "\t-s: disable SEGMEXEC\t\t-S: enable SEGMEXEC\n\n"
    "\t-v: view flags\t\t\t-z: restore default flags\n"
    "\t-q: suppress error messages\t-Q: report flags in short format flags\n"
  );
  exit(EXIT_SUCCESS);
}

static int is_elf32(const void * const map, struct pax_state * const state)
{
  const Elf32_Ehdr * const ehdr = (const Elf32_Ehdr *)map;

  if (memcmp(ehdr->e_ident, ELFMAG, SELFMAG))
    return 0;
  if (ehdr->e_ehsize != sizeof(Elf32_Ehdr))
    return 0;
  if ((ehdr->e_version != EV_CURRENT) || (ehdr->e_ident[EI_CLASS] != ELFCLASS32))
    return 0;
  if ((ehdr->e_type != ET_EXEC) && (ehdr->e_type != ET_DYN))
    return 0;
  if (!ehdr->e_phoff || !ehdr->e_phnum || sizeof(Elf32_Phdr) != ehdr->e_phentsize)
    return 0;
  if (ehdr->e_phnum > 65536U / ehdr->e_phentsize)
    return 0;
  if (ehdr->e_phoff > ehdr->e_phoff + ehdr->e_phentsize * ehdr->e_phnum)
    return 0;

  state->ops = &elf32;
  state->ops->phoff._32 = ehdr->e_phoff;
  state->ops->phentsize._32 = ehdr->e_phentsize;
  state->ops->phnum._32 = ehdr->e_phnum;

  return 1;
}

static int is_elf64(const void * const map, struct pax_state * const state)
{
  const Elf64_Ehdr * const ehdr = (const Elf64_Ehdr *)map;

  if (memcmp(ehdr->e_ident, ELFMAG, SELFMAG))
    return 0;
  if (ehdr->e_ehsize != sizeof(Elf64_Ehdr))
    return 0;
  if ((ehdr->e_version != EV_CURRENT) || (ehdr->e_ident[EI_CLASS] != ELFCLASS64))
    return 0;
  if ((ehdr->e_type != ET_EXEC) && (ehdr->e_type != ET_DYN))
    return 0;
  if (!ehdr->e_phoff || !ehdr->e_phnum || sizeof(Elf64_Phdr) != ehdr->e_phentsize)
    return 0;
  if (ehdr->e_phnum > 65536U / ehdr->e_phentsize)
    return 0;
  if (ehdr->e_phoff > ehdr->e_phoff + ehdr->e_phentsize * ehdr->e_phnum)
    return 0;

  state->ops = &elf64;
  state->ops->phoff._64 = ehdr->e_phoff;
  state->ops->phentsize._64 = ehdr->e_phentsize;
  state->ops->phnum._64 = ehdr->e_phnum;

  return 1;
}

static int pax_verify_file(struct pax_state * const state)
{
  int fd;
  const void * map;
  size_t size = sizeof(Elf64_Ehdr);

  fd = open(state->argv[state->files], O_RDONLY);
  if (0 > fd) {
    if (!state->quiet)
      perror("open: ");
    return -1;
  }

  map = mmap(NULL, size, PROT_READ, MAP_PRIVATE, fd, (off_t)0);
  close(fd);
  if (MAP_FAILED == map) {
    if (!state->quiet)
      perror("mmap: ");
    return -1;
  }

  if (!is_elf32(map, state) && !is_elf64(map, state)) {
    munmap((void *)map, size);
    if (!state->quiet)
      fprintf(stderr, "file %s is not a valid ELF executable\n", state->argv[state->files]);
    return -1;
  }

  munmap((void *)map, size);
  return 0;
}

static int pax_modify_file(const struct pax_state * const state)
{
  int fd;
  void * map;

  /* open file */
  fd = open(state->argv[state->files], (state->flags_on | state->flags_off) ? O_RDWR : O_RDONLY);
  if (0 > fd) {
    if (!state->quiet)
      perror("open: ");
    return -1;
  }

  /* mmap file */
  map = state->ops->map_phdr(fd, state);
  close(fd);
  if (MAP_FAILED == map) {
    if (!state->quiet)
      perror("mmap: ");
    return -1;
  }

  /* report/modify flags */
  state->ops->modify_phdr(map, state);
  return 0;
}

static int pax_process_file(struct pax_state * const state)
{
  /* get/verify ELF header */
  if (0 > pax_verify_file(state))
    return -1;

  /* report/modify program header */
  if (0 > pax_modify_file(state))
    return -1;

  return 0;
}

static int pax_process_files(struct pax_state * const state)
{
  while (state->argv[state->files]) {
    pax_process_file(state);
    ++state->files;
  }

  return 0;
}

static int pax_parse_args(int argc, struct pax_state * const state)
{
  while (1) {
    switch(getopt(argc, state->argv, "pPsSmMeErRxXvqQz")) {
    case -1:
      state->files = optind;
      return optind < argc ? 0 : -1;

    case '?':
      return -1;

#define parse_flag(option1, option2, flag)	\
  case option1:					\
    state->flags_on &= ~PF_##flag;		\
    state->flags_on |= PF_NO##flag;		\
    state->flags_off &= ~PF_NO##flag;		\
    state->flags_off |= PF_##flag;		\
    break;					\
  case option2:					\
    state->flags_on &= ~PF_NO##flag;		\
    state->flags_on |= PF_##flag;		\
    state->flags_off &= ~PF_##flag;		\
    state->flags_off |= PF_NO##flag;		\
    break;

   parse_flag('p', 'P', PAGEEXEC);
   parse_flag('s', 'S', SEGMEXEC);
   parse_flag('m', 'M', MPROTECT);
   parse_flag('e', 'E', EMUTRAMP);
   parse_flag('r', 'R', RANDMMAP);
   parse_flag('x', 'X', RANDEXEC);

#undef parse_flag

    case 'v':
      state->view = 1;
      break;

    case 'q':
      state->quiet = 1;
      break;

    case 'Q':
      state->shortonly = 1;
      break;

    case 'z':
      state->flags_on = 0U;
      state->flags_off = PF_PAX_MASK;
      break;
    }
  }
}

int main(int argc, char * argv[])
{
  struct pax_state state = {
    .argv = argv,
    .flags_on = 0U,
    .flags_off = 0U,
    .files = 0U,
    .quiet = 0,
    .shortonly = 0,
    .view = 0,
    .ops = NULL,
  };

  if (3 > argc)
    usage();

  /* parse arguments */
  if (0 > pax_parse_args(argc, &state))
    return -1;

  if (state.view)
    banner();

  /* process files */
  if (0 > pax_process_files(&state))
    return -2;

  return 0;
}