#include <assert.h>
#include <err.h>
#include <termios.h>
#include <time.h>

// include the GNU extension of asprintf
#define _GNU_SOURCE
#include <stdio.h>

#include "libcreddy_common.h"

#define KEY_SUFFIX  "_private.pem"
#define CERT_SUFFIX "_ID.pem"
/* Size of password memory allocation */
#define PWLEN 128

//
// ID object
//
struct _creddy_id_t {
    char *keyid;
    char *cn;
    certificate_t *cert;
    private_key_t *key;

    int refcount;
};

/* Callback configuration */
struct cb_opts {
    bool use_prompt;    /* Print a prompt to stderr */
    bool use_echo;      /* If true, turn off input echo on stdin */
    unsigned int tries; /* Number of attempts allowed */
    char prompt[20];    /* The prompt to display if use_echo is true */
};

static char *_get_keyid(certificate_t *cert);
static chunk_t _passphrase_callback(void *user, int try);
static private_key_t *_generate_key(void);
static certificate_t *_generate_cert(private_key_t *private, char *cn, int validity);
static void _encode_base64(FILE *out, chunk_t encoding);

/**
 * Helper function for building a creddy ID from a cert. Used by
 * creddy_id_from_*
 */
creddy_id_t *_creddy_id_from_cert(certificate_t *cert) {
    creddy_id_t *id = creddy_xmalloc(sizeof(creddy_id_t));
    id->keyid = NULL;
    id->cn = NULL;
    id->cert = cert;
    id->key = NULL;

    id->keyid = _get_keyid(id->cert);

    // get the CN from the cert
    id_part_t type;
    chunk_t data;

    identification_t *cert_id = id->cert->get_subject(id->cert);
    enumerator_t *id_enum = cert_id->create_part_enumerator(cert_id);
    while (id_enum->enumerate(id_enum, &type, &data))
        if (type == ID_PART_RDN_CN) {
            id->cn = creddy_xmalloc(data.len + 1);
            memcpy(id->cn, data.ptr, data.len);
            id->cn[data.len] = 0;
        }
    id_enum->destroy(id_enum);

    id->refcount = 1;

    return id;
}

/**
 * Load an ID cert from a file.
 */
creddy_id_t *creddy_id_from_file(char *filename) {
    libabac_init();

    certificate_t *cert = lib->creds->create(lib->creds,
        CRED_CERTIFICATE, CERT_X509,
        BUILD_FROM_FILE, filename,
        BUILD_X509_FLAG, X509_AA,
        BUILD_END
    );

    if (cert == NULL)
        return NULL;

    return _creddy_id_from_cert(cert);
}

/**
 * Load an ID cert from a chunk.
 */
creddy_id_t *creddy_id_from_chunk(abac_chunk_t achunk) {
    chunk_t chunk = { .ptr = achunk.ptr, .len = achunk.len };

    libabac_init();

    certificate_t *cert = lib->creds->create(lib->creds,
        CRED_CERTIFICATE, CERT_X509,
        BUILD_BLOB_ASN1_DER, chunk,
        BUILD_X509_FLAG, X509_AA,
        BUILD_END
    );

    if (cert == NULL)
        return NULL;

    return _creddy_id_from_cert(cert);
}

/**
 * Load private key for a cert.
 */
int creddy_id_load_privkey(creddy_id_t *id, char *filename) {
    struct cb_opts c_opts = { 1, 0, 3, "Key password:" };

    assert(id != NULL);

    libabac_init();

    // load signer key
    private_key_t *key = lib->creds->create(lib->creds,
        CRED_PRIVATE_KEY, KEY_RSA,
        BUILD_FROM_FILE, filename,
        /* Ask for password if the key's encrypted */
        BUILD_PASSPHRASE_CALLBACK, _passphrase_callback, &c_opts,
        BUILD_END
    );
    if (key == NULL)
        return 0;

    id->key = key;
    return 1;
}

/**
 * Generate an ID with the specified CN and validity.
 *
 * validity is measured in seconds (as of 0.2.0)
 */
int creddy_id_generate(creddy_id_t **ret, char *cn, int validity) {
    if (cn == NULL || !creddy_clean_name(cn))
        return CREDDY_GENERATE_INVALID_CN;

    if (validity < 0)
        return CREDDY_GENERATE_INVALID_VALIDITY;

    creddy_id_t *id = creddy_xmalloc(sizeof(creddy_id_t));

    id->cn = creddy_xstrdup(cn);
    id->key = _generate_key();
    id->cert = _generate_cert(id->key, cn, validity);
    id->keyid = _get_keyid(id->cert);

    id->refcount = 1;

    *ret = id;
    return CREDDY_SUCCESS;
}

char *creddy_id_keyid(creddy_id_t *id) {
    assert(id != NULL);

    return id->keyid;
}

/**
 * Get the issuer of an ID cert.
 * Returns a malloc'd string that must be free'd.
 */
