#include <string.h>
#include <assert.h>

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

// shamelessly copied from:
// http://www.roman10.net/how-to-calculate-iptcpudp-checksumpart-2-implementation/
unsigned short compute_checksum(unsigned short *addr, unsigned int count) {

  unsigned long sum = 0;

  while (count > 1) {
    sum += * addr++;
    count -= 2;
  }

  //if any bytes left, pad the bytes and add
  if(count > 0) {
    sum += ((*addr)&htons(0xFF00));
  }

  //Fold sum to 16 bits: add carrier to result
  while (sum>>16) {
      sum = (sum & 0xffff) + (sum >> 16);
  }

  //one's complement
  sum = ~sum;

  return ((unsigned short)sum);

}

void compute_ip_checksum(struct iphdr *pkt)
{
  pkt->check = 0x0000;
  pkt->check = compute_checksum( (unsigned short*) pkt, pkt->ihl * 4 );
}

int build_wrapped_ipv4_packet( struct peer_context *pctx, struct packet *pkt, struct rsp_data *out )
{
    struct iphdr *wrap_hdr = (struct iphdr *) out->scratch;
    uint16_t wrap_hdr_size = sizeof( struct iphdr );
    
    ssize_t enc_size;
    size_t orig_data_size = ntohs( pkt->hdr.ip.tot_len );
    size_t bytes_to_encrypt = orig_data_size > 512 ? 512 : orig_data_size;
    
    debug( "Wrapping an IPv4 packet" );
    debug( "wrap_hdr_size: %u, orig_data_size: %zu, bytes_to_encrypt: %zu", wrap_hdr_size, orig_data_size, bytes_to_encrypt );

    // Areas in scratch we'll be using later
    // We use two bytes to store the size of the encrypted blob
    uint16_t *pkt_enc_size = (uint16_t *) (out->scratch + wrap_hdr_size );
    unsigned char * pkt_enc_data = out->scratch + wrap_hdr_size + 2;
    
    // Keep track of the total size of the data in out as we go
    uint16_t out_len = 0;
    
    out->count = 0;
    
    memset( wrap_hdr, 0, wrap_hdr_size );
    
    wrap_hdr->version  = 0x04;
    wrap_hdr->ihl      = wrap_hdr_size / 4;
    wrap_hdr->ttl      = IPDEFTTL;
    wrap_hdr->protocol = IPPROTO_HIDE_EID;
    wrap_hdr->frag_off = htons( IP_DF ); // DF bit set    
    wrap_hdr->saddr    = pctx->x->addr.ip4.s_addr;
    wrap_hdr->daddr    = pctx->y->addr.ip4.s_addr;
    // FIXME: Do we need to set an ID ?
    
    // iovec 0: encapsulating IP header.
    out->iovs[0].iov_base = wrap_hdr;
    out->iovs[0].iov_len = wrap_hdr_size;
    out_len += wrap_hdr_size;
    out->count++;
    

    // Encrypt the first 512 or so bytes of the data. FIXME: introspect and
    // calculate exactly how many bytes for TCP, UDP, etc. to do as little work
    // as we can get away with, here. fragments > 0 don't need encrypting at all
    enc_size = rlocs_encrypt(
        pctx,
        (unsigned char *)&pkt->hdr, bytes_to_encrypt,
        pkt_enc_data, IP_MAXPACKET - wrap_hdr_size - 2
    );

    if ( enc_size  < 0 ) { 
        debug( "Failed to encrypt, dropping packet" );
        return 0;
    }
    debug( "enc_size: %li", enc_size );
    
    *pkt_enc_size = htons( enc_size );
    
    // iovec 1: encrypted portion of encpasulated packet
    out->iovs[1].iov_base = pkt_enc_size;
    out->iovs[1].iov_len = enc_size + 2;
    out_len += enc_size + 2;
    out->count++;
    
    debug( "iovs[0]: %p, %zu", out->iovs[0].iov_base, out->iovs[0].iov_len );
    debug( "iovs[1]: %p, %zu", out->iovs[1].iov_base, out->iovs[1].iov_len );
    
    // iovec 2: unencrypted remains of encapsulated packet, if present
    if ( bytes_to_encrypt < orig_data_size ) {
        out->iovs[2].iov_base = ((char *) pkt) + bytes_to_encrypt;
        out->iovs[2].iov_len  = orig_data_size - bytes_to_encrypt;
        out_len += orig_data_size - bytes_to_encrypt;
        out->count++;
        debug( "iovs[2]: %p, %zu", out->iovs[2].iov_base, out->iovs[2].iov_len );
        debug( "pkt: %p", pkt );
    }
    
    wrap_hdr->tot_len = htons( out_len );
    compute_ip_checksum( wrap_hdr );
    
    debug( "Finished wrapping IPv4 packet" );
    return 1;
}

