/*
 * Copyright (c) 2021 Clementine Computing LLC.
 *
 * This file is part of PopuFare.
 *
 * PopuFare is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * PopuFare 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 Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with PopuFare.  If not, see <https://www.gnu.org/licenses/>.
 *
 */

#include "fareqr.h"

// AES CBC, sha256 message digest, with salt.
// Functions modified from https://wiki.openssl.org/index.php/EVP_Symmetric_Encryption_and_Decryption
//
// openssl example commands to encrypt/decrypt:
// msg=";123456789012345?"
// key="oovevobie8woid1Ou3iu6aboochei2AeKoceeCh1iePheuRae4Yai1dahtheegi7"
// openssl enc -aes-256-cbc -in <( echo -n "$msg" ) -out /dev/stdout -pass pass:"$key" -e -base64 -md sha256
//
// ---
//
// to descypt with openssl
//
// openssl aes-256-cbc -d -a  -in <( echo "U2FsdGVkX19gpKzk+Y5SGgnQr1uwpUYUpea6nu77wJfG6bD4GZpRneSaxaallz0n" ) -out /dev/stdout -pass pass:"$key" -md sha256
//


static int aes_encrypt(unsigned char *plaintext, int plaintext_len, unsigned char *key, unsigned char *iv, unsigned char *ciphertext) {
  EVP_CIPHER_CTX *ctx;

  int len;

  int ciphertext_len;

  // Create and initialise the context
  if(!(ctx = EVP_CIPHER_CTX_new())) { return -1; }

  // Initialise the encryption operation. IMPORTANT - ensure you use a key
  // and IV size appropriate for your cipher
  // In this example we are using 256 bit AES (i.e. a 256 bit key). The
  // IV size for *most* modes is the same as the block size. For AES this
  // is 128 bits
  //
  if(1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv)) { return -2; }

  // Provide the message to be encrypted, and obtain the encrypted output.
  // EVP_EncryptUpdate can be called multiple times if necessary
  //
  if(1 != EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len)) { return -3; }
  ciphertext_len = len;

  // Finalise the encryption. Further ciphertext bytes may be written at
  // this stage.
  //
  if(1 != EVP_EncryptFinal_ex(ctx, ciphertext + len, &len)) { return -4; }
  ciphertext_len += len;

  // Clean up
  //
  EVP_CIPHER_CTX_free(ctx);

  return ciphertext_len;
}

static int aes_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *key, unsigned char *iv, unsigned char *plaintext) {   
  EVP_CIPHER_CTX *ctx;
  int plaintext_len, len;

  // Create and initialise the context
  //
  if(!(ctx = EVP_CIPHER_CTX_new())) { return -1; }

  // Initialise the decryption operation. IMPORTANT - ensure you use a key
  // and IV size appropriate for your cipher
  // In this example we are using 256 bit AES (i.e. a 256 bit key). The
  // IV size for *most* modes is the same as the block size. For AES this
  // is 128 bits
  //
  if(1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv)) { return -2; }

  // Provide the message to be decrypted, and obtain the plaintext output.
  // EVP_DecryptUpdate can be called multiple times if necessary.
  //
  if(1 != EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len)) { return -3; }
  plaintext_len = len;

  // Finalise the decryption. Further plaintext bytes may be written at
  // this stage.
  //
  if(1 != EVP_DecryptFinal_ex(ctx, plaintext + len, &len)) { return -4; }
  plaintext_len += len;

  // Clean up
  //
  EVP_CIPHER_CTX_free(ctx);

  return plaintext_len;
}

