/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2023 Rubicon Communications, LLC (Netgate)
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *    - Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *    - Redistributions in binary form must reproduce the above
 *      copyright notice, this list of conditions and the following
 *      disclaimer in the documentation and/or other materials provided
 *      with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 */
#include <sys/cdefs.h>

#include <err.h>
#include <errno.h>
#include <netdb.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <net/pflow.h>

#include <netlink/netlink.h>
#include <netlink/netlink_generic.h>
#include <netlink/netlink_snl.h>
#include <netlink/netlink_snl_generic.h>
#include <netlink/netlink_snl_route.h>

static int get(int id);

static bool verbose = false;

extern char *__progname;

static void
usage(void)
{
        fprintf(stderr,
"usage: %s [-lc] [-d id] [-s id ...] [-v]\n",
            __progname);

        exit(1);
}

static int
pflow_to_id(const char *name)
{
        int ret, id;

        ret = sscanf(name, "pflow%d", &id);
        if (ret == 1)
                return (id);

        ret = sscanf(name, "%d", &id);
        if (ret == 1)
                return (id);

        return (-1);
}

struct pflowctl_list {
        int id;
};
#define _IN(_field)     offsetof(struct genlmsghdr, _field)
#define _OUT(_field)    offsetof(struct pflowctl_list, _field)
static struct snl_attr_parser ap_list[] = {
        { .type = PFLOWNL_L_ID, .off = _OUT(id), .cb = snl_attr_get_int32 },
};
static struct snl_field_parser fp_list[] = {};
#undef _IN
#undef _OUT
SNL_DECLARE_PARSER(list_parser, struct genlmsghdr, fp_list, ap_list);

static int
list(void)
{
        struct snl_state ss = {};
        struct snl_errmsg_data e = {};
        struct pflowctl_list l = {};
        struct snl_writer nw;
        struct nlmsghdr *hdr;
        uint32_t seq_id;
        int family_id;

        snl_init(&ss, NETLINK_GENERIC);
        family_id = snl_get_genl_family(&ss, PFLOWNL_FAMILY_NAME);
        if (family_id == 0)
                errx(1, "pflow.ko is not loaded.");

        snl_init_writer(&ss, &nw);
        hdr = snl_create_genl_msg_request(&nw, family_id, PFLOWNL_CMD_LIST);

        hdr = snl_finalize_msg(&nw);
        if (hdr == NULL)
                return (ENOMEM);
        seq_id = hdr->nlmsg_seq;

        snl_send_message(&ss, hdr);

        while ((hdr = snl_read_reply_multi(&ss, seq_id, &e)) != NULL) {
                if (! snl_parse_nlmsg(&ss, hdr, &list_parser, &l))
                        continue;

                get(l.id);
        }

        if (e.error)
                errc(1, e.error, "failed to list");

        return (0);
}

struct pflowctl_create {
        int id;
};
#define _IN(_field)     offsetof(struct genlmsghsdr, _field)
#define _OUT(_field)    offsetof(struct pflowctl_create, _field)
static struct snl_attr_parser ap_create[] = {
        { .type = PFLOWNL_CREATE_ID, .off = _OUT(id), .cb = snl_attr_get_int32 },
};
static struct snl_field_parser pf_create[] = {};
#undef _IN
#undef _OUT
SNL_DECLARE_PARSER(create_parser, struct genlmsghdr, pf_create, ap_create);

static int
create(void)
{
        struct snl_state ss = {};
        struct snl_errmsg_data e = {};
        struct pflowctl_create c = {};
        struct snl_writer nw;
        struct nlmsghdr *hdr;
        uint32_t seq_id;
        int family_id;

        snl_init(&ss, NETLINK_GENERIC);
        family_id = snl_get_genl_family(&ss, PFLOWNL_FAMILY_NAME);
        if (family_id == 0)
                errx(1, "pflow.ko is not loaded.");

        snl_init_writer(&ss, &nw);
        hdr = snl_create_genl_msg_request(&nw, family_id, PFLOWNL_CMD_CREATE);

        hdr = snl_finalize_msg(&nw);
        if (hdr == NULL)
                return (ENOMEM);
        seq_id = hdr->nlmsg_seq;

        snl_send_message(&ss, hdr);

        while ((hdr = snl_read_reply_multi(&ss, seq_id, &e)) != NULL) {
                if (! snl_parse_nlmsg(&ss, hdr, &create_parser, &c))
                        continue;

                printf("pflow%d\n", c.id);
        }

        if (e.error)
                errc(1, e.error, "failed to create");

        return (0);
}

