/*
 * nsntrace is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by the
 * Free Software Foundation; either version 2 of the License, or (at your
 * option) any later version.
 *
 * nsntrace is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
 * for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with nsntraces; if not, see <https://www.gnu.org/licenses/>.
 *
 */
#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <ftw.h>
#include <grp.h>
#include <getopt.h>
#include <linux/limits.h>
#include <sched.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>
#include <pwd.h>

#include "capture.h"
#include "net.h"

/*
 * This application attempts to trace the network traffic of a single process.
 * It does so by creating a new network namespace. In Linux a network namespace
 * partitions the use of the network effectively virtualizing the network,
 * devices, addresses, ports, routes firewall rules, etc into separate boxes.
 *
 * So what this application does is use the clone syscall to create a new
 * network namespace (CLONE_NEWNET) and from that namespace launch the
 * requested process as well as start a trace. This will ensure that all
 * the packets we trace come from the process.
 *
 * The problem we are left with is that the process is isolated in the
 * namespace and cannot reach any other network. We get around that by
 * creating virtual network interfaces. We keep one of them in the
 * root network namespace and but the other one in the newly created one where
 * our tracing takes place. We set the root namespaced one as the default gw
 * of the trace namespaced virtual device.
 *
 * And then to make sure we can reach our indented net, we use ip
 * tables and NAT to forward all traffic from the virtual device to our
 * default network interface.
 *
 * This will allow us to capture the packets from a single process while
 * it is communicating with our default network. A limitation is that our
 * ip address will be the NAT one of the virtual device.
 *
 * Another limitation is, that since we are using iptables and since
 * we are tracing raw sockets. This application needs to be run as root.
 *
 */

#define STACK_SIZE (1024 * 64) /* 64 kB stack */
#define DEFAULT_OUTFILE "nsntrace.pcap"
#define NETNS_RUN_DIR "/run/netns"
#define PROC_SELF_NETNS_PATH "/proc/self/ns/net"

static FILE *where;

struct nsntrace_options {
	char *outfile;
	char *device;
	int   use_public_dns;
	char *user;
	char *filter;
	char * const *args;
};

struct nsntrace_common {
	struct nsntrace_options *options;
	struct nsntrace_if_info if_info;
};

#define PUBLIC_DNS 1000
static const char *short_opt = "+o:d:u:f:h";
static struct option long_opt[] = {
	{ "outfile", required_argument, NULL, 'o' },
	{ "device",  required_argument, NULL, 'd' },
	{ "use-public-dns", no_argument, NULL, PUBLIC_DNS },
	{ "user",    required_argument, NULL, 'u' },
	{ "filter",  required_argument, NULL, 'f' },
	{ "help",    required_argument, NULL, 'h' },
	{ NULL,	     0,			NULL,  0 }
};

static char child_stack[STACK_SIZE];
static pid_t child_pid;

/*
 * We will attempt to catch the signals that can make us exit since
 * we want to remove the temporary network configurations when we leave
 *
 * This is hard. Which signals should we catch? This is what we have
 * now:
 *
 * SIGHUP:  Disconnects a process from the parent process.
 * SIGINT:  Same as pressing ctrl-c.
 * SIGABRT: The abort signal, most often used on one self.
 * SIGTERM: A request to a process to stop running (kill $pid),
 * SIGQUIT: Similar to SIGINT but also a request to core dump.
 * SIGSEGV: Generally sent to process by the kernel when the
 *	    process is accessing memory incorrectly. There are
 *	    no guarantees for what we can do when this happens.
 *	    But let's try to clean up!
 */
const int nsntrace_signals[] = {
	SIGHUP,
	SIGINT,
	SIGABRT,
	SIGTERM,
	SIGQUIT,
	SIGSEGV,
};

static void
_nsntrace_handle_signals(void (*handler)(int))
{
	struct sigaction action = { 0 };
	int s;

	action.sa_handler = handler;
	for (s = 0; s < sizeof(nsntrace_signals) / sizeof(int); s++) {
		sigaction(nsntrace_signals[s], &action, NULL);
	}
}