// Helper function to take secret key/pass (pass_key) and the encoded
// base64 string and decrypt it with AES CBC 256, sha256 message digest.
// 
// return the number of bytes decoded on success
// return non-positive number on error
//
static int decode_b64(unsigned char *decrypt_text,
                      unsigned char *pass_key,
                      unsigned char *enc_b64) {
  int i=0, _ret = 0, r=0;
  unsigned char *enc_b = NULL, salt[8];
  unsigned char *key = NULL, *iv = NULL, *ciphertext_b = NULL;
  size_t enc_len=0, enc_size=0, ciphertext_len;

  if ((!decrypt_text) || (!pass_key) || (!enc_b64)) { return -9; }

  enc_size = Base64decode_len((char *)enc_b64);
  if (enc_size < 1) {
    _ret=-1;
    goto _decode_b64_cleanup;
  }

  enc_b = (unsigned char *)calloc(enc_size, sizeof(char));
  if (!enc_b) {
    _ret=-2;
    goto _decode_b64_cleanup;
  }
  Base64decode((char *)enc_b, (char *)enc_b64);
  enc_len = enc_size-1;

  // Check exmpanded string is well formed (has "Salted__" prefix, etc.)
  // and extract the salt
  //
  if (enc_len < 16) {
    _ret = -3;
    goto _decode_b64_cleanup;
  }
  if (strncmp((char *)enc_b, "Salted__", 8)!=0) {
    _ret = -4;
    goto _decode_b64_cleanup;
  }
  for (i=0; i<8; i++) { salt[i] = enc_b[i+8]; }

  // Move past the salt prefix 8 bytes) and 8 bytes of salt
  //
  ciphertext_b    = enc_b   + 16;
  ciphertext_len  = enc_len - 16;

  // Allocate iv, key and extract them from the salt and pass phrase
  //
  iv  = (unsigned char *)calloc(256, sizeof(char));
  key = (unsigned char *)calloc(256, sizeof(char));
  r = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha256(), salt, pass_key, strlen((const char *)pass_key), 1, key, iv);
  if (!r) {
    _ret = -5;
    goto _decode_b64_cleanup;
  }


  r = aes_decrypt(ciphertext_b, ciphertext_len, key, iv, decrypt_text);
  if (r<=0) {
    _ret = -6;
    goto _decode_b64_cleanup;
  }

  decrypt_text[r] = '\0';
  _ret = r;

_decode_b64_cleanup:

  if (enc_b) { free(enc_b); }
  if (key) { free(key); }
  if (iv) { free(iv); }

  return _ret;
}

static int encode_b64(unsigned char *enc_b64,
                      unsigned char *pass_key,
                      unsigned char *msg) {
  int i=0, _ret = 0, r=0;
  unsigned char *enc_b = NULL;
  unsigned char *key = NULL, *iv = NULL;
  size_t b64_sz=0, enc_len=0;

  unsigned char salt[8];
  char _pfx[] = "Salted__";

  if ((!enc_b64) || (!pass_key) || (!msg)) { return -9; }

  for (i=0; i<8; i++) { salt[i] = (unsigned char)(rand() % 256); }

  iv  = (unsigned char *)calloc(256, sizeof(char));
  key = (unsigned char *)calloc(256, sizeof(char));
  r = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha256(), salt, pass_key, strlen((char *)pass_key), 1, key, iv);
  if (!r) {
    _ret = -5;
    goto _encode_b64_cleanup;
  }

  enc_b = (unsigned char *)calloc(2*LINE_BUFFER_SIZE, sizeof(char));
  if (!enc_b) {
    _ret = -4;
    goto _encode_b64_cleanup;
  }

  for (i=0; i<8; i++) {
    enc_b[i] = _pfx[i];
    enc_b[i+8] = salt[i];
  }

  enc_len = aes_encrypt(msg, strlen((char *)msg), key, iv, enc_b+16);
  if (enc_len<=0) {
    _ret = -6;
    goto _encode_b64_cleanup;
  }
  enc_len += 16;

  b64_sz = Base64encode_len(enc_len);
  Base64encode((char *)enc_b64, (char *)enc_b, enc_len);
  _ret = b64_sz;

_encode_b64_cleanup:

  if (enc_b) { free(enc_b); }
  if (key) { free(key); }
  if (iv) { free(iv); }

  return _ret;
}



