/*
 * Secure eXtended Message Passing framework
 * Secure eXtended Transport layer implementation: (libsxt)
 * - very similar to SSH2/TLS
 * - using already proven and tested crypto algos
 * - better than TLS for message passing
 *
 * PublicPrivateKeyPairs operation API
 *
 * (c) Alexander Vdolainen 2016 <avdolainen@zoho.com>
 *
 * libsxmp is free software: you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * libsxmp is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.";
 *
 */

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <stdarg.h>
#include <errno.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <fcntl.h>

#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/engine.h>

#include <sxt/errno.h>
#include <sxt/safebuf.h>
#include <sxt/rdb.h>
#include <sxt/crypt.h>
#include <sxt/sxtkey.h>
#include <sxt/ed25519.h>
#include <sxt/ciphers.h>
#include <sxt/base64.h>

/* locals */
static int sxt_key_export_priv_ed25519(const sxtkey_t *key, const char *pass, sxtsafebuffer_t **bin);
static int sxt_public_ed25519_2sb(const sxtkey_t *key, sxtsafebuffer_t **buf);
static int sxt_public_ed25519_2rdb(const sxtkey_t *key, sxtrdb_t **out);

int sxt_key_generate(sxtkey_t *key, int type, int opt)
{
  int r = 0;

  if(!key) return SXT_EINVAL;

  /* set defaults */
  key->type = type;
  key->flags = SXT_PPKP_PRIVATE | SXT_PPKP_PUBLIC;

  switch(type) {
  case PPKP_ED25519:
    key->pubkey = malloc(sizeof(ed25519_pubkey));
    if(!key->pubkey) {
      r = SXT_ENOMEM;
      goto __fall;
    }
    key->privkey = malloc(sizeof(ed25519_privkey));
    if(!key->privkey) {
      r = SXT_ENOMEM;
      goto __safefall0;
    }

    r = crypto_sign_ed25519_keypair(*key->pubkey, *key->privkey);
    if(r) {
      r = SXT_ECRYPTO;
      goto __safefall0;
    }
    break;
  default:
    return SXT_EINVAL;
  }

  if(!r) return SXT_SUCCESS;

 __safefall0:
  if(key->pubkey) free(key->pubkey);
  if(key->privkey) free(key->privkey);

 __fall:
  return r;
}

sxtkey_t *sxt_key_alloc(void)
{
  sxtkey_t *key = malloc(sizeof(sxtkey_t));

  if(!key) return NULL;

  return key;
}

void sxt_key_burn(sxtkey_t *key)
{
  if(!key) return;

  key->priv = NULL;

  switch(key->type) {
  case PPKP_ED25519:
    if(key->pubkey) memset(key->pubkey, 0, sizeof(ed25519_pubkey));
    if(key->privkey) memset(key->privkey, 0, sizeof(ed25519_privkey));
    break;
  default:
    return;
  }

  key->type = 0;
  key->flags = 0;

  return;
}

void sxt_key_free(sxtkey_t *key)
{
  if(!key) return;

  switch(key->type) {
  case PPKP_ED25519:
    if(key->pubkey) free(key->pubkey);
    if(key->privkey) free(key->privkey);
    break;
  default: return; /* cannot free unrecognized key due to the
                    * potential memleak
                    */
  }

  free(key);
  return;
}

const char *sxt_key_name(int type)
{
  switch(type) {
  case PPKP_ED25519:
    return "ppkp-ed25519";
  default: return NULL;
  }

  return NULL;
}

uint8_t sxt_key_type_fname(const char *name)
{
  if(!name) return 0;

  if(!strcmp(name, "ppkp-ed25519")) return PPKP_ED25519;
  else return 0;
}

int sxt_key_public(const sxtkey_t *key)
{
  if(!key) return 0;

  return (key->flags & SXT_PPKP_PUBLIC) == SXT_PPKP_PUBLIC;
}

int sxt_key_private(const sxtkey_t *key)
{
  if(!key) return 0;

  return (key->flags & SXT_PPKP_PRIVATE) == SXT_PPKP_PRIVATE;
}

int sxt_key_assign_hash(sxtkey_t *key, uint64_t hash)
{
  if(!key) return SXT_EINVAL;

  if(!sxt_key_private(key)) return SXT_EINVAL;

  if(key->hash) return SXT_EKEY;
  else key->hash = hash;

  return SXT_SUCCESS;
}

uint64_t sxt_key_hash(const sxtkey_t *key)
{
  if(!key) return 0;

  return key->hash;
}

