/*
 * auditd-plugin-clickhouse is an auditd plugin for sending auditd data
 * to clickhouse DB.
 * Copyright (C) 2019-2020 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 "auditd-record.hpp"
#include "utils.hpp"

#include <functional>
#include <set>

#include <clickhouse-cpp/columns/array.h>
#include <clickhouse-cpp/columns/nullable.h>
#include <clickhouse-cpp/columns/numeric.h>
#include <clickhouse-cpp/columns/string.h>

namespace {

std::map<auparse_type_t, std::function<bool(const std::string &)> >& get_audit_type_check_instance()
{
	static std::map<auparse_type_t, std::function<bool(const std::string &)> > instance;
	return instance;
}

class AuditTypeCheckRegister
{
public:
	AuditTypeCheckRegister(auparse_type_t type, const std::function<bool(const std::string &)> &func)
	{
		get_audit_type_check_instance()[type] = func;
	}
};

std::function<bool(const std::string &)> integer_record_type_check_function()
{
	return [](const std::string &value) -> bool {
		std::set<std::string> allowed_values = {
			"integer"
		};

		return (allowed_values.find(value) != allowed_values.end());
	};
}

std::function<bool(const std::string &)> interpreted_string_record_type_check_function()
{
	return [](const std::string &value) -> bool {
		std::set<std::string> allowed_values = {
			"string",
			"string_array"
		};

		return (allowed_values.find(value) != allowed_values.end());
	};
}

std::function<bool(const std::string &)> any_record_type_check_function()
{
	return [](const std::string &value) -> bool {
		return true;
	};
}

} // unnamed namespace

#define register_record_type(audit_type, check_function) \
	static const AuditTypeCheckRegister audit_type_check_register_##audit_type(audit_type, check_function)

register_record_type(AUPARSE_TYPE_UNCLASSIFIED, any_record_type_check_function());
register_record_type(AUPARSE_TYPE_UID,          integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_GID,          integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_SYSCALL,      integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_ARCH,         interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_EXIT,         integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_ESCAPED,      interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_PERM,         integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_MODE,         integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_SOCKADDR,     interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_FLAGS,        integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_PROMISC,      integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_CAPABILITY,   integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_SUCCESS,      interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_A0,           interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_A1,           interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_A2,           interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_A3,           interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_SIGNAL,       integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_LIST,         integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_TTY_DATA,     interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_SESSION,      integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_CAP_BITMAP,   interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_NFPROTO,      integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_ICMPTYPE,     integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_PROTOCOL,     integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_ADDR,         interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_PERSONALITY,  integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_SECCOMP,      integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_OFLAG,        integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_MMAP,         integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_MODE_SHORT,   integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_MAC_LABEL,    interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_PROCTITLE,    interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_HOOK,         integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_NETACTION,    integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_MACPROTO,     integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_IOCTL_REQ,    integer_record_type_check_function());
register_record_type(AUPARSE_TYPE_ESCAPED_KEY,  interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_ESCAPED_FILE, interpreted_string_record_type_check_function());
register_record_type(AUPARSE_TYPE_FANOTIFY,     interpreted_string_record_type_check_function());

#undef register_record_type

bool check_field_type(auparse_type_t field_type, const std::string &database_type, const std::string &database_field_name)
{
	const auto &check_instance = get_audit_type_check_instance();

	auto check_type_function = check_instance.find(field_type);
	if (check_type_function != check_instance.end())
	{
		return check_type_function->second(database_type);
	}
	else
	{
		fprintf(stderr, "Warning: unknown record type for record name \"%s\"\n", database_field_name.c_str());
		return interpreted_string_record_type_check_function()(database_type);
	}
}

AbstractRecordField::AbstractRecordField(const std::string &name)
	: m_name(name)
{
}

CommonStringRecordField::CommonStringRecordField(const std::string &name)
	: AbstractRecordField(name)
{
}

std::vector<AbstractRecordField::Column> CommonStringRecordField::generateColumnsAndNames() const
{
	return std::vector<AbstractRecordField::Column> {
		Column { m_name, std::make_shared<clickhouse::ColumnNullable>(std::make_shared<clickhouse::ColumnString>(), std::make_shared<clickhouse::ColumnUInt8>()) }
	};
}

void CommonStringRecordField::addToColumn(const std::vector<Column> &columns) const
{
	if ((columns.size() != 1) || (!columns[0].value))
	{
		throw std::runtime_error("CommonStringRecordField::addToColumn: invalid columns argument");
	}

	auto nullable = columns[0].value->As<clickhouse::ColumnNullable>();
	if (!nullable)
	{
		throw std::runtime_error("Invalid column type: not ColumnNullable");
	}

	auto nested = nullable->Nested();
	if ((!nested) || (nested->Type()->GetCode() != clickhouse::Type::String))
	{
		throw std::runtime_error("Invalid nested column type: not String");
	}
	
	auto data_column = std::make_shared<clickhouse::ColumnString>();
	auto null_column = std::make_shared<clickhouse::ColumnUInt8>();

	if (m_value)
	{
		data_column->Append(*m_value);
		null_column->Append(0);
	}
	else
	{
		data_column->Append(std::string());
		null_column->Append(1);
	}

	columns[0].value->Append(std::make_shared<clickhouse::ColumnNullable>(data_column, null_column));
}

std::shared_ptr<AbstractRecordField> StringRecordField::createRecord(const std::string &name)
{
	std::shared_ptr<AbstractRecordField> result;
	result.reset(new StringRecordField(name));
	return result;
}

void StringRecordField::addOrUpdateValue(auparse_state_t *record)
{
	if (record)
	{
		m_value = string_or_null(auparse_get_field_str(record));
	}
	else
	{
		m_value = boost::none;
	}
}

StringRecordField::StringRecordField(const std::string &name)
	: CommonStringRecordField(name)
{
}

AbstractRecordField::Type StringRecordField::getType() const
{
	return AbstractRecordField::Type::String;
}

std::shared_ptr<AbstractRecordField> InterpretedStringRecordField::createRecord(const std::string &name)
{
	std::shared_ptr<AbstractRecordField> result;
	result.reset(new InterpretedStringRecordField(name));
	return result;
}

InterpretedStringRecordField::InterpretedStringRecordField(const std::string &name)
	: CommonStringRecordField(name)
{
}

void InterpretedStringRecordField::addOrUpdateValue(auparse_state_t *record)
{
	if (record)
	{
		m_value = string_or_null(auparse_interpret_field(record));
	}
	else
	{
		m_value = boost::none;
	}
}

AbstractRecordField::Type InterpretedStringRecordField::getType() const
{
	return AbstractRecordField::Type::InterpretedString;
}

std::shared_ptr<AbstractRecordField> IntegerRecordField::createRecord(const std::string &name)
{
	std::shared_ptr<AbstractRecordField> result;
	result.reset(new IntegerRecordField(name));
	return result;
}

IntegerRecordField::IntegerRecordField(const std::string &name)
	: InterpretedStringRecordField(name)
{
}

void IntegerRecordField::addOrUpdateValue(auparse_state_t *record)
{
	InterpretedStringRecordField::addOrUpdateValue(record);
	
	if (record)
	{
		m_int_value = auparse_get_field_int(record);
	}
	else
	{
		m_int_value = boost::none;
	}
}

std::vector<AbstractRecordField::Column> IntegerRecordField::generateColumnsAndNames() const
{
	return std::vector<AbstractRecordField::Column> {
		Column { m_name + "_IntValue", std::make_shared<clickhouse::ColumnNullable>(std::make_shared<clickhouse::ColumnInt64>(), std::make_shared<clickhouse::ColumnUInt8>()) },
		Column { m_name + "_InterpretedValue", std::make_shared<clickhouse::ColumnNullable>(std::make_shared<clickhouse::ColumnString>(), std::make_shared<clickhouse::ColumnUInt8>()) }
	};
}

void IntegerRecordField::addToColumn(const std::vector<Column> &columns) const
{
	if ((columns.size() != 2) || (!columns[0].value) || (!columns[1].value))
	{
		throw std::runtime_error("IntegerRecordField::addToColumn: invalid columns argument");
	}

	auto nullable = columns[0].value->As<clickhouse::ColumnNullable>();
	if (!nullable)
	{
		throw std::runtime_error("Invalid column type: not ColumnNullable");
	}

	auto nested = nullable->Nested();
	if ((!nested) || (nested->Type()->GetCode() != clickhouse::Type::Int64))
	{
		throw std::runtime_error("Invalid nested column type: not Int64");
	}
	
	auto data_column = std::make_shared<clickhouse::ColumnInt64>();
	auto null_column = std::make_shared<clickhouse::ColumnUInt8>();

	if (m_int_value)
	{
		data_column->Append(*m_int_value);
		null_column->Append(0);
	}
	else
	{
		data_column->Append(0);
		null_column->Append(1);
	}

	std::vector<Column> string_columns;
	string_columns.push_back(columns[1]);

	columns[0].value->Append(std::make_shared<clickhouse::ColumnNullable>(data_column, null_column));

	// now also add string value
	InterpretedStringRecordField::addToColumn(string_columns);
}

AbstractRecordField::Type IntegerRecordField::getType() const
{
	return AbstractRecordField::Type::Int;
}

std::shared_ptr<AbstractRecordField> InterpretedStringArrayRecordField::createRecord(const std::string &name)
{
	std::shared_ptr<AbstractRecordField> result;
	result.reset(new InterpretedStringArrayRecordField(name));
	return result;
}

void InterpretedStringArrayRecordField::addOrUpdateValue(auparse_state_t *record)
{
	if (record)
	{
		m_names_array.push_back(string_or_null(auparse_get_field_name(record)));
		m_values_array.push_back(string_or_null(auparse_interpret_field(record)));
	}
}

std::vector<AbstractRecordField::Column> InterpretedStringArrayRecordField::generateColumnsAndNames() const
{
	return std::vector<AbstractRecordField::Column> {
		Column { m_name + "_Name", std::make_shared<clickhouse::ColumnArray>(std::make_shared<clickhouse::ColumnString>()) },
		Column { m_name + "_Value", std::make_shared<clickhouse::ColumnArray>(std::make_shared<clickhouse::ColumnString>()) }
	};
}

void InterpretedStringArrayRecordField::addToColumn(const std::vector<Column> &columns) const
{
	if ((columns.size() != 2) || (!columns[0].value) || (!columns[1].value))
	{
		throw std::runtime_error("InterpretedStringArrayRecordField::addToColumn: invalid columns argument");
	}

	auto array_names = columns[0].value->As<clickhouse::ColumnArray>();
	if (!array_names)
	{
		throw std::runtime_error("Invalid column 0 type: not ColumnArray");
	}

	auto array_values = columns[1].value->As<clickhouse::ColumnArray>();
	if (!array_values)
	{
		throw std::runtime_error("Invalid column 1 type: not ColumnArray");
	}

	auto name_column = std::make_shared<clickhouse::ColumnString>();
	for (auto iter = m_names_array.begin(); iter != m_names_array.end(); ++iter)
	{
		name_column->Append(*iter);
	}

	auto value_column = std::make_shared<clickhouse::ColumnString>();
	for (auto iter = m_values_array.begin(); iter != m_values_array.end(); ++iter)
	{
		value_column->Append(*iter);
	}

	array_names->AppendAsColumn(name_column);
	array_values->AppendAsColumn(value_column);
}

InterpretedStringArrayRecordField::InterpretedStringArrayRecordField(const std::string &name)
	: AbstractRecordField(name)
{
}

AbstractRecordField::Type InterpretedStringArrayRecordField::getType() const
{
	return AbstractRecordField::Type::InterpretedStringArray;
}