// `fqreqr_lookup_decode` uses the public key portion of `fareqr_s`
// to lookup the private pass/key in the file `seedfn`.
// If found, it then proceeds to decrypt and store the encrypted
// credential in `dst_cred`
//
// returns 0 on success
// returns non zero on error
// 
int fareqr_lookup_decode(char *seedfn, char *fareqr_s, char *dst_cred) {
  int i, n, r, _ret = 0;;
  char *p=NULL, *stop_tok=NULL;
  char *enc_str = NULL, *dec_str = NULL, *plain_str=NULL,
    *pub_key = NULL, *priv_key = NULL;

  if ((!seedfn) || (!fareqr_s) || (!dst_cred)) { return -1; }
  if (fareqr_s[0] != '@') { return -2; }
  stop_tok = strchr(fareqr_s, '%');
  if (!stop_tok) { return -3; }

  pub_key = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
  priv_key = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
  enc_str = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
  dec_str = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
  plain_str = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
  pub_key[0] = priv_key[0] = enc_str[0] = dec_str[0] = plain_str[0] = '\0';

  for (i=0, p = (fareqr_s+1); (i<(LINE_BUFFER_SIZE-1)) && (p<stop_tok) && (*p); p++, i++) {
    pub_key[i] = *p;
  }
  pub_key[i] = '\0';

  for (n=0, p = (stop_tok+1); (n<(LINE_BUFFER_SIZE-1)) && (*p) && ((*p) != '$'); p++, n++) {
    enc_str[n] = *p;
  }
  enc_str[n] = '\0';

  r = fareqr_lookup_seed_secret(seedfn, pub_key, priv_key);
  if (r<0) { _ret = r; }
  else {
    //r = fareqr_decode(fareqr_s, pub_key, priv_key, dst_cred);
    r = decode_b64((unsigned char *)dst_cred, (unsigned char *)priv_key, (unsigned char *)enc_str);
    if (r<0) { _ret = r; }
  }

  if (pub_key) { free(pub_key); }
  if (priv_key) { free(priv_key); }
  if (enc_str) { free(enc_str); }
  if (dec_str) { free(dec_str); }
  if (plain_str) { free(plain_str); }

  return _ret;
}

int fareqr_encode(char *tok_public, char *tok_secret, char *tok_cred, char *fareqr_str) {
  int r=0, _ret=0;
  uint8_t *dst_data = NULL;

  if ((!tok_public) || (!tok_secret) || (!tok_cred) || (!fareqr_str)) { return -9; }

  dst_data = (uint8_t *)malloc(sizeof(uint8_t)*LINE_BUFFER_SIZE);

  r = encode_b64(dst_data, (unsigned char *)tok_secret, (unsigned char *)tok_cred);
  if (r < 0) {
    _ret = -1;
    goto _fareqr_encode_cleanup;
  }

  if (fareqr_str) {
    snprintf(fareqr_str, LINE_BUFFER_SIZE-1, "@%s%%%s$", tok_public, dst_data);
  }

_fareqr_encode_cleanup:

  if (dst_data) { free(dst_data); }
  return _ret;
}


// depreicated (doing away with one-time pads in favor of AES above)
//
int fareqr_decode(char *fareqr_s, char *check_pub, char *tok_secret, char *dst_cred) {
  int n=0, r=0, _ret=0;
  char *p, *stop_tok;
  int pub_tok_read_len = 0, check_pub_len = 0;

  char *enc_str = NULL;

  if ((!fareqr_s) || (!dst_cred)) { return -9; }

  if (fareqr_s[0] != '@') { return -1; }
  stop_tok = strchr(fareqr_s, '%');
  if (!stop_tok) { return -2; }
  if (!tok_secret) { return -3; }

  if (check_pub) {
    check_pub_len = strlen(check_pub);
    for ( p = (fareqr_s+1); p < stop_tok; p++) {
      if (pub_tok_read_len >= check_pub_len) { return -3; }
      if ( (*p) != check_pub[pub_tok_read_len] ) { return -4; }
      pub_tok_read_len++;
    }
  }

  enc_str = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
  for (n=0, p = (stop_tok+1); (*p) && ((*p) != '$'); p++, n++) {
    enc_str[n] = *p;
  }
  enc_str[n] = '\0';

  r = decode_b64((unsigned char *)dst_cred, (unsigned char *)tok_secret, (unsigned char *)enc_str);
  if (r < 0) {
    _ret = -1;
    goto _fareqr_decode_cleanup;
  }

_fareqr_decode_cleanup:

  if (enc_str) { free(enc_str); }
  return _ret;
}

// return negative on error or not found
// 0 on success (found)
//
int fareqr_lookup_seed_secret(char *seedfn, char *pub, char *priv) {
  int i;
  FILE *fp;
  char buf[LINE_BUFFER_SIZE] = {0};
  int pos = 0, ch=0, line_no=0;
  char *tok0_ptr=NULL,
       *tok1_ptr=NULL;

  fp = fopen(seedfn, "r");
  if (!fp) { perror(seedfn); return -2;}
  while (!feof(fp)) {
    ch = fgetc(fp);

    // process line if we reach a newline or eof
    //
    if (feof(fp) || (ch == '\n')) {

      // if the line is empty or a comment, skip
      //
      if ((pos==0) || (buf[0] == '#')) {

      }

      // get tokens out of line, using ' ' as the
      // delimeter
      //
      else {
        tok0_ptr = buf;
        tok1_ptr = strchr(buf, ' ');
        if (tok1_ptr) {
          *tok1_ptr = '\0';
          tok1_ptr++;


          // Cechk it against our supplied 'pulbic' key
          //
          if (strcmp(pub, tok0_ptr)==0) {

            // If we've found it, copy it over to the `priv`
            // above and return
            //
            for (i=0; tok1_ptr[i]; i++) {
              priv[i] = tok1_ptr[i];
            }
            priv[i]='\0';
            return 0;

          }
        }

        line_no++;
      }

      pos=0;
      buf[0]='\0';
      continue;
    }

    buf[pos] = ch;
    pos++;
    buf[pos] = '\0';
  }
  fclose(fp);

  return -1;
}

