// Copyright (c) 2013-2015 IONU Security, Inc. All rights reserved.
//
// Keyring and Key classes for secure storage and management of keys

#ifndef eyering_h
#define eyering_h

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <openssl/evp.h>

#include "eyeconstants.h"
#include "eyejson.h"
#include "eyeencrypteddb.h"

#ifndef LIBEYE_DLL
#ifdef WIN32
#  ifdef DLLEXPORT
#    define LIBEYE_DLL __declspec(dllexport)
#  else
#    define LIBEYE_DLL __declspec(dllimport)
#  endif
#else
#define LIBEYE_DLL
#endif
#endif

namespace sequencelogic {

#define SL_KEYRING_DIGEST_LEN 20  // sha1
#define SL_ECDH_HMAC_LEN 32       // sha256
#define SL_ECDH_SECRET_LEN 64     // sha512


class Key {
 friend class Keyring;
 public:
    enum KEY_TYPE {
        RSA,
        RSA_PUBLIC,
        AES,
        EC,
        EC_PUBLIC,
        STRING,
        AESHMAC,
        GENERIC
    };
    enum KEY_FORMAT {
        PEM,
        DER,
        PKCS8,
        PKCS12,
        UNKNOWN_FORMAT
    };

    /**
     * @constructor
     * @param name of the key
     * @param description of the key
     * @param length of the key
     * @param key bytes
     * @param type of the key
     */
    LIBEYE_DLL Key (const std::string& name, const std::string& desc, size_t len, const unsigned char* key, KEY_TYPE type);

    /**
     * @constructor
     * @param name of the key
     * @param description of the key
     * @param key bytes
     * @param type of the key
     */
    LIBEYE_DLL Key (const std::string& name, const std::string& desc, const std::string& key, KEY_TYPE type);

    /**
     * @constructor
     * @param name of the key
     * @param description of the key
     * @param filename of key (PEM or DER format)
     * @param password or "" if not protected by PKCS#8
     */
    LIBEYE_DLL Key (const std::string& name, const std::string& desc, const std::string& filename, const std::string& password);

    /**
     * @constructor
     * @param filename of key (PEM or DER format)
     * @param password or "" if not protected by PKCS#8
     */
    LIBEYE_DLL Key (const std::string& filename, const std::string& password);

    /**
     * @constructor
     */
    LIBEYE_DLL Key();

    /**
     * @destructor
     */
    LIBEYE_DLL ~Key();

    /**
     * Check if the key is valid
     * @return true or false
     */
    LIBEYE_DLL bool IsValid() const;

    /**
     * AES-CBC encrypt with 128 or 256 bit key
     * @param input plaintext
     * @param input length
     * @param output length
     * @param initial value (16 bytes)
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* SymmetricDecryptBuffer (const unsigned char* in, size_t inlen, size_t* outlen,
                                                      const unsigned char* iv) const;
    /**
     * AES-CBC decrypt with 128 or 256 bit key
     * @param input ciphertext
     * @param input length
     * @param output length
     * @param initial value (16 bytes)
     * @return plaintext or nullptr
     */
    LIBEYE_DLL unsigned char* SymmetricEncryptBuffer (const unsigned char* in, size_t inlen, size_t* outlen,
                                                      const unsigned char* iv) const;
    /**
     * AES-CBC encrypt file with 128 or 256 bit key
     * @param plaintext filename
     * @param encrypted filename
     * @param initial value (16 bytes)
     * @return true for success
     */
    LIBEYE_DLL bool SymmetricEncryptFile (const std::string& plainfile, const std::string& encryptedfile,
                                                      const unsigned char* iv) const;
    /**
     * AES-CBC decrypt file with 128 or 256 bit key
     * @param encrypted filename
     * @param plaintext filename
     * @param initial value (16 bytes)
     * @return true for success
     */
    LIBEYE_DLL bool SymmetricDecryptFile (const std::string& encryptedfile, const std::string& plainfile,
                                                      const unsigned char* iv) const;