void build_icmp_too_big( uint16_t max_mtu, struct in_addr *rloc_src, struct packet *pkt, struct rsp_data *out )
{
    struct iphdr *ip = (struct iphdr*) out->scratch;
    struct icmphdr *icmp = (struct icmphdr *) ( out->scratch + sizeof( struct iphdr ) );
    uint16_t icmp_size = sizeof( struct icmphdr ) + ( pkt->hdr.ip.ihl * 4 ) + 8;
     
    debug( "Building ICMP Too Big packet" );

    memset( out->scratch, 0, sizeof( struct iphdr ) + sizeof( struct icmphdr ) );

    ip->version = 4;
    ip->ihl = sizeof( struct iphdr ) / 4;
    ip->ttl = IPDEFTTL;
    ip->protocol = IPPROTO_ICMP;
    ip->saddr = rloc_src->s_addr;
    ip->daddr = pkt->hdr.ip.saddr;

    icmp->type = ICMP_DEST_UNREACH;
    icmp->code = ICMP_FRAG_NEEDED;
    icmp->un.frag.mtu = htons( max_mtu );
    
    out->iovs[0].iov_base = out->scratch;
    out->iovs[0].iov_len  = sizeof( struct iphdr ) + icmp_size;
    
    memcpy( ((char *)icmp) + sizeof( struct icmphdr ), pkt, icmp_size - sizeof( struct icmphdr ) );
    
    //out->iovs[1].iov_base = pkt;
    //out->iovs[1].iov_len = ( pkt->hdr.ip.ihl * 4 ) + 8;        
    
    ip->tot_len = htons( out->iovs[0].iov_len );

    icmp->checksum = compute_checksum( (unsigned short *)icmp, icmp_size );
    compute_ip_checksum( ip );
    out->count = 1;

    return;
}

struct peer_context *packet_peer_context( struct rlocs *reg, struct packet *pkt, int wrapping )
{
    char bad_eid_txt[128] = {0};
    struct rloc *src_rloc = NULL, *dst_rloc = NULL;
    struct in_addr *ip4_src = (struct in_addr*) &pkt->hdr.ip.saddr;
    struct in_addr *ip4_dst = (struct in_addr*) &pkt->hdr.ip.daddr;
    struct in6_addr *ip6_src = &pkt->hdr.ip6.ip6_src;
    struct in6_addr *ip6_dst = &pkt->hdr.ip6.ip6_dst;

    

    switch( pkt->hdr.ip.version ) {
        case 0x04: // ipv4 
            ip6_src = NULL;
            ip6_dst = NULL;
            if ( wrapping ) { // source and destination are eids
                src_rloc = rloc_find_for_ipv4( reg, ip4_src );
                dst_rloc = rloc_find_for_ipv4( reg, ip4_dst );
            }
            break;
        case 0x06: // ipv6
            ip4_src = NULL;
            ip6_src = NULL;
            if ( wrapping ) { // source and destination are eids
                src_rloc = rloc_find_for_ipv6( reg, ip6_src );
                dst_rloc = rloc_find_for_ipv6( reg, ip6_dst );
            }
            break;
    }

    if ( !wrapping ) { // source and destination are rlocs
        src_rloc = rloc_find_by_address( reg, ip4_src, ip6_src );
        dst_rloc = rloc_find_by_address( reg, ip4_dst, ip6_dst );
    }


    if ( src_rloc == NULL ) {
        if ( ip4_src != NULL ) {
            inet_ntop( AF_INET, ip4_src, bad_eid_txt, 128 );
        } else if ( ip6_src != NULL ) {
            inet_ntop( AF_INET6, &pkt->hdr.ip6.ip6_src, bad_eid_txt, 128 );
        }
        warn( "Couldn't find source RLOC from %s", bad_eid_txt );
    }