/* will duplicate a key depends on it's kind */
int sxt_key_dup(const sxtkey_t *key, sxtkey_t **dest)
{
  int r = SXT_EKEY;

  if(!key) return SXT_EINVAL;

  if(sxt_key_public(key)) return sxt_key_dup_public(key, dest);
  else if(sxt_key_private(key)) return sxt_key_dup_private(key, dest);
  else return sxt_key_dup_private(key, dest);

  /* never riched */
  return r;
}

/* will duplicate public key, if key was private or pair - it becomes public */
int sxt_key_dup_public(const sxtkey_t *key, sxtkey_t **dest)
{
  sxtkey_t *nkey = NULL;
  int r = SXT_SUCCESS;

  if(!key) return SXT_EINVAL;

  if(!sxt_key_private(key)) return SXT_EKEY;

  if(!(nkey = sxt_key_alloc())) return SXT_ENOMEM;
  else nkey->flags = 0;

  switch(key->type) {
  case PPKP_ED25519: /* actually the same still, but in case of other key types will differ */
    /* check first */
    if(!key->pubkey || !key->privkey) {
      r = SXT_EKEY;
      goto __failed;
    }

    /* allocate new values */
    if(!(nkey->pubkey = malloc(ED25519_PK_LEN))) {
      r = SXT_ENOMEM;
      goto __failed;
    }
    if(!(nkey->privkey = malloc(ED25519_SK_LEN))) {
      free(nkey->pubkey);
      r = SXT_ENOMEM;
      goto __failed;
    }

    /* copy that */
    memcpy(nkey->pubkey, key->pubkey, ED25519_PK_LEN);
    memcpy(nkey->privkey, key->privkey, ED25519_SK_LEN);
    break;
  default: /* don't now about it */
    sxt_key_free(nkey);
    return SXT_EKEY;
  }

  /* copy other values */
  nkey->flags |= SXT_PPKP_PUBLIC;
  nkey->type = key->type;
  nkey->hash = key->hash;
  nkey->priv = key->priv;

  *dest = nkey;

  return r;

 __failed:
  sxt_key_free(nkey);
  return r;
}

/* will duplicate private key, if key public error will returns */
int sxt_key_dup_private(const sxtkey_t *key, sxtkey_t **dest)
{
  sxtkey_t *nkey = NULL;
  int r = SXT_SUCCESS;

  if(!key) return SXT_EINVAL;
  //if(sxt_key_public(key)) return SXT_EKEY; /* cannot duplicate private key from public */

  if(!(nkey = sxt_key_alloc())) return SXT_ENOMEM;
  else nkey->flags = 0;

  switch(key->type) {
  case PPKP_ED25519:
    /* check first */
    if(!key->pubkey || !key->privkey) {
      r = SXT_EKEY;
      goto __failed;
    }

    /* allocate new values */
    if(!(nkey->pubkey = malloc(ED25519_PK_LEN))) {
      r = SXT_ENOMEM;
      goto __failed;
    }
    if(!(nkey->privkey = malloc(ED25519_SK_LEN))) {
      free(nkey->pubkey);
      r = SXT_ENOMEM;
      goto __failed;
    }

    /* copy that */
    memcpy(nkey->pubkey, key->pubkey, ED25519_PK_LEN);
    memcpy(nkey->privkey, key->privkey, ED25519_SK_LEN);
    break;
  default: /* don't now about it */
    sxt_key_free(nkey);
    return SXT_EKEY;
  }

  /* copy other values */
  nkey->flags = SXT_PPKP_PRIVATE;
  nkey->type = key->type;
  nkey->hash = key->hash;
  nkey->priv = key->priv;

  *dest = nkey;

  return r;

 __failed:
  sxt_key_free(nkey);
  return r;
}

