#define _GNU_SOURCE
#include <errno.h>
#include <stdbool.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <xtables.h>
#include <linux/netfilter/xt_connlabel.h>
#include <libnetfilter_conntrack/libnetfilter_conntrack.h>

enum {
	O_LABEL = 0,
	O_SET = 1,
};

static struct nfct_labelmap *map;

static void connlabel_mt_help(void)
{
	puts(
"connlabel match options:\n"
"[!] --label name     Match if label has been set on connection\n"
"    --set            Set label on connection");
}

static const struct xt_option_entry connlabel_mt_opts[] = {
	{.name = "label", .id = O_LABEL, .type = XTTYPE_STRING,
	 .min = 1, .flags = XTOPT_MAND|XTOPT_INVERT},
	{.name = "set", .id = O_SET, .type = XTTYPE_NONE},
	XTOPT_TABLEEND,
};

/* cannot do this via _init, else static builds might spew error message
 * for every iptables invocation.
 */
static int connlabel_open(void)
{
	const char *fname;

	if (map)
		return 0;

	map = nfct_labelmap_new(NULL);
	if (map != NULL)
		return 0;

	fname = nfct_labels_get_path();
	if (errno) {
		fprintf(stderr, "Warning: cannot open %s: %s\n",
			fname, strerror(errno));
	} else {
		xtables_error(RESOURCE_PROBLEM,
			"cannot parse %s: no labels found", fname);
	}
	return 1;
}

static int connlabel_value_parse(const char *in)
{
	char *end;
	unsigned long value = strtoul(in, &end, 0);

	if (in[0] == '\0' || *end != '\0')
		return -1;

	return value;
}

static void connlabel_mt_parse(struct xt_option_call *cb)
{
	struct xt_connlabel_mtinfo *info = cb->data;
	bool have_labelmap = !connlabel_open();
	int tmp;

	xtables_option_parse(cb);

	switch (cb->entry->id) {
	case O_LABEL:
		if (have_labelmap)
			tmp = nfct_labelmap_get_bit(map, cb->arg);
		else
			tmp = connlabel_value_parse(cb->arg);

		if (tmp < 0)
			xtables_error(PARAMETER_PROBLEM,
				      "label '%s' not found or invalid value",
				      cb->arg);

		info->bit = tmp;
		if (cb->invert)
			info->options |= XT_CONNLABEL_OP_INVERT;
		break;
	case O_SET:
		info->options |= XT_CONNLABEL_OP_SET;
		break;
	}

}

static const char *connlabel_get_name(int b)
{
	const char *name;

	if (connlabel_open())
		return NULL;

	name = nfct_labelmap_get_name(map, b);
	if (name && strcmp(name, ""))
		return name;
	return NULL;
}

static void
connlabel_mt_print_op(const struct xt_connlabel_mtinfo *info, const char *prefix)
{
	if (info->options & XT_CONNLABEL_OP_SET)
		printf(" %sset", prefix);
}

static void
connlabel_mt_print(const void *ip, const struct xt_entry_match *match, int numeric)
{
	const struct xt_connlabel_mtinfo *info = (const void *)match->data;
	const char *name = connlabel_get_name(info->bit);

	printf(" connlabel");
	if (info->options & XT_CONNLABEL_OP_INVERT)
		printf(" !");
	if (numeric || name == NULL) {
		printf(" %u", info->bit);
	} else {
		printf(" '%s'", name);
	}
	connlabel_mt_print_op(info, "");
}

static void
connlabel_mt_save(const void *ip, const struct xt_entry_match *match)
{
	const struct xt_connlabel_mtinfo *info = (const void *)match->data;
	const char *name = connlabel_get_name(info->bit);

	if (info->options & XT_CONNLABEL_OP_INVERT)
		printf(" !");
	if (name)
		printf(" --label \"%s\"", name);
	else
		printf(" --label \"%u\"", info->bit);
	connlabel_mt_print_op(info, "--");
}

static int connlabel_mt_xlate(struct xt_xlate *xl,
			      const struct xt_xlate_mt_params *params)
{
	const struct xt_connlabel_mtinfo *info =
		(const void *)params->match->data;
	const char *name = connlabel_get_name(info->bit);
	char *valbuf = NULL;

	if (name == NULL) {
		if (asprintf(&valbuf, "%u", info->bit) < 0)
			return 0;
		name = valbuf;
	}

	if (info->options & XT_CONNLABEL_OP_SET)
		xt_xlate_add(xl, "ct label set %s ", name);

	xt_xlate_add(xl, "ct label ");
	if (info->options & XT_CONNLABEL_OP_INVERT)
		xt_xlate_add(xl, "and %s != ", name);
	xt_xlate_add(xl, "%s", name);

	free(valbuf);
	return 1;
}

static struct xtables_match connlabel_mt_reg = {
	.family        = NFPROTO_UNSPEC,
	.name          = "connlabel",
	.version       = XTABLES_VERSION,
	.size          = XT_ALIGN(sizeof(struct xt_connlabel_mtinfo)),
	.userspacesize = offsetof(struct xt_connlabel_mtinfo, bit),
	.help          = connlabel_mt_help,
	.print         = connlabel_mt_print,
	.save          = connlabel_mt_save,
	.x6_parse      = connlabel_mt_parse,
	.x6_options    = connlabel_mt_opts,
	.xlate	       = connlabel_mt_xlate,
};

void _init(void)
{
	xtables_register_match(&connlabel_mt_reg);
}