/*
 * gre.c - userspace GRE tunnel
 *
 * Copyright (C) 2015 - 2017, Xiaoxiao <i@pxx.io>
 * Copyright (C) 2019 - 2020, Mikael Magnusson <mikma@users.sourceforge.net>
 *
 * 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 <http://www.gnu.org/licenses/>.
 */

#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <linux/if_tun.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <netdb.h>
#include <netinet/in.h>
#include <pwd.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

static int tun;
static int sock;
static struct sockaddr_storage remote;
static size_t remote_len;
static int remote_any;
static short type = IFF_TUN;
static uint8_t mac[6];
static int foreground;

uint8_t buf[65536];

static void gre_cb(void);
static int gre_ipv4(uint8_t **buf, int *n);
static int gre_ipv6(uint8_t **buf, int *n, const struct sockaddr_in6 *src);
static int gre_any(const uint8_t *buf, int n);
static int tun_cb(void);
static int tun_new(short type, const char *dev);
static int setnonblock(int fd);
static int runas(const char *user);
static int daemonize(void);
static int inet_addr_storage(const char *cp, struct sockaddr_storage *sp, size_t *sp_len);
static int is_any(const struct sockaddr_storage *ss);

static int usage(const char *arg0)
{
    printf("usage: %s [-n|--tun|-p|--tap] <tun> [remote] local\n", arg0);
    exit(EXIT_FAILURE);
}


int main(int argc, char **argv)
{
    fd_set readset;

    while (1)
    {
        int option_index = 0;
        static struct option long_options[] = {
            {"foreground", no_argument, 0, 'f' },
            {"help", no_argument, 0, 'h' },
            {"tun", no_argument, 0,  'n' },
            {"tap", no_argument, 0,  'p' },
            {0,     0,           0,  0 }
        };

        char c = getopt_long(argc, argv, "fnp",
                             long_options, &option_index);
        if (c == -1)
            break;

        switch (c) {
        case 'f':
            foreground = 1;
            break;

        case 'h':
            usage(argv[0]);
            break;

        case 'n':
            printf("tun\n");
            type = IFF_TUN;
            break;

        case 'p':
            printf("tap\n");
            type = IFF_TAP;
            break;

        case '?':
            break;

        default:
            printf("?? getopt returned character code 0%o ??\n", c);
        }
    }

    int optnum = argc - optind;

    if (optnum < 2 || optnum  > 3)
    {
        usage(argv[0]);
    }

    const char *dev = argv[optind++];
    const char *remote_opt = NULL;
    const char *local_opt = NULL;

    if (optnum == 3)
        remote_opt = argv[optind++];

    local_opt = argv[optind++];

    tun = tun_new(type, dev);
    if (tun < 0)
    {
        printf("failed to init %s device\n", type == IFF_TAP ? "tap" : "tun");
        return EXIT_FAILURE;
    }

    struct sockaddr_storage local;
    size_t local_len = 0;
    if (inet_addr_storage(local_opt, &local, &local_len))
    {
        fprintf(stderr, "bad local address\n");
        return EXIT_FAILURE;
    }

    sock = socket(local.ss_family, SOCK_RAW, IPPROTO_GRE);
    if (sock < 0)
    {
        perror("socket");
        return EXIT_FAILURE;
    }

    if (bind(sock, (struct sockaddr *)&local, local_len) != 0)
    {
        perror("bind");
        return EXIT_FAILURE;
    }

    if (!remote_opt)
    {
        if (local.ss_family == AF_INET)
            remote_opt = "0.0.0.0";
        else
            remote_opt = "::";
    }

    if (inet_addr_storage(remote_opt, &remote, &remote_len))
    {
        fprintf(stderr, "bad remote address\n");
        return EXIT_FAILURE;
    }

    remote_any = is_any(&remote);
    if (remote_any)
    {
        fprintf(stderr, "remote any\n");

        if (type == IFF_TUN)
        {
            fprintf(stderr, "tun required specifying remote\n");
            return EXIT_FAILURE;
        }

        if (local.ss_family != AF_INET)
        {
            fprintf(stderr, "IPv4 required for mGRE\n");
            return EXIT_FAILURE;
        }
    }

    setnonblock(sock);
    setnonblock(tun);
    runas("nobody");

    if (!foreground)
        daemonize();

    int maxfd = (tun > sock ? tun : sock) + 1;
    while (1)
    {
        FD_ZERO(&readset);
        FD_SET(tun, &readset);
        FD_SET(sock, &readset);

        int r = select(maxfd, &readset, NULL, NULL, NULL);
        if (r < 0)
        {
            if (errno == EINTR)
            {
                continue;
            }
            else
            {
                perror("select");
                break;
            }
        }

        if (FD_ISSET(sock, &readset))
        {
            gre_cb();
        }

        if (FD_ISSET(tun, &readset))
        {
            if (tun_cb() < 0)
                return 0;
        }
    }

    return 0;
}