    /**
     * AES-GCM encrypt with 256 bit key, authenticated with optional additional authenticated data
     * @param input plaintext
     * @param input length
     * @param output length
     * @param initial value (16 bytes)
     * @param additional authenticated data
     * @param length of additional authenticated data
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* SymmetricEncryptAuthenticated (const unsigned char *plaintext, size_t inlen, size_t* outlen,
                                                   const unsigned char *iv,
                                                   const unsigned char *aad = nullptr, size_t aadlen = 0) const;
    /**
     * AES-GCM decrypt with 256 bit key, authenticated with optional additional authenticated data
     * @param input ciphertext
     * @param input length
     * @param output length
     * @param initial value (16 bytes)
     * @param additional authenticated data
     * @param length of additional authenticated data
     * @return plaintext or nullptr
     */
    LIBEYE_DLL unsigned char* SymmetricDecryptAuthenticated (const unsigned char *ciphertext, size_t inlen, size_t* outlen,
                                                   const unsigned char *iv,
                                                   const unsigned char *aad = nullptr, size_t aadlen = 0) const;
    /**
     * AES-CBC encrypt with 256 bit key and keyed HMAC for authentication
     * @param input plaintext
     * @param input length
     * @param output length
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* HMACEncrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * AES-CBC decrypt with 256 bit key and keyed HMAC for authentication
     * @param input ciphertext
     * @param input length
     * @param output length
     * @return plaintext
     */
    LIBEYE_DLL unsigned char* HMACDecrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * RSA encryption, plaintext length is 214 byte maximum with EME-OAEP padding
     * @param input plaintext
     * @param input length
     * @param output length
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* PublicKeyEncrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * RSA decryption
     * @param input ciphertext
     * @param input length
     * @param output length
     * @return plaintext or nullptr
     */
    LIBEYE_DLL unsigned char* PrivateKeyDecrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * RSA sign
     * @param input to be signed
     * @param input length
     * @param signature length
     * @return signature or nullptr
     */
    LIBEYE_DLL unsigned char* PrivateKeySign (const unsigned char* tbs, size_t tbslen, size_t* siglen) const;

    /**
     * RSA verify
     * @parameter signature to verify
     * @param signature length
     * @param input to be signed
     * @param input length
     * @return 0 for success, non-zero for failure
     */
    LIBEYE_DLL int PublicKeyVerify (const unsigned char* sig, size_t siglen, const unsigned char* tbs, size_t tbslen) const;

    // These are unsecure and misnamed, provided for compatibility with the libionu version for CG token exchange
    LIBEYE_DLL unsigned char* RSAPrivateKeyEncrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;
    LIBEYE_DLL unsigned char* RSAPublicKeyDecrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * ECIES based encryption, uses ephemeral EC key
     * @param input plaintext
     * @param input length
     * @param output length
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* ECIESEncrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * ECIES based decryption, uses ephemeral EC key
     * @param input ciphertext
     * @param input length
     * @param output length
     * @return plaintext or nullptr
     */
    LIBEYE_DLL unsigned char* ECIESDecrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * RSA IES based encryption, uses ephemeral AES key to extend size
     * @param input plaintext
     * @param input length
     * @param output length
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* RSAIESEncrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    /**
     * RSA IES based decryption, uses ephemeral AES key to extend size
     * @param input ciphertext
     * @param input length
     * @param output length
     * @return plaintext or nullptr
     */
    LIBEYE_DLL unsigned char* RSAIESDecrypt (const unsigned char* in, size_t inlen, size_t* outlen) const;

    // Generate ECDH shared secret with peer's public key
    LIBEYE_DLL unsigned char* ECDHEncrypt (const Key* peer, const unsigned char* in, size_t inlen, size_t* outlen) const;
    LIBEYE_DLL unsigned char* ECDHDecrypt (const Key* peer, const unsigned char* in, size_t inlen, size_t* outlen) const;
    LIBEYE_DLL unsigned char* GenerateECDHSecret (const Key* peer, size_t* secret_len) const;

