/* See LICENSE for license details. */
#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <linux/landlock.h>
#include <linux/prctl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/prctl.h>
#include <sys/syscall.h>
#include <unistd.h>

static int parse_path_access(const char *str);
static int parse_port_access(const char *str);

static int  landlock_init       (char *argv0, struct landlock_ruleset_attr *attr);
static void landlock_add_rule   (char **argv, int ruleset_fd);

static char *argv0;

#define LANDLOCK_ACCESS_FS_ALL \
		(LANDLOCK_ACCESS_FS_EXECUTE \
		| LANDLOCK_ACCESS_FS_IOCTL_DEV \
		| LANDLOCK_ACCESS_FS_MAKE_BLOCK \
		| LANDLOCK_ACCESS_FS_MAKE_CHAR \
		| LANDLOCK_ACCESS_FS_MAKE_DIR \
		| LANDLOCK_ACCESS_FS_MAKE_FIFO \
		| LANDLOCK_ACCESS_FS_MAKE_REG \
		| LANDLOCK_ACCESS_FS_MAKE_SOCK \
		| LANDLOCK_ACCESS_FS_MAKE_SYM \
		| LANDLOCK_ACCESS_FS_READ_DIR \
		| LANDLOCK_ACCESS_FS_READ_FILE \
		| LANDLOCK_ACCESS_FS_REFER \
		| LANDLOCK_ACCESS_FS_REMOVE_DIR \
		| LANDLOCK_ACCESS_FS_REMOVE_FILE \
		| LANDLOCK_ACCESS_FS_TRUNCATE \
		| LANDLOCK_ACCESS_FS_WRITE_FILE)

#define LANDLOCK_ACCESS_NET_ALL (LANDLOCK_ACCESS_NET_BIND_TCP | LANDLOCK_ACCESS_NET_CONNECT_TCP)

#ifndef RUN_LANDLOCK_VERSION
#define RUN_LANDLOCK_VERSION "devel"
#endif

int main(int argc, char **argv)
{
	int ruleset = 0, opt;
	struct landlock_ruleset_attr attr = {0};

	attr.handled_access_fs  = LANDLOCK_ACCESS_FS_ALL;
	attr.handled_access_net = LANDLOCK_ACCESS_NET_ALL;

	argv0 = argv[0];
	opterr = 0;

	// First argument parse: denied accesss
	while ((opt = getopt(argc, argv, "+hrv")) != -1) {
		switch (opt) {
			case 'r':
				if (ruleset == 0) {
					ruleset = landlock_init(argv0, &attr);
				}

				landlock_add_rule(argv, ruleset);
				break;

			case 'h':
				printf("Usage: %s [[-r type subject actions] ...] command [args ...]\n", argv0);
				return EXIT_SUCCESS;

			case 'v':
				printf("run_landlock " RUN_LANDLOCK_VERSION "\n");
				return EXIT_SUCCESS;

			default:
				fprintf(stderr, "%s: invalid option -- '%c'\n", argv0, optopt);
				return EXIT_FAILURE;
		}
	}

	if (ruleset == 0) {
		ruleset = landlock_init(argv0, &attr);
	}

	if (ruleset != -1) {
		prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0);

		if (syscall(SYS_landlock_restrict_self, ruleset, 0) == -1) {
			fprintf(stderr, "%s: failed to apply landlock (%s)\n", argv0, strerror(errno));
			return EXIT_FAILURE;
		}
	}

	if (argv[optind] == NULL) {
		fprintf(stderr, "%s: missing command argument\n", argv0);
		return EXIT_FAILURE;
	}

	execvp(argv[optind], &argv[optind]);
	fprintf(stderr, "%s: failed to exec '%s' (%s)\n", argv0, argv[optind], strerror(errno));
	return EXIT_FAILURE;
}


#define map(name, access) if (strcmp(name, str) == 0) return access;

static int parse_path_access(const char *str)
{
	map("fs_execute",      LANDLOCK_ACCESS_FS_EXECUTE)
	map("fs_ioctl_dev",    LANDLOCK_ACCESS_FS_IOCTL_DEV)
	map("fs_make_block",   LANDLOCK_ACCESS_FS_MAKE_BLOCK)
	map("fs_make_char",    LANDLOCK_ACCESS_FS_MAKE_CHAR)
	map("fs_make_dir",     LANDLOCK_ACCESS_FS_MAKE_DIR)
	map("fs_make_fifo",    LANDLOCK_ACCESS_FS_MAKE_FIFO)
	map("fs_make_reg",     LANDLOCK_ACCESS_FS_MAKE_REG)
	map("fs_make_sock",    LANDLOCK_ACCESS_FS_MAKE_SOCK)
	map("fs_make_sym",     LANDLOCK_ACCESS_FS_MAKE_SYM)
	map("fs_read_dir",     LANDLOCK_ACCESS_FS_READ_DIR)
	map("fs_read_file",    LANDLOCK_ACCESS_FS_READ_FILE)
	map("fs_refer",        LANDLOCK_ACCESS_FS_REFER)
	map("fs_remove_dir",   LANDLOCK_ACCESS_FS_REMOVE_DIR)
	map("fs_remove_file",  LANDLOCK_ACCESS_FS_REMOVE_FILE)
	map("fs_truncate",     LANDLOCK_ACCESS_FS_TRUNCATE)
	map("fs_write_file",   LANDLOCK_ACCESS_FS_WRITE_FILE)
	map("fs_all",          LANDLOCK_ACCESS_FS_ALL)
	return -1;
}