#if 0
#define dbg_mark  printf("%s:%d\n", __FUNCTION__, __LINE__)
#endif
int sxt_key_export_priv_file(const sxtkey_t *key, const char *file, const char *passkey,
                             int (*ask_passkey)(char *pkbuf, size_t length, int confirm, void *priv),
                             void *priv)
{
  FILE *out = NULL;
  char *pass = NULL;
  sxtsafebuffer_t *bin = NULL;
  int r = 0;

  if(!key) return SXT_EINVAL;
  if(!sxt_key_private(key)) return SXT_EKEY;

  if((out = fopen(file, "wb")) == NULL) return SXT_EIO;

  /* passkey challenge */
  if(passkey) pass = (char *)passkey;
  else if(ask_passkey) {
    if((pass = malloc(sizeof(char)*64)) == NULL) {
      r = SXT_ENOMEM;
      goto __failed;
    }
    r = ask_passkey(pass, 64, 1, priv);
    if(r != SXT_SUCCESS)      goto __failed;
  }

  /* ok, get the base64 encoded data for the key */
  switch(key->type) {
  case PPKP_ED25519:
    r = sxt_key_export_priv_ed25519(key, pass, &bin);
    if(!bin) goto __failed;
    break;
  default:
    r = SXT_EKEY;
    goto __failed;
  }

  /* sxt key container as follows: (<key-type> "<base64 encoded content of the key>")\n */
  fprintf(out, "(%s \"%s\")\n", sxt_key_name(key->type), (char *)sxtsafebuffer_getdata(bin));

  sxtsafebuffer_destroy(bin);

 __failed:
  if(pass && pass != passkey) {
    memset(pass, 0, sizeof(char)*64); /* clean to disallow passkeys found in coredumps */
    free(pass);
  }
  fclose(out);

  return r;
}

static int sxt_key_publickey2blob_whash(const sxtkey_t *key, sxtsafebuffer_t **o)
{
  sxtsafebuffer_t *buf = NULL, *pk = NULL;
  sxtrdb_t *b = NULL;
  int r = SXT_SUCCESS;

  if(!key) return SXT_EINVAL;

  /* get a public key blob without hash */
  switch(key->type) {
  case PPKP_ED25519:
    if((r = sxt_public_ed25519_2sb(key, &pk)) != SXT_SUCCESS)
      return r;
    break;
  default: return SXT_EKEY;
  }

  if(!(b = sxtrdb_new())) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  sxtrdb_print(b, "pq", sxtsafebuffer_length(pk), sxtsafebuffer_getdata(pk),
               key->hash);
  if(!(buf = sxtsafebuffer_new(sxtrdb_length(b)))) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  memcpy(sxtsafebuffer_getdata(buf), sxtrdb_rdata(b), sxtrdb_length(b));

  *o = buf;

 __failed:
  if(pk) sxtsafebuffer_destroy(pk);
  if(b) sxtrdb_free(b);

  return r;
}

static int sxt_key_hash2blob(const sxtkey_t *key, sxtsafebuffer_t **o)
{
  sxtsafebuffer_t *buf = NULL;
  sxtrdb_t *b = NULL;
  int r = SXT_SUCCESS;

  if(!key) return SXT_EINVAL;

  if(!(b = sxtrdb_new())) return SXT_ENOMEM;

  sxtrdb_print(b, "q", key->hash);
  if(!(buf = sxtsafebuffer_new(sxtrdb_length(b)))) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  memcpy(sxtsafebuffer_getdata(buf), sxtrdb_rdata(b), sxtrdb_length(b));

  *o = buf;

 __failed:
  if(b) sxtrdb_free(b);

  return r;
}

int sxt_key_export_public_file(const sxtkey_t *key, const char *file)
{
  FILE *out = NULL;
  sxtsafebuffer_t *hashbuf = NULL, *pubkeybuf = NULL;
  sxtsafebuffer_t *pubkey_sb = NULL, *hash_sb = NULL;
  char *pckey = NULL, *pchash = NULL;
  int r = SXT_SUCCESS;

  if(!key) return SXT_EINVAL;
  if(!sxt_key_public(key)) return SXT_EKEY;

  if((out = fopen(file, "wb")) == NULL) return SXT_EIO;

  /*
   * format described as follows:
   * (<key-type> 'public "<pubkey><hash>" "<hash>")\n
   * all scoped values are base64 encoded
   */

  /* get a public key with hash blob */
  if((r = sxt_key_publickey2blob_whash(key, &pubkey_sb)) != SXT_SUCCESS) goto __failed;

  /* get key hash blob */
  if((r = sxt_key_hash2blob(key, &hash_sb)) != SXT_SUCCESS) goto __failed;

  /* allocate buffers for base64 encoded data */
  if(!(pubkeybuf = sxtsafebuffer_new(sxt_rawlen2b64len(sxtsafebuffer_length(pubkey_sb))))) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  if(!(hashbuf = sxtsafebuffer_new(sxt_rawlen2b64len(sxtsafebuffer_length(hash_sb))))) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  /* do base64 encode */
  sxt_b64encode_in(sxtsafebuffer_getdata(pubkey_sb), sxtsafebuffer_getdata(pubkeybuf),
                   sxtsafebuffer_length(pubkey_sb)); /* pubkey */
  sxt_b64encode_in(sxtsafebuffer_getdata(hash_sb), sxtsafebuffer_getdata(hashbuf),
                   sxtsafebuffer_length(hash_sb)); /* hash */

  /* get cstring to output */
  if(!(pckey = sxtsafebuffer_getcstr(pubkeybuf))) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  if(!(pchash = sxtsafebuffer_getcstr(hashbuf))) {
    r = SXT_ENOMEM;
    sxtsafebuffer_freecstr(pckey);
    goto __failed;
  }

  /* output this */
  fprintf(out, "(%s 'public \"%s\" \"%s\")\n", sxt_key_name(key->type), pckey, pchash);

  /* free all stuff */
  sxtsafebuffer_freecstr(pckey);
  sxtsafebuffer_freecstr(pchash);

 __failed:
  if(pubkey_sb) sxtsafebuffer_destroy(pubkey_sb);
  if(hash_sb) sxtsafebuffer_destroy(hash_sb);
  if(hashbuf) sxtsafebuffer_destroy(hashbuf);
  if(pubkeybuf) sxtsafebuffer_destroy(pubkeybuf);
  fclose(out);

  return r;
}

