/*
    mtr  --  a network diagnostic tool
    Copyright (C) 2016  Matt Kimball

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License version 2 as
    published by the Free Software Foundation.

    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, write to the Free Software
    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/

#include <netinet/in.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <unistd.h>

#include "packet/protocols.h"

#define MAX_PACKET_SIZE 9000

/*
    The first probe sent by mtr-packet will have this sequence number,
    so wait for an ICMP packet with this sequence ID.
*/
#define SEQUENCE_NUM 33000

/*
    Check to see if the packet we've recieved is intended for this test
    process.  We expected the ICMP sequence number to be equal to our
    process ID.
*/
bool is_packet_for_us4(
    char *packet,
    int packet_size)
{
    int ip_icmp_size = sizeof(struct IPHeader) + sizeof(struct ICMPHeader);
    int expected_sequence;
    struct IPHeader *ip;
    struct ICMPHeader *icmp;

    if (packet_size < ip_icmp_size) {
        return false;
    }

    ip = (struct IPHeader *)packet;
    icmp = (struct ICMPHeader *)(ip + 1);

    expected_sequence = htons(SEQUENCE_NUM);
    if (icmp->sequence == expected_sequence) {
        return true;
    }

    return false;
}

/*
    Check to see if the ICMPv6 packet is for us.
    Unlike ICMPv4 packets, ICMPv6 packets don't include the IP header.
*/
bool is_packet_for_us6(
    char *packet,
    int packet_size)
{
    int expected_sequence;
    struct ICMPHeader *icmp;

    if (packet_size < sizeof(struct ICMPHeader)) {
        return false;
    }

    icmp = (struct ICMPHeader *)packet;

    expected_sequence = htons(SEQUENCE_NUM);
    if (icmp->sequence == expected_sequence) {
        return true;
    }

    return false;
}

/*
    Check that all the bytes in the body of the packet have the same value.
    If so, return that value.  If not, return -1.
*/
int get_packet_pattern(
    unsigned char *packet,
    int packet_size)
{
    int fill_value;
    int i;

    if (packet_size <= 0) {
        return -1;
    }

    fill_value = packet[0];
    for (i = 1; i < packet_size; i++) {
        if (packet[i] != fill_value) {
            return -1;
        }
    }

    return fill_value;
}

/*  Print information about the ICMPv4 packet we received  */
void dump_packet_info4(
    char *packet,
    int packet_size)
{
    int ip_icmp_size = sizeof(struct IPHeader) + sizeof(struct ICMPHeader);
    int pattern;
    struct IPHeader *ip;
    struct ICMPHeader *icmp;
    unsigned char *body;
    int body_size;

    ip = (struct IPHeader *)packet;
    icmp = (struct ICMPHeader *)(ip + 1);
    body = (unsigned char *)(icmp + 1);
    body_size = packet_size - ip_icmp_size;

    printf("size %d\n", packet_size);
    printf("tos %d\n", ip->tos);

    pattern = get_packet_pattern(body, body_size);
    if (pattern < 0) {
        printf("bitpattern none\n");
    } else {
        printf("bitpattern %d\n", pattern);
    }
}

/*  Print information about an ICMPv6 packet  */
void dump_packet_info6(
    char *packet,
    int packet_size)
{
    int pattern;
    struct ICMPHeader *icmp;
    unsigned char *body;
    int body_size;
    int total_size;

    icmp = (struct ICMPHeader *)packet;
    body = (unsigned char *)(icmp + 1);
    body_size = packet_size - sizeof(struct ICMPHeader);

    total_size = packet_size + sizeof(struct IP6Header);
    printf("size %d\n", total_size);

    pattern = get_packet_pattern(body, body_size);
    if (pattern < 0) {
        printf("bitpattern none\n");
    } else {
        printf("bitpattern %d\n", pattern);
    }
}

/*  Receive ICMP packets until we get one intended for this test process  */
void loop_on_receive(
    int icmp_socket,
    int ip_version)
{
    int packet_size;
    char packet[MAX_PACKET_SIZE];

    while (true) {
        packet_size = recv(icmp_socket, packet, MAX_PACKET_SIZE, 0);
        if (packet_size < -1) {
            perror("Failure during receive");
            exit(EXIT_FAILURE);
        }

        if (ip_version == 6) {
            if (is_packet_for_us6(packet, packet_size)) {
                dump_packet_info6(packet, packet_size);
                return;
            }
        } else {
            if (is_packet_for_us4(packet, packet_size)) {
                dump_packet_info4(packet, packet_size);
                return;
            }
        }
    }
}

/*  Parse the commandline arguments  */
void parse_cmdline(
    int argc,
    char **argv,
    int *ip_version)
{
    int opt;

    *ip_version = 4;

    while ((opt = getopt(argc, argv, "46")) != -1) {
        if (opt == '4') {
            *ip_version = 4;
        }

        if (opt == '6') {
            *ip_version = 6;
        }
    }
}

/*
    A helper for mtr-packet testing which waits for an ICMP packet
    intended for this test process, and then prints information about
    it.
*/
int main(
    int argc,
    char **argv)
{
    int icmp_socket;
    int ip_version;

    parse_cmdline(argc, argv, &ip_version);

    if (ip_version == 6) {
        icmp_socket = socket(AF_INET6, SOCK_RAW, IPPROTO_ICMPV6);
    } else {
        icmp_socket = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    }
    if (icmp_socket < 0) {
        perror("Failure opening listening socket");
        exit(EXIT_FAILURE);
    }

    printf("status listening\n");
    fflush(stdout);

    loop_on_receive(icmp_socket, ip_version);

    return EXIT_SUCCESS;
}