/* Registry dependencies */

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#include <errno.h>
#include <fcntl.h>

/* We use RSA public-key encryption to begin with. We expect the config file
 * to give us PEM-encoded data.
 */
#include <openssl/ssl.h>
#include <openssl/pem.h>
#include <openssl/rsa.h>
#include <openssl/bio.h>
#include <openssl/rand.h>
#include <openssl/err.h>

#include <openssl/evp.h>
#include <openssl/ec.h>
#include <openssl/obj_mac.h>

#include <json/json_tokener.h>
#include <json/linkhash.h>

#include <arpa/inet.h>

//#include "util.h"
#include "rlocs.h"

#define JSON_TYPE "rloc-registry"

void show_ssl_errors(void)
{
    ERR_print_errors_fp(stderr);
    return;
}

void rlocs_init(void)
{
    uint32_t seed;
    SSL_load_error_strings();                /* readable error messages */
    SSL_library_init();                      /* initialize library */
    OpenSSL_add_all_algorithms();
    
    // set the random seed to something good
    if ( RAND_bytes( (unsigned char*)&seed, sizeof( uint32_t ) ) != 1 ) {
        warn( "Failed to initialize random seed" );
        show_ssl_errors();
        exit(1);
    }
    srand( seed );

    return;
}


int rlocs_rloc_from_json(struct rlocs* reg, const char* key, json_object* val )
{
    ssize_t pem_data_size;
    const char* pem_data;
    struct rloc *rloc;
    EVP_PKEY *evp_key = NULL;
    EC_KEY *ec_key = NULL;
    BIO *bp;

    if ( !json_object_is_type( val, json_type_string ) ) {
        warn( "Non-string value in 'keys' object" );
        return 0;
    }
    
    if ( reg->num_entries >= MAX_RLOCS ) {
        warn( "Too many RLOCs in config-file!" );
        return 0;
    }
    
    rloc = &reg->entries[reg->num_entries];
    
    if ( strchr( key, ':' ) == NULL ) { /* IPv4 */
        rloc->family = AF_INET;
    } else { /* IPv6 */
        rloc->family = AF_INET6;
    }
    
    if ( inet_pton( rloc->family, key, (void*)&rloc->addr ) != 1 ) {
        warn( "Couldn't parse %s as an IP address", key );
        return 0;
    }
    
    // Next, we need to read and store the PEM-encoded RSA public key
    pem_data = json_object_get_string( val );
    pem_data_size = json_object_get_string_len( val );
    
    if ( ( evp_key = EVP_PKEY_new() ) == NULL ) {
        warn( "couldn't allocate new EVP_pkey for rloc %s", key );
        goto fail;
    }
    
    if ( !( bp = BIO_new_mem_buf( (void*)pem_data, pem_data_size ) ) ) {
        warn( "couldn't allocate new BIO to interpret key for rloc %s", key );
        goto fail;
    }
    
    ec_key = PEM_read_bio_EC_PUBKEY( bp, &ec_key, NULL, NULL );    
    BIO_vfree( bp );
    
    if ( ec_key == NULL ) {
        warn( "Couldn't parse key for rloc %s as an EC secp160r2 public key", key );
        goto fail;
    }
    
    // assign so we don't have to free the EC_KEY object any more
    if ( !EVP_PKEY_assign_EC_KEY( evp_key, ec_key ) ) {
        warn( "couldn't assign key for rloc %s to container", key );
        goto fail;
    }
    
    rloc->key = evp_key;
    rloc->context_id = reg->num_entries;
    reg->num_entries++;
    return 1;

fail:
    show_ssl_errors();

    if ( evp_key != NULL ) {
        EVP_PKEY_free( evp_key );
    }

    if ( ec_key != NULL ) {
        EC_KEY_free( ec_key );
    }

    return 0;
}