    /**
     * Get PEM format public key for RSA or EC private key
     * @return PEM formatted public key or nullptr
     */
    LIBEYE_DLL char* GetPubKey() const;

    /**
     * Get the public key fingerprint (e.g. '38:52:77:78:11:24:e6:8b:33:ea:37:1c:9a:10:f3:1d:5d:dc:d4:5d')
     * @return fingerprint
     */
    LIBEYE_DLL std::string PublicKeyFingerPrint () const;

    /**
     * Export to PEM or DER format as determined by extension
     * @param filename of key (PEM or DER format)
     * @param password or "" if not protected by PKCS#8
     * @return true for success
     */
    LIBEYE_DLL bool Export (const std::string& filename, const std::string& password) const;

    /**
     * Import key from file
     * @param filename of key (PEM or DER format)
     * @param password or "" if not protected by PKCS#8
     * @return true for success
     */
    LIBEYE_DLL bool Import (const std::string& filename, const std::string& password);

    /**
     * Import key from buffer
     * @param buffer containing key data
     * @param length of buffer
     * @param password or "" if not protected by PKCS#8
     * @param format of key (PEM or DER)
     * @return true for success
     */
    LIBEYE_DLL bool Import (const char* buffer, size_t bytes, const std::string& password, KEY_FORMAT format);

    // Duplicate a key, rekey
    void Duplicate (const Key* dupkey);
    void SetType (KEY_TYPE type) {_type = type; }
    void SetKey (const unsigned char* key, size_t bytes);
    void SetIv (const unsigned char* iv, size_t bytes);

    const std::string GetName() const         { return _name; }
    const std::string GetDescription() const  { return _desc; }
    size_t GetLength() const                  { return _len; }
    LIBEYE_DLL const unsigned char* GetKey() const       { return _key; }
    const unsigned char* GetIv() const        { return _iv; }
    KEY_TYPE GetType() const                  { return _type; }
    LIBEYE_DLL const std::string GetStringKey() const;

    LIBEYE_DLL bool Equals (const Key& other) const;
    LIBEYE_DLL void Dump() const;

 private:
    KEY_TYPE _type;            // Key type - RSA, AES, TGI, STRING, GENERIC
    std::string _name;         // Name of the key
    std::string _desc;         // Descriptive text [optional]
    size_t _len;               // Length of key in bytes
    unsigned char* _key;       // Actual key data
    unsigned char* _iv;        // Initial value or nonce, may or may not be present and is not saved
    unsigned char _obf;        // Random obfuscation byte

    unsigned char* HMACCrypt (bool encrypt, const unsigned char* in, size_t inlen, size_t* outlen) const;
    unsigned char* PublicKeyCrypt (bool encrypt, const unsigned char* in, size_t inlen, size_t* outlen) const;
    unsigned char* RSAKeyCrypt (bool encrypt, const unsigned char* in, size_t inlen, size_t* outlen) const;
    unsigned char* ECDHKeyCrypt (const Key* peer, bool encrypt, const unsigned char* in, size_t inlen, size_t* outlen) const;

    // From the file extension determine the format to read/write
    KEY_FORMAT GetKeyFormat (const std::string& filename) const;
    std::string GetKeyTypeStr() const;

    // Encrypt/Decrypt the key
    void Encrypt (EVP_CIPHER_CTX* ctx, const unsigned char* key, const unsigned char* iv);
    void Decrypt (EVP_CIPHER_CTX* ctx, const unsigned char* key, const unsigned char* iv);

