/*
 * Logserver
 * Copyright (C) 2017-2025 Joel Reardon
 *
 * This program 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 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#ifndef __RUN__H__
#define __RUN__H__

#include <cstdint>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <string>
#include <vector>

#include <sys/socket.h>
#include <sys/time.h>
#include <sys/unistd.h>
#include <sys/wait.h>

#include "tokenizer.h"

using namespace std;

/* Wrapper around running a command. Takes a |-separated sequence of commands
 * and runs them in order, feeding the output as input as would be expected.
 * Allows the result to be read back in at the end */
class Run {
	/* stores a pair of FDs for a pipe and allows closing either side */
	struct PipePair {
		PipePair() {
			int r = pipe(&read_end);
			if (r) {
				read_end = 0;
				write_end = 0;
			}
		}
		~PipePair() {
			close();
		}
		void close() {
			if (read_end) ::close(read_end);
			if (write_end) ::close(write_end);
			read_end = 0;
			write_end = 0;
		}
		void set_write() {
			::close(read_end);
			read_end = 0;
		}
		void set_read() {
			::close(write_end);
			write_end = 0;
		}

		int read_end;
		int write_end;

		/* casting to int returns the read_end */
		operator int*() {
			return (int*) this;
		}
	};

public:
	Run() : Run("", "") {}
	explicit Run(const string& cmd) : Run(cmd, "") {}

	/* constructor takes command and input to use for stdin. it separates
	 * the |-separated commands and assembles the pipeline
	 */
	Run(const string& cmd, const string& input)
		: _cmd(cmd), _started(false), _read(false),
		  _done(false), _status(0) {
		vector<string_view> pipeline;
		_pipes.push_back(nullptr);
		_pipes.back().reset(new PipePair());

		Tokenizer::split_mind_quote(_cmd, "|", &pipeline);
		for (const auto& x: pipeline) {
			connect(x);
		}
		write_input(input);
		safety();
	}

	virtual ~Run() {}

	virtual string command() const { return _cmd; }

	/* function invoke operator will execute the pipeline */
	void operator()() {
		_started = true;
		size_t pos = 0;
		while (pos != _argvs.size()) {
			_pids.push_back(execute(pos));
			if (pos == 0) _pipes[0]->set_write();
			++pos;
		}
		_pipes.back()->set_read();
	}

	/* get the return value of the last command in the pipeline */
	int result() {
		if (!_read) read();
		return _status;
	}

	/* TODO: move read and redirect to a single stream based runner */
	int redirect(const string &filename) {
		ofstream fout(filename);
		const size_t SIZE = 4096;
		assert(_pipes.back()->read_end);
		int r = 0;
		size_t pid_pos = 0;
		char buf[SIZE];
		while (true) {
			r = ::read(_pipes.back()->read_end, buf, SIZE);
			if (r == 0 || (r < 0 && errno == EAGAIN)) {
				if (pid_pos == _pids.size()) {
					break;
				}
				waitpid(_pids[pid_pos++], &_status, 0);
				continue;
			}
			if (r < 0) {
				return -1;
			}
			fout.write(buf, r);
		}
		return 0;
	}

	/* read output as a vector */
	void read(vector<string>* out) {
		stringstream ss;
		ss << read(1000);
		string s;
		while (getline(ss, s)) {
			out->push_back(s);
		}
	}

	/* returns stdout of last process */
	string read() {
		return read(1000);
	}

	/* runs through the pipeline closing stdin and waiting for the process
	 * to finish */
	void finish() {
		size_t pid_pos = 0;
		while (pid_pos < _pids.size()) {
			_pipes[pid_pos]->close();
			waitpid(_pids[pid_pos++], &_status, 0);
		}
	}

	/* read stdout of last process with a timeout */
	string read(int runtime) {
		const size_t SIZE = 4096;
		stringstream ss;
		char buf[SIZE];
		size_t pid_pos = 0;
		thread kill_thread(bind(&Run::abort_run, this, runtime));
		int r = 0;
		prepare_read();
		while (true) {
			r = ::read(_pipes.back()->read_end, buf, SIZE);
			if (r == 0 || (r < 0 && errno == EAGAIN)) {
				if (pid_pos == _pids.size()) {
					break;
				}
				waitpid(_pids[pid_pos++], &_status, 0);
				continue;
			}
			if (r < 0) {
				_done = true;
				kill_thread.join();
				return "";
			}
			ss.write(buf, r);
		}
		_done = true;
		kill_thread.join();
		return ss.str();
	}

	/* gets the FD for the last command's stdout */
	int read_fd() {
		prepare_read();
		return _pipes.back()->read_end;
	}

	/* gets the FD for the first command's stdin */
	int write_fd() { return _pipes.front()->write_end; }

	/* closes the write end of the pipelien */
	void close_write() {
		size_t pos = 0;
		while (pos + 1 < _pipes.size()) {
			_pipes[pos]->close();
			++pos;
		}
	}

protected:
	/* adds a new parameter cmd as a command to the end of the current
	 * pipeline */
	void connect(const string_view& cmd) {
		_argvs.push_back(vector<string>());
		_pipes.push_back(nullptr);
		_pipes.back().reset(new PipePair());
		Tokenizer::split_mind_quote_and_copy(cmd, " ", &_argvs.back());
	}

