/* $OpenBSD: tcpdrop.c,v 1.4 2004/05/22 23:55:22 deraadt Exp $ */

/*-
 * Copyright (c) 2009 Juli Mallett <jmallett@FreeBSD.org>
 * Copyright (c) 2004 Markus Friedl <markus@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/cdefs.h>
__FBSDID("$FreeBSD$");

#include <sys/param.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/socketvar.h>
#include <sys/sysctl.h>

#include <netinet/in.h>
#include <netinet/in_pcb.h>
#define TCPSTATES
#include <netinet/tcp_fsm.h>
#include <netinet/tcp_var.h>

#include <err.h>
#include <netdb.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define TCPDROP_FOREIGN         0
#define TCPDROP_LOCAL           1

struct host_service {
        char hs_host[NI_MAXHOST];
        char hs_service[NI_MAXSERV];
};

static bool tcpdrop_list_commands = false;

static char *findport(const char *);
static struct xtcpcb *tcpcblist(const char *, size_t *);
static void sockinfo(const struct sockaddr *, struct host_service *);
static bool tcpdrop(const struct sockaddr *, const struct sockaddr *);
static bool tcpdropall(void);
static bool tcpdropbyname(const char *, const char *, const char *,
    const char *);
static bool tcpdropconn(const struct in_conninfo *);
static void usage(void);

/*
 * Drop a tcp connection.
 */
int
main(int argc, char *argv[])
{
        char *lport, *fport;
        bool dropall;
        int ch;

        dropall = false;

        while ((ch = getopt(argc, argv, "al")) != -1) {
                switch (ch) {
                case 'a':
                        dropall = true;
                        break;
                case 'l':
                        tcpdrop_list_commands = true;
                        break;
                default:
                        usage();
                }
        }
        argc -= optind;
        argv += optind;

        if (dropall) {
                if (argc != 0)
                        usage();
                if (!tcpdropall())
                        exit(1);
                exit(0);
        }

        if ((argc != 2 && argc != 4) || tcpdrop_list_commands)
                usage();

        if (argc == 2) {
                lport = findport(argv[0]);
                fport = findport(argv[1]);
                if (lport == NULL || lport[1] == '\0' || fport == NULL ||
                    fport[1] == '\0')
                        usage();
                *lport++ = '\0';
                *fport++ = '\0';
                if (!tcpdropbyname(argv[0], lport, argv[1], fport))
                        exit(1);
        } else if (!tcpdropbyname(argv[0], argv[1], argv[2], argv[3]))
                exit(1);

        exit(0);
}

static char *
findport(const char *arg)
{
        char *dot, *colon;

        /* A strrspn() or strrpbrk() would be nice. */
        dot = strrchr(arg, '.');
        colon = strrchr(arg, ':');
        if (dot == NULL)
                return (colon);
        if (colon == NULL)
                return (dot);
        if (dot < colon)
                return (colon);
        else
                return (dot);
}

static struct xtcpcb *
tcpcblist(const char *name, size_t *len0)
{
        struct xtcpcb *tcp;
        size_t len;
        int rv;

        len = 0;
        rv = sysctlbyname(name, NULL, &len, NULL, 0);
        if (rv == -1)
                err(1, "sysctlbyname %s", name);

        if (len == 0)
                errx(1, "%s is empty", name);

        tcp = malloc(len);
        if (tcp == NULL)
                errx(1, "malloc failed");

        rv = sysctlbyname(name, tcp, &len, NULL, 0);
        if (rv == -1)
                err(1, "sysctlbyname %s", name);

        *len0 = len;
        return (tcp);
}

static void
sockinfo(const struct sockaddr *sa, struct host_service *hs)
{
        static const int flags = NI_NUMERICHOST | NI_NUMERICSERV;
        int rv;

        rv = getnameinfo(sa, sa->sa_len, hs->hs_host, sizeof hs->hs_host,
            hs->hs_service, sizeof hs->hs_service, flags);
        if (rv == -1)
                err(1, "getnameinfo");
}

static bool
tcpdrop(const struct sockaddr *lsa, const struct sockaddr *fsa)
{
        struct host_service local, foreign;
        struct sockaddr_storage addrs[2];
        int rv;

        memcpy(&addrs[TCPDROP_FOREIGN], fsa, fsa->sa_len);
        memcpy(&addrs[TCPDROP_LOCAL], lsa, lsa->sa_len);

        sockinfo(lsa, &local);
        sockinfo(fsa, &foreign);

        if (tcpdrop_list_commands) {
                printf("tcpdrop %s %s %s %s\n", local.hs_host, local.hs_service,
                    foreign.hs_host, foreign.hs_service);
                return (true);
        }

        rv = sysctlbyname("net.inet.tcp.drop", NULL, NULL, &addrs,
            sizeof addrs);
        if (rv == -1) {
                warn("%s %s %s %s", local.hs_host, local.hs_service,
                    foreign.hs_host, foreign.hs_service);
                return (false);
        }
        printf("%s %s %s %s: dropped\n", local.hs_host, local.hs_service,
            foreign.hs_host, foreign.hs_service);
        return (true);
}

