#include "LIEF/LIEF.h"

#include <err.h>
#include <stdio.h>
#include <errno.h>
#include <signal.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <seccomp.h>
#include <inttypes.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <sys/prctl.h>
#include <sys/syscall.h>

#define MAX_SIZE        0x1000
#define MARKER          "deadbeef"
#define MARKER_SIZE     (sizeof(MARKER) - 1)

static int readall(int fd, void *buf, size_t count)
{
  ssize_t n;
  size_t i;
  char *p;

  p = buf;
  i = count;
  while (i > 0) {
    n = read(fd, p, i);
    if (n == 0) {
      warnx("read failed");
      return -1;
    } else if (n == -1) {
      if (errno == EINTR)
        continue;
      warn("read failed");
      return -1;
    }
    i -= n;
    p += n;

    if (count - i >= MARKER_SIZE) {
      if (memcmp(p - MARKER_SIZE, MARKER, MARKER_SIZE) == 0) {
        printf("[*] received %lu bytes\n", count - i);
        break;
      }
    }
  }

  return count - i;
}

static int writeall(int fd, void *buf, size_t count)
{
  ssize_t n;
  size_t i;
  char *p;

  p = buf;
  i = count;
  while (i > 0) {
    n = write(fd, p, i);
    if (n == 0) {
      warnx("write failed");
      return -1;
    } else if (n == -1) {
      if (errno == EINTR)
        continue;
      warn("write failed");
      return -1;
    }
    i -= n;
    p += n;
  }

  return 0;
}

static int recvfile(int infd, char *template)
{
  char buf[MAX_SIZE];
  ssize_t count;
  int outfd;

  count = readall(infd, buf, sizeof(buf));
  if (count == -1)
    return -1;

  outfd = mkstemp(template);
  if (outfd == -1) {
    warn("mkstemp");
    return -1;
  }

  if (writeall(outfd, buf, count) != 0) {
    close(outfd);
    return -1;
  }

  if (fchmod(outfd, 0500) != 0) {
    warn("fchmod");
    close(outfd);
    return -1;
  }

  close(outfd);

  return 0;
}

static unsigned long get_mmap_min_addr(unsigned long *mmap_min_addr)
{
  FILE *fp;
  int ret;

  fp = fopen("/proc/sys/vm/mmap_min_addr", "r");
  if (fp == NULL) {
    warn("failed to open \"/proc/sys/vm/mmap_min_addr\"");
    return -1;
  }

  ret = (fscanf(fp, "%lu", mmap_min_addr) != 1);
  fclose(fp);

  return ret;
}

static int install_syscall_filter(unsigned long mmap_min_addr)
{
  int allowed_syscall[] = {
    SCMP_SYS(rt_sigreturn),
    SCMP_SYS(rt_sigaction),
    SCMP_SYS(rt_sigprocmask),
    SCMP_SYS(sigreturn),
    SCMP_SYS(exit_group),
    SCMP_SYS(exit),
    SCMP_SYS(brk),
    SCMP_SYS(access),
    SCMP_SYS(fstat),
    SCMP_SYS(write),
    SCMP_SYS(close),
    SCMP_SYS(mprotect),
    SCMP_SYS(arch_prctl),
    SCMP_SYS(munmap),
    SCMP_SYS(fstat),
    SCMP_SYS(readlink),
    SCMP_SYS(uname),
  };
  scmp_filter_ctx ctx;
  unsigned int i;
  int ret;

  ctx = seccomp_init(SCMP_ACT_KILL);
  if (ctx == NULL) {
    warn("seccomp_init");
    return -1;
  }

  for (i = 0; i < sizeof(allowed_syscall) / sizeof(int); i++) {
    if (seccomp_rule_add(ctx, SCMP_ACT_ALLOW, allowed_syscall[i], 0) != 0) {
      warn("seccomp_rule_add");
      ret = -1;
      goto out;
    }
  }

  /* prevent mmap to map mmap_min_addr */
  if (seccomp_rule_add(ctx, SCMP_ACT_ALLOW, SCMP_SYS(mmap), 1,
                       SCMP_A0(SCMP_CMP_GE, mmap_min_addr + PAGE_SIZE)) != 0) {
    warn("seccomp_rule_add");
    ret = -1;
    goto out;
  }

  /* first execve argument (filename) must be mapped at mmap_min_addr */
  if (seccomp_rule_add(ctx, SCMP_ACT_ALLOW, SCMP_SYS(execve), 1,
                       SCMP_A0(SCMP_CMP_EQ, mmap_min_addr)) != 0) {
    warn("seccomp_rule_add");
    ret = -1;
    goto out;
  }

  puts("[*] seccomp-bpf filters installed");

  ret = seccomp_load(ctx);
  if (ret != 0)
    warn("seccomp_load");

out:
  seccomp_release(ctx);
  return ret;
}