static int
del(char *idstr)
{
        struct snl_state ss = {};
        struct snl_errmsg_data e = {};
        struct snl_writer nw;
        struct nlmsghdr *hdr;
        int family_id;
        int id;

        id = pflow_to_id(idstr);
        if (id < 0)
                return (EINVAL);

        snl_init(&ss, NETLINK_GENERIC);
        family_id = snl_get_genl_family(&ss, PFLOWNL_FAMILY_NAME);
        if (family_id == 0)
                errx(1, "pflow.ko is not loaded.");

        snl_init_writer(&ss, &nw);
        hdr = snl_create_genl_msg_request(&nw, family_id, PFLOWNL_CMD_DEL);

        snl_add_msg_attr_s32(&nw, PFLOWNL_DEL_ID, id);

        hdr = snl_finalize_msg(&nw);
        if (hdr == NULL)
                return (ENOMEM);

        snl_send_message(&ss, hdr);
        snl_read_reply_code(&ss, hdr->nlmsg_seq, &e);

        if (e.error)
                errc(1, e.error, "failed to delete");

        return (0);
}

struct pflowctl_sockaddr {
        union {
                struct sockaddr_in in;
                struct sockaddr_in6 in6;
                struct sockaddr_storage storage;
        };
};
static bool
pflowctl_post_sockaddr(struct snl_state* ss __unused, void *target)
{
        struct pflowctl_sockaddr *s = (struct pflowctl_sockaddr *)target;

        if (s->storage.ss_family == AF_INET)
                s->storage.ss_len = sizeof(struct sockaddr_in);
        else if (s->storage.ss_family == AF_INET6)
                s->storage.ss_len = sizeof(struct sockaddr_in6);
        else
                return (false);

        return (true);
}
#define _OUT(_field)    offsetof(struct pflowctl_sockaddr, _field)
static struct snl_attr_parser nla_p_sockaddr[] = {
        { .type = PFLOWNL_ADDR_FAMILY, .off = _OUT(in.sin_family), .cb = snl_attr_get_uint8 },
        { .type = PFLOWNL_ADDR_PORT, .off = _OUT(in.sin_port), .cb = snl_attr_get_uint16 },
        { .type = PFLOWNL_ADDR_IP, .off = _OUT(in.sin_addr), .cb = snl_attr_get_in_addr },
        { .type = PFLOWNL_ADDR_IP6, .off = _OUT(in6.sin6_addr), .cb = snl_attr_get_in6_addr },
};
SNL_DECLARE_ATTR_PARSER_EXT(sockaddr_parser, 0, nla_p_sockaddr, pflowctl_post_sockaddr);
#undef _OUT

struct pflowctl_get {
        int id;
        int version;
        struct pflowctl_sockaddr src;
        struct pflowctl_sockaddr dst;
        uint32_t obs_dom;
        uint8_t so_status;
};
#define _IN(_field)     offsetof(struct genlmsghdr, _field)
#define _OUT(_field)    offsetof(struct pflowctl_get, _field)
static struct snl_attr_parser ap_get[] = {
        { .type = PFLOWNL_GET_ID, .off = _OUT(id), .cb = snl_attr_get_int32 },
        { .type = PFLOWNL_GET_VERSION, .off = _OUT(version), .cb = snl_attr_get_int16 },
        { .type = PFLOWNL_GET_SRC, .off = _OUT(src), .arg = &sockaddr_parser, .cb = snl_attr_get_nested },
        { .type = PFLOWNL_GET_DST, .off = _OUT(dst), .arg = &sockaddr_parser, .cb = snl_attr_get_nested },
        { .type = PFLOWNL_GET_OBSERVATION_DOMAIN, .off = _OUT(obs_dom), .cb = snl_attr_get_uint32 },
        { .type = PFLOWNL_GET_SOCKET_STATUS, .off = _OUT(so_status), .cb = snl_attr_get_uint8 },
};
static struct snl_field_parser fp_get[] = {};
#undef _IN
#undef _OUT
SNL_DECLARE_PARSER(get_parser, struct genlmsghdr, fp_get, ap_get);