    // Get the value parameters as a json string, load from JSON
    std::string GetValueJSON() const;
    bool LoadValueJSON (const std::string& name, const std::string& json);
    EyeJSONObject* GetJSONObject() const;
    std::string ToJSON () const;
};


// Keyring holds any number of keys and linked keyrings
// It is always encrypted on disk, and until unlocked also in memory
class Keyring {
 public:
   enum RING_STATUS {
      VALID,
      INVALID,       // Not a valid keyring, can't open file
      CORRUPT        // keyring is corrupt
   };

    /**
     * Keyring constructor
     * @param name of the keyring
     * @param filename of the SQLite database file
     * @param key to lock/unlock the keyring
     */
    LIBEYE_DLL Keyring (const std::string& name, const std::string& database, const Key& key);

    /**
     * Keyring constructor
     * @param name of the keyring
     * @param filename of the SQLite database file
     * @param password to lock/unlock the keyring
     */
    LIBEYE_DLL Keyring (const std::string& name, const std::string& database, const std::string& password);

    /**
     * Keyring constructor
     * @param name of the keyring
     * @param description of the keyring
     */
    LIBEYE_DLL Keyring (const char* name, const char* desc);

    /**
     * Keyring constructor
     * @param filename for the keyring
     */
    LIBEYE_DLL Keyring (const char* filename);

    /**
     * Keyring destructor
     */
    LIBEYE_DLL ~Keyring();
    
    LIBEYE_DLL bool IsLocked () { return _isLocked; }
    LIBEYE_DLL RING_STATUS GetStatus () { return _status; }

    /**
     * Lock the keyring with a password
     * @param password
     * @return true for success
     */
    LIBEYE_DLL bool Lock (const char* password);

    /**
     * Unlock the keyring with a password
     * @param password
     * @return true for success
     */
    LIBEYE_DLL bool Unlock (const char* password);

    /**
     * Lock the keyring with a 256 bit AES key
     * @param key to lock keyring
     * @return true for success
     */
    LIBEYE_DLL bool Lock (const unsigned char* key);

    /**
     * Unlock the keyring with a 256 bit AES key
     * @param key to unlock keyring
     * @return true for success
     */
    LIBEYE_DLL bool Unlock (const unsigned char* key);

    /**
     * Save the keyring to disk
     * @return true for success
     */
    LIBEYE_DLL bool Save ();

    /**
     * Save the keyring to disk
     * @param filename
     * @return true for success
     */
    LIBEYE_DLL bool SaveAs (const char* filename);

    /**
     * Add key to keyring, replaces existing key of same name
     * @param key to add
     * @return true for success
     */
    LIBEYE_DLL bool AddKey (Key* key);

    /**
     * Add key to keyring, replaces existing key of same name
     * @param name of the key
     * @param description of the key
     * @param length of the key
     * @param key bytes
     * @param type of the key
     */
    LIBEYE_DLL bool AddKey (const std::string& name, const std::string& desc, size_t len, const unsigned char* key, Key::KEY_TYPE type);

    /**
     * Get key from keyring
     * @param name of the key
     * @return key or nullptr if not found
     */
    LIBEYE_DLL const Key* GetKey (const std::string& keyname) const;

    /**
     * Get key from keyring
     * @param name of the key
     * @param key to fill in
     * @return true if found
     */
    LIBEYE_DLL bool GetKey (const std::string& name, Key& key) const;

    /**
     * Get public key from keyring
     * @param name of the key
     * @return true if found and RSA or EC key
     */
    LIBEYE_DLL char* GetPubKey (const char* keyname);

    /**
     * Remove key from keyring
     * @param name of the key
     * @return true if removed
     */
    LIBEYE_DLL bool RemoveKey (const std::string& keyname);

    /**
     * Update key in keyring
     * @param name of the key
     * @return true if updated
     */
    LIBEYE_DLL bool UpdateKey (const Key* key);

    /**
     * Generate a random 2048 bit RSA keypair and add private key to keyring
     * @param name of the key
     * @param description of the key
     * @return true for success
     */
    LIBEYE_DLL bool GenerateRSAKey (const char* name, const char* desc);

