module:set_global();

local audit_log_limit = module:get_option_number("audit_log_limit", 10000);

local cleanup_after = module:get_option_string("audit_log_expires_after", "2w");

local attach_ips = module:get_option_boolean("audit_log_ips", true);

local attach_ipv4_prefix = module:get_option_number("audit_log_ipv4_prefix", nil);

local attach_ipv6_prefix = module:get_option_number("audit_log_ipv6_prefix", nil);

local have_geoip, geoip = pcall(require, "geoip.country");

local attach_location = have_geoip and module:get_option_boolean("audit_log_location", true);

local geoip4_country, geoip6_country;

if have_geoip and attach_location then

	geoip4_country = geoip.open(module:get_option_string("geoip_ipv4_country", "/usr/share/GeoIP/GeoIP.dat"));

	geoip6_country = geoip.open(module:get_option_string("geoip_ipv6_country", "/usr/share/GeoIP/GeoIPv6.dat"));

end

local time_now = os.time;

local ip = require "util.ip";

local st = require "util.stanza";

local moduleapi = require "core.moduleapi";

local host_wide_user = "@";

local stores = {};

local function get_store(self, host)

	local store = rawget(self, host);

	if store then

		return store

	end

	store = module:context(host):open_store("audit", "archive");

	rawset(self, host, store);

	return store;

end

setmetatable(stores, { __index = get_store });

local function get_ip_network(ip_addr)

	local _ip = ip.new_ip(ip_addr);

	local proto = _ip.proto;

	local network;

	if proto == "IPv4" and attach_ipv4_prefix then

		network = ip.truncate(_ip, attach_ipv4_prefix).normal.."/"..attach_ipv4_prefix;

	elseif proto == "IPv6" and attach_ipv6_prefix then

		network = ip.truncate(_ip, attach_ipv6_prefix).normal.."/"..attach_ipv6_prefix;

	end

	return network;

end

local function session_extra(session)

	local attr = {

		xmlns = "xmpp:prosody.im/audit",

	};

	if session.id then

		attr.id = session.id;

	end

	if session.type then

		attr.type = session.type;

	end

	local stanza = st.stanza("session", attr);

	if attach_ips and session.ip then

		local remote_ip, network = session.ip;

		if attach_ipv4_prefix or attach_ipv6_prefix then

			network = get_ip_network(remote_ip);

		end

		stanza:text_tag("remote-ip", network or remote_ip);

	end

	if attach_location and session.ip then

		local remote_ip = ip.new(session.ip);

		local geoip_country = ip.proto == "IPv6" and geoip6_country or geoip4_country;

		stanza:tag("location", {

			country = geoip_country:query_by_addr(remote_ip.normal);

		}):up();

	end

	if session.client_id then

		stanza:text_tag("client", session.client_id);

	end

	return stanza

end

local function audit(host, user, source, event_type, extra)

	if not host or host == "*" then

		error("cannot log audit events for global");

	end

	local user_key = user or host_wide_user;

	local attr = {

		["source"] = source,

		["type"] = event_type,

	};

	if user_key ~= host_wide_user then

		attr.user = user_key;

	end

	local stanza = st.stanza("audit-event", attr);

	if extra ~= nil then

		if extra.session then

			local child = session_extra(extra.session);

			if child then

				stanza:add_child(child);

			end

		end

		if extra.custom then

			for _, child in extra.custom do

				if not st.is_stanza(child) then

					error("all extra.custom items must be stanzas")

				end

				stanza:add_child(child);

			end

		end

	end

	local id, err = stores[host]:append(nil, nil, stanza, time_now(), user_key);

	if err then

		module:log("error", "failed to persist audit event: %s", err);

		return

	else

		module:log("debug", "persisted audit event %s as %s", stanza:top_tag(), id);

	end

end

function moduleapi.audit(module, user, event_type, extra)

	audit(module.host, user, "mod_" .. module:get_name(), event_type, extra);

end

function module.command(_arg)

	local arg = require "util.argparse".parse(_arg, { value_params = { "limit", "user" } });

	local host = arg[1];

	if not host then

		print("EE: Please supply the host for which you want to show events");

		return 1;

	elseif not prosody.hosts[host] then

		print("EE: Unknown host: "..host);

		return 1;

	end

	require "core.storagemanager".initialize_host(host);

	local store = stores[host];

	local c = 0;

	local results, err = store:find(nil, {

		with = arg.user;

		limit = arg.limit and tonumber(arg.limit) or nil;

	})

	if not results then

		print("EE: Failed to query audit log: "..tostring(err));

		return 1;

	end

	local colspec = {

		{ title = "Date", key = "when", width = 19, mapper = function (when) return os.date("%Y-%m-%d %R:%S", when); end };

		{ title = "Source", key = "source", width = 18 };

		{ title = "Event", key = "event_type", width = 22 };

	};

	if not arg.global then

		table.insert(colspec, {

			title = "User", key = "username", width = 30,

			mapper = function (user)

				if user == "@" then return ""; end

				if user:sub(-#host-1, -1) == ("@"..host) then

					return (user:gsub("@.+$", ""));

				end

			end;

		});

		if attach_ips then

			table.insert(colspec, {

				title = "IP", key = "ip", width = "28";

			});

		end

		if attach_location then

			table.insert(colspec, {

				title = "Location", key = "country", width = 2;

			});

		end

	end

	local width = tonumber(os.getenv("COLUMNS")) or 80;

	local row = require "util.human.io".table(colspec, width);

	print(string.rep("-", width));

	print(row());

	print(string.rep("-", width));

	for _, entry, when, user in results do

		c = c + 1;

		print(row({

			when = when;

			source = entry.attr.source;

			event_type = entry.attr.type:gsub("%-", " ");

			username = user;

			ip = entry:get_child_text("remote-ip");

			location = entry:find("location@country");

		}));

	end

	print(string.rep("-", width));

	print(("%d records displayed"):format(c));

end