int rlocs_eid_from_json( struct rlocs *reg, json_object *map)
{
    json_object *entry;
    const char *str;
    int mask;

    if ( !json_object_is_type( map, json_type_object ) ) {
        warn( "eid_rloc_maps array in config contains non-object" );
        return 0;
    }
        
    /* Initial processing of the netmask entry */
    entry = json_object_object_get( map, "netmask" );
    if ( !json_object_is_type( entry, json_type_int ) ) {
        warn( "No netmask specified" );
        return 0;        
    }
    
    mask = json_object_get_int( entry );
    if ( mask < 0 ) {
        warn( "Bad netmask specified (%i)", mask );
        return 0;
    }
    
    /* Work out the family and finish off processing */
    entry = json_object_object_get( map, "family" );
    if ( !json_object_is_type( entry, json_type_string ) ) {
        warn( "No ip address family specified" );
        return 0;
    }
    str = json_object_get_string( entry );
    
    if ( strcmp( "ipv4", str ) == 0 ) {
        struct ip4_eid_map_entry *eid_mapping;
        struct in_addr addr;

        if ( reg->num_ip4_map_entries >= MAX_EID_MAPPINGS ) {
            warn( "Limit of %i IPv4 EID mappings reached!", MAX_EID_MAPPINGS );
            return 0;
        }

        if ( mask > 32 ) {
            warn( "Bad netmask specified for IPv4 (%i)", mask );
            return 0;
        }
        
        eid_mapping = &reg->ip4_mappings[reg->num_ip4_map_entries];
        eid_mapping->mask = (unsigned int) mask;

        // parse our network address
        entry = json_object_object_get( map, "network" );
        if (!json_object_is_type( entry, json_type_string ) ) {
            warn( "No network specified" );
            return 0;
        }
        str = json_object_get_string( entry );
        if ( inet_pton( AF_INET, str, &eid_mapping->network ) != 1 ) {
            warn( "Invalid network specified: %s", str );
            return 0;
        }
        
        // Remember, in_addrs are stored in network byte order
        eid_mapping->broadcast.s_addr =
            htonl(  ntohl( eid_mapping->network.s_addr ) | (( 1 << (32 - mask )) -1 ) );

        // Get a pointer to the rloc these eids are mapped to
        entry = json_object_object_get( map, "rloc" );
        if ( !json_object_is_type( entry, json_type_string ) ) {
            warn( "no rloc specified" );
            return 0;
        }
        str = json_object_get_string( entry );
        if ( inet_pton( AF_INET, str, &addr ) != 1 ) {
            warn( "Invalid rloc specified: %s", str );
            return 0;
        }
        
        eid_mapping->rloc = rloc_find_by_address( reg, &addr, NULL );
        if ( eid_mapping->rloc == NULL ) {
            warn( "Unknown rloc: %s", str );
            return 0;
        }
        
        // success!
        reg->num_ip4_map_entries++;
    } else if ( strcmp( "ipv6", str ) == 0 ) {
        struct ip6_eid_map_entry *eid_mapping;
        struct in6_addr addr;

        if ( reg->num_ip6_map_entries >= MAX_EID_MAPPINGS ) {
            warn( "Limit of %i IPv6 EID mappings reached!", MAX_EID_MAPPINGS );
            return 0;
        }


        if ( mask > 128 ) {
            warn( "Bad netmask specified for IPv6 (%i)", mask );
            return 0;
        }

        eid_mapping = &reg->ip6_mappings[reg->num_ip6_map_entries];
        eid_mapping->mask = (unsigned int) mask;

        // parse our network address
        entry = json_object_object_get( map, "network" );
        if (!json_object_is_type( entry, json_type_string ) ) {
            warn( "No network specified" );
            return 0;
        }
        str = json_object_get_string( entry );
        if ( inet_pton( AF_INET6, str, &eid_mapping->network ) != 1 ) {
            warn( "Invalid network specified: %s", str );
            return 0;
        }
        warn( "setting eid_mapping->network with %s", str );
        
        memcpy( &eid_mapping->broadcast, &eid_mapping->network, sizeof( struct in6_addr ) );
        // FIXME: calculate ipv6 broadcast here. ipv6 is broken until we do
         warn( "IPv6 broadcast not calculated yet" );
        
        
        // Get a pointer to the rloc these eids are mapped to
        entry = json_object_object_get( map, "rloc" );
        if ( !json_object_is_type( entry, json_type_string ) ) {
            warn( "no rloc specified" );
            return 0;
        }
        str = json_object_get_string( entry );
        if ( inet_pton( AF_INET6, str, &addr ) != 1 ) {
            warn( "Invalid rloc specified: %s", str );
            return 0;
        }
        
        eid_mapping->rloc = rloc_find_by_address( reg, NULL, &addr );
        if ( eid_mapping->rloc == NULL ) {
            warn( "Unknown rloc: %s", str );
            return 0;
        }
        
        // success!
        reg->num_ip6_map_entries++;        
    } else {
        warn( "Unknown family specified: %s", str );
        return 0;
    }
    
    return 1;

}

