#include "../module/module.h"
#include "../util.h"
#include "state.h"

#include <jail.h>
#include <sys/param.h>
#include <sys/jail.h>

#include <netgraph.h>
#include <netgraph/ng_message.h>
#include <netgraph/ng_eiface.h>

#include <sys/mount.h>
#include <sys/ioctl.h>
#include <fs/devfs/devfs.h>

char const *Usage = "usage: %1$s start        [name] [configuration]\n"
                    "       %1$s stop         [name] [configuration] [data]\n"
                    "       %1$s get-endpoint [name] [configuration] [data] [interface]\n"
                    "       %1$s cmd          [name] [configuration] [data] sh\n"
                    "       %1$s cmd          [name] [configuration] [data] wireshark [interface]\n"
                    "unsupported commands: mod\n";

static configuration Configuration;
static data          Data;
static bool          ConfigurationLoaded;
static bool          DataLoaded;

// HELPERS

static void
Free_all(void)
{
	if(ConfigurationLoaded == true) Free_configuration(&Configuration);
	if(DataLoaded == true) Free_data(&Data);
}

static void
Load_configuration(ucl_object_t *root)
{
	char *Error = Parse_configuration(&Configuration, root, "");
	if(Error != NULL)
	{
		fprintf(stderr, "%s: configuration error: %s\n", Arg0, Error);
		free(Error);
		Free_all();
		exit(-1);
	}
	ConfigurationLoaded = true;
}

static void
Load_data(ucl_object_t *root)
{
	char *Error = Parse_data(&Data, root, "");
	if(Error != NULL)
	{
		fprintf(stderr, "%s: data error: %s\n", Arg0, Error);
		free(Error);
		Free_all();
		exit(-1);
	}
	DataLoaded = true;
	// TODO: check that data matches configuration
}

static int
start_interface(interface *E, interface_configuration *C, fd Control, char const *JailName)
{
	E->Type = C->Type;
	switch(C->Type)
	{
		case interface_type_eiface:
		{
			if(Create_eiface(&E->ID, Control) == -1) return -1;
			char *InterfaceName = getifname(E->ID, Control);
			if(InterfaceName == NULL)
			{
				DestroyNetgraphNode(E->ID, Control);
				return -1;
			}
			// NOTE: This first renames the interface and only then vnets it.
			int Result =
			  command("ifconfig", "ifconfig", InterfaceName, "name", C->Name, "vnet", JailName, NULL);
			free(InterfaceName);
			return Result;
		}
		case interface_type_steal:
		{
			return command("ifconfig", "ifconfig", C->steal.Interface, "name", C->Name, "vnet", JailName, NULL);
		}
		case interface_type_COUNT:
		default: __builtin_unreachable();
	}
}

static int
stop_interface(interface *E, interface_configuration *C, fd Control, char const *JailName)
{
	switch(C->Type)
	{
		case interface_type_eiface: return DestroyNetgraphNode(E->ID, Control);
		case interface_type_steal:
		{
			return command("ifconfig",
			               "ifconfig",
			               C->Name,
			               "-vnet",
			               JailName,
			               "name",
			               C->steal.Interface,
			               NULL);
		}
		case interface_type_COUNT:
		default: __builtin_unreachable();
	}
}

static int
stop_interfaces(interfaces *E, interfaces_configuration *C, fd Control, char const *JailName)
{
	for(size_t i = 0; i < E->Count; ++i)
	{
		if(stop_interface(E->_ + i, C->_ + i, Control, JailName) == -1)
		{
			fprintf(stderr, "failed to stop interface '%s'\n", C->_[i].Name);
			return -1;
		}
	}
	return 0;
}

static int
start_interfaces(interfaces *E, interfaces_configuration *C, fd Control, char const *Name)
{
	Data.Interfaces._ = calloc(sizeof(interface), Configuration.Interfaces.Count);
	for(E->Count = 0; E->Count < C->Count; ++E->Count)
	{
		interface               *E_ = E->_ + E->Count;
		interface_configuration *C_ = C->_ + E->Count;
		if(start_interface(E_, C_, Control, Name) == -1) goto error;
	}
	return 0;
error:
	stop_interfaces(E, C, Control, Name);
	Free_interfaces(E);
	return -1;
}

