#include <CLI/CLI.hpp>
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <linux/landlock.h>
#include <linux/prctl.h>
#include <linux/seccomp.h>
#include <sys/syscall.h>
#include <string>
#include <string_view>
#include <unistd.h>
#include <vector>

#ifdef SUPPORT_SECCOMP
#include <seccomp.h>
#include <sys/prctl.h>
#endif

#ifndef landlock_create_ruleset
static inline int
landlock_create_ruleset (const struct landlock_ruleset_attr *const attr,
                         const size_t size, const uint32_t flags)
{
  return syscall (__NR_landlock_create_ruleset, attr, size, flags);
}
#endif


#ifndef landlock_create_ruleset
static inline int
landlock_add_rule (const int ruleset_fd,
                   const enum landlock_rule_type rule_type,
                   const void *const rule_attr,
                   const uint32_t flags)
{
  return syscall (__NR_landlock_add_rule, ruleset_fd,
                  rule_type, rule_attr, flags);
}
#endif

#ifndef landlock_create_ruleset
static inline int
landlock_restrict_self (const int ruleset_fd,
                        const uint32_t flags)
{
  return syscall (__NR_landlock_restrict_self, ruleset_fd, flags);
}
#endif

enum class SyscallFilterMode : uint8_t
{
  None = 0,
  Whitelist,
  Blacklist
};

struct Options
{
  std::vector<std::string> fs_ro;
  std::vector<std::string> fs_rw;
  std::vector<int> tcp_bind;
  std::vector<int> tcp_connect;
  std::vector<std::string> scopes;
#ifdef SUPPORT_SECCOMP
  std::vector<std::string> syscalls;
  SyscallFilterMode syscall_filter_mode = SyscallFilterMode::None;
#endif
  /* Includes program arguments.  */
  std::vector<std::string> command;
};

static bool
init_seccomp (
#ifndef SUPPORT_SECCOMP
              [[maybe_unused]]
#endif
              const Options &opts)
{
#ifdef SUPPORT_SECCOMP
  using std::cerr;
  scmp_filter_ctx ctx;

#define SECCOMP_INIT_FINAL_RULESET(VIOLATE_ACTION, RULESET_ACTION)                      \
  do                                                                                    \
    {                                                                                   \
      ctx = seccomp_init (VIOLATE_ACTION);                                              \
      if (ctx == nullptr)                                                               \
        {                                                                               \
          cerr << "Failed to initialize seccomp context, bailing.\n";                   \
          return false;                                                                 \
        }                                                                               \
      for (const auto &syscall : opts.syscalls)                                         \
        {                                                                               \
          const int resolved_syscall =                                                  \
            seccomp_syscall_resolve_name(syscall.c_str ());                             \
          if (resolved_syscall == __NR_SCMP_ERROR)                                      \
            {                                                                           \
              cerr << "Failed to resolve syscall '" << syscall << "', bailing.\n";      \
              seccomp_release (ctx);                                                    \
              return false;                                                             \
            }                                                                           \
          if (seccomp_rule_add_exact (ctx, RULESET_ACTION, resolved_syscall, 0) > 0)    \
            {                                                                           \
              cerr << "Failed to add seccomp rule " #RULESET_ACTION " for syscall '"    \
                   << syscall << "', bailing.\n";                                       \
              seccomp_release (ctx);                                                    \
              return false;                                                             \
            }                                                                           \
        }                                                                               \
    } while (0);
  
  switch (opts.syscall_filter_mode)
    {
    case SyscallFilterMode::Whitelist:
      SECCOMP_INIT_FINAL_RULESET (SCMP_ACT_KILL, SCMP_ACT_ALLOW)
      break;
    case SyscallFilterMode::Blacklist:
      SECCOMP_INIT_FINAL_RULESET (SCMP_ACT_ALLOW, SCMP_ACT_KILL)
      break;
    case SyscallFilterMode::None:
      return true;
    default:
      /* Unreachable.  */
      __builtin_unreachable ();
    }
#undef SECCOMP_INIT_FINAL_RULESET

  if (seccomp_load (ctx) < 0)
    {
      cerr << "Failed to load seccomp ruleset, bailing.\n";
      seccomp_release (ctx);
      return false;
    }

  seccomp_release (ctx);
  
  /* No new privs for the calling thread, preserved across execve.  */
  if (prctl (PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) < 0)
    {
      cerr << "Failed to restrict calling thread to no new privileges: "
           << strerror (errno) << '\n';
      return false;
    }
#endif
  
  return true;
}