static int _nsntrace_unlink_cb(const char *fpath,
			       const struct stat *sb,
			       int typeflag,
			       struct FTW *ftwbuf)
{
    int rv = remove(fpath);
    if (rv)
        perror(fpath);
    return rv;
}

static void
_nsntrace_remove_rundir(pid_t pid)
{
	char path[PATH_MAX] = { 0, };

	snprintf(path, PATH_MAX, "%s/%d", NSNTRACE_RUN_DIR, getppid());
	nftw(path, _nsntrace_unlink_cb, 64, FTW_DEPTH | FTW_PHYS);
}

static void _nsntrace_remove_ns()
{
	char netns_path[PATH_MAX];

	snprintf(netns_path, sizeof(netns_path), "%s/nsntrace-%d", NETNS_RUN_DIR, getpid());
	umount2(netns_path, MNT_DETACH);
	if (unlink(netns_path) < 0) {
		perror("unlink");
	}

	/* clean up run-time files */
	_nsntrace_remove_rundir(getppid());
}

static void
_nsntrace_cleanup_ns()
{
	kill(child_pid, SIGTERM);
	waitpid(child_pid, NULL, 0);

	_nsntrace_remove_ns();

	fprintf(where, "Finished capturing %lu packets.\n",
		nsntrace_capture_packet_count());
	nsntrace_capture_flush();
}

static void
_nsntrace_cleanup_ns_signal_callback(int sig)
{
	_nsntrace_cleanup_ns();
	exit(EXIT_SUCCESS);
}

static void
_nsntrace_cleanup(int sig) {
	/*
	 * Make sure we do not just die when we receive our
	 * terminating signals. We need to clean up after
	 * our children.
	 */
	wait(NULL);
}

/*
 * This will give the network namespace a name in the eyes of tools like
 * ip-netns. We need to make sure we clean up after us, since a bind mount
 * of /proc/<pid>/ns/net will keep the network namespace alive after the
 * process exits.
 */
static void
_nsntrace_set_name()
{
	char netns_path[PATH_MAX];
	int fd;
	int ret = 0;

	/* Create the base netns directory if it doesn't exist */
	if ((ret = mkdir(NETNS_RUN_DIR, S_IRWXU|S_IRGRP|S_IXGRP|S_IROTH|S_IXOTH)) < 0) {
		if (errno != EEXIST) {
			goto out;
		}
	}

	snprintf(netns_path, PATH_MAX, "%s/nsntrace-%d", NETNS_RUN_DIR, getpid());
	fd = open(netns_path, O_RDONLY|O_CREAT|O_EXCL, 0);
	if (fd < 0) {
		ret = fd;
		goto out;
	}
	close(fd);

	ret = mount(PROC_SELF_NETNS_PATH, netns_path, "none", MS_BIND, NULL);
out:
	if (ret < 0) {
		fprintf(stderr, "failed to set namespace name\n");
	}
}

static void
_nsntrace_start_tracer(struct nsntrace_options *options,
					   struct nsntrace_if_info *info)
{
	int ret;
	/*
	 * If outfile is a lone dash ("-"), the user wants us to output to STDOUT,
	 * in that case we avoid printing anything else to STDOUT and use STDERR
	 * instead.
	 */
	FILE *fp;
	if (options->outfile[0] == '-' && options->outfile[1] == '\0') {
		where = stderr;
		fp = stdout;
	} else {
		where = stdout;
		fp = fopen(options->outfile, "w");
		if (!fp) {
			perror("fopen");
			exit(EXIT_FAILURE);
		}
	}

	fprintf(where, "Starting network trace of '%s' on interface %s.\n"
		"Your IP address in this trace is %s.\n"
		"Use ctrl-c to end at any time.\n\n",
		options->args[0], options->device, info->ns_ip);
	ret = nsntrace_capture_start(info->ns_if, options->filter, fp);
	if (ret != 0) {
		exit(ret);
	}
}