static int
start_jail(jid_t *JID, char const *Name, char const *Path)
{
	*JID = jail_setv(JAIL_CREATE, "name", Name, "vnet", "new", "persist", "true", "path", Path, NULL);
	if(*JID == -1)
	{
		fprintf(stderr, "jail creation failed: %s\n", strerror(errno));
		return -1;
	}
	return 0;
}

static int
stop_jail(jid_t JID)
{
	if(jail_remove(JID) == -1)
	{
		fprintf(stderr, "jail removal failed: %s\n", strerror(errno));
		Free_all();
		return -1;
	}
	return 0;
}

NAMED_ENUM(mapping_type, unionfs_below, unionfs, devfs, nullfs);

static int
mapping_mount(mapping *M, char const *TargetPrefix, mapping_type Type)
{
	char Target[MAXPATHLEN];
	snprintf(Target, sizeof(Target), "%s/%s", TargetPrefix, M->Target);
	if(mkdir(Target, 755) == -1)
	{
		if(errno != EEXIST)
		{
			fprintf(stderr, "mkdir %s: %s\n", Target, strerror(errno));
			return -1;
		}
	}
	switch(Type)
	{
		case mapping_type_unionfs_below: return mount_unionfs(Target, M->Source, true);
		case mapping_type_unionfs: return mount_unionfs(Target, M->Source, false);
		case mapping_type_devfs: return mount_devfs(Target);
		case mapping_type_nullfs: return mount_nullfs(Target, M->Source);
		case mapping_type_COUNT:
		default: __builtin_unreachable();
	}
}

static int
mapping_umount(char const *TargetPrefix, char const *Target_)
{
	char Target[MAXPATHLEN];
	snprintf(Target, sizeof(Target), "%s/%s", TargetPrefix, Target_);
	// NOTE: unmount fails because some processes somehow still hold a reference to the jail FS
	while(unmount(Target, MNT_FORCE) == -1)
	{
		fprintf(stderr, "unmount %s failed: %s\n", Target, strerror(errno));
		if(errno != EBUSY) return -1;
		sleep(1);
	}
	return 0;
}

static int
rm(char const *Path)
{
	pid_t PID = fork();
	if(PID == -1) return -1;
	if(PID == 0)
	{
		close(STDIN_FILENO);
		dup2(STDERR_FILENO, STDOUT_FILENO);
		close_all();
		fprintf(stderr, "rm -r %s\n", Path);
		execlp("rm", "rm", "-r", Path, NULL);
		exit(-1);
	}
	int Exit;
	waitpid(PID, &Exit, 0);
	if(!WIFEXITED(Exit) || (WEXITSTATUS(Exit) != 0)) return -1;
	return 0;
}

static int
ApplyDevfsRuleset(char const *Prefix, s64 Ruleset)
{
	if(Ruleset == 0) return 0;
	devfs_rsnum RS = (devfs_rsnum)Ruleset;
	if(RS != Ruleset)
	{
		fprintf(stderr, "Ruleset invalid: %ld\n", Ruleset);
		return -1;
	}
	char Path[MAXPATHLEN];
	snprintf(Path, sizeof(Path), "%s/dev", Prefix);
	fd DevFD = open(Path, O_RDONLY);
	if(DevFD == -1)
	{
		fprintf(stderr, "open '%s' failed: %s\n", Path, strerror(errno));
		return -1;
	}
	if(ioctl(DevFD, DEVFSIO_SAPPLY, &RS) == -1)
	{
		fprintf(stderr, "ruleset %hu apply failed: %s\n", RS, strerror(errno));
		close(DevFD);
		return -1;
	}
	close(DevFD);
	return 0;
}

