#include "util.h"
#include "rlocs.h"
#include "packet.h"

#include <unistd.h>
#include <errno.h>

#include <sys/uio.h>


int process_icmpv4_rloc_update( struct rlocs *reg, struct packet *packet  )
{
    uint16_t hdr_len = packet->hdr.ip.ihl * 4;
    uint16_t inner_ip_hdr_offset = hdr_len + sizeof( struct icmphdr );
    if ( ntohs( packet->hdr.ip.tot_len ) < inner_ip_hdr_offset + sizeof( struct iphdr ) ) {
        debug( "Truncated ICMP packet is unidentifiable" );
        return 0;
    }

    struct rloc *s_rloc;
    struct rloc *d_rloc;
    
    struct icmphdr *icmp = (struct icmphdr *) packet->payload + ( packet->hdr.ip.ihl * 4 );

    if ( icmp->type != ICMP_DEST_UNREACH && icmp->code != ICMP_FRAG_NEEDED ) {
        return 0; // It may be going elsewhere
    }
    
    // Be careful with this - some of payload may be past allocated memory
    struct packet *inner_ip = (struct packet *) ((char *) packet + inner_ip_hdr_offset );
    
    //  Not much we can do with this case right now
    if ( inner_ip->hdr.ip.protocol != IPPROTO_HIDE_EID ) {
        debug( "ICMP Too Big response to an unwrapped packet. Peculiar." );
        return -1;
    }
    
    if ( !rlocs_find_two_ipv4( reg, &s_rloc, (struct in_addr *)&inner_ip->hdr.ip.saddr, &d_rloc, (struct in_addr *)&inner_ip->hdr.ip.daddr ) ) {
            return 0;
    }

    // All we're interested in is setting path mtu
    uint16_t new_mtu = ntohs( icmp->un.frag.mtu );
    rlocs_set_path_mtu( reg, s_rloc, d_rloc, new_mtu );
    debug( "Set MTU for %s <-> %s to %u", s_rloc->presentation, d_rloc->presentation, new_mtu );
    return 1;
}


int process_icmpv6_rloc_update( struct rlocs *reg, struct packet *packet )
{
    warn( "STUB: process_icmpv6_rloc_update" );
    return 0;
}

int process_packet( struct rlocs *reg, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2 )
{
    int protocol = packet_find_protocol( pkt );

    if ( protocol == -1 ) {
        warn( "Couldn't work out version / protocol of received packet, discarding" );
        return 0;
    }

    if ( protocol == IPPROTO_HIDE_EID ) {
        return unwrap_packet( reg, pkt, frag1 );
    }

    if ( protocol == IPPROTO_ICMP && process_icmpv4_rloc_update( reg, pkt ) ) {
        return 0;
    }

    if ( protocol == IPPROTO_ICMPV6 && process_icmpv6_rloc_update( reg, pkt ) ) {
        return 0;
    }

    return wrap_packet( reg, pkt, frag1, frag2 );
}

/*
 * Entry point. Expects an invocation like:
 * hide-eid <filename of rloc database> <listen_ifname> <output_ifname>
 * 
 * hide-eid wraps packets that come to it unwrapped, and unwraps packets that
 * come to it wrapped. This makes it handy for all-in-one boxes
 */ 
int main(int argc, char** argv)
{
    struct session session;
    struct packet recv_pkt;
    
    struct rsp_data frag1;
    struct rsp_data frag2;
    ssize_t count;
    
    if ( argc < 4 ) {
        warn( "Usage: %s <rloc database> <listen_ifname> <output_ifname> [ <rloc> <keyfile> ]n", argv[0] );
        return 1;
    }
    
    rlocs_init();
    if ( !session_setup( &session, argv[1], argv[2], argv[3] ) ) {
        warn( "Failed to set up session, exiting" );
        return 1;
    }

    if  (argc > 4 ) {
        if ( !session_upgrade_rlocs( &session, argc - 4, argv + 4 ) ) {
            warn( "Failed to upgrade rlocs for session, exiting" );
            session_teardown( &session );
            return 1;
        }    
    }    
    
    memset( &recv_pkt, 0, sizeof( struct packet ) );
    memset( &frag1, 0, sizeof( struct rsp_data ) );
    memset( &frag2, 0, sizeof( struct rsp_data ) );

    warn( "TODO: Write BGP interventions to file" );

    info( "Listening for packets" );
    while( 1 ) {
        int rsp_count;
        
        if ( ( count = read( session.listen_if, &recv_pkt, sizeof( struct packet ) ) ) < 0 ) {
            warn( "Failed to get a packet (%s)", strerror( errno ) );
            break;
        }
        
        if ( count == 0 ) {
            warn( "Got EOF" );
            break;
            
        }

        if ( ( rsp_count = process_packet( session.rlocs, &recv_pkt, &frag1, &frag2 ) ) < 0 ) {
            debug( "Error processing packet, dropping" );
        }
        
        if ( rsp_count > 0 ) {
            if ( ( count = writev( session.output_if, frag1.iovs, frag1.count ) ) < 0 ) {
                debug( "Error writing processed packet to output: %s", strerror(errno) );
            }
        }
        
        if ( rsp_count == 2 ) {
            if ( ( count = writev( session.output_if, frag2.iovs, frag2.count ) ) < 0 ) {
                debug( "Error writing second processed packet to output: %s", strerror(errno) );
            }
        }
        
    }
    
    info( "Finished, cleaning up" );
    session_teardown( &session );

    return 0;
}