static int parse_port_access(const char *str)
{
	map("net_bind_tcp",    LANDLOCK_ACCESS_NET_BIND_TCP)
	map("net_connect_tcp", LANDLOCK_ACCESS_NET_CONNECT_TCP)
	map("net_all",         LANDLOCK_ACCESS_NET_ALL)
	return -1;
}

enum landlock_rule_type parse_rule_type(const char *str)
{
	map("path", LANDLOCK_RULE_PATH_BENEATH);
	map("port", LANDLOCK_RULE_NET_PORT);
	return -1;
}

union landlock_rule_attr {
	struct landlock_path_beneath_attr path_beneath_attr;
	struct landlock_net_port_attr     net_port_attr;
};

static void landlock_rule_attr_add_path_beneath_access(union landlock_rule_attr *rule_attr, int access)
{
	rule_attr->path_beneath_attr.allowed_access |= access;
}

static void landlock_rule_attr_add_net_port_access(union landlock_rule_attr *rule_attr, int access)
{
	rule_attr->net_port_attr.allowed_access |= access;
}

static int landlock_init(char *argv0, struct landlock_ruleset_attr *attr)
{
	int ruleset;

	errno = 0;

	ruleset = syscall(SYS_landlock_create_ruleset, attr, sizeof(struct landlock_ruleset_attr), 0);

	if (ruleset != -1)
		return ruleset;

	switch (errno) {
		case 0:
			return ruleset;

		case EOPNOTSUPP:
			fprintf(stderr, "%s: landlock disabled, running without\n", argv0);
			return -1;

		case ENOSYS:
			fprintf(stderr, "%s: landlock not supported, running without\n", argv0);
			return -1;

		case ENOMSG:
			fprintf(stderr, "%s: no access specified\n", argv0);
			exit(EXIT_FAILURE);

		default:
			fprintf(stderr, "%s: failed to initialize landlock (%s)\n", argv0, strerror(errno));
			exit(EXIT_FAILURE);
	}
}

static int get_port(const char *str)
{
	char *endptr;
	long port;

	port = strtol(str, &endptr, 10);

	if (*endptr != '\0' || port < 0) {
		fprintf(stderr, "%s: invalid port: %s\n", argv0, str);
		exit(EXIT_FAILURE);
	}

	return port;
}

static int get_path_fd(const char *path)
{
	int fd;

	if ((fd = open(path, O_PATH)) == -1) {
		fprintf(stderr, "%s: failed to open '%s' (%s)\n", argv0, path, strerror(errno));
		exit(EXIT_FAILURE);
	}

	return fd;
}

static void landlock_add_rule(char **argv, int ruleset_fd)
{
	char *arg_type;
	char *arg_subject;
	char *arg_actions;
	char *access_name;
	enum landlock_rule_type rule_type;
	int  access;
	int  (*parse_access)(const char *);
	void (*allow_access)(union landlock_rule_attr*, int);
	union landlock_rule_attr rule_attr = {0};

	if ((arg_type = argv[optind ++]) == NULL) {
		fprintf(stderr, "%s: missing 'type' argument", argv0);
		exit(EXIT_FAILURE);
	}

	if ((arg_subject = argv[optind ++]) == NULL) {
		fprintf(stderr, "%s: missing 'subject' argument", argv0);
		exit(EXIT_FAILURE);
	}

	if ((arg_actions = argv[optind ++]) == NULL) {
		fprintf(stderr, "%s: missing 'actions' argument", argv0);
		exit(EXIT_FAILURE);
	}

	switch ((rule_type = parse_rule_type(arg_type))) {
		case LANDLOCK_RULE_PATH_BENEATH:
			parse_access = parse_path_access;
			allow_access = landlock_rule_attr_add_path_beneath_access;
			rule_attr.path_beneath_attr.parent_fd = get_path_fd(arg_subject);
			break;

		case LANDLOCK_RULE_NET_PORT:
			parse_access = parse_port_access;
			allow_access = landlock_rule_attr_add_net_port_access;
			rule_attr.net_port_attr.port = get_port(arg_subject);
			break;

		default:
			fprintf(stderr, "%s: invalid 'type' argument: '%s'\n", argv0, arg_type);
			exit(EXIT_FAILURE);
	}


	access_name = strtok(arg_actions, ",");
	while (access_name) {
		if ((access = parse_access(access_name)) == -1) {
			fprintf(stderr, "run_landlock: unexpected access argument on object '%s': '%s'\n", access_name, arg_subject);
			exit(EXIT_FAILURE);
		}

		allow_access(&rule_attr, access);

		access_name = strtok(NULL, ",");
	}

	if (ruleset_fd != -1) {
		if (syscall(SYS_landlock_add_rule, ruleset_fd, rule_type, &rule_attr, 0) == -1) {

			if (errno == EINVAL)
				fprintf(stderr, "%s: can't apply one of these access to '%s': %s\n", argv0, arg_subject, arg_actions);
			else
				fprintf(stderr, "%s: failed to add landlock rule (%s)\n", argv0, strerror(errno));
			exit(EXIT_FAILURE);
		}
	}
}