static void
_nsntrace_start_tracee(struct nsntrace_options *options)
{
	uid_t uid;
	gid_t gid;

	/*
	 * The getpwnam() function will return a pointer to fields in
	 * the password database that matches the username given. This can
	 * be the local password file or NIS, LDAP etc.
	 *
	 * If the user supplied a username to run the traced process as, then
	 * we look up the uid and gid and set those ids on our newly forked
	 * process. Along with other user specific goodies.
	 */
	if (options->user) {
		struct passwd* pwd;

		if (sscanf(options->user, "%i", &uid) == 1) {
			pwd = getpwuid(uid);
		} else {
			pwd = getpwnam(options->user);
		}
		if (!pwd) {
			fprintf(stderr,"Cannot find user '%s'\n",
				options->user);
			exit(EXIT_FAILURE);
		}
		uid = pwd->pw_uid;
		gid = pwd->pw_gid;
		/*
		 * Set common environment variables that application looks
		 * at and expects to reflect the user.
		 *
		 * Any more that should be set?
		 */
		setenv("HOME", pwd->pw_dir, 1);
		setenv("USER", pwd->pw_name, 1);
		setenv("USERNAME", pwd->pw_name, 1);
		/*
		 * The initgroups() function needs to be run before we
		 * lose privileges (setuid).
		 */
		if (initgroups(options->user, gid) < 0) {
			fprintf(stderr, "Failed to initialize groups\n");
			exit(EXIT_FAILURE);
		}
	} else {
		uid = getuid();
		gid = getgid();
	}

	if (setgid(gid) < 0) {
		fprintf(stderr, "Unable to set process group ID\n");
		exit(EXIT_FAILURE);
	}

	if (setuid(uid) < 0) {
		fprintf(stderr, "Unable to set process user ID\n");
		exit(EXIT_FAILURE);
	}

	/*
	 * If outfile is a lone dash we output only PCAP data on STDOUT.
	 * We use the dup2 syscall to redirect the tracee output to STDERR.
	 */
	if (options->outfile[0] == '-' && options->outfile[1] == '\0') {
		dup2(STDERR_FILENO, STDOUT_FILENO);
	}

	/* launch the application to trace */
	if (execvp(options->args[0], options->args) < 0) {
		fprintf(stderr, "Unable to start '%s'\n", options->args[0]);
		exit(EXIT_FAILURE);
	}
}

static int
netns_main(void *arg) {
	int status;
	int ret = EXIT_SUCCESS;
	struct nsntrace_common *common = (struct nsntrace_common *) arg;
	struct nsntrace_options *options = common->options;

	if (nsntrace_net_ns_init(options->use_public_dns, &common->if_info) < 0) {
		fprintf(stderr, "failed to setup network environment\n");
		return EXIT_FAILURE;
	}

	_nsntrace_handle_signals(_nsntrace_cleanup_ns_signal_callback);
	_nsntrace_start_tracer(common->options, &common->if_info);

	_nsntrace_set_name();

	child_pid = fork();
	if (child_pid < 0) {
		return EXIT_FAILURE;
	} else if (child_pid > 0) { /* parent - tracer */
		struct timespec timeout = { 0 };

		waitpid(child_pid, &status, 0);
		if (WIFEXITED(status)) {
			ret = WEXITSTATUS(status);
		} else {
			ret = EXIT_FAILURE;
		}

		/*
		 * Sleep so that all packets can be processed.
		 * Do not bother with error handling for this,
		 * it is best effort to catch all packets.
		 */
		timeout.tv_sec = 2; /* 2 seconds timeout */
		nanosleep(&timeout, NULL);

		/* the tracee exited, we waited, stop capture */
		nsntrace_capture_stop();

		/* exited capture loop, clean up */
		_nsntrace_cleanup_ns();
		exit(ret);
	} else { /* child - tracee */
		_nsntrace_start_tracee(common->options);
	}

	return ret;
}

static void
_nsntrace_usage()
{
	fprintf(stderr, "usage: nsntrace [options] program [arguments]\n"
	       "Perform network trace of a single process by using "
	       "network namespaces.\n\n"
	       "Options:\n"
	       "-o file\t\t\tsend trace output to file (default nsntrace.pcap), use '-' for stdout\n"
	       "-d device\t\tthe network device to trace\n"
	       "-f filter\t\tan optional capture filter\n"
	       "-u username\t\trun PROG as username/uid\n"
	       "--use-public-dns\toverride resolv.conf to use public nameservers from\n"
	       "\t\t\tQuad9, Cloudflare, Google and OpenDNS\n");
}