int
main (int argc, char **argv)
{
  using std::cout, std::cerr;
  /* Argument parsing.  */
  std::string_view progname = argv[0];
  CLI::App app{"Initalise and launch a restricted program environment"};
  Options opts;

  app.set_help_flag ("--help", "Display this help and exit");
  app.add_option ("--fs-ro", opts.fs_ro, "Add PATH to read-only allowed paths")
    ->type_name ("PATH");
  app.add_option ("--fs-rw", opts.fs_rw, "Add PATH to read-write allowed paths")
    ->type_name ("PATH");
  app.add_option ("--tcp-bind", opts.tcp_bind, "Allow binding to PORT")
    ->type_name ("PORT");
  app.add_option ("--tcp-connect", opts.tcp_connect, "Allow connecting to PORT")
    ->type_name ("PORT");
  app.add_option ("--scope", opts.scopes)
    ->type_name ("ACTION")
    ->check (CLI::IsMember ({"abstract-socket", "send-signal", "a", "s"}))
    ->transform (CLI::CheckedTransformer ({
          {"a", "abstract-socket"},
          {"s", "send-signal"}}))
    ->description (R"(Deny ACTION outside of the landlock domain
  - abstract-socket: Restrict opening abstract unix sockets
  - send-signal: Restrict sending signals)");
#ifdef SUPPORT_SECCOMP
  auto syscall_whitelist = app.add_option ("--syscall-whitelist", opts.syscalls,
                                           "Whitelist syscall")
    ->type_name ("SYSCALL");
  auto syscall_blacklist = app.add_option ("--syscall-blacklist", opts.syscalls,
                                           "Blacklist syscall")
    ->type_name ("SYSCALL");

  syscall_whitelist->excludes (syscall_blacklist);
  syscall_blacklist->excludes (syscall_whitelist);
#endif
  app.add_option ("COMMAND", opts.command, "Command to run (with arguments)")
    ->required ()
    ->allow_extra_args ();
  
  CLI11_PARSE (app, argc, argv);
  
  /* Initalise ruleset.  */
  int abi = landlock_create_ruleset (nullptr, 0,
                                     LANDLOCK_CREATE_RULESET_VERSION);
  if (abi < 0)
    {
      switch (errno)
        {
        case ENOSYS:
          /* Landlock not supported. */
          cerr << "The running kernel does not support Landlock, bailing.\n";
          break;
        case EOPNOTSUPP:
          /* Landlock disabled. */
          cerr << "The running kernel has Landlock supported disabled, bailing.\n";
          break;
        }
      return EXIT_FAILURE;
    }

#ifdef SUPPORT_SECCOMP
  if (syscall_whitelist->count () > 0)
    opts.syscall_filter_mode = SyscallFilterMode::Whitelist;
  else if (syscall_blacklist->count () > 0)
    opts.syscall_filter_mode = SyscallFilterMode::Blacklist;
#endif

  if (!init_seccomp(opts))
    return EXIT_FAILURE;


  std::vector<char *> args;
  for (const auto &arg : opts.command)
    args.push_back (const_cast<char *>(arg.c_str ()));
  args.push_back (nullptr);

  execvp(args.at (0), args.data ());
  /* If the command launched succesfully, this entire section should
     become unreachable.  */
  cerr << opts.command.at (0) << ": " << strerror (errno) << '\n';
  return EXIT_FAILURE;
}