// to compile:
// gcc -D__FAREQR_MAIN__ fareqr.c b64.c -o fareqr
//

#ifdef __FAREQR_MAIN__

#define FAREQR_VERSION "0.1.0"

void show_version(FILE *ofp) {
  fprintf(ofp, "fareqr version %s\n", FAREQR_VERSION);
}

void show_help(FILE *ofp) {
  show_version(ofp);
  fprintf(ofp, "\nusage:\n\n");
  fprintf(ofp, "  fareqr encode <pubkey> <privatekey> <str>\n");
  fprintf(ofp, "  fareqr decode <privatekey> <encstr>\n");
  fprintf(ofp, "  fareqr dbdecode <qrseedfile> <fareqr>\n");
  fprintf(ofp, "  fareqr version\n");
  fprintf(ofp, "  fareqr help\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "fareqr is a program to help with encoding and decoding 'fareqr' strings.\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "A fareqr string is of the form:\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "  @<pubkey>%%<aes(privkey,credential)>$\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "Where `<aes(privkey,credential)>` is the base64 encoded AES encrypted string\n");
  fprintf(ofp, "of the credential, encrypted with the privkey.\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "The <qrseedfile> is a text file of <pubkey> <privkey> pairs.\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "Here is some example usage:\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "  $ fareqr encode 'wu9XouSh' 'ohNgizahkephain3aosoh2AeH1aethoo4cie6oiSaezimaighai2eiVaefahfien' ';123456789060535?'\n");
  fprintf(ofp, "  @wu9XouSh%%U2FsdGVkX19nxmlzUf9K7GAgplhbmU4tcFYYa/Xz6oq0XrxSYTbBOA2yffAi7A0z$\n");
  fprintf(ofp, "  $ fareqr decode 'ohNgizahkephain3aosoh2AeH1aethoo4cie6oiSaezimaighai2eiVaefahfien' '@wu9XouSh%%U2FsdGVkX19nxmlzUf9K7GAgplhbmU4tcFYYa/Xz6oq0XrxSYTbBOA2yffAi7A0z$'\n");
  fprintf(ofp, "  ;123456789060535?\n");
  fprintf(ofp, "  $ echo 'wu9XouSh ohNgizahkephain3aosoh2AeH1aethoo4cie6oiSaezimaighai2eiVaefahfien' > ./qr.seed\n");
  fprintf(ofp, "  $ fareqr dbdecode ./qr.seed '@wu9XouSh%%U2FsdGVkX19nxmlzUf9K7GAgplhbmU4tcFYYa/Xz6oq0XrxSYTbBOA2yffAi7A0z$'\n");
  fprintf(ofp, "  ;123456789060535?\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "Where 'wu9XouSh' is the public key, 'ohNgizahkephain3aosoh2AeH1aethoo4cie6oiSaezimaighai2eiVaefahfien' is the private key and\n");
  fprintf(ofp, "';123456789060535?' is the credential information to be encoded.\n");
  fprintf(ofp, "\n");
  fprintf(ofp, "\n");
  fflush(ofp);
}

/* 
 *
 * quick test/start:
 *
 * $ fareqr encode 'wu9XouSh' 'ohNgizahkephain3aosoh2AeH1aethoo4cie6oiSaezimaighai2eiVaefahfien' ';123456789060535?'
 * @wu9XouSh%VFl8VF1PV19TXEBeUVxdBl4=
 *
 * $ fareqr decode 'ohNgizahkephain3aosoh2AeH1aethoo4cie6oiSaezimaighai2eiVaefahfien' '@wu9XouSh%VFl8VF1PV19TXEBeUVxdBl4=$' 
 * ;123456789060535?
 *
 * $ fareqr dbdecode '@wu9XouSh%VFl8VF1PV19TXEBeUVxdBl4=$'
 * ;123456789060535?
 *
 */

int main(int argc, char **argv) {
  int i, r;
  char *tok_public = NULL,
       *tok_secret = NULL,
       *tok_cred = NULL,
       *fareqr_str = NULL;

  uint8_t x,y,z;

  uint8_t *src_data = NULL, *dst_data = NULL;
  int src_data_n = 0, dst_data_n = 0;

  char *fn = NULL;

  if (argc <= 1) {
    show_help(stderr);
    exit(1);
  }

  if (strcmp(argv[1], "encode")==0) {
    if (argc>2) {
      tok_public = strdup(argv[2]);
      if (argc>3) {
        tok_secret = strdup(argv[3]);
        if (argc>4) {
          tok_cred = strdup(argv[4]);
        }
      }
    }

    if ((!tok_public) || (!tok_secret) || (!tok_cred)) {
      show_help(stderr);
      if (tok_public) { free(tok_public); }
      if (tok_secret) { free(tok_secret); }
      if (tok_cred) { free(tok_cred); }
      exit(2);
    }

    fareqr_str = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
    r = fareqr_encode(tok_public, tok_secret, tok_cred, fareqr_str);
    if (r==0) {
      printf("%s\n", fareqr_str);
    }
  }

  else if (strcmp(argv[1], "decode")==0) {
    if (argc>2) {
      tok_secret = strdup(argv[2]);
      if (argc>3) {
        fareqr_str = strdup(argv[3]);
      }
    }
    tok_cred = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
    tok_cred[0] = '0';
    r = fareqr_decode(fareqr_str, tok_public, tok_secret, tok_cred);
    if (r<0) {
      fprintf(stderr, "error, failed to decode qr fare string (%i)\n", r);
    }
    else {
      printf("%s\n", tok_cred);
    }
  }

  else if (strcmp(argv[1], "dbdecode")==0) {
    if (argc>2) {
      fn = strdup(argv[2]);
      if (argc>3) {
        fareqr_str = strdup(argv[3]);
      }
    }

    tok_cred = (char *)malloc(sizeof(char)*LINE_BUFFER_SIZE);
    tok_cred[0] = '0';
    r = fareqr_lookup_decode(fn, fareqr_str, tok_cred);
    if (r<0) {
      fprintf(stderr, "could not decode '%s' with db file '%s', exiting (got %i)\n",
          fareqr_str, fn, r);
    }
    else {
      printf("%s\n", tok_cred);
    }

  }

  // test out our AES encryption
  //
  else if (strcmp(argv[1], "aes-decode")==0) {
    if (argc>2) {
      tok_secret = strdup(argv[2]);
      if (argc>3) {
        fareqr_str = strdup(argv[3]);
      }
    }

    tok_cred = (unsigned char *)calloc(LINE_BUFFER_SIZE, sizeof(char));
    r = decode_b64(tok_cred, tok_secret, fareqr_str);
    if (r<0) {
      fprintf(stderr, "could not decode aes base64 string (AES CBC 256, sha256 md) (got %i)\n", r);
    }
    else {
      printf("%s\n", tok_cred);
    }

  }

  // test out our AES encryption
  //
  else if (strcmp(argv[1], "aes-encode")==0) {
    if (argc>2) {
      tok_secret = strdup(argv[2]);
      if (argc>3) {
        tok_cred = strdup(argv[3]);
      }
    }

    fareqr_str = (unsigned char *)calloc(2*LINE_BUFFER_SIZE, sizeof(char));
    r = encode_b64(fareqr_str, tok_secret, tok_cred);
    if (r<0) {
      fprintf(stderr, "could not encode aes base64 string (AES CBC 256, sha256 md) (got %i)\n", r);
    }
    else {
      printf("%s\n", fareqr_str);
    }

  }

  else if (strcmp(argv[1], "help")==0) {
    show_help(stdout);
    exit(0);
  }

  else if (strcmp(argv[1], "version")==0) {
    show_version(stdout);
    exit(0);
  }

  else {
    fprintf(stderr, "unknown operation '%s'\n", argv[1]);
    show_help(stderr);
    exit(3);
  }

  if (tok_public) { free(tok_public); }
  if (tok_secret) { free(tok_secret); }
  if (tok_cred) { free(tok_cred); }
  if (fareqr_str) { free(fareqr_str); }

  if (src_data) { free(src_data); }
  if (dst_data) { free(dst_data); }
  if (fn) { free(fn); }
}

#endif