static void gre_cb(void)
{
    int res;
    int n;
    struct sockaddr_storage src;
    socklen_t src_len = sizeof(src);
    uint8_t offset = 0; //type == IFF_TAP ? 6 + 6 + 2 : 0;

    memset(&src, 0, src_len);
    n = recvfrom(sock, buf + offset, sizeof(buf)-offset, 0, (struct sockaddr*)&src, &src_len);
    if (n < 0)
    {
        perror("recv");
        return;
    }

    uint8_t *frame = buf + offset;

    switch (remote.ss_family) {
        case AF_INET: res = gre_ipv4(&frame, &n); break;
        case AF_INET6: res = gre_ipv6(&frame, &n, (const struct sockaddr_in6*)&src); break;
        default: printf("Unknown family\n"); return;
    }

    if (res < 0)
    {
        printf("Failed GRE decode\n");
        return;
    }

    if ( 0 ) //type == IFF_TAP )
    {
        /* Set tuntap header */
        memset(frame - offset, 0, 2);
        memcpy(frame - offset + 2, frame + 2, 2);
        /* Set MAC destination */
        memcpy(frame - offset + 4, mac, sizeof(mac));
        /* Set MAC source */
        memset(frame - offset + 4 + 6, 2, 6);
        /* Protocol reused from tuntap header. */
        memcpy(frame - offset + 4 + 6 + 6, frame - offset + 2, 2);
        gre_any(frame - offset, n + offset);
    } else {
        gre_any(frame, n);
    }
}

static int gre_ipv4(uint8_t **bufp, int *np)
{
    uint8_t *buf = *bufp;
    int n = *np;
    int ihl;    // IP header length

    ihl = 4 * (buf[0] & 0x0f);
    if (ihl > 60 || ihl < 20)
    {
        printf("IPv4 header too long\n");
        return -1;
    }
    // check source IPv4 address
    const struct sockaddr_in *remote_in = (const struct sockaddr_in *)&remote;
    if (!remote_any && *(uint32_t *)(buf + 12) != remote_in->sin_addr.s_addr)
    {
        printf("Wrong IPv4 source\n");
        return -1;
    }

    *bufp = buf + ihl;
    *np = n - ihl;
    return 0;
}

static int gre_ipv6(uint8_t **bufp, int *np, const struct sockaddr_in6 *src)
{
    /* uint8_t *buf = *bufp; */
    int n = *np;

    if (n < 40)
    {
        return -1;
    }
    // check source IPv6 address
    const struct sockaddr_in6 *remote_in6 = (const struct sockaddr_in6 *)&remote;
    if (memcmp(src->sin6_addr.s6_addr, remote_in6->sin6_addr.s6_addr, 16) != 0)
    {
        return -1;
    }

    /* Unchanged */
    /* *bufp = buf; */
    /* *np = n; */
    return 0;
}

static int gre_any(const uint8_t *buf, int n)
{
    //uint8_t offset = type == IFF_TAP ? 6 + 6 + 2 : 0;
    uint8_t offset = 0;

    // parse GRE header
    uint16_t *hdr = (uint16_t *)(buf+offset);
    if (hdr[0] != 0)
    {
        printf("Bad GRE header: %04x %04x\n", ntohs(hdr[0]), ntohs(hdr[1]));
        return -1;
    }
    uint16_t protocol = ntohs(*(uint16_t *)(buf + 2));
    if (type == IFF_TUN || remote_any)
    {
        switch (protocol) {
        case ETHERTYPE_IP:
        case ETHERTYPE_IPV6:
            break;
        default:
            printf("Wrong proto: %04x\n", protocol);
            return -1;
        }
    } else {
        if (protocol != ETH_P_TEB)
        {
            printf("Wrong proto: %04x\n", protocol);
            return -1;
        }
    }

    if (remote_any)
    {
        // TODO Add header
    }

    return write(tun, buf, n);
}