struct rlocs* rlocs_new( char* filename )
{
    struct rlocs* result;
    struct stat filedata;
    json_object* config;
    ssize_t have_read = 0;
    char* json_text;
    int fd;
    
    fd = open( filename, O_RDONLY );
    if ( errno < 0 ) {
        warn( "Error %s (%i) opening %s", strerror(errno), errno, filename );
        return NULL;
    }
    
    fstat( fd, &filedata );
    if ( errno < 0 ) {
        warn( "Error %s (%i) getting size of %s", strerror(errno), errno, filename );
        return NULL;        
    }
    
    /* Make sure we're null-terminated */
    json_text = xmalloc( filedata.st_size + 1 );
    
    while ( have_read < filedata.st_size ) {
        ssize_t bytes = read( fd, json_text + have_read, filedata.st_size - have_read );
        if ( bytes < 0 ) {
            warn( "Error %s (%i) reading %s", strerror(errno), errno, filename ) ;
            free( json_text );
            close( fd );
            return NULL;
        }

        have_read += bytes;

        //EOF
        if ( bytes == 0 ) {
            warn( "short read of %s: %zu instead of %zu", filename, have_read, filedata.st_size );
            filedata.st_size = have_read;
            
        }
    }
    
    //info( "Current config:\n %s", json_text );

    close( fd );
    config = json_tokener_parse( json_text );
    free( json_text );
    
    if ( !json_object_is_type( config, json_type_object ) ) {
        warn( "Config does not parse to a JSON object" );
        return NULL;
    }
    
    json_object* type_field = json_object_object_get( config, "type" );
    const char* type_data = json_object_get_string( type_field );
    if ( strcmp( JSON_TYPE, type_data ) != 0 ) {
        warn( "Expected type to be %s, not %s", JSON_TYPE, type_data );
        return NULL;
    }
    
    json_object * keys_field = json_object_object_get( config, "keys" );
    
    if ( !json_object_is_type( keys_field, json_type_object ) ) {
        warn( "keys field is not present in config" );
        return NULL;
    }
    
    result = xmalloc( sizeof( struct rlocs ) );
    result->config = config;
    
    /* populate our array of rloc objects */
    json_object_object_foreach( keys_field, key, val ) {
        if ( !rlocs_rloc_from_json( result, key, val ) ) {
            warn( "Failed to process rloc %s", key );
            goto fail;
        }
    }
    

    json_object * maps_field = json_object_object_get( config, "eid_rloc_map" );
    if ( !json_object_is_type( maps_field, json_type_array ) ) {
        warn( "Couldn't find eid_rloc_map field in config" );
        goto fail;   
    }
    
    int i, len = json_object_array_length( maps_field );
    for ( i = 0; i < len ; i++ ) {
        json_object *map_field = json_object_array_get_idx( maps_field, i );

        if ( !rlocs_eid_from_json( result, map_field ) ) {
            warn( "Failed to process eid map %i", i );
            goto fail;
        }
    }

    return result;
    
fail:
    rlocs_free( result ); // deallocate any struct rloc we already created
    return NULL;
}

// FIXME: these really do need to be as fast as possible. This is a 2-minute effort.
struct rloc *rloc_find_for_ipv4( struct rlocs *reg, struct in_addr *eid )
{
    struct ip4_eid_map_entry *current;
    uint32_t host_eid = ntohl( eid->s_addr );
    int i;
    
    for ( i = 0 ; i < reg->num_ip4_map_entries ; i++ ) {
        current = &reg->ip4_mappings[i];

        uint32_t host_start = ntohl( current->network.s_addr );
        uint32_t host_end   = ntohl( current->broadcast.s_addr );
        
        if ( host_start <= host_eid && host_end >= host_eid ) {
            break;
        }

    }
    
    return i < reg->num_ip4_map_entries ? current->rloc : NULL;
}

struct rloc *rloc_find_for_ipv6( struct rlocs *reg, struct in6_addr *eid )
{
    warn( "Searching for IPv6 RLOCs not done yet" );
    return NULL;
}