static int sxt_key_export_priv_ed25519(const sxtkey_t *key, const char *pass, sxtsafebuffer_t **bin)
{
  int r = 0;
  sxtsafebuffer_t *b64 = NULL, *pubkey_sb = NULL;
  sxtrdb_t *keybuf = NULL, *privbuf = NULL;
  sxtrdb_t *kdfopts = NULL;
  uint32_t urand;
  uint8_t flags = 0;

  if(!key) return SXT_EINVAL;
  if(key->type != PPKP_ED25519) return SXT_EKEY;

  /* allocate a buffer for the key at all */
  if(!(keybuf = sxtrdb_new())) return SXT_ENOMEM;
  else sxtrdb_setflags(keybuf, SXTRDB_BURN);

  /* ok firstly get private-key-data according to SXT key container doc */
  if((r = sxt_public_ed25519_2sb(key, &pubkey_sb)) != SXT_SUCCESS) goto __failed;

  /* now we need random number */
  sxt_get_random(&urand, sizeof(uint32_t), 1);

  if(!(privbuf = sxtrdb_new())) {
    r = SXT_ENOMEM;
    goto __failed;
  } else sxtrdb_setflags(privbuf, SXTRDB_BURN);

  /* two random numbers */
  if(!sxtrdb_print(privbuf, "dd", urand, urand)) {
    r = SXT_ENOMEM;
    goto __failed;
  }

  /* pub, priv, hash, zero padding */
  if(!sxtrdb_print(privbuf, "dpdpqb", (uint32_t)ED25519_PK_LEN, (size_t)ED25519_PK_LEN,
                   key->pubkey, (uint32_t)ED25519_SK_LEN, (size_t)ED25519_SK_LEN,
                   key->privkey, key->hash, (uint8_t)'\0')) {
    r = SXT_ENOMEM;
    goto __failed;
  }

  /* if hash is non-zero - include it i.e. set flag */
  if(key->hash) flags |= SXT_PPKP_IHASH;

  if(pass) { /* we need to encrypt the key i.e. privbuf *must* be encrypted */
    /* to avoid yet another wheel invention,
     * the following was taken from ssh implementation, i mean
     * kdf and crypto in use to encrypt the key
     */
    uint8_t pad = 1;
    size_t keystub_len;
    sxtsafebuffer_t *salt;
    sxt_cipher_t *cipher = sxt_cipher_get("aes128-cbc");
    uint8_t keystub[128];

    /* set flags */
    flags |= SXT_PPKP_ENCRYPT;

    if(!cipher) {
      /* it might be in case of forgotten initialization function */
      r = SXT_EKEY;
      goto __failed;
    }

    /* get the salt */
    if(!(salt = sxtsafebuffer_new(16))) {
    __enomem1:
      sxt_cipher_free(cipher);
      r = SXT_ENOMEM;
      goto __failed;
    } else sxt_get_random(sxtsafebuffer_getdata(salt), 16, 1);

    /* generate kdf description <salt><rounds> */
    if(!(kdfopts = sxtrdb_new())) {
      sxtsafebuffer_destroy(salt);
      goto __enomem1;
    }
    sxtrdb_print(kdfopts, "pd", 16, sxtsafebuffer_getdata(salt), (uint32_t)32);

    /* ok, now we need to get a padding, just to be sure to be aligned for the
       cipher block size */
    while(sxtrdb_length(privbuf) % cipher->blksize != 0) {
      if(sxtrdb_write_u8(privbuf, pad) != sizeof(uint8_t)) {
        sxt_cipher_free(cipher);
        r = SXT_ENOMEM;
        goto __failed;
      }
      pad++;
    }

    /* key stub */
    keystub_len = cipher->keysize/8 + cipher->blksize;
    if(keystub_len > sizeof(keystub)) {
      r = SXT_EKEY;
      sxt_cipher_free(cipher);
      goto __failed;
    }

    /* bcrypt those stuff */
    if(bcrypt_pbkdf(pass, strlen(pass), sxtsafebuffer_getdata(salt),
                    sxtsafebuffer_length(salt), keystub, keystub_len, 32)) {
      r = SXT_ERROR;
      sxt_cipher_free(cipher);
      goto __failed;
    }

    /* encrypt with cipher */
    sxt_cipher_set_encrypt_key(cipher, keystub, keystub + cipher->keysize/8);
    sxt_cipher_encrypt(cipher, sxtrdb_rdata(privbuf), sxtrdb_rdata(privbuf), sxtrdb_length(privbuf));

    /* clean up all */
    memset(keystub, 0, sizeof(keystub));
    sxt_cipher_free(cipher);
    sxtsafebuffer_destroy(salt);
  }

  /* in case of non-crypted key we anyways should include encrypt options */
  if(!kdfopts) {
    if(!(kdfopts = sxtrdb_new())) {
      r = SXT_ENOMEM;
      goto __failed;
    }
    sxtrdb_print(kdfopts, "pd", 3, "nil", (uint32_t)0);
  }

  /* there are a time to print all required content to the buffer */
  /* <MAGIC VERSION SIGN><8bit flags><key-type>
   * <cipher-name><kdf-name><kdf-opts>
   * <public-key>
   * <private-key-data>
   * according to the documentation
   */
  sxtrdb_print(keybuf, "sbsssppp", PPKP_MAGIC, flags, "ppkp-ed25519",
               pass ? "aes128-cbc" : "nil", pass ? "bcrypt" : "nil",
               sxtrdb_length(kdfopts), sxtrdb_rdata(kdfopts),
               sxtsafebuffer_length(pubkey_sb), sxtsafebuffer_getdata(pubkey_sb),
               sxtrdb_length(privbuf), sxtrdb_rdata(privbuf));

  /* as described - encode all with base64 */
  b64 = sxtsafebuffer_new(sxt_rawlen2b64len(sxtrdb_length(keybuf)));
  if(!b64) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  /* encode buffer */
  sxt_b64encode_in((const char *)sxtrdb_rdata(keybuf), (char *)sxtsafebuffer_getdata(b64),
                   sxtrdb_length(keybuf));

  /* the last ... */
  *bin = b64;
  r = SXT_SUCCESS;

 __failed:
  if(keybuf) sxtrdb_free(keybuf); /* container free */
  if(pubkey_sb) sxtsafebuffer_destroy(pubkey_sb); /* public key bin free */
  if(privbuf) sxtrdb_free(privbuf); /* private key buffer free */

  return r;
}