static bool
tcpdropall(void)
{
        struct xtcpcb *tcp;
        size_t len;
        bool ok;
        int cnt, i;

        ok = true;

        tcp = tcpcblist("net.inet.tcp.pcblist", &len);
        cnt = len / sizeof(*tcp);

        for (i = 0; i < cnt; ++i) {
                struct xtcpcb *xpcb;
                struct tcpcb *tp;
                struct inpcb *inp;

                xpcb = &tcp[i];
                if (xpcb->xt_len != sizeof(*xpcb))
                        break;

                tp = &xpcb->xt_tp;
                inp = &xpcb->xt_inp;

                /*
                 * XXX
                 * Check protocol, support just v4 or v6, etc.
                 */

                /* Skip listening sockets.  */
                if (tp->t_state == TCPS_LISTEN)
                        continue;

                if (!tcpdropconn(&inp->inp_inc))
                        ok = false;
        }
        free(tcp);

        return (ok);
}

static bool
tcpdropbyname(const char *lhost, const char *lport, const char *fhost,
    const char *fport)
{
        static const struct addrinfo hints = {
                /*
                 * Look for streams in all domains.
                 */
                .ai_family = AF_UNSPEC,
                .ai_socktype = SOCK_STREAM,
        };
        struct addrinfo *ail, *local, *aif, *foreign;
        int error;
        bool ok, infamily;

        error = getaddrinfo(lhost, lport, &hints, &local);
        if (error != 0)
                errx(1, "getaddrinfo: %s port %s: %s", lhost, lport,
                    gai_strerror(error));

        error = getaddrinfo(fhost, fport, &hints, &foreign);
        if (error != 0) {
                freeaddrinfo(local); /* XXX gratuitous */
                errx(1, "getaddrinfo: %s port %s: %s", fhost, fport,
                    gai_strerror(error));
        }

        ok = true;
        infamily = false;

        /*
         * Try every combination of local and foreign address pairs.
         */
        for (ail = local; ail != NULL; ail = ail->ai_next) {
                for (aif = foreign; aif != NULL; aif = aif->ai_next) {
                        if (ail->ai_family != aif->ai_family)
                                continue;
                        infamily = true;
                        if (!tcpdrop(ail->ai_addr, aif->ai_addr))
                                ok = false;
                }
        }

        if (!infamily) {
                warnx("%s %s %s %s: different address families", lhost, lport,
                    fhost, fport);
                ok = false;
        }

        freeaddrinfo(local);
        freeaddrinfo(foreign);

        return (ok);
}

static bool
tcpdropconn(const struct in_conninfo *inc)
{
        struct sockaddr *local, *foreign;
        struct sockaddr_in6 sin6[2];
        struct sockaddr_in sin4[2];

        if (inc->inc_isipv6) {
                memset(sin6, 0, sizeof sin6);

                sin6[TCPDROP_LOCAL].sin6_len = sizeof sin6[TCPDROP_LOCAL];
                sin6[TCPDROP_LOCAL].sin6_family = AF_INET6;
                sin6[TCPDROP_LOCAL].sin6_port = inc->inc_lport;
                memcpy(&sin6[TCPDROP_LOCAL].sin6_addr, &inc->inc6_laddr,
                    sizeof inc->inc6_laddr);
                local = (struct sockaddr *)&sin6[TCPDROP_LOCAL];

                sin6[TCPDROP_FOREIGN].sin6_len = sizeof sin6[TCPDROP_FOREIGN];
                sin6[TCPDROP_FOREIGN].sin6_family = AF_INET6;
                sin6[TCPDROP_FOREIGN].sin6_port = inc->inc_fport;
                memcpy(&sin6[TCPDROP_FOREIGN].sin6_addr, &inc->inc6_faddr,
                    sizeof inc->inc6_faddr);
                foreign = (struct sockaddr *)&sin6[TCPDROP_FOREIGN];
        } else {
                memset(&sin4[TCPDROP_LOCAL], 0, sizeof sin4[TCPDROP_LOCAL]);

                sin4[TCPDROP_LOCAL].sin_len = sizeof sin4[TCPDROP_LOCAL];
                sin4[TCPDROP_LOCAL].sin_family = AF_INET;
                sin4[TCPDROP_LOCAL].sin_port = inc->inc_lport;
                memcpy(&sin4[TCPDROP_LOCAL].sin_addr, &inc->inc_laddr,
                    sizeof inc->inc_laddr);
                local = (struct sockaddr *)&sin4[TCPDROP_LOCAL];

                sin4[TCPDROP_FOREIGN].sin_len = sizeof sin4[TCPDROP_FOREIGN];
                sin4[TCPDROP_FOREIGN].sin_family = AF_INET;
                sin4[TCPDROP_FOREIGN].sin_port = inc->inc_fport;
                memcpy(&sin4[TCPDROP_FOREIGN].sin_addr, &inc->inc_faddr,
                    sizeof inc->inc_faddr);
                foreign = (struct sockaddr *)&sin4[TCPDROP_FOREIGN];
        }

        return (tcpdrop(local, foreign));
}

static void
usage(void)
{
        fprintf(stderr,
"usage: tcpdrop local-address local-port foreign-address foreign-port\n"
"       tcpdrop local-address:local-port foreign-address:foreign-port\n"
"       tcpdrop local-address.local-port foreign-address.foreign-port\n"
"       tcpdrop [-l] -a\n");
        exit(1);
}