static int tun_cb(void)
{
    int n;
    struct sockaddr_storage dest;

    if (!remote_any)
    {
        memcpy(&dest, &remote, sizeof(struct sockaddr_storage));
    } else {
        memset(&dest, 0, sizeof(struct sockaddr_storage));
    }

    n = read(tun, buf, sizeof(buf));
    if (n < 0)
    {
        int err = errno;
        perror("read");
        if (err == EBADFD)
            return -1;

        return 0;
    }
    uint16_t proto = ntohs(*(uint16_t *)(buf + 2));
    uint8_t *frame = buf;

    if (remote_any)
    {
        printf("%02x:%02x:%02x:%02x:%02x:%02x\n",
          frame[4 + 0],
          frame[4 + 1],
          frame[4 + 2],
          frame[4 + 3],
          frame[4 + 4],
          frame[4 + 5]
        );

        if (frame[4] != 0x2e || frame[5] != 0x2e)
        {
            printf("Bad magic number!\n");
            return 0;
        }

        struct sockaddr_in *sin = (struct sockaddr_in *)&dest;
        sin->sin_family = AF_INET;
        memcpy(&sin->sin_addr.s_addr, frame + 4 + 2, 4);

        int offset = 2 + 6 + 6;
        frame += offset;
        n -= offset;
    }
    else if (type == IFF_TAP)
    {
        *(uint16_t*)(frame + 2) = htons(ETH_P_TEB);
    }

    frame[0] = 0;
    frame[1] = 0;
    switch (proto) {
    case ETHERTYPE_IP:
    case ETHERTYPE_IPV6:
        break;
    case ETHERTYPE_ARP:
        if (type == IFF_TAP && !remote_any)
            break;
        return 0;
    default:
        return 0;
    }

    if (is_any(&dest))
    {
        printf("Remote unset!\n");
        return -1;
    }

    ssize_t res = sendto(sock, frame, n, 0, (struct sockaddr *)&dest, remote_len);
    if (res < 0)
    {
        perror("sendto");
    }
    return res;
}

static int tun_new(short type, const char *dev)
{
    struct ifreq ifr;
    int fd, err;

    fd = open("/dev/net/tun", O_RDWR);
    if (fd < 0)
    {
        return -1;
    }

    bzero(&ifr, sizeof(struct ifreq));

    ifr.ifr_flags = type;
    if (*dev != '\0')
    {
        strncpy(ifr.ifr_name, dev, IFNAMSIZ);
    }

    err = ioctl(fd, TUNSETIFF, (void *)&ifr);
    if (err < 0)
    {
        return err;
    }

    bzero(&ifr.ifr_hwaddr, sizeof(struct sockaddr));

    err = ioctl(fd, SIOCGIFHWADDR, (void *)&ifr);
    if (err < 0)
    {
        return err;
    }

    memcpy(mac, ifr.ifr_hwaddr.sa_data, 6);

    return fd;
}

static int setnonblock(int fd)
{
    int flags;
    flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1)
    {
        return -1;
    }
    if (-1 == fcntl(fd, F_SETFL, flags | O_NONBLOCK))
    {
        return -1;
    }
    return 0;
}

static int runas(const char *user)
{
    struct passwd *pw_ent = getpwnam(user);

    if (pw_ent != NULL)
    {
        if (setegid(pw_ent->pw_gid) != 0)
        {
            return -1;
        }
        if (seteuid(pw_ent->pw_uid) != 0)
        {
            return -1;
        }
    }

    return 0;
}

static int daemonize(void)
{
    pid_t pid;

    pid = fork();
    if (pid < 0)
    {
        perror("fork");
        return -1;
    }

    if (pid > 0)
    {
        exit(0);
    }

    umask(0);

    if (setsid() < 0)
    {
        perror("setsid");
        return -1;
    }

    return 0;
}

static int inet_addr_storage(const char *cp, struct sockaddr_storage *sp, size_t *sp_len)
{
    struct addrinfo hints;
    struct addrinfo *result = NULL;
    int res;

    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_flags = AI_NUMERICHOST | AI_ADDRCONFIG;
    res = getaddrinfo(cp, NULL, &hints, &result);
    if (res != 0) {
        fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(res));
        return -1;
    }

    memcpy(sp, result->ai_addr, result->ai_addrlen);
    *sp_len = result->ai_addrlen;

    freeaddrinfo(result);
    result = NULL;

    if (sp->ss_family == AF_INET) {
        struct sockaddr_in *sin = (struct sockaddr_in *)sp;
        sin->sin_port = htons(IPPROTO_GRE);
    } else if (sp->ss_family == AF_INET6) {
        struct sockaddr_in6 *sin = (struct sockaddr_in6 *)sp;
        sin->sin6_port = htons(IPPROTO_GRE);
    }

    return 0;
}

static int is_any(const struct sockaddr_storage *ss)
{
    switch(ss->ss_family) {
    case AF_INET: {
        const struct sockaddr_in *sin = (const struct sockaddr_in*)ss;
        return sin->sin_addr.s_addr == INADDR_ANY;
    }
    case AF_INET6: {
        const struct sockaddr_in6 *sin6 = (const struct sockaddr_in6*)ss;
        return !memcmp(&in6addr_any, sin6->sin6_addr.s6_addr, sizeof(struct in6_addr));
    }
    default:
        return -1;
    }
}