auditd-plugin-clickhouse/auditd-plugin-clickhouse.cpp

/*
 * auditd-plugin-clickhouse is an auditd plugin for sending auditd data
 * to clickhouse DB.
 * Copyright (C) 2019 Aleksei Nikiforov <darktemplar@basealt.ru>
 *
 * 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/>.
 *
 */

#include <signal.h>
#include <poll.h>
#include <unistd.h>
#include <errno.h>

#include <string>
#include <map>
#include <fstream>
#include <vector>

#include <auparse.h>

#include <clickhouse-cpp/client.h>

#include <boost/scope_exit.hpp>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/ini_parser.hpp>

int runpipes[2] = { -1, -1 };
volatile bool running = true;

static void term_handler(int sig)
{
	if (runpipes[1] != -1)
	{
		write(runpipes[1], "1", 1);
	}

	running = false;
}

bool is_valid_table_name(const std::string &value)
{
	return (std::find_if_not(value.begin(), value.end(), [] (const char c) { return (std::isalnum(c) || (strchr("_", c) != NULL)); } ) == value.end());
}

template <typename T>
bool always_true(const T &)
{
	return true;
}

template <typename T>
void optional_set(T &value, const boost::property_tree::ptree &tree, const char *element_name, bool (*check_function)(const T &))
{
	auto element = tree.get_child_optional(element_name);
	if (element)
	{
		if (check_function(element->get_value<T>()))
		{
			value = element->get_value<T>();
		}
		else
		{
			throw std::runtime_error(std::string("Invalid value for option '") + std::string(element_name) + std::string("'"));
		}
	}
}

template <typename T>
void optional_set(T &value, const boost::property_tree::ptree &tree, const char *element_name)
{
	optional_set(value, tree, element_name, &always_true<T>);
}

void auparse_callback(auparse_state_t *au, auparse_cb_event_t cb_event_type, void *user_data)
{
}

int main(int argc, char **argv)
{
	if (argc != 2)
	{
		fprintf(stderr, "Error: USAGE: %s config\n", argv[0]);
		return -1;
	}

	try
	{
		clickhouse::ClientOptions client_options;

		size_t buffer_size = 4096;

		if (pipe(runpipes) < 0)
		{
			throw std::runtime_error("Failed to create pipes");
		}

		BOOST_SCOPE_EXIT(&runpipes)
		{
			close(runpipes[0]);
			close(runpipes[1]);
			runpipes[0] = -1;
			runpipes[1] = -1;
		} BOOST_SCOPE_EXIT_END

		{
			struct sigaction sa;

			sa.sa_flags = 0;
			sigemptyset(&sa.sa_mask);

			sa.sa_handler = term_handler;
			if (sigaction(SIGTERM, &sa, NULL) < 0)
			{
				throw std::runtime_error("Failed to set sigterm handler");
			}
		}

		/* First read config */
		{
			boost::property_tree::ptree clickhouse_config_tree;

			std::ifstream stream(argv[1]);
			boost::property_tree::read_ini(stream, clickhouse_config_tree);

			auto general = clickhouse_config_tree.get_child_optional("General");
			if (general)
			{
				optional_set(buffer_size, *general, "BufferSize");
			}

			auto client_connection = clickhouse_config_tree.get_child_optional("Connection");
			if (client_connection)
			{
				optional_set(client_options.host, *client_connection, "Hostname");
				optional_set(client_options.port, *client_connection, "Port");
				optional_set(client_options.default_database, *client_connection, "DefaultDatabase");
				optional_set(client_options.user, *client_connection, "Username");
				optional_set(client_options.password, *client_connection, "Password");
				optional_set(client_options.ping_before_query, *client_connection, "PingBeforeQuery");
				optional_set(client_options.send_retries, *client_connection, "SendRetries");

				{
					auto element = client_connection->get_child_optional("RetryTimeout");
					if (element)
					{
						client_options.retry_timeout = std::chrono::seconds(element->get_value<int>());
					}
				}

				{
					auto element = client_connection->get_child_optional("CompressionMethod");
					if (element)
					{
						const auto value = element->get_value<std::string>();

						const std::map<std::string, clickhouse::CompressionMethod> compression_method_map =
						{
							{ "None", clickhouse::CompressionMethod::None },
							{ "LZ4", clickhouse::CompressionMethod::LZ4 },
						};

						auto iter = compression_method_map.find(value);
						if (iter != compression_method_map.end())
						{
							client_options.compression_method = iter->second;
						}
						else
						{
							client_options.compression_method = clickhouse::CompressionMethod::None;
						}
					}
				}
			}
		}

		/* Now connect to clickhouse */
		clickhouse::Client client(client_options);

		std::unique_ptr<auparse_state_t, void(*)(auparse_state_t*)> au(auparse_init(AUSOURCE_FEED, 0), [](auparse_state_t *obj){
			if (obj)
			{
				auparse_flush_feed(obj);
				auparse_destroy(obj);
			}
		});

		if (!au)
		{
			throw std::runtime_error("Failed to initialize audit");
		}

		auparse_add_callback(au.get(), auparse_callback, NULL, NULL);

		std::vector<char> data;
		data.resize(buffer_size);

		struct pollfd pollfds[2];

		pollfds[0].fd = STDIN_FILENO;
		pollfds[1].fd = runpipes[0];

		while (running)
		{
			pollfds[0].events = POLLIN | POLLHUP;
			pollfds[0].revents = 0;
			pollfds[1].events = POLLIN | POLLHUP;
			pollfds[1].revents = 0;

			int res = poll(pollfds, sizeof(pollfds) / sizeof(pollfds[0]), -1);
			if (res < 0)
			{
				// error occured, finish running
				fprintf(stderr, "Poll returned error: %d\n", errno);
				running = false;
			}
			else if (res == 0)
			{
				// timeout, do nothing
			}
			else if (pollfds[0].revents & POLLIN)
			{
				ssize_t readsize = read(STDIN_FILENO, data.data(), data.size());
				if (readsize > 0)
				{
					auparse_feed(au.get(), data.data(), readsize);
				}
				else
				{
					// error occured, finish running
					fprintf(stderr, "Read returned error: %d\n", errno);
					running = false;
				}
			}
			else if (pollfds[0].revents & POLLHUP)
			{
				// stdin closed, no more data, finish running
				running = false;
			}
		}
	}
	catch (const std::exception &e)
	{
		fprintf(stderr, "Caught exception: %s\n", e.what());
		return -1;
	}
	catch (...)
	{
		fprintf(stderr, "Caught unknown exception\n");
		return -1;
	}

	return 0;
}