    if ( dst_rloc == NULL ) {
        if ( ip4_dst != NULL ) {
            inet_ntop( AF_INET, ip4_dst, bad_eid_txt, 128 );
        } else if ( ip6_dst != NULL ) {
            inet_ntop( AF_INET6, &pkt->hdr.ip6.ip6_dst, bad_eid_txt, 128 );
        }
        warn( "Couldn't find destination RLOC from %s", bad_eid_txt );
    }

    if ( src_rloc == NULL || dst_rloc == NULL ) {
        return NULL;
    }

    if ( wrapping ) { // when wrapping, src=x, dst=y
        return rlocs_get_peer_ctx( reg, src_rloc, dst_rloc );
    } else { // when unwrapping, src=y, dst=x
        return rlocs_get_peer_ctx( reg, dst_rloc, src_rloc );
    }
}

int wrap_ipv4_packet_in_ipv4( struct peer_context *pctx, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2 )
{

    uint16_t max_size = pctx->path_mtu;
    uint16_t pkt_tot_len = ntohs( pkt->hdr.ip.tot_len );
    uint16_t pkt_hdr_len = pkt->hdr.ip.ihl * 4;
    
    int num_packets = 1;
    
    // fragmentation is needed.
    if ( pkt_tot_len > max_size - WRAP_OVERHEAD ) {
        debug( "Packet needs fragmenting" );
        // DF bit set, so return ICMP Too Big
        if ( ntohs( pkt->hdr.ip.frag_off ) & IP_DF ) {
            build_icmp_too_big( max_size, &pctx->x->addr.ip4, pkt, frag1 );
            return 1;
        }
        
        num_packets = 2;

        // good enough, it's getting encrypted and only needs to be unique for
        // a short period of time
        uint16_t frag_id = (uint16_t) rand();

        // Must be an 8-byte offset
        uint16_t frag_off = ( pkt_tot_len - pkt_hdr_len ) / 2;
        frag_off += frag_off%8;
        uint16_t frag2_size = pkt_tot_len - pkt_hdr_len - frag_off;
                
        if ( pkt_hdr_len > sizeof( struct iphdr ) ) {
            warn( "FIXME: options specified with IP header are not handled correctly during fragmentation yet" );
        }

        // wrap_ipv4_packet only touches scratch upto IP_MAXPACKET. We allocate
        // double that.
        struct packet *pkt2 = (struct packet *) frag2->scratch + IP_MAXPACKET;

        pkt->hdr.ip.tot_len = htons( pkt_hdr_len + frag_off );
        pkt->hdr.ip.id = htons( frag_id );
        pkt->hdr.ip.frag_off = htons( 0 | IP_MF );
        
        memcpy( pkt2, pkt, pkt_hdr_len );
        
        pkt2->hdr.ip.tot_len = frag2_size;
        pkt2->hdr.ip.frag_off = htons( frag_off / 8 );
        memcpy( ((char *)pkt2) + pkt_hdr_len, ((char*)pkt)+pkt_hdr_len, frag2_size );
        
        // Need to recompute these
        compute_ip_checksum( &pkt->hdr.ip );
        compute_ip_checksum( &pkt2->hdr.ip );
        
        if ( !build_wrapped_ipv4_packet( pctx, pkt2, frag2 ) ) {
            debug( "Couldn't wrap packet 2 of 2 ");
            return 0;
        }
    }

    if ( !build_wrapped_ipv4_packet( pctx, pkt, frag1 ) ) {
        debug( "Couldn't wrap packet 1 of %i", num_packets );
        return 0;
    }
    
    return num_packets;
}

int wrap_ipv6_packet_in_ipv4(struct peer_context *pctx, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2)
{
    warn( "STUB: wrap_ipv6_packet_in_ipv4" );
    return 0;   
}

int wrap_ipv4_packet_in_ipv6(struct peer_context *pctx, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2)
{
    warn( "STUB: wrap_ipv4_packet_in_ipv6" );
    return 0;   
}

int wrap_ipv6_packet_in_ipv6(struct peer_context *pctx, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2)
{
    warn( "STUB: wrap_ipv6_packet_in_ipv6" );
    return 0;   
}