struct rloc *rloc_find_by_address( struct rlocs *reg, struct in_addr *ipv4, struct in6_addr *ipv6 )
{
    struct rloc *current;
    int i, cmp;

    for( i = 0 ; i < reg->num_entries; i++ ) {
        current = &reg->entries[i];
        cmp = 1;
        switch( current->family ) {
            case AF_INET:
                if ( ipv4 != NULL ) {
                    cmp = memcmp( &current->addr, ipv4, sizeof( struct in_addr ) );
                }
                break;
            case AF_INET6:
                if ( ipv6 != NULL ) {
                    cmp = memcmp( &current->addr, ipv6, sizeof( struct in6_addr ) );
                }
                break;
        }
        if ( cmp == 0 ) {
            break;
        }
    }
    
    return i < reg->num_entries ? current : NULL;
}

/* Replaces the public key in the rloc struct with a private key so we can
 * unwrap, as well as wrap, packets.
 */
int rloc_add_private_key( struct rloc *rloc, char *filename )
{
    BIO *key_data = BIO_new_file( filename, "r" );    
    EC_KEY* key = PEM_read_bio_ECPrivateKey( key_data, NULL, NULL, NULL );
    
    if ( key == NULL ) {
        warn( "Failed to add private key %s", filename );
        goto fail;
    }
    
    if ( !EVP_PKEY_assign_EC_KEY( rloc->key, key ) ) {
        warn( "Failed to assign private key in %s to rloc", filename );
        goto fail;
    }
    
    return 1;

fail:
    show_ssl_errors();

    if ( key != NULL ) {
        EC_KEY_free( key );
    }

    return 0;    
}

int rlocs_update_key_context(struct rlocs *reg, struct rloc *x, struct rloc *y)
{
    struct key_context *entry = &reg->key_contexts[x->context_id][y->context_id];

    unsigned char secret[1024]; // Should be enough buffer space
    size_t secret_len = 1024;
    
    unsigned char secret_hash[SHA256_DIGEST_LENGTH];
    
    EVP_PKEY_CTX *ctx = NULL;
    int result;
    
    /* Create the context for the shared secret derivation */
    if( ( ctx = EVP_PKEY_CTX_new( x->key, NULL ) ) == NULL) {
        warn( "EVP_PKEY_CTX_new failed" );
        goto fail;
        
    }

    /* Initialise */
    if( ( result = EVP_PKEY_derive_init( ctx ) ) != 1 ) {
        warn( "EVP_PKEY_derive_init failed: %i", result );
        goto fail;
    }

    /* Provide the peer public key */
    if( ( result = EVP_PKEY_derive_set_peer( ctx, y->key ) ) != 1 ) {
        warn( "EV_PKEY_derive_set_peer failed: %i", result );
        goto fail;
    }

    /* Derive the shared secret */
    if( ( result = EVP_PKEY_derive( ctx, &secret[0], &secret_len ) ) != 1 ) {
        warn( "EVP_PKEY_derive failed: %i", result );
        goto fail;
    }
    
    /* SHA256 hash of the shared secret. TODO: Add a salt if we want one... */
    if ( !sha256sum( &secret[0], secret_len, secret_hash ) ) {
        warn( "SHA256( secret ) failed!" );
        goto fail;
    }
    memcpy( &entry->secret[0], &secret_hash[0], SHA256_DIGEST_LENGTH );
    
    EVP_PKEY_CTX_free( ctx );

    /* Now we finally have our secret key, we can set up an AES256
     * symmetric key cipher with it, which is what will be used to encrypt /
     * decrypt packet data. */

    EVP_CIPHER_CTX_init( &entry->ctx );
    entry->in_use = 1;

    return 1;

    
fail:
    show_ssl_errors();
    if ( ctx != NULL) {
        EVP_PKEY_CTX_free( ctx );
    }
    return 0;
}