char *creddy_id_issuer(creddy_id_t *id) {
    char *ret;
    int rv = asprintf(&ret, "%Y", id->cert->get_issuer(id->cert));

    if (rv < 0)
        err(1, "couldn't malloc string for issuer\n");

    return ret;
}

/**
 * Gets the subject DN of an ID cert.
 * Returns a malloc'd string that must be free'd.
 */
char *creddy_id_subject(creddy_id_t *id) {
    char *ret;
    int rv = asprintf(&ret, "%Y", id->cert->get_subject(id->cert));

    if (rv < 0)
        err(1, "couldn't malloc string for subject\n");

    return ret;
}

/**
 * Get the validity period.
 */
void creddy_id_validity(creddy_id_t *id, time_t *not_before, time_t *not_after) {
    id->cert->get_validity(id->cert, NULL, not_before, not_after);
}

certificate_t *creddy_id_cert(creddy_id_t *id) {
    assert(id != NULL);

    return id->cert;
}

private_key_t *creddy_id_privkey(creddy_id_t *id) {
    assert(id != NULL);

    return id->key;
}

/**
 * Get the default filename for the cert. Value must be freed by caller.
 */
char *creddy_id_cert_filename(creddy_id_t *id) {
    assert(id != NULL);
    assert(id->cn != NULL);

    // malloc the filename
    int len = strlen(id->cn) + strlen(CERT_SUFFIX) + 1;
    char *filename = creddy_xmalloc(len);
    sprintf(filename, "%s" CERT_SUFFIX, id->cn);

    return filename;
}

/**
 * Write the ID cert to an open file pointer.
 */
void creddy_id_write_cert(creddy_id_t *id, FILE *out) {
    assert(id != NULL);

    chunk_t encoding = id->cert->get_encoding(id->cert);
    _encode_base64(out, encoding);
    free(encoding.ptr);
}

/**
 * Default private key filename. Value must be freed by caller.
 */
char *creddy_id_privkey_filename(creddy_id_t *id) {
    assert(id != NULL);
    assert(id->cn != NULL);

    // malloc the filename
    int len = strlen(id->cn) + strlen(KEY_SUFFIX) + 1;
    char *filename = creddy_xmalloc(len);
    sprintf(filename, "%s" KEY_SUFFIX, id->cn);

    return filename;
}

/**
 * Write the private key to a file.
 * Returns false if there's no private key loaded
 */
int creddy_id_write_privkey(creddy_id_t *id, FILE *out) {
    int ret;
    chunk_t encoding;

    assert(id != NULL);

    if (id->key == NULL)
        return 0;

    ret = id->key->get_encoding(id->key, KEY_PRIV_PEM, &encoding);
    if (!ret)
        errx(1, "Couldn't encode private key");

    fwrite(encoding.ptr, encoding.len, 1, out);

    free(encoding.ptr);
    return 1;
}

/**
 * Get a DER-encoded chunk representing the cert.
 */
abac_chunk_t creddy_id_cert_chunk(creddy_id_t *id) {
    chunk_t encoding = id->cert->get_encoding(id->cert);
    abac_chunk_t ret = { encoding.ptr, encoding.len };
    return ret;
}

/**
 * Copy a creddy ID. Actually just increases its reference count.
 */
creddy_id_t *creddy_id_dup(creddy_id_t *id) {
    ++id->refcount;
    return id;
}

void creddy_id_free(creddy_id_t *id) {
    if (id == NULL)
        return;

    --id->refcount;
    if (id->refcount > 0)
        return;

    // free once the reference count reaches 0
    DESTROY_IF(id->cert);
    DESTROY_IF(id->key);

    free(id->keyid);
    free(id);
}

//
// Helper functions below
//

static char *_get_keyid(certificate_t *cert) {
    // get the keyid
    x509_t *x509 = (x509_t *)cert;
    chunk_t keyid = x509->get_subjectKeyIdentifier(x509);
    chunk_t string = chunk_to_hex(keyid, NULL, 0);
    return (char *)string.ptr;
}