static char *setup_sandbox(unsigned long mmap_min_addr)
{
  int flags, prot;
  void *p;

  flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED;
  prot = PROT_READ | PROT_WRITE;
  p = mmap((void *)mmap_min_addr, PAGE_SIZE, prot, flags, -1, 0);
  if (p == MAP_FAILED) {
    warn("mmap");
    return NULL;
  }

  if (install_syscall_filter(mmap_min_addr) != 0) {
    if (munmap(p, PAGE_SIZE) != 0)
      warn("munmap");
    return NULL;
  }

  return (char *)p;
}

static void execute(char *filename, unsigned long mmap_min_addr)
{
  if (prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0) != 0)
    err(EXIT_FAILURE, "prctl(PR_SET_PDEATHSIG)");

  char *p = setup_sandbox(mmap_min_addr);
  if (p == NULL)
    exit(EXIT_FAILURE);

  strncpy(p, filename, PAGE_SIZE);
  p[PAGE_SIZE-1] = '\x00';

  printf("[*] let's execute \"%s\" in that sandbox\n", filename);
  if (execl(p, filename, NULL) != 0)
    err(EXIT_FAILURE, "execl of \"%s\" failed", p);

  /* never reached */
  exit(EXIT_SUCCESS);
}

static int elf_check(const char *filename, unsigned long mmap_min_addr)
{
  Elf_Binary_t *elf_binary = elf_parse(filename);
  if (elf_binary == NULL) {
    warnx("failed to parse ELF binary");
    return -1;
  }

  Elf_Header_t header = elf_binary->header;
  uint8_t *identity = header.identity;
  int ret = 0;

  if (identity[EI_CLASS] != ELFCLASS64) {
    warnx("invalid ELF class \"%s\"", ELF_CLASS_to_string(identity[EI_CLASS]));
    ret = -1;
    goto out;
  }

  mmap_min_addr += PAGE_SIZE;

  unsigned int i;
  Elf_Section_t** sections = elf_binary->sections;
  for (i = 0; i < header.numberof_sections; ++i) {
    Elf_Section_t* section = sections[i];
    if (section == NULL) {
      warnx("invalid section %d", i);
      ret = -1;
      goto out;
    }
    if (section->virtual_address != 0 &&
        section->virtual_address < mmap_min_addr + PAGE_SIZE) {
      warnx("invalid section \"%s\" (0x%lx)", section->name,
            section->virtual_address);
      ret = -1;
      goto out;
    }
  }

  Elf_Segment_t** segments = elf_binary->segments;
  for (i = 0; segments[i] != NULL; ++i) {
    Elf_Segment_t* segment = segments[i];
    if (segment->virtual_address != 0 &&
        segment->virtual_address < mmap_min_addr + PAGE_SIZE) {
      warnx("invalid segment (0x%lx)", segment->virtual_address);
      ret = -1;
      goto out;
    }
  }

out:
  elf_binary_destroy(elf_binary);
  return ret;
}

int main(void)
{
  unsigned long mmap_min_addr;
  char filename[PAGE_SIZE];
  int ret, status;
  pid_t pid;

  setbuf(stdin, NULL);
  setbuf(stdout, NULL);
  setbuf(stderr, NULL);

  system("/bin/ls -l .");

  if (get_mmap_min_addr(&mmap_min_addr) != 0)
    return EXIT_FAILURE;

  strncpy(filename, "/tmp/execve-sandbox-XXXXXX", sizeof(filename));

  puts("[*] waiting for an ELF binary...");
  if (recvfile(STDIN_FILENO, filename) != 0)
    return EXIT_FAILURE;

  if (elf_check(filename, mmap_min_addr) != 0) {
    ret = EXIT_FAILURE;
    goto out;
  }

  pid = fork();
  if (pid == -1) {
    warn("fork");
    ret = EXIT_FAILURE;
    goto out;
  }

  /* child */
  if (pid == 0)
    execute(filename, mmap_min_addr);

  /* parent */
  if (waitpid(pid, &status, 0) == -1) {
    warn("waitpid");
  } else {
    if (WIFSIGNALED(status) && WTERMSIG(status) == SIGSYS)
      puts("[*] \xc2\xaf\x5c\x5f\x28\xe3\x83\x84\x29\x5f\x2f\xc2\xaf");
  }

  ret = EXIT_SUCCESS;

out:
  if (unlink(filename) == -1)
    warn("unlink \"%s\" failed", filename);

  return ret;
}