static void
_nsntrace_parse_options(struct nsntrace_options *options,
			int argc, char **argv)
{
	int c;

	opterr = 0;
	while ((c = getopt_long(argc, argv, short_opt, long_opt, NULL)) > 0) {
		switch(c) {
		case -1:
		case 0:
			break;

		case 'o':
			options->outfile = strdup(optarg);
			break;

		case 'd':
			options->device = strdup(optarg);
			break;

		case 'u':
			options->user = strdup(optarg);
			break;

		case 'f':
			options->filter = strdup(optarg);
			break;

		case PUBLIC_DNS:
			options->use_public_dns = 1;
			break;

		case 'h':
			_nsntrace_usage();
			exit(EXIT_SUCCESS);
			break;

		default:
			fprintf(stderr, "Invalid option '%c'\n", c);
			_nsntrace_usage();
			exit(EXIT_FAILURE);
		}
	}

	if (!options->device) {
		options->device = nsntrace_capture_default_device();
	}
	if (!options->outfile) {
		options->outfile = strdup(DEFAULT_OUTFILE);
	}
	options->args = argv + optind; /* the arguments after options parsed */
	if (!options->args[0]) {
		_nsntrace_usage();
		exit(EXIT_FAILURE);
	}
}

/*
 * Create the pid based run dir for this nsntrace instance.
 * This is for storing run-time files and to be able to cleanup
 * easily without removing other nsntrace instances files.
 */
static void
_nsntrace_mkrundir()
{
	char path[PATH_MAX] = { 0, };

	if (mkdir(NSNTRACE_RUN_DIR, 0644) < 0) {
		if (errno != EEXIST) {
			perror("mkdir");
		}
	}
	snprintf(path, PATH_MAX, "%s/%d", NSNTRACE_RUN_DIR, getpid());
	if (mkdir(path, 0644) < 0) {
		if (errno != EEXIST) {
			perror("mkdir");
		}
	}
}

int
main(int argc, char **argv)
{
	struct nsntrace_options options = { 0 };
	struct nsntrace_common common = { &options,  };
	pid_t pid = 0;
	int status;
	int ret = EXIT_SUCCESS;

	_nsntrace_parse_options(&options, argc, argv);

	/* geteuid() returns the effective user ID, 0 if root */
	if (geteuid() != 0) {
		fprintf(stderr,
			"You need root privileges to run this application\n");
		exit(EXIT_FAILURE);
	}

	if (!nsntrace_net_ip_forward_enabled()) {
		fprintf(stderr, "Please enable IP forwarding:\n");
		if( NULL == getenv("SUDO_UID") ) {
			fprintf(stderr, "# sysctl net.ipv4.ip_forward=1\n");
		} else {
			fprintf(stderr, "$ sudo sysctl net.ipv4.ip_forward=1\n");
		}
		exit(EXIT_FAILURE);
	}

	_nsntrace_mkrundir();

	if (nsntrace_net_get_if_info(getpid(), &common.if_info) < 0) {
		ret = EXIT_FAILURE;
		goto out;
	}

	/* here we create a new process in a new network namespace */
	pid = clone(netns_main, child_stack + STACK_SIZE,
		    CLONE_NEWNET | SIGCHLD, &common);
	if (pid < 0) {
		perror("clone");
		ret = EXIT_FAILURE;
		goto out;
	}

	_nsntrace_handle_signals(_nsntrace_cleanup);

	if ((ret = nsntrace_net_init(pid, options.device, &common.if_info)) < 0 ||
	    (ret = nsntrace_capture_check_device(options.device))) {
		fprintf(stderr, "Failed to setup networking environment\n");
		kill(pid, SIGKILL);
		goto out;
	}

	/* wait here until our traced process exists or the user aborts */
	waitpid(pid, &status, 0);
	if (WIFEXITED(status)) {
		ret = WEXITSTATUS(status);
	} else {
		ret = EXIT_FAILURE;
	}

out:
	if (pid != 0) {
		nsntrace_net_deinit(pid, options.device, &common.if_info);
	}
	_nsntrace_remove_rundir(getpid());
	return ret;
}