static void
print_sockaddr(const char *prefix, const struct sockaddr_storage *s)
{
        char buf[INET6_ADDRSTRLEN];
        int error;

        if (s->ss_family != AF_INET && s->ss_family != AF_INET6)
                return;

        if (s->ss_family == AF_INET ||
            s->ss_family == AF_INET6) {
                error = getnameinfo((const struct sockaddr *)s,
                    s->ss_len, buf, sizeof(buf), NULL, 0,
                    NI_NUMERICHOST);
                if (error)
                        err(1, "sender: %s", gai_strerror(error));
        }

        printf("%s", prefix);
        switch (s->ss_family) {
        case AF_INET: {
                const struct sockaddr_in *sin = (const struct sockaddr_in *)s;
                if (sin->sin_addr.s_addr != INADDR_ANY) {
                        printf("%s", buf);
                        if (sin->sin_port != 0)
                                printf(":%u", ntohs(sin->sin_port));
                }
                break;
        }
        case AF_INET6: {
                const struct sockaddr_in6 *sin6 = (const struct sockaddr_in6 *)s;
                if (!IN6_IS_ADDR_UNSPECIFIED(&sin6->sin6_addr)) {
                        printf("[%s]", buf);
                        if (sin6->sin6_port != 0)
                                printf(":%u", ntohs(sin6->sin6_port));
                }
                break;
        }
        }
}

static int
get(int id)
{
        struct snl_state ss = {};
        struct snl_errmsg_data e = {};
        struct pflowctl_get g = {};
        struct snl_writer nw;
        struct nlmsghdr *hdr;
        uint32_t seq_id;
        int family_id;

        snl_init(&ss, NETLINK_GENERIC);
        family_id = snl_get_genl_family(&ss, PFLOWNL_FAMILY_NAME);
        if (family_id == 0)
                errx(1, "pflow.ko is not loaded.");

        snl_init_writer(&ss, &nw);
        hdr = snl_create_genl_msg_request(&nw, family_id, PFLOWNL_CMD_GET);
        snl_add_msg_attr_s32(&nw, PFLOWNL_GET_ID, id);

        hdr = snl_finalize_msg(&nw);
        if (hdr == NULL)
                return (ENOMEM);
        seq_id = hdr->nlmsg_seq;

        snl_send_message(&ss, hdr);

        while ((hdr = snl_read_reply_multi(&ss, seq_id, &e)) != NULL) {
                if (! snl_parse_nlmsg(&ss, hdr, &get_parser, &g))
                        continue;

                printf("pflow%d: version %d domain %u", g.id, g.version, g.obs_dom);
                print_sockaddr(" src ", &g.src.storage);
                print_sockaddr(" dst ", &g.dst.storage);
                printf("\n");
                if (verbose) {
                        printf("\tsocket: %s\n",
                            g.so_status ? "connected" : "disconnected");
                }
        }

        if (e.error)
                errc(1, e.error, "failed to get");

        return (0);
}

struct pflowctl_set {
        int id;
        uint16_t version;
        struct sockaddr_storage src;
        struct sockaddr_storage dst;
        uint32_t obs_dom;
};
static inline bool
snl_add_msg_attr_sockaddr(struct snl_writer *nw, int attrtype, struct sockaddr_storage *s)
{
        int off = snl_add_msg_attr_nested(nw, attrtype);

        snl_add_msg_attr_u8(nw, PFLOWNL_ADDR_FAMILY, s->ss_family);

        switch (s->ss_family) {
        case AF_INET: {
                const struct sockaddr_in *in = (const struct sockaddr_in *)s;
                snl_add_msg_attr_u16(nw, PFLOWNL_ADDR_PORT, in->sin_port);
                snl_add_msg_attr_ip4(nw, PFLOWNL_ADDR_IP, &in->sin_addr);
                break;
        }
        case AF_INET6: {
                const struct sockaddr_in6 *in6 = (const struct sockaddr_in6 *)s;
                snl_add_msg_attr_u16(nw, PFLOWNL_ADDR_PORT, in6->sin6_port);
                snl_add_msg_attr_ip6(nw, PFLOWNL_ADDR_IP6, &in6->sin6_addr);
                break;
        }
        default:
                return (false);
        }
        snl_end_attr_nested(nw, off);

        return (true);
}

static int
do_set(struct pflowctl_set *s)
{
        struct snl_state ss = {};
        struct snl_errmsg_data e = {};
        struct snl_writer nw;
        struct nlmsghdr *hdr;
        int family_id;

        snl_init(&ss, NETLINK_GENERIC);
        family_id = snl_get_genl_family(&ss, PFLOWNL_FAMILY_NAME);
        if (family_id == 0)
                errx(1, "pflow.ko is not loaded.");

        snl_init_writer(&ss, &nw);
        snl_create_genl_msg_request(&nw, family_id, PFLOWNL_CMD_SET);

        snl_add_msg_attr_s32(&nw, PFLOWNL_SET_ID, s->id);
        if (s->version != 0)
                snl_add_msg_attr_u16(&nw, PFLOWNL_SET_VERSION, s->version);
        if (s->src.ss_len != 0)
                snl_add_msg_attr_sockaddr(&nw, PFLOWNL_SET_SRC, &s->src);
        if (s->dst.ss_len != 0)
                snl_add_msg_attr_sockaddr(&nw, PFLOWNL_SET_DST, &s->dst);
        if (s->obs_dom != 0)
                snl_add_msg_attr_u32(&nw, PFLOWNL_SET_OBSERVATION_DOMAIN, s->obs_dom);

        hdr = snl_finalize_msg(&nw);
        if (hdr == NULL)
                return (1);

        snl_send_message(&ss, hdr);
        snl_read_reply_code(&ss, hdr->nlmsg_seq, &e);

        if (e.error)
                errc(1, e.error, "failed to set");

        return (0);
}