    /**
     * Generate a random 256 bit AES key and add to keyring
     * @param name of the key
     * @param description of the key
     * @return true for success
     */
    LIBEYE_DLL bool GenerateAESKey (const char* name, const char* desc);

    /**
     * Generate a random 256 bit AES plus 256 bit HMAC key and add to keyring
     * @param name of the key
     * @param description of the key
     * @return true for success
     */
    LIBEYE_DLL bool GenerateAESHMACKey (const char* name, const char* desc);

    /**
     * Generate random 521 bit EC key using ANSI X9.62 Prime 521 curve and add to the keyring
     * @param name of the key
     * @param description of the key
     * @return true for success
     */
    LIBEYE_DLL bool GenerateECKey (const char* name, const char* desc);

    /**
     * Import the public key from PEM encoded X509 certificate
     * @param filename for certificate
     * @return true for success
     */
    LIBEYE_DLL bool ImportPubKeyFromCert (const std::string& filename);

    /**
     * Generate 2048 RSA keypair using a secret string to seed a DBRG, such that
     * given same secret will produce the same key, and add the private key to keyring
     * @param name of the key
     * @param description of the key
     * @param seed for DRBG
     * @return true for success
     */
    LIBEYE_DLL bool GenerateRSAKey (const std::string& keyname, const std::string& desc, const std::string& seed);
    
    /**
     * Provision a challenge based account key pair
     * @param keyname    – name of the key to be added to keyring
     * @param challenges – challenge answer set
     * @param salt       – empty for initial provision, cgsalt for subsequent
     * @return salt      - base64 encoded salt
     */
    LIBEYE_DLL bool ProvisionChallengeRSAKey (const std::string& keyname, const std::string& challenges, std::string& salt);

    /**
     * Provision an account code based key pair
     * @param keyname    – name of the key to be added to keyring
     * @param code       – account/recovery code, may be empty for initial provision
     * @param salt       – empty for initial provision, cgsalt for subsequent
     * @param challenges – new challenge answer set
     * @return code,salt - code,base64 encoded salt
     */
    LIBEYE_DLL bool ProvisionAccountRSAKey (const std::string& keyname, std::string& code, std::string& salt, const std::string& challenges);

    /**
     * Generate deterministic 2048 bit RSA keypair, add private key to keyring
     *  - old style challenge only, call with challenges: secret and salt == ""
     *  - new style pre-provision only, call with secret: challenges and salt == ""
     *  - provision and get salt, call with challenges and secret
     *  - new style salted seed, call with challenges and salt: secret == ""
     *  - recovery/provision, call with secret and salt
     * @param name of the key
     * @param description of the key
     * @param challenge answer string
     * @param secret provision/recovery code, input and output
     * @param salt as hex, base64 or byte string
     * @return true for success
     */
    LIBEYE_DLL bool GenerateRSAKey (const std::string& name, const std::string& desc, const std::string& challenges, const std::string& secret, std::string& salt);

    /**
     * Generate 521 bit EC keypair using a json challenge string to seed a DBRG, such that
     * given same challenges will produce the same key, and add the private key to keyring
     * @param name of the key
     * @param description of the key
     * @param challenge answer string to seed DRBG
     * @return true for success
     */
    LIBEYE_DLL bool GenerateECKey (const char* name, const char* desc, const char* challenges);

    /**
     * Generate deterministic 521 bit EC keypair, add private key to keyring
     *  - old style challenge only, call with challenges: secret and salt == ""
     *  - new style pre-provision only, call with secret: challenges and salt == ""
     *  - provision and get salt, call with challenges and secret
     *  - new style salted seed, call with challenges and salt: secret == ""
     * @param name of the key
     * @param description of the key
     * @param challenge answer string
     * @param secret string
     * @param salt as hex, base64 or byte string
     * @return true for success
     */
    LIBEYE_DLL bool GenerateECKey (const std::string& name, const std::string& desc, const std::string& challenges, const std::string& secret, std::string& salt);