static int
start_fs(filesystem_configuration *E, char const *Path)
{
	bool   CreatedPath       = false;
	size_t LayerAt           = 0;
	size_t VolumeAt          = 0;
	bool   Mounted_devfs     = false;
	bool   Mounted_temporary = false;
	if(mkdir(Path, 0755) == -1)
	{
		fprintf(stderr, "mkdir %s: %s\n", Path, strerror(errno));
		goto error;
	}
	CreatedPath = true;
	for(LayerAt = 0; LayerAt < E->LayerCount; ++LayerAt)
	{
		if(mapping_mount(&(mapping){ .Source = E->Layer[LayerAt], .Target = "/" },
		                 Path,
		                 mapping_type_unionfs_below) == -1)
			goto error;
	}
	if(E->Temporary != NULL)
	{
		if(mkdir(E->Temporary, 755) == -1)
		{
			if(errno != EEXIST)
			{
				fprintf(stderr, "mkdir %s: %s\n", E->Temporary, strerror(errno));
				goto error;
			}
		}
		if(mapping_mount(&(mapping){ .Source = E->Temporary, .Target = "/" }, Path, mapping_type_unionfs) == -1)
			goto error;
	}
	Mounted_temporary = true;
	if(mapping_mount(&(mapping){ .Target = "/dev" }, Path, mapping_type_devfs) == -1) goto error;
	if(ApplyDevfsRuleset(Path, E->DevfsRuleset) == -1) goto error;
	for(VolumeAt = 0; VolumeAt < E->VolumeCount; ++VolumeAt)
	{
		if(mapping_mount(E->Volume + VolumeAt, Path, mapping_type_nullfs) == -1) goto error;
	}
	return 0;
error:
	for(size_t i = 0; i < VolumeAt; ++i) mapping_umount(Path, E->Volume[E->VolumeCount - 1 - i].Target);
	if(Mounted_devfs) mapping_umount(Path, "/dev");
	if(Mounted_temporary) mapping_umount(Path, "/");
	for(size_t i = 0; i < LayerAt; ++i) mapping_umount(Path, "/");
	if(CreatedPath) rm(Path);
	return -1;
}

static int
stop_fs(filesystem_configuration *E, char const *Path)
{
	for(size_t i = 0; i < E->VolumeCount; ++i) mapping_umount(Path, E->Volume[E->VolumeCount - 1 - i].Target);
	mapping_umount(Path, "/dev");
	if(E->Temporary != NULL) mapping_umount(Path, "/");
	for(size_t i = 0; i < E->LayerCount; ++i) mapping_umount(Path, "/");
	if(rm(Path) == -1)
	{
		fprintf(stderr, "rm -r %s: %s\n", Path, strerror(errno));
		goto error;
	}
	return 0;
error:
	return -1;
}

static bool
ValidJailName(char const *Name)
{
	for(char const *At = Name; *At; ++At)
	{
		if(isnumber(*At) == false) return true;
	}
	return true;
}

// COMMANDS

static char const *FilesystemPrefix = "/tmp/sf";

int
start(char const *Name, ucl_object_t *configuration)
{
	if(ValidJailName(Name) == false) return -1;
	Load_configuration(configuration);
	fd Control = -1;
	Data.JID   = -1;
	char Path[PATH_MAX];
	snprintf(Path, sizeof(Path), "%s/%s", FilesystemPrefix, Name);
	if(mkdir(FilesystemPrefix, 0755) == -1)
	{
		if(errno != EEXIST)
		{
			fprintf(stderr, "mkdir %s: %s\n", FilesystemPrefix, strerror(errno));
			goto error;
		}
	}
	if(start_fs(&Configuration.Filesystem, Path) == -1) goto error;
	if(start_jail(&Data.JID, Name, Path) == -1) goto error;
	if(NgMkSockNode(NULL, &Control, NULL) == -1)
	{
		fprintf(stderr, "Failed to create netgraph socket: %s\n", strerror(errno));
		goto error;
	}
	if(start_interfaces(&Data.Interfaces, &Configuration.Interfaces, Control, Name) == -1)
	{
		fprintf(stderr, "failed to start interfaces\n");
		goto error;
	}
	DataLoaded = true;
	close(Control);
	exec_string_array(&Configuration.Init, false, Data.JID);
	jprint_state S;
	bzero(&S, sizeof(S));
	S.F = stdout;
	Save_data(&S, &Data);
	Free_all();
	return 0;
error:
	if(Control != -1) close(Control);
	if(Data.JID != -1) stop_jail(Data.JID);
	Free_all();
	return -1;
}

int
stop(char const *Name, ucl_object_t *configuration, ucl_object_t *data)
{
	if(ValidJailName(Name) == false) return -1;
	Load_configuration(configuration);
	Load_data(data);
	exec_string_array(&Configuration.Shutdown, true, Data.JID);
	fd Control = -1;
	if(NgMkSockNode(NULL, &Control, NULL) == -1)
	{
		fprintf(stderr, "Failed to create netgraph socket: %s\n", strerror(errno));
		goto error;
	}
	if(stop_interfaces(&Data.Interfaces, &Configuration.Interfaces, Control, Name) == -1) goto error;
	if(stop_jail(Data.JID) == -1) goto error;
	char Path[PATH_MAX];
	snprintf(Path, sizeof(Path), "%s/%s", FilesystemPrefix, Name);
	if(stop_fs(&Configuration.Filesystem, Path) == -1) goto error;
	Free_all();
	return 0;
error:
	Free_all();
	return -1;
}