static chunk_t _passphrase_callback(void *user, int try) {
    /* Get a password from stdin and return it as a chunk_t.  If too many tries
     * have occurred or there is any other problem, return an empty chunk_t,
     * which libstrongswan takes as giving up.  The chunk is alloated here
     * (inside getline), and presumably freed by libstrongswan. User points to
     * a cb_opts struct, which affects this routine in the obvious ways.
     */
    /* Configuration options */
    struct cb_opts *opts = (struct cb_opts *) user;
    chunk_t rv = chunk_empty;   /* Return value, starts empty */

    if (try -1 < opts->tries ) {
        struct termios t;   /* Terminal settings */
        size_t len = 0;     /* Length of string from getline */
        tcflag_t orig = 0;  /* Holds the original local flags (echo in here) */

        if (!opts->use_echo) {
            /* Use tc{get,set}attr to turn echo off and restore the intial
             * echo settings */
            if (!tcgetattr(0, &t)) {
                orig = t.c_lflag;

                t.c_lflag &= ~ECHO;
                if ( tcsetattr(0, TCSANOW, &t) ) {
                    perror("Cannot turn off echo");
                    return rv;
                }
            }
            else {
                perror("Cannot turn get attributes to off echo");
                return rv;
            }
        }
        if (opts->use_prompt) printf("%s", opts->prompt);

        /* Because rv.ptr starts as NULL, getline allocates memory.  The size
         * of the allocation returns in rv.len and the size of the string
         * (including newline and NUL) is in len.  */
        if ((rv.ptr = (u_char *) malloc(rv.len = PWLEN))) {
            if ( fgets(rv.ptr, rv.len, stdin) ) {
                /* Readjust the chunk_t's len field to the size of the string
                 * w/o the newline or NUL */
                /* would prefer strnlen, but no such luck in FBSD7 or earlier*/
                size_t len = strlen(rv.ptr);

                if (rv.ptr[len-2] == '\n') rv.len = len-2;
                else rv.len = len -1;
            }
            else {
                /* Read failed.  Deallocate and clear rv */
                free(rv.ptr);
                rv = chunk_empty;
            }
        }
        else {
            /* Failed malloc.  Restore rv to empty and return it */
            perror("malloc");
            rv = chunk_empty;
            return rv;
        }

        if (!opts->use_echo ) {
            /* Pop echo beck to its original setting. */
            t.c_lflag = orig;

            if ( tcsetattr(0, TCSANOW, &t) )
                perror("Cannot restore echo setting?");

            if (opts->use_prompt) printf("\n");
        }
    }
    else fprintf(stderr, "Too many tries (%d)", try-1);
    return rv;
}

/**
 * Generate a private key.
 */
static private_key_t *_generate_key(void) {
    private_key_t *key;
    libabac_init();

    // generate the key
    key = lib->creds->create(
        lib->creds,
        CRED_PRIVATE_KEY, KEY_RSA,
        BUILD_KEY_SIZE, 2048,
        BUILD_END
    );
    if (key == NULL)
        errx(1, "Key generation failed");

    return key;
}

static char *_create_dn(char *cn) {

#define DN "cn="

    char *dn = creddy_xmalloc(sizeof(DN) + strlen(cn));
    memcpy(dn, DN, sizeof(DN));
    strcat(dn, cn);

    return dn;
}

/**
 * Generate certificate.
 *
 * validity is measured in seconds (as of 0.2.0)
 */
static certificate_t *_generate_cert(private_key_t *private, char *cn, int validity) {
    // build the DN
    char *dn_string = _create_dn(cn);
    libabac_init();

    identification_t *id = identification_create_from_string(dn_string);
    if (id == NULL)
        errx(1, "couldn't create ID from DN %s", dn_string);
    free(dn_string);

    // get the public key
    public_key_t *public = private->get_public_key(private);
    if (public == NULL)
        errx(1, "couldn't get public key from private key");

    // create a serial (stolen from strongswan pki)
    rng_t *rng = lib->crypto->create_rng(lib->crypto, RNG_WEAK);
    if (!rng)
        errx(1, "no random number generator");

    // random serial
    chunk_t serial = creddy_generate_serial();

    // validity period
    time_t not_before = time(NULL);
    time_t not_after = not_before + validity;

    // create!
    certificate_t *cert = lib->creds->create(lib->creds,
        CRED_CERTIFICATE, CERT_X509,
        BUILD_SIGNING_KEY, private,
        BUILD_PUBLIC_KEY, public,
        BUILD_SUBJECT, id,
        BUILD_NOT_BEFORE_TIME, not_before,
        BUILD_NOT_AFTER_TIME, not_after,
        BUILD_SERIAL, serial,
        BUILD_DIGEST_ALG, HASH_SHA1,
        BUILD_X509_FLAG, X509_CA,
        BUILD_PATHLEN, X509_NO_PATH_LEN_CONSTRAINT,
        BUILD_END
    );
    if (cert == NULL)
        errx(1, "couldn't build cert :(");

    DESTROY_IF(id);
    DESTROY_IF(public);
    free(serial.ptr);

    return cert;
}

#define BYTES_PER_LINE 64

// thx libstrongswan
static void _encode_base64(FILE *out, chunk_t encoding) {
    int start;

    chunk_t b64 = chunk_to_base64(encoding, NULL);

    fprintf(out, "-----BEGIN CERTIFICATE-----\n");

    for (start = 0; start < b64.len; start += BYTES_PER_LINE) {
        int left = b64.len - start;
        int len = left < BYTES_PER_LINE ? left : BYTES_PER_LINE;
        fwrite(b64.ptr + start, len, 1, out);
        fprintf(out, "\n");
    }

    fprintf(out, "-----END CERTIFICATE-----\n");

    free(b64.ptr);
}