	// TODO: make this a config file
	/* idea here is to preven running commands that may do real things on
	 * the FS like rm */
	static inline const set<string> ALLOWED = {
		"cat", "sort", "uniq", "ls", "grep",
		"cut", "tr", "sed", "awk", "csv_mr",
		"fgrep", "/bin/which", "whoami",
	 	"base64", "echo", "file",
		"packetgrep", "wc", "xsel", "mplayer"
	};

	/* checks all commands are on the allow list and throws if not */
	void safety() {
		for (const auto &x : _argvs) {
			if (!ALLOWED.count(string(x[0]))) {
				_argvs.clear();
				_pipes.clear();
				throw logic_error(_cmd + " not allowed");
			}
		}
	}

	/* runs the pipeline command at parameter pos's position */
	pid_t execute(size_t pos) {
		const char &c = _argvs[pos][0][0];
		if (c != '/' && c != '~' && c != '.') {
			string which_path;
			char *path_env = getenv("PATH");
			if (!path_env) return 0;
			char *delete_path = strdup(path_env);
			if (!delete_path) return 0;

			char *dir = strtok(delete_path, ":");
			while (dir) {
				which_path = string(dir) + "/" +
					_argvs[pos][0];
				if (access(which_path.c_str(), X_OK) == 0) {
					break;
				}
				dir = strtok(NULL, ":");
			}
			free(delete_path);
			if (which_path.empty())
				throw logic_error(string("fork(): failed"));
			_argvs[pos][0] = which_path;
		}
		pid_t pid = fork();
		if (pid == -1) {
			throw logic_error(string("fork(): failed"));
		}
		if (pid == 0) {
			// I am child
			int r = 0;
			for (size_t i = 0; i < _pipes.size(); ++i) {
				if (i == pos || i == pos+1) continue;
				_pipes[i]->close();
			}
			_pipes[pos]->set_read();
			_pipes[pos+1]->set_write();
			int keep_read = _pipes[pos]->read_end;
			int keep_write = _pipes[pos + 1]->write_end;

			int fdlimit = (int)sysconf(_SC_OPEN_MAX);
			for (int i = 3; i < fdlimit; i++) {
				if (i == keep_read || i == keep_write) continue;
				close(i);
			}

			char** argv = get_c_args(pos);
			r = dup2(_pipes[pos]->read_end, STDIN_FILENO);
			if (r == -1) {
				exit(-1);
			}
			close(STDERR_FILENO);
			// TODO: decide what to do with STDERR_FILENO
			r = dup2(_pipes[pos + 1]->write_end, STDOUT_FILENO);
			if (r == -1) {
				exit(-1);
			}
			execv(argv[0], (char * const *) argv);

			_pipes[pos]->close();
			_pipes[pos + 1]->close();
			char** p = argv;
			while (!*p++) delete[](*p);
			delete[] argv;
			exit(-1);
		}
		return pid;

	}

	/* used to separate check if command is still running and after a
	 * timeout will kill the processes */
	void abort_run(int when) {
		struct timeval tv;
		gettimeofday(&tv, nullptr);
		int now = tv.tv_sec;
		while (!_done && (now + when > tv.tv_sec)) {
			usleep(1000);
			gettimeofday(&tv, nullptr);
		}
		if (_done) return;

		for (auto &x : _pids) {
			kill(x, SIGKILL);
		}

	}

	void prepare_read() {
                _read = true;
                assert(_pipes.size());
                assert(_pids.size());
                assert(_pipes.back()->read_end);
                _done = false;
	}

	/* writes data to the front of the pipeline */
	void write_input(const string& data) {
		size_t todo = data.length();
		while (todo) {
			int r = ::write(_pipes.front()->write_end, data.c_str(), data.length());
			if (r < 0) throw logic_error(string("write(): failed."));
			size_t written = (size_t) r;
			if (written == 0) throw logic_error(string("write(): failed."));
			todo -= written;
		}
	}

	/* converts the pipeline command at parameter pos's position into a
	 * argv-style char** for running */
	char** get_c_args(int pos) {
		char** retval = new char*[_argvs[pos].size() + 1];

		for (size_t i = 0; i < _argvs[pos].size(); ++i) {
			retval[i] = new char[_argvs[pos][i].length() + 1];
			string arg = _argvs[pos][i];
			assert(arg.size());
			if (arg.size() > 2 && arg[0] == '"' && arg[arg.size() - 1] == '"') {
				strncpy(retval[i], arg.data() + 1, arg.size() - 2);
				retval[i][arg.size() - 2] = '\0';
			} else {
				strncpy(retval[i], arg.data(), arg.size());
				retval[i][arg.size()] = '\0';
			}

		}
		retval[_argvs[pos].size()] = nullptr;
		return retval;
	}

	/* vector of each command, consisting of vector of args */
	vector<vector<string>> _argvs;
	/* pipes connecting the processes */
	vector<unique_ptr<PipePair>> _pipes;
	/* pids for the processes */
	vector<pid_t> _pids;
	/* the command we are given */
	string _cmd;
	/* has execution started */
	bool _started;
	/* has reading happened */
	bool _read;
	/* is execution done */
	atomic<bool> _done;
	/* what is return val of last command */
	int _status;
};

#endif  // __RUN__H__