ssize_t rlocs_encrypt( struct rlocs *reg, struct rloc *x, struct rloc *y, unsigned char *data, size_t data_len, unsigned char *dest, size_t dest_len )
{
    struct key_context *entry = &reg->key_contexts[x->context_id][y->context_id];

    if ( !entry->in_use && !rlocs_update_key_context( reg, x, y ) ) {
        warn( "Couldn't build an encryption context for these rlocs" );
        return -1;
    }
    
    // FIXME: I don't know that this is good. It could be terrible.
    // We use a PRNG to generate an IV per message, and put it at the start of
    // the encrypted blob. Since we're using aes256 (hardcoded), IV is 128 bits
    if ( dest_len < 48 ) {
        warn( "dest_len must be at least 48 bytes, maybe more");
        goto fail;
    }

    unsigned char *iv = dest;
    unsigned char *secret = (unsigned char *)&entry->secret[0];
    size_t written = 16;
    int outl = dest_len - written;

    RAND_pseudo_bytes( iv, 16 );
    
    if ( !EVP_EncryptInit_ex( &entry->ctx, EVP_aes_256_cbc(), NULL, secret, iv ) ) {
        warn( "EVP_EncryptInit_ex() failed" );
        goto fail;
    }
    
    if ( !EVP_EncryptUpdate( &entry->ctx, dest + written, &outl, data, data_len ) ) {
        warn( "EVP_EncryptUpdate() failed" );
        goto fail;
    }
    
    written += outl;
    outl = dest_len - written;

    if ( !EVP_EncryptFinal_ex( &entry->ctx, dest + written, &outl ) ) {
        warn( "EVP_EncryptFinal_ex() failed" );
        goto fail;
    }

    return written + outl;

fail:
    show_ssl_errors();
    return -1;
}

ssize_t rlocs_decrypt( struct rlocs *reg, struct rloc *x, struct rloc *y, unsigned char *data, size_t data_len, unsigned char *dest, size_t dest_len )
{
    struct key_context *entry = &reg->key_contexts[x->context_id][y->context_id];

    if ( !entry->in_use && !rlocs_update_key_context( reg, x, y ) ) {
        warn( "Couldn't build a decryption context for these rlocs" );
        return -1;
    }
    
    unsigned char *iv = data;
    unsigned char *secret = (unsigned char *)&entry->secret[0];
    size_t written = 0;
    int outl = dest_len;
    
    if ( !EVP_DecryptInit_ex( &entry->ctx, EVP_aes_256_cbc(), NULL, secret, iv ) ) {
        warn( "EVP_DecrypttInit_ex() failed" );
        goto fail;
    }
    
    if ( !EVP_DecryptUpdate( &entry->ctx, dest, &outl, data + 16, data_len - 16 ) ) {
        warn( "EVP_DecryptUpdate() failed" );
        goto fail;
    }
    
    written += outl;
    outl = dest_len - written;

    if ( !EVP_DecryptFinal_ex( &entry->ctx, dest + written, &outl ) ) {
        warn( "EVP_DecryptFinal_ex() failed" );
        goto fail;
    }

    return written + outl;

fail:
    show_ssl_errors();
    return -1;
}


void rlocs_debug_output( struct rlocs *reg )
{
    int i;
    char addr1[128], addr2[128];
    
    info( "Configuration has parsed as:" );
    for( i = 0 ; i < reg->num_entries; i++ ) {
        struct rloc *current = &reg->entries[i];
        inet_ntop( current->family, &current->addr, &addr1[0], 128 );
        
        info( "RLOC %i: family %i, address %s", i, current->family, &addr1[0] );
    }
    
    for ( i = 0 ; i < reg->num_ip4_map_entries ; i++ ) {
        struct ip4_eid_map_entry *current = &reg->ip4_mappings[i];
        inet_ntop( AF_INET, &current->network, &addr1[0], 128 );
        inet_ntop( AF_INET, &current->rloc->addr, &addr2[0], 128 );
        info( "IPv4 EID mapping %i: %s/%u => %s", i, addr1, current->mask, addr2 );
    }

    for ( i = 0 ; i < reg->num_ip6_map_entries ; i++ ) {
        struct ip6_eid_map_entry *current = &reg->ip6_mappings[i];
        inet_ntop( AF_INET6, &current->network, &addr1[0], 128 );
        inet_ntop( AF_INET6, &current->rloc->addr, &addr2[0], 128 );
        info( "IPv6 EID mapping %i: %s/%u => %s", i, addr1, current->mask, addr2 );
    }    

    return;
}


void rlocs_free( struct rlocs* registry )
{
    int i;
    
    for( i = 0; i < registry->num_entries; i++ ) {
        if ( registry->entries[i].key != NULL ) {
            EVP_PKEY_free( registry->entries[i].key );
        }
    }
    // No need to do json_object_put() here.
    free( registry );
    return;
}