int wrap_packet( struct rlocs *reg, struct packet *pkt, struct rsp_data *frag1, struct rsp_data *frag2 )
{
    struct peer_context *pctx = packet_peer_context( reg, pkt, 1 );

    if ( pctx == NULL ) {
        return 0;
    }

    debug( "Wrapping, x=%s, y=%s", pctx->x->presentation, pctx->y->presentation );

    int result = 0;

    switch ( pctx->x->family ) {
        case AF_INET:
            switch ( pkt->hdr.ip.version ) {
                case 0x04: // ipv4
                    result = wrap_ipv4_packet_in_ipv4( pctx, pkt, frag1, frag2 );
                    break;
                case 0x06: // ipv6
                    result = wrap_ipv6_packet_in_ipv4( pctx, pkt, frag1, frag2 );
                    break;
            }
            break;
        case AF_INET6:
            switch ( pkt->hdr.ip.version ) {
                case 0x04: // ipv4
                    result = wrap_ipv4_packet_in_ipv6( pctx, pkt, frag1, frag2 );
                    break;
                case 0x06: // ipv6
                    result = wrap_ipv6_packet_in_ipv6( pctx, pkt, frag1, frag2 );
                    break;
        }
        break;
        default:
            warn( "Unknown family of peer context: %i", pctx->x->family );
    }

    return result;
}



int unwrap_ipv4_packet( struct peer_context *pctx, struct packet *pkt, struct rsp_data *out )
{
    out->count = 2;
    assert( out->count < MAX_IOVS );
    
    // first, check this is actually a hide-eid packet.
    if ( pkt->hdr.ip.protocol != IPPROTO_HIDE_EID ) {
        warn( "expected IP protocol %u, not %u", IPPROTO_HIDE_EID, pkt->hdr.ip.protocol );
        return 0;
    }

    uint16_t hdr_size = pkt->hdr.ip.ihl * 4;
    uint16_t encrypted_size = ntohs( *((uint16_t*)pkt + ( hdr_size / 2 )) );
    info( "encrypted_size: %u", encrypted_size );

    // iovec 0: decrypted data. This should be an IP header.
    unsigned char *encrypted_data = ((unsigned char *)pkt) + hdr_size + 2;
    unsigned char *scratch = &out->scratch[0];
    
    int decrypted_size = rlocs_decrypt( pctx, encrypted_data, encrypted_size, scratch, IP_MAXPACKET );
    
    if ( decrypted_size < 0 ) {
        warn( "Failed to decrypt packet!" );
        return 0;        
    }
    info( "decrypted_size: %u", decrypted_size );
   
    out->iovs[0].iov_base = scratch;
    out->iovs[0].iov_len = decrypted_size;
       
    // iovec 1: never-encrypted part
    out->iovs[1].iov_base = encrypted_data + encrypted_size;
    out->iovs[1].iov_len = ntohs( pkt->hdr.ip.tot_len ) - hdr_size - encrypted_size;
    
    if ( out->iovs[0].iov_len + out->iovs[1].iov_len > IP_MAXPACKET ) {
        warn( "Unwrapped packet is too large, dropping it" );
        warn( "iovs[0] is %zu, iovs[1] is %zu", out->iovs[0].iov_len, out->iovs[1].iov_len );
        warn( "hdr_size = %u, encrypted_size = %u, tot_len = %u", hdr_size, encrypted_size, ntohs( pkt->hdr.ip.tot_len ) );
        return 0;
    }

    debug( "Finished unwrapping IPv4 packet" );
    return 1;
}

int unwrap_ipv6_packet(struct peer_context *pctx, struct packet *pkt, struct rsp_data *out)
{
    warn( "STUB: unwrap_ipv6_packet" );

    return 0;   
}


int unwrap_packet( struct rlocs *reg, struct packet *pkt, struct rsp_data *out )
{
    int result;
    struct peer_context *pctx = packet_peer_context( reg, pkt, 0 );

    if ( pctx == NULL ) {
        return 0;
    } 

    debug( "Unwrapping, x=%s, y=%s", pctx->x->presentation, pctx->y->presentation );

    switch ( pkt->hdr.ip.version ) {
        case 0x04:
            result = unwrap_ipv4_packet( pctx, pkt, out );
            break;
        case 0x06:
            result = unwrap_ipv6_packet( pctx, pkt, out );
            break;
        default:
            warn( "Couldn't unwrap packet with version %i, discarding", pkt->hdr.ip.version );
            result = 0;
    }

    return result;
}