    /**
     * Pre-provision a set of keys from a van and list of email addresses, keyname=email, desc=secret
     * @param domain string, van, company name...
     * @param list of users to create keys for, could be email, username...
     * @return true for success
     */
    LIBEYE_DLL bool GenerateRSAKeys (const std::string& van, const std::vector<std::string>& emailist);

    /**
     * Encrypt plaintext using specifed key (can be AES, AESHMAC, EC, RSA)
     * @param key to use for encryption
     * @param input plaintext
     * @param input length
     * @param output length
     * @param initialization vector
     * @return ciphertext or nullptr
     */
    LIBEYE_DLL unsigned char* Encrypt (const std::string keyname, const unsigned char *plaintext,
                                       size_t inlen, size_t* outlen, const unsigned char *iv = nullptr) const;
    /**
     * Decrypt ciphertext using specifed key (can be AES, AESHMAC, EC, RSA)
     * @param key to use for encryption
     * @param input ciphertext
     * @param input length
     * @param output length
     * @param initialization vector
     * @return plaintext or nullptr
     */
    LIBEYE_DLL unsigned char* Decrypt (const std::string keyname, const unsigned char *ciphertext,
                                       size_t inlen, size_t* outlen, const unsigned char *iv = nullptr) const;
    /**
     * Add keyring and transfer ownership
     * @param keyring to add
     * @return true for success
     */
    LIBEYE_DLL bool AddKeyring (Keyring* keyring);

    /**
     * Remove and delete keyring
     * @param keyring to remove
     * @return true for success
     */
    LIBEYE_DLL bool RemoveKeyring (const char* ringname);

    /**
     * Get the named keyring
     * @param keyring to get
     * @return keyring or nullptr
     */
    LIBEYE_DLL const Keyring* GetKeyring (const char* ringname) const;

    /**
     * Load a keyring from JSON, deletes any exisiting keys
     * @param json string to load keyring with
     * @return true for success
     */
    LIBEYE_DLL bool LoadFromJSON (const std::string& json);

    /**
     * Get key names
     * @return JSON array of key names
     */
    LIBEYE_DLL std::string GetKeyNames();
       
    /**
     * Convert the format to JSON string
     * @param key to lock keyring
     * @return JSON string representation of Keyring
     */
    LIBEYE_DLL std::string Stringify (const Key& key);

    LIBEYE_DLL void Dump();

    const char* GetName()         { return _name.c_str(); }
    const char* GetDescription()  { return _desc.c_str(); }
    size_t GetNumKeys()           { return _keys.size(); }
    size_t GetNumRings()          { return _keyrings.size(); }

 private:
    RING_STATUS _status;               // Status of keyring
    char* _filename;                   // Current filename
    EyeEncryptedDB* _db;               // Database (if used)
    std::string _version;              // Version of this key ring
    bool _isLocked;                    // Has the Keyring been unlocked
    bool _sync;                        // Should this keyring be sync'd
    unsigned char _attempts;           // Number of attempts to unlock
    std::string _name;                 // Name of the keyring
    std::string _desc;                 // Descriptive text [optional]
    unsigned char _iv[SL_AES_BLOCK_LEN]; // Initialization vector
    unsigned char _digest[SL_KEYRING_DIGEST_LEN]; // Message digest of key data
    std::vector<Key*> _keys;           // Array of keys
    std::vector<Keyring*> _keyrings;   // Array of linked keyrings
    EVP_CIPHER_CTX _ctx;               // Encryption context
    
    static int GetRounds (const char* password);
    bool DigestKeys (unsigned char* md_value);
    EyeJSONObject* GetJSONObject() const;
    EyeJSONObject* GetJSONObject (const Key& key);
    bool AddJSONKeyring (EyeJSONObject* jring);
    bool AddJSONKey (EyeJSONObject* jkey);
};


} //namespce ionu
#endif