// we need this so sched.h exports unshare and CLONE_*
#define _GNU_SOURCE
#include <sched.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <signal.h>

pid_t pid_child;

void drop_root(void) {
  uid_t uid = getuid();
  // Drop root privileges
  if (setresuid(-1,uid,uid) == -1)
    {
      int err = errno;
      printf("Failed to drop root privileges with setresuid (%d)\n", err);
      exit(err);
    }

  gid_t gid = getgid();
  if (setresgid(-1,gid,gid) == -1)
    {
      int err = errno;
      printf("Failed to drop root privileges with setegid (%d)\n", err);
      exit(err);
    }

  // sanity check
  if(seteuid(0) != -1) {
    printf("Sanity check failed. I was able to regain root.\n");
    exit(1);
  }
}

void forward_signal(int sig)
{
  if(kill(pid_child, sig) == -1) {
    printf("Unable to forward signal %d to child\n", sig);
    if(sig == SIGTERM)
      exit(1);
  }
}

char** argdup(int argc, const char** argv)
{
  char** newargs = malloc(sizeof(char*) * (argc+1));
  for(size_t i = 0; i < argc; i++)
    {
      newargs[i] = strdup(argv[i]);
    }
  newargs[argc] = NULL;
  return newargs;
}

int main(int argc, const char** argv)
{
  if(argc == 1) {
    printf("Usage: pidjail PROGRAM ARGUMENTS...\n"
           "Run command within its own pid namespace. Integrated init process.\n");
    return 0;
  }

  // next fork shall be in a new pid namespace
  if (unshare(CLONE_NEWPID) != 0)
    {
      int err = errno;
      printf("Failed to unshare pid namespace (%d)\n", err);
      return err;
    }

  pid_t pid = fork();

  if (pid == -1)
    {
      int err = errno;
      printf("Failed to fork (%d)\n", err);
      return err;
    }

  // Drop root privileges, we only needed those for the unshare call and fork above.
  drop_root();

  if (pid != 0)
    {

      // Setup signal handler to forward SIGTERM
      pid_child = pid;
      if(signal(SIGTERM, forward_signal) == SIG_ERR) {
        printf("Unable to setup signal handler in head\n");
      }
      // parent waits for child then exits
      int status;
      if(waitpid(pid, &status, 0) == -1)
        {
          int err = errno;
          printf("Failed to wait (%d)\n", err);
          return err;
        }

      return WEXITSTATUS(status);
    }
  else
    {
      // Child should be in new pid namespace and
      // functions as the the init process
      // it needs to fork again then wait for any child.
      // if the forked child exits then exit.

      pid = fork();
      if (pid != 0)
        {
          // Init process wait for anything and exit if first child exits.
          pid_t first_child = pid;
          pid_t exited_child;
          int child_status;
          int err;

          // Setup forward for SIGTERM
          pid_child = first_child;
          if(signal(SIGTERM, forward_signal) == SIG_ERR) {
            printf("Unable to setup signal forward in init. Aborting.\n");
            return 1;
          }

          do {
            exited_child = wait(&child_status);
            err = errno;
          } while(exited_child != first_child && exited_child != -1);

          if (exited_child == -1)
            {
              return err;
            }
          else
            {
              int exit_code = WEXITSTATUS(child_status);
              return exit_code;
            }
        }
      else
        {
          // First child of init process. do exec here
          // use cli arguments for subprocess. skip 0 as it's our programs name.

          char** newargs = argdup(argc-1, &argv[1]);

          if (execvp(newargs[0], newargs) == -1)
            {
              printf("Failed to exec (%d)\n", errno);
              return errno;
            }
        }
    }
}