static int sxt_public_ed25519_2sb(const sxtkey_t *key, sxtsafebuffer_t **buf)
{
  sxtrdb_t *rbuf = NULL;
  sxtsafebuffer_t *o = NULL;
  int r = SXT_SUCCESS;

  if(!key->pubkey) return SXT_EKEY;

  r = sxt_public_ed25519_2rdb(key, &rbuf);
  if(!rbuf) return r;

  if(!(o = sxtsafebuffer_new((size_t)sxtrdb_length(rbuf)))) {
    r = SXT_ENOMEM;
    goto __failed;
  }
  sxtsafebuffer_setdata(o, sxtrdb_rdata(rbuf), (size_t)sxtrdb_length(rbuf));
  *buf = o;

 __failed:
  sxtrdb_free(rbuf);

  return r;
}

static int sxt_public_ed25519_2rdb(const sxtkey_t *key, sxtrdb_t **out)
{
  sxtrdb_t *buf = NULL;

  if(!key->pubkey) return SXT_EKEY;

  if(!(buf = sxtrdb_new())) return SXT_ENOMEM;
  else sxtrdb_setflags(buf, SXTRDB_BURN);

  if(!sxtrdb_print(buf, "dp", (uint32_t)ED25519_PK_LEN, (size_t)ED25519_PK_LEN,
                   key->pubkey)) {
    sxtrdb_free(buf);
    return SXT_EKEY;
  }

  *out = buf;

  return SXT_SUCCESS;
}