#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <libnet.h>
#include <netinet/igmp.h>

#define IPOPT_RA                148             /* router alert */
#define IGMP_V3_MEMBERSHIP_REPORT 0x22

/* Structure to represent IGMP extra information */
struct igmp_extra {
    u_int8_t igmp_version;
    char *igmp_tag;
    u_int8_t igmp_type;             /* IGMP type */
    u_int8_t igmp_code;             /* routing code */
    char *igmp_dst;
    int group_override_dst;
} g_igmp_pkts[] = {
    /* name,  type (or version+type), code */
    { 1, "query",  0x11, 0, "224.0.0.1", 0 },
    { 1, "report", 0x12, 0, "224.0.0.1", 1 },
    { 1, "dvmrp",  0x13, 0, "224.0.0.1", 0 },
    { 2, "query",  0x11, 1, "224.0.0.1", 0 },
    { 2, "report", 0x16, 1, "224.0.0.2", 1 },
    { 2, "leave",  0x17, 1, "224.0.0.2", 0 },
    { 3, "report", 0x22, 1, "224.0.0.22", 0 },

    /* Note: end of list (please keep) */
    { 0, 0, 0, 0 },
};

void usage(char *);

int main(int argc, char **argv)
{
    u_int32_t ip_src = 0;
    u_int32_t ip_dst = 0;
    u_int32_t igmp_group = 0;      
    char *igmp_group_str = NULL;
    char *ip_src_str = NULL;
    char *ip_dst_str = NULL;

    /* ports */
    u_short src_prt = 3141;
    u_short dst_prt = 5926;   

    /* igmp stuff */
    u_char igmp_type = 0;
    u_char igmp_code = 0;
    u_int8_t igmp_version = 0;
    int igmp_override;
    struct igmp_extra *pkt_ptr;

    int ip_src_test = 1;
    int ip_dst_test = 1;
    int igmp_group_test =1;

    /* libnet stuff */
    char neterr[LIBNET_ERRBUF_SIZE];
    libnet_ptag_t ptag;
    libnet_t *netcontext = NULL;

    /* misc */
    char *device = "eth0";
    int c;                      
    char *cp = NULL; 
    int found=0;

    printf("IGMP packet generator\n\n");

    /* Parsing command line */
    while((c = getopt(argc, argv, "i:t:g:d:s:")) != EOF)
    {
        switch (c)
        {
            case 'i':
                printf("  Net interface = [%s]\n",optarg);
                device = optarg;
                break;

            case 't':
                if (!(cp = strrchr(optarg, '.'))) {
                    usage(argv[0]);
                    exit(1);
                }
                *cp++ = 0;

                // igmp_version = (u_short)atoi(cp);
                igmp_version = (u_short)atoi(optarg);

                pkt_ptr = g_igmp_pkts;
                while(pkt_ptr->igmp_version || pkt_ptr->igmp_tag){
                    // if ((strcasecmp(pkt_ptr->igmp_tag, optarg) == 0) 
                    if ((strcasecmp(pkt_ptr->igmp_tag, cp) == 0) 
                            && (igmp_version == pkt_ptr->igmp_version)){
                        found = 1;
                        igmp_type = pkt_ptr->igmp_type;
                        igmp_code = pkt_ptr->igmp_code;
                        ip_dst_str = pkt_ptr->igmp_dst;
                        igmp_override = pkt_ptr->group_override_dst;

                        break;
                    }
                    pkt_ptr++; 
                }
                if (found){
                    printf("  Packet = [%s] version [%d]\n", pkt_ptr->igmp_tag, igmp_version);
                } else {
                    usage(argv[0]);
                    exit(1);
                }
                break;

            case 'g':
                // Group argument handling logic
                printf("  Group = [%s]\n", optarg);
                igmp_group_str = optarg;
                break;

            case 'n':
                // delay between packets in sec
                break;

                /*
                 *  We expect the input to be of the form `ip.ip.ip.ip.port`.  We
                 *  point cp to the last dot of the IP address/port string and
                 *  then seperate them with a NULL byte.  The optarg now points to
                 *  just the IP address, and cp points to the port.
                 */
            case 'd':
                if ((cp = strrchr(optarg, ':')))
                {
                    *cp++ = 0;
                    dst_prt = (u_short)atoi(cp);
                } else {
                    dst_prt = 0;
                }
                ip_dst_str = strdup(optarg);
                printf("  Destination = ip [%s] port [%d]\n", ip_dst_str, dst_prt);
                break;

            case 's':
                if ((cp = strrchr(optarg, ':')))
                {
                    *cp++ = 0;
                    src_prt = (u_short)atoi(cp);
                } else {
                    src_prt = 0;
                }
                ip_src_str = strdup(optarg);
                printf("  Source = ip [%s] port [%d]\n", ip_src_str, src_prt);
                break;

            default:
                usage(argv[0]);
                exit(1);
        }
    }
    if (!igmp_group_str || !device)
    {
        usage(argv[0]);
        exit(EXIT_FAILURE);
    }

    printf("done\n");

    /* Memory initialization */
    printf("Initializing libnet context...\n");
    netcontext = libnet_init(LIBNET_RAW4, device, neterr);
    if (!netcontext) {
        fprintf(stderr, "%s\n", neterr);
        exit(1);
    }
    libnet_clear_packet(netcontext);

    if (!ip_src_str) {
        ip_src = libnet_get_ipaddr4(netcontext);
        ip_src_str = libnet_addr2name4(ip_src, LIBNET_DONT_RESOLVE);
    }

    printf("Packet construction...\n");

    /* Packet construction: IGMP */
    printf("  Building IGMP content...\n");
    ptag = libnet_build_igmp(
        igmp_type,        // IGMP type
        igmp_code,        // IGMP code
        0,                // checksum
        0,                // group address
        NULL,             // payload
        0,                // payload size
        netcontext,       // libnet context
        0                 // ptag
    );
    if (ptag == -1) {
        fprintf(stderr, "Error building IGMP header: %s\n", libnet_geterror(netcontext));
        libnet_destroy(netcontext);
        exit(EXIT_FAILURE);
    }

    printf("  done\n");

    /* Send packet */
    if (libnet_write(netcontext) == -1) {
        fprintf(stderr, "Error sending packet: %s\n", libnet_geterror(netcontext));
    } else {
        printf("Packet sent successfully.\n");
    }

    printf("Packet construction...\n");

    /*
     *  Packet construction : IGMP
     */
    printf("  Building IGMP content...\n");

    // Override dest with group info when needed
    if (igmp_override) {
        ip_dst_str = igmp_group_str;
    }

    switch ( igmp_type ) {
        case 0x11: // IGMP_MEMBERSHIP_QUERY
            if (igmp_version == 2){
                // group specific
            } else if (igmp_version == 1) {
                // general
                igmp_group = 0;
                igmp_group_str = "0.0.0.0";
                igmp_group_test = 0;
            }
            // DST = 224.0.0.1
            //igmp_group = 0;
            //igmp_group_test = 0;
            break;
        case 0x12: //
        default:
            break;
    }

    if (ip_src_test
            && !(ip_src = libnet_name2addr4(netcontext, ip_src_str, LIBNET_RESOLVE)))
    {
        fprintf(stderr,"Bad source IP address: %s\n", ip_src_str);
        fprintf(stderr,"%s\n", libnet_geterror(netcontext));
        exit(1);
    }

    if (ip_dst_test 
            && !(ip_dst = libnet_name2addr4(netcontext, ip_dst_str, LIBNET_RESOLVE)))
    {
        fprintf(stderr,"Bad destination IP address: %s\n", ip_dst_str);
        fprintf(stderr,"%s\n", libnet_geterror(netcontext));
        exit(1);
    }

    if (igmp_group_test && !(igmp_group = libnet_name2addr4(netcontext, igmp_group_str, LIBNET_RESOLVE)))
    {
        fprintf(stderr,"Bad group IP address: %s\n", optarg);
        fprintf(stderr,"%s\n", libnet_geterror(netcontext));
        exit(1);
    }

    printf("    IP SRC IP  = %s\n", ip_src_str);
    printf("    IP DEST IP = %s\n", ip_dst_str);
    printf("    IGMP TYPE  = 0x%02x\n", igmp_type);
    printf("    IGMP CODE  = 0x%02x\n", igmp_code);
    printf("    IGMP GROUP = %s\n", igmp_group_str);
    printf("    GROUP TEST = %d\n", igmp_group_test);

    /*
       int test = 1;
       if ( *(char *) test == 1) {
    // little endian
    printf("little endian\n");
    } else {
    // big endian
    printf("big endian\n");
    }
    */

    ptag = libnet_build_igmp(
            igmp_type,              // IGMP type 
            igmp_code,              // IGMP code or TTL 
            0,                      // checksum
            htonl(igmp_group),                // ip addr 
            NULL,                      // Ptr to packet data (or null) 
            0,                      // Packet payload length 
            netcontext,             // Ptr to libnet context
            0);                     // Build a new packet header
    if (ptag < 0)
    { 
        fprintf(stderr, "Error while building IGMP header data (err %d)\n", ptag);
        fprintf(stderr,"%s\n", libnet_geterror(netcontext));
        exit(1);
    }


    u_int8_t ipopt[2];
    ipopt[0] = IPOPT_RA;
    ipopt[1] = 4;
    ptag = libnet_build_ipv4_options(
            ipopt,       // options
            2,           // options_s (length of option string)
            netcontext,  // libnet context
            0         // ptag
            );                
    if (ptag < 0)
    { 
        fprintf(stderr, "Error while building IPv4 options data (err %d)\n", ptag);
    }

    ptag = libnet_build_ipv4(LIBNET_IPV4_H + LIBNET_IGMP_H,  // pkg size + data
            0,         // IP tos
            getpid(),  // IP ID
            IP_DF,         // frag flags and offset
            1,         // TTL
            IPPROTO_IGMP,           // transport protocol
            0,                      // checksum
            ip_src,                 // source IP
            ip_dst,                 // destination IP
            NULL,                   // payload (none)
            0,                      // payload length
            netcontext,             // libnet context
            0                       // build a new header
            );                
    if (ptag < 0)
    { 
        fprintf(stderr, "Error while building IPv4 header data (err %d)\n", ptag);
        fprintf(stderr,"%s\n", libnet_geterror(netcontext));
        exit(1);
    }

    printf("  done\n"); //IGMP content

    c = libnet_write(netcontext);
    if (c < 0)
    {
        fprintf(stderr,"%s\n", libnet_geterror(netcontext));
    }
    else
    {
        printf("construction and injection completed, wrote all %d bytes\n", c);
    } 

    /*
     *  Free packet memory.
     */
    libnet_diag_dump_pblock(netcontext);
    libnet_destroy(netcontext);
    exit(0);

    return (c == -1 ? EXIT_FAILURE : EXIT_SUCCESS);
}

void usage(char *name)
{
    struct igmp_extra *pkt_ptr;

    fprintf(stderr, "Usage: %s [options]\n", name);
    fprintf(stderr, "\nOptions:\n");
    fprintf(stderr, "  -i <interface>   Specify the network interface (e.g., eth0)\n");
    fprintf(stderr, "  -t <type>        Specify the IGMP packet type and version (e.g., 1.query)\n");
    fprintf(stderr, "  -g <group>       Specify the IGMP group (e.g., 224.0.0.1)\n");
    fprintf(stderr, "  -s <source>      Specify the source IP and port (e.g., 192.168.1.1:1234)\n");
    fprintf(stderr, "  -d <destination> Specify the destination IP and port (e.g., 224.0.0.2:5678)\n");
    fprintf(stderr, "  -n <number>      Specify delay between packets in seconds (optional)\n");

    fprintf(stderr, "\nAvailable IGMP packet types:\n");
    pkt_ptr = g_igmp_pkts;
    while(pkt_ptr->igmp_version || pkt_ptr->igmp_tag){
        fprintf(stderr," - %d.%s\n", pkt_ptr->igmp_version, pkt_ptr->igmp_tag);
        pkt_ptr++; 
    }
}

// vim: ts=4 sts=4 sw=4 et