static void
pflowctl_addr(const char *val, struct sockaddr_storage *ss)
{
        struct addrinfo *res0;
        int error;
        bool flag;
        char *ip, *port;
        char buf[sysconf(_SC_HOST_NAME_MAX) + 1 + sizeof(":65535")];
        struct addrinfo hints = {
                .ai_family = AF_UNSPEC,
                .ai_socktype = SOCK_DGRAM, /*dummy*/
                .ai_flags = AI_NUMERICHOST,
        };

        if (strlcpy(buf, val, sizeof(buf)) >= sizeof(buf))
                errx(1, "%s bad value", val);

        port = NULL;
        flag = *buf == '[';

        for (char *cp = buf; *cp; ++cp) {
                if (*cp == ']' && *(cp + 1) == ':' && flag) {
                        *cp = '\0';
                        *(cp + 1) = '\0';
                        port = cp + 2;
                        break;
                }
                if (*cp == ']' && *(cp + 1) == '\0' && flag) {
                        *cp = '\0';
                        port = NULL;
                        break;
                }
                if (*cp == ':' && !flag) {
                        *cp = '\0';
                        port = cp + 1;
                        break;
                }
        }

        ip = buf;
        if (flag)
                ip++;

        if ((error = getaddrinfo(ip, port, &hints, &res0)) != 0)
                errx(1, "error in parsing address string: %s",
                    gai_strerror(error));

        memcpy(ss, res0->ai_addr, res0->ai_addr->sa_len);
        freeaddrinfo(res0);
}

static int
set(char *idstr, int argc, char *argv[])
{
        struct pflowctl_set s = {};

        s.id = pflow_to_id(idstr);
        if (s.id < 0)
                return (EINVAL);

        while (argc > 0) {
                if (strcmp(argv[0], "src") == 0) {
                        if (argc < 2)
                                usage();

                        pflowctl_addr(argv[1], &s.src);

                        argc -= 2;
                        argv += 2;
                } else if (strcmp(argv[0], "dst") == 0) {
                        if (argc < 2)
                                usage();

                        pflowctl_addr(argv[1], &s.dst);

                        argc -= 2;
                        argv += 2;
                } else if (strcmp(argv[0], "proto") == 0) {
                        if (argc < 2)
                                usage();

                        s.version = strtol(argv[1], NULL, 10);

                        argc -= 2;
                        argv += 2;
                } else if (strcmp(argv[0], "domain") == 0) {
                        if (argc < 2)
                                usage();

                        s.obs_dom = strtol(argv[1], NULL, 10);

                        argc -= 2;
                        argv += 2;
                } else {
                        usage();
                }
        }

        return (do_set(&s));
}

static const struct snl_hdr_parser *all_parsers[] = {
        &list_parser,
        &get_parser,
};

enum pflowctl_op_t {
        OP_HELP,
        OP_LIST,
        OP_CREATE,
        OP_DELETE,
        OP_SET,
};
int
main(int argc, char *argv[])
{
        int ch;
        enum pflowctl_op_t op = OP_HELP;
        char **set_args = NULL;
        size_t set_arg_count = 0;

        SNL_VERIFY_PARSERS(all_parsers);

        if (argc < 2)
                usage();

        while ((ch = getopt(argc, argv,
            "lcd:s:v")) != -1) {
                switch (ch) {
                case 'l':
                        op = OP_LIST;
                        break;
                case 'c':
                        op = OP_CREATE;
                        break;
                case 'd':
                        op = OP_DELETE;
                        break;
                case 's':
                        op = OP_SET;
                        set_arg_count = argc - optind;
                        set_args = argv + optind;
                        break;
                case 'v':
                        verbose = true;
                        break;
                }
        }

        switch (op) {
        case OP_LIST:
                return (list());
        case OP_CREATE:
                return (create());
        case OP_DELETE:
                return (del(optarg));
        case OP_SET:
                return (set(optarg, set_arg_count, set_args));
        case OP_HELP:
                usage();
                break;
        }

        return (0);
}