int
get_endpoint(char const *Name, ucl_object_t *configuration, ucl_object_t *data, char const *Interface)
{
	if(ValidJailName(Name) == false) return -1;
	Load_configuration(configuration);
	Load_data(data);
	for(size_t i = 0; i < Configuration.Interfaces.Count; ++i)
	{
		interface_configuration *C = Configuration.Interfaces._ + i;
		if(strcmp(C->Name, Interface) == 0)
		{
			interface   *E = Data.Interfaces._ + i;
			jprint_state S;
			bzero(&S, sizeof(S));
			S.F = stdout;
			JPrintObjectBegin(&S);
			JPrintMember(&S, "address");
			char Address[NG_PATHSIZ];
			snprintf(Address, sizeof(Address), "[%x]:", E->ID);
			JPrint_string(&S, Address);
			JPrintMember(&S, "hook");
			JPrint_string(&S, "ether");
			JPrintObjectEnd(&S);
			Free_all();
			return 0;
		}
	}
	Free_all();
	return -1;
}

int
cmd(char const *Name, ucl_object_t *configuration, ucl_object_t *data, size_t ArgCount, char **Arg)
{
	if(ValidJailName(Name) == false) return -1;
	Load_configuration(configuration);
	Load_data(data);
	if(ArgCount < 1) usage();
	char const *Command = Arg[0];
	--ArgCount;
	++Arg;
	if(strcmp(Command, "sh") == 0)
	{
		size_t ExtraArguments = 2;
		char **Args           = alloca(sizeof(char *) * (ArgCount + ExtraArguments + 1));
		Args[0]               = "jexec";
		Args[1]               = strdup(Name);
		memcpy(Args + ExtraArguments, Arg, sizeof(char *) * ArgCount);
		Args[ArgCount + ExtraArguments] = NULL;
		execvp("jexec", Args);
		fprintf(stderr, "execvp failed: %s\n", strerror(errno));
		exit(-1);
	} else if(strcmp(Command, "wireshark") == 0)
	{
		if(ArgCount < 1) usage();
		char const *Interface = Arg[0];
		--ArgCount;
		++Arg;
		fd Pipe[2];
		if(pipe(Pipe) == -1)
		{
			fprintf(stderr, "pipe: %s\n", strerror(errno));
			return -1;
		}
		pid_t PID = fork();
		if(PID == -1)
		{
			fprintf(stderr, "fork: %s\n", strerror(errno));
			return -1;
		} else if(PID == 0)
		{
			dup2(Pipe[1], STDOUT_FILENO);
			close(Pipe[0]);
			close(Pipe[1]);
			close_all();
			if(jail_attach(Data.JID))
			{
				fprintf(stderr, "jail_attach(%d): %s\n", Data.JID, strerror(errno));
				exit(-1);
			}
			execlp("tcpdump", "tcpdump", "-s", "-0", "-U", "-w", "-", "-i", Interface, NULL);
			fprintf(stderr, "execlp failed: %s\n", strerror(errno));
			exit(-1);
		} else
		{
			dup2(Pipe[0], STDIN_FILENO);
			close(Pipe[0]);
			close(Pipe[1]);
			close_all();
			char Title[64];
			snprintf(Title, sizeof(Title), "gui.window_title:%s:%s", Name, Interface);
			execlp("wireshark", "wireshark", "-o", Title, "-k", "-i", "-", NULL);
			fprintf(stderr, "execlp failed: %s\n", strerror(errno));
			exit(-1);
		}
	} else
	{
		fprintf(stderr, "unknown command %s\n", Command);
		usage();
	}
}

int
mod(char const *Name, ucl_object_t *configuration, ucl_object_t *data, size_t ArgCount, char **Arg)
{
	if(ValidJailName(Name) == false) return -1;
	Load_configuration(configuration);
	Load_data(data);
	UNUSED(ArgCount, Arg);
	usage();
}