// Package goenc contains functions for working with encryption
package goenc

// work is derived from many sources:
//
// http://stackoverflow.com/questions/21151714/go-generate-an-ssh-public-key
// https://golang.org/pkg/crypto/cipher/
// https://leanpub.com/gocrypto/read#leanpub-auto-aes-cbc
// https://github.com/hashicorp/memberlist/blob/master/security.go

import (
    "crypto/rand"
    "encoding/binary"
    "io"

    "io/ioutil"

    "os"

    "github.com/alistanis/goenc/aes/cbc"
    "github.com/alistanis/goenc/aes/cfb"
    "github.com/alistanis/goenc/aes/ctr"
    "github.com/alistanis/goenc/aes/gcm"
    "github.com/alistanis/goenc/encerrors"
    "github.com/alistanis/goenc/generate"
    "github.com/alistanis/goenc/nacl"
    "golang.org/x/crypto/nacl/box"
    "golang.org/x/crypto/nacl/secretbox"
    "golang.org/x/crypto/scrypt"
)

/*
    TODO(cmc): Verify this isn't horrifically insecure and have this reviewed by a(n) expert(s) before publishing
*/

// BlockCipher represents a cipher that encodes and decodes chunks of data at a time
type BlockCipher interface {
    Encrypt(key, plaintext []byte) ([]byte, error)
    Decrypt(key, ciphertext []byte) ([]byte, error)
    KeySize() int
}

//---------------------------------------------------------------------------
// BlockCipherInterface Functions - these should not be used with large files
//--------------------------------------------------------------------------------

// EncryptAndSaveWithPerms encrypts data and saves it to a file with the given permissions using the given key
func EncryptAndSaveWithPerms(cipher BlockCipher, key, plaintext []byte, path string, perm os.FileMode) error {
    data, err := cipher.Encrypt(key, plaintext)
    if err != nil {
        return err
    }
    return ioutil.WriteFile(path, data, perm)
}

// EncryptAndSave encrypts data and saves it to a file with the permissions 0644
func EncryptAndSave(cipher BlockCipher, key, plaintext []byte, path string) error {
    return EncryptAndSaveWithPerms(cipher, key, plaintext, path, 0644)
}

// ReadEncryptedFile reads a file a path and attempts to decrypt the data there with the given key
func ReadEncryptedFile(cipher BlockCipher, key []byte, path string) ([]byte, error) {
    ciphertext, err := ioutil.ReadFile(path)
    if err != nil {
        return nil, err
    }
    plaintext, err := cipher.Decrypt(key, ciphertext)
    return plaintext, err
}

// CipherKind represents what kind of cipher to use
type CipherKind int

// CipherKind constants
const (
    CBC CipherKind = iota
    CFB
    CTR
    GCM
    NaCL

    Mock
)

const (
    // SaltSize sets a generic salt size
    SaltSize = 64
)

// Cipher is a struct that contains a BlockCipher interface and stores a DerivedKey Complexity number
type Cipher struct {
    BlockCipher
    DerivedKeyN int
}

// NewCipher returns a new Cipher containing a BlockCipher interface based on the CipherKind
func NewCipher(kind CipherKind, derivedKeyN int, args ...[]byte) (*Cipher, error) {
    c := &Cipher{DerivedKeyN: derivedKeyN}
    switch kind {
    case GCM:
        c.BlockCipher = gcm.New()
    case NaCL:
        // special case, we need to define a pad for nacl
        if len(args) == 0 {
            return nil, encerrors.ErrNoPadProvided
        }
        n := &nacl.Cipher{}
        n.Pad = args[0]
        c.BlockCipher = n
    case CFB:
        c.BlockCipher = cfb.New()
    case CBC:
        c.BlockCipher = cbc.New()
    case CTR:
        c.BlockCipher = ctr.New()
    case Mock:
        c.BlockCipher = &MockBlockCipher{}
    default:
        return nil, encerrors.ErrInvalidCipherKind
    }
    return c, nil
}

// Encrypt takes a password, plaintext, and derives a key based on that password,
// then encrypting that data with the underlying block cipher
func (c *Cipher) Encrypt(password, plaintext []byte) ([]byte, error) {
    salt, err := generate.RandBytes(SaltSize)
    if err != nil {
        return nil, err
    }

    key, err := DeriveKey(password, salt, c.DerivedKeyN, c.BlockCipher.KeySize())
    if err != nil {
        return nil, err
    }

    out, err := c.BlockCipher.Encrypt(key, plaintext)
    Zero(key)
    if err != nil {
        return nil, err
    }

    out = append(salt, out...)
    return out, nil
}

// Overhead is the amount of Overhead contained in the ciphertext
const Overhead = SaltSize + secretbox.Overhead + generate.NonceSize

// Decrypt takes a password and ciphertext, derives a key, and attempts to decrypt that data
func (c *Cipher) Decrypt(password, ciphertext []byte) ([]byte, error) {
    if len(ciphertext) < Overhead {
        return nil, encerrors.ErrInvalidMessageLength
    }

    key, err := DeriveKey(password, ciphertext[:SaltSize], c.DerivedKeyN, c.KeySize())
    if err != nil {
        return nil, err
    }

    out, err := c.BlockCipher.Decrypt(key, ciphertext[SaltSize:])
    Zero(key)
    if err != nil {
        return nil, err
    }

    return out, nil
}

// MockBlockCipher implements BlockCipher but does nothing
type MockBlockCipher struct{}

// Encrypt in this case is only implementing the BlockCipher interface, it doesn't do anything
func (m *MockBlockCipher) Encrypt(key, plaintext []byte) ([]byte, error) {
    return plaintext, nil
}

// Decrypt in this case is only implementing the BlockCipher interface, it doesn't do anything
func (m *MockBlockCipher) Decrypt(key, ciphertext []byte) ([]byte, error) {
    return ciphertext, nil
}

// KeySize is a mock key size to use with the mock cipher
func (m *MockBlockCipher) KeySize() int {
    return 32
}

// Message represents a message being passed, and contains its contents and a sequence number
type Message struct {
    Number   uint32
    Contents []byte
}

// NewMessage returns a new message
func NewMessage(in []byte, num uint32) *Message {
    return &Message{Contents: in, Number: num}
}

// Marshal encodes a sequence number into the data that we wish to send
func (m *Message) Marshal() []byte {
    out := make([]byte, 4, len(m.Contents)+4)
    binary.BigEndian.PutUint32(out[:4], m.Number)
    return append(out, m.Contents...)
}

// UnmarshalMessage decodes bytes into a message pointer
func UnmarshalMessage(in []byte) (*Message, error) {
    m := &Message{}
    if len(in) <= 4 {
        return m, encerrors.ErrInvalidMessageLength
    }

    m.Number = binary.BigEndian.Uint32(in[:4])
    m.Contents = in[4:]
    return m, nil
}

// Channel is a typed io.ReadWriter used for communicating securely
type Channel io.ReadWriter

// Session represents a session that can be used to pass messages over a secure channel
type Session struct {
    Cipher *Cipher
    Channel
    lastSent uint32
    lastRecv uint32
    sendKey  *[32]byte
    recvKey  *[32]byte
}

// LastSent returns the last sent message id
func (s *Session) LastSent() uint32 {
    return s.lastSent
}

// LastRecv returns the last received message id
func (s *Session) LastRecv() uint32 {
    return s.lastRecv
}

// Encrypt encrypts a message with an embedded message id
func (s *Session) Encrypt(message []byte) ([]byte, error) {
    if len(message) == 0 {
        return nil, encerrors.ErrInvalidMessageLength
    }

    s.lastSent++
    m := NewMessage(message, s.lastSent)
    return s.Cipher.Encrypt(s.sendKey[:], m.Marshal())
}

// Decrypt decrypts a message and checks that its message id is valid
func (s *Session) Decrypt(message []byte) ([]byte, error) {
    out, err := s.Cipher.Decrypt(s.recvKey[:], message)
    if err != nil {
        return nil, err
    }

    m, err := UnmarshalMessage(out)
    if err != nil {
        return nil, err
    }

    // if this number is less than or equal to the last received message, this is a replay and we bail
    if m.Number <= s.lastRecv {
        return nil, encerrors.ErrInvalidMessageID
    }

    s.lastRecv = m.Number

    return m.Contents, nil
}

// Send encrypts the message and sends it out over the channel.
func (s *Session) Send(message []byte) error {
    m, err := s.Encrypt(message)
    if err != nil {
        return err
    }

    err = binary.Write(s.Channel, binary.BigEndian, uint32(len(m)))
    if err != nil {
        return err
    }

    _, err = s.Channel.Write(m)
    return err
}

// Receive listens for a new message on the channel.
func (s *Session) Receive() ([]byte, error) {
    var mlen uint32
    err := binary.Read(s.Channel, binary.BigEndian, &mlen)
    if err != nil {
        return nil, err
    }

    message := make([]byte, int(mlen))
    _, err = io.ReadFull(s.Channel, message)
    if err != nil {
        return nil, err
    }

    return s.Decrypt(message)
}

// GenerateKeyPair generates a new key pair. This can be used to get a
// new key pair for setting up a rekeying operation during the session.
func GenerateKeyPair() (pub *[64]byte, priv *[64]byte, err error) {
    pub = new([64]byte)
    priv = new([64]byte)

    recvPub, recvPriv, err := box.GenerateKey(rand.Reader)
    if err != nil {
        return nil, nil, err
    }
    copy(pub[:], recvPub[:])
    copy(priv[:], recvPriv[:])

    sendPub, sendPriv, err := box.GenerateKey(rand.Reader)
    if err != nil {
        return nil, nil, err
    }
    copy(pub[32:], sendPub[:])
    copy(priv[32:], sendPriv[:])
    return pub, priv, err
}

// Close zeroises the keys in the session. Once a session is closed,
// the traffic that was sent over the channel can no longer be decrypted
// and any attempts at sending or receiving messages over the channel
// will fail.
func (s *Session) Close() error {
    Zero(s.sendKey[:])
    Zero(s.recvKey[:])
    return nil
}

// keyExchange is a convenience function that takes keys as byte slices,
// copying them into the appropriate arrays.
func keyExchange(shared *[32]byte, priv, pub []byte) {
    // Copy the private key and wipe it, as it will no longer be needed.
    var kexPriv [32]byte
    copy(kexPriv[:], priv)
    Zero(priv)

    var kexPub [32]byte
    copy(kexPub[:], pub)

    box.Precompute(shared, &kexPub, &kexPriv)
    Zero(kexPriv[:])
}

// NewSession returns a new *Session
func NewSession(ch Channel, c *Cipher) *Session {
    return &Session{
        Cipher:  c,
        Channel: ch,
        recvKey: new([32]byte),
        sendKey: new([32]byte),
    }
}

// Dial sets up a new session over the channel by generating a new pair
// of Curve25519 keypairs, sending its public keys to the peer, and
// reading the peer's public keys back.
func Dial(ch Channel, c *Cipher) (*Session, error) {
    var peer [64]byte
    pub, priv, err := GenerateKeyPair()
    if err != nil {
        return nil, err
    }

    _, err = ch.Write(pub[:])
    if err != nil {
        return nil, err
    }

    // Make sure the entire public key is read.
    _, err = io.ReadFull(ch, peer[:])
    if err != nil {
        return nil, err
    }

    s := NewSession(ch, c)

    s.KeyExchange(priv, &peer, true)
    return s, nil
}

// Listen waits for a peer to Dial in, then sets up a key exchange
// and session.
func Listen(ch Channel, c *Cipher) (*Session, error) {
    var peer [64]byte
    pub, priv, err := GenerateKeyPair()
    if err != nil {
        return nil, err
    }

    // Ensure the entire peer key is read.
    _, err = io.ReadFull(ch, peer[:])
    if err != nil {
        return nil, err
    }

    _, err = ch.Write(pub[:])
    if err != nil {
        return nil, err
    }

    s := NewSession(ch, c)

    s.KeyExchange(priv, &peer, false)
    return s, nil
}

// KeyExchange - Rekey is used to perform the key exchange once both sides have
// exchanged their public keys. The underlying message protocol will
// need to actually initiate and carry out the key exchange, and call
// this once that is finished. The private key will be zeroised after
// calling this function. If the session is on the side that initiated
// the key exchange (e.g. by calling Dial), it should set the dialer
// argument to true. This will also reset the message counters for the
// session, as it will cause the session to use a new key.
func (s *Session) KeyExchange(priv, peer *[64]byte, dialer bool) {
    // This function denotes the dialer, who initiates the session,
    // as A. The listener is denoted as B. A is started using Dial,
    // and B is started using Listen.
    if dialer {
        // The first 32 bytes are the A->B link, where A is the
        // dialer. This key material should be used to set up the
        // A send key.
        keyExchange(s.sendKey, priv[:32], peer[:32])

        // The last 32 bytes are the B->A link, where A is the
        // dialer. This key material should be used to set up the A
        // receive key.
        keyExchange(s.recvKey, priv[32:], peer[32:])
    } else {
        // The first 32 bytes are the A->B link, where A is the
        // dialer. This key material should be used to set up the
        // B receive key.
        keyExchange(s.recvKey, priv[:32], peer[:32])

        // The last 32 bytes are the B->A link, where A is the
        // dialer. This key material should be used to set up the
        // B send key.
        keyExchange(s.sendKey, priv[32:], peer[32:])
    }
    s.lastSent = 0
    s.lastRecv = 0
}

const (
    // testComplexity is unexported because we don't want to use such a weak key in the wild
    testComplexity = 1 << (iota + 7)
)

const (
    // N Complexity in powers of 2 for key Derivation

    // InteractiveComplexity - recommended complexity for interactive sessions
    InteractiveComplexity = 1 << (iota + 14)
    // Complexity15 is 2^15
    Complexity15
    // Complexity16 is 2^16
    Complexity16
    // Complexity17 is 2^17
    Complexity17
    // Complexity18 is 2^18
    Complexity18
    // Complexity19 is 2^19
    Complexity19
    // AggressiveComplexity is 2^20 (don't use this unless you have incredibly strong CPU power
    AggressiveComplexity
)

// DeriveKey generates a new NaCl key from a passphrase and salt.
// This is a costly operation.
func DeriveKey(pass, salt []byte, N, keySize int) ([]byte, error) {
    var naclKey = make([]byte, keySize)
    key, err := scrypt.Key(pass, salt, N, 8, 1, keySize)
    if err != nil {
        return nil, err
    }

    copy(naclKey, key)
    Zero(key)
    return naclKey, nil
}

// Zero zeroes out bytes of data so that it does not stay in memory any longer than necessary
func Zero(data []byte) {
    for i := 0; i < len(data); i++ {
        data[i] = 0
    }
}

// Package cbc supports cbc encryption
package cbc

// https://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Cipher_Block_Chaining_.28CBC.29

import (
    "crypto/aes"
    "crypto/cipher"
    "crypto/hmac"
    "crypto/rand"
    "crypto/sha256"

    "io"

    "github.com/alistanis/goenc/encerrors"
    "github.com/alistanis/goenc/generate"
)

const (
    // NonceSize to use for nonces
    NonceSize = aes.BlockSize
    // MACSize is the output size of HMAC-SHA-256
    MACSize = 32
    // CKeySize - Cipher key size - AES-256
    CKeySize = 32
    // MKeySize - HMAC key size - HMAC-SHA-256
    MKeySize = 32
    // KeySize is the key size for CBC
    KeySize = CKeySize + MKeySize
)

// pad pads input to match the correct size
func pad(in []byte) []byte {
    padding := 16 - (len(in) % 16)
    for i := 0; i < padding; i++ {
        in = append(in, byte(padding))
    }
    return in
}

// unpad removes unnecessary bytes that were added during initial padding
func unpad(in []byte) []byte {
    if len(in) == 0 {
        return nil
    }

    padding := in[len(in)-1]
    if int(padding) > len(in) || padding > aes.BlockSize {
        return nil
    } else if padding == 0 {
        return nil
    }

    for i := len(in) - 1; i > len(in)-int(padding)-1; i-- {
        if in[i] != padding {
            return nil
        }
    }
    return in[:len(in)-int(padding)]
}

// Cipher implements the BlockCipher interface
type Cipher struct{}

// Encrypt implements the BlockCipher interface
func (c *Cipher) Encrypt(key, plaintext []byte) ([]byte, error) {
    return Encrypt(key, plaintext)
}

// Decrypt implements the BlockCipher interface
func (c *Cipher) Decrypt(key, ciphertext []byte) ([]byte, error) {
    return Decrypt(key, ciphertext)
}

// KeySize returns CBC KeySize and implements the BlockCipher interface
func (c *Cipher) KeySize() int {
    return KeySize
}

// New returns a new cbc cipher
func New() *Cipher {
    return &Cipher{}
}

// Key returns a random key as a pointer to an array of bytes specified by KeySize
func Key() (*[KeySize]byte, error) {
    key := new([KeySize]byte)
    _, err := io.ReadFull(rand.Reader, key[:])
    return key, err
}

// Encrypt encrypts plaintext using the given key with CBC encryption
func Encrypt(key, plaintext []byte) ([]byte, error) {
    if len(key) != KeySize {
        return nil, encerrors.ErrInvalidKeyLength
    }

    iv, err := generate.RandBytes(NonceSize)
    if err != nil {
        return nil, err
    }

    pmessage := pad(plaintext)
    ct := make([]byte, len(pmessage))

    // NewCipher only returns an error with an invalid key size,
    // but the key size was checked at the beginning of the function.
    c, _ := aes.NewCipher(key[:CKeySize])
    ctr := cipher.NewCBCEncrypter(c, iv)
    ctr.CryptBlocks(ct, pmessage)

    h := hmac.New(sha256.New, key[CKeySize:])
    ct = append(iv, ct...)
    h.Write(ct)
    ct = h.Sum(ct)
    return ct, nil
}

// Decrypt decrypts ciphertext using the given key
func Decrypt(key, ciphertext []byte) ([]byte, error) {
    if len(key) != KeySize {
        return nil, encerrors.ErrInvalidKeyLength
    }

    // HMAC-SHA-256 returns a MAC that is also a multiple of the
    // block size.
    if (len(ciphertext) % aes.BlockSize) != 0 {
        return nil, encerrors.ErrInvalidMessageLength
    }

    // A ciphertext must have at least an IV block, a ciphertext block,
    // and two blocks of HMAC.
    if len(ciphertext) < (4 * aes.BlockSize) {
        return nil, encerrors.ErrInvalidMessageLength
    }

    macStart := len(ciphertext) - MACSize
    tag := ciphertext[macStart:]
    out := make([]byte, macStart-NonceSize)
    ciphertext = ciphertext[:macStart]

    h := hmac.New(sha256.New, key[CKeySize:])
    h.Write(ciphertext)
    mac := h.Sum(nil)
    if !hmac.Equal(mac, tag) {
        return nil, encerrors.ErrInvalidSum
    }

    // NewCipher only returns an error with an invalid key size,
    // but the key size was checked at the beginning of the function.
    c, _ := aes.NewCipher(key[:CKeySize])
    ctr := cipher.NewCBCDecrypter(c, ciphertext[:NonceSize])
    ctr.CryptBlocks(out, ciphertext[NonceSize:])

    pt := unpad(out)
    if pt == nil {
        return nil, encerrors.ErrInvalidPadding
    }

    return pt, nil
}

// Package cfb supports basic cfb encryption with NO HMAC
package cfb

// https://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Cipher_Feedback_.28CFB.29

import (
    "crypto/aes"
    "crypto/cipher"
    "crypto/rand"
    "io"

    "github.com/alistanis/goenc/encerrors"
    "github.com/alistanis/goenc/generate"
)

// KeySize for CFB uses the generic key size
const KeySize = generate.KeySize

// Cipher to use for implementing the BlockCipher interface
type Cipher struct {
}

// New returns a new cfb cipher
func New() *Cipher {
    return &Cipher{}
}

// Encrypt implements the BlockCipher interface
func (c *Cipher) Encrypt(key, plaintext []byte) ([]byte, error) {
    return Encrypt(key, plaintext)
}

// Decrypt implements the BlockCipher interface
func (c *Cipher) Decrypt(key, ciphertext []byte) ([]byte, error) {
    return Decrypt(key, ciphertext)
}

// KeySize implements the BlockCipher interface
func (c *Cipher) KeySize() int {
    return KeySize
}

// Decrypt decrypts ciphertext using the given key
func Decrypt(key, ciphertext []byte) ([]byte, error) {

    // Create the AES cipher
    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }

    if len(ciphertext) < aes.BlockSize {
        return nil, encerrors.ErrInvalidMessageLength
    }

    // get first 16 bytes from ciphertext
    iv := ciphertext[:aes.BlockSize]

    // Remove the IV from the ciphertext
    ciphertext = ciphertext[aes.BlockSize:]

    // Return a decrypted stream
    stream := cipher.NewCFBDecrypter(block, iv)

    // SimpleDecrypt bytes from ciphertext
    stream.XORKeyStream(ciphertext, ciphertext)

    return ciphertext, nil
}

// Encrypt encrypts ciphertext using the given key.
// NOTE: This is not secure without being authenticated (crypto/hmac)
func Encrypt(key, plaintext []byte) ([]byte, error) {
    // Create the AES cipher
    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }

    // Empty array of 16 + plaintext length
    // Include the IV at the beginning
    ciphertext := make([]byte, aes.BlockSize+len(plaintext))

    // Slice of first 16 bytes
    iv := ciphertext[:aes.BlockSize]

    // Write 16 rand bytes to fill iv
    if _, err := io.ReadFull(rand.Reader, iv); err != nil {
        return nil, err
    }

    // Return an encrypted stream
    stream := cipher.NewCFBEncrypter(block, iv)

    // SimpleEncrypt bytes from plaintext to ciphertext
    stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)
    return ciphertext, nil
}

// DecryptString decrypts ciphertext using the given key
func DecryptString(key, ciphertext string) (string, error) {
    b, err := Decrypt([]byte(key), []byte(ciphertext))
    return string(b), err
}

// EncryptString encrypts ciphertext using the given key
func EncryptString(key, plaintext string) (string, error) {
    b, err := Encrypt([]byte(key), []byte(plaintext))
    return string(b), err
}

// Package ctr supports ctr encryption
package ctr

// https://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Counter_.28CTR.29

import (
    "crypto/aes"
    "crypto/cipher"
    "crypto/hmac"
    "crypto/rand"
    "crypto/sha256"

    "io"

    "github.com/alistanis/goenc/encerrors"
    "github.com/alistanis/goenc/generate"
)

const (
    // NonceSize to use for nonces
    NonceSize = aes.BlockSize
    // MACSize is the output size of HMAC-SHA-256
    MACSize = 32
    // CKeySize - Cipher key size - AES-256
    CKeySize = 32
    // MKeySize - HMAC key size - HMAC-SHA-256
    MKeySize = 32
    // KeySize to use for keys, 64 bytes
    KeySize = CKeySize + MKeySize
)

// Cipher to implement the BlockCipher interface
type Cipher struct {
}

// New returns a new ctr cipher
func New() *Cipher {
    return &Cipher{}
}

// Encrypt implements the BlockCipher interface
func (c *Cipher) Encrypt(key, plaintext []byte) ([]byte, error) {
    return Encrypt(key, plaintext)
}

// Decrypt implements the BlockCipher interface
func (c *Cipher) Decrypt(key, ciphertext []byte) ([]byte, error) {
    return Decrypt(key, ciphertext)
}

// KeySize implements the BlockCipher interface
func (c *Cipher) KeySize() int {
    return KeySize
}

// Key returns a pointer to an array of bytes with the given KeySize
func Key() (*[KeySize]byte, error) {
    key := new([KeySize]byte)
    _, err := io.ReadFull(rand.Reader, key[:])
    return key, err
}

// Encrypt encrypts plaintext using the given key with CTR encryption
func Encrypt(key, plaintext []byte) ([]byte, error) {
    if len(key) != KeySize {
        return nil, encerrors.ErrInvalidKeyLength
    }

    nonce, err := generate.RandBytes(NonceSize)
    if err != nil {
        return nil, err
    }

    ct := make([]byte, len(plaintext))

    // NewCipher only returns an error with an invalid key size,
    // but the key size was checked at the beginning of the function.
    c, _ := aes.NewCipher(key[:CKeySize])
    ctr := cipher.NewCTR(c, nonce)
    ctr.XORKeyStream(ct, plaintext)

    h := hmac.New(sha256.New, key[CKeySize:])
    ct = append(nonce, ct...)
    h.Write(ct)
    ct = h.Sum(ct)
    return ct, nil
}

// Decrypt decrypts ciphertext using the given key
func Decrypt(key, ciphertext []byte) ([]byte, error) {
    if len(key) != KeySize {
        return nil, encerrors.ErrInvalidKeyLength
    }

    if len(ciphertext) <= (NonceSize + MACSize) {
        return nil, encerrors.ErrInvalidMessageLength
    }

    macStart := len(ciphertext) - MACSize
    tag := ciphertext[macStart:]
    out := make([]byte, macStart-NonceSize)
    ciphertext = ciphertext[:macStart]

    h := hmac.New(sha256.New, key[CKeySize:])
    h.Write(ciphertext)
    mac := h.Sum(nil)
    if !hmac.Equal(mac, tag) {
        return nil, encerrors.ErrInvalidSum
    }

    c, _ := aes.NewCipher(key[:CKeySize])
    ctr := cipher.NewCTR(c, ciphertext[:NonceSize])
    ctr.XORKeyStream(out, ciphertext[NonceSize:])
    return out, nil
}

// Package gcm supports gcm encryption
package gcm

// https://en.wikipedia.org/wiki/Galois/Counter_Mode

import (
    "crypto/aes"
    "crypto/cipher"
    "encoding/binary"

    "github.com/alistanis/goenc/encerrors"
    "github.com/alistanis/goenc/generate"
)

// NonceSize // generic NonceSize
const NonceSize = generate.NonceSize

// KeySize // generic KeySize
const KeySize = generate.KeySize

// Cipher to implement the BlockCipher interface
type Cipher struct {
}

// New returns a new GCM cipher
func New() *Cipher {
    return &Cipher{}
}

// Encrypt implements the BlockCipher interface
func (c *Cipher) Encrypt(key, plaintext []byte) ([]byte, error) {
    return Encrypt(key, plaintext)
}

// Decrypt implements the BlockCipher interface
func (c *Cipher) Decrypt(key, ciphertext []byte) ([]byte, error) {
    return Decrypt(key, ciphertext)
}

// KeySize returns the GCM key size
func (c *Cipher) KeySize() int {
    return KeySize
}

// Encrypt secures a message using AES-GCM.
func Encrypt(key, plaintext []byte) ([]byte, error) {
    c, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }

    gcm, err := cipher.NewGCMWithNonceSize(c, NonceSize)
    if err != nil {
        return nil, err
    }

    nonce, err := generate.Nonce()
    if err != nil {
        return nil, err
    }

    // Seal will append the output to the first argument; the usage
    // here appends the ciphertext to the nonce. The final parameter
    // is any additional data to be authenticated.
    out := gcm.Seal(nonce[:], nonce[:], plaintext, nil)
    return out, nil
}

// EncryptString is a convenience function for working with strings
func EncryptString(key, plaintext string) (string, error) {
    data, err := Encrypt([]byte(key), []byte(plaintext))
    return string(data), err
}

// Decrypt decrypts data using AES-GCM
func Decrypt(key, ciphertext []byte) ([]byte, error) {
    // Create the AES cipher
    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }
    gcm, err := cipher.NewGCMWithNonceSize(block, NonceSize)
    if err != nil {
        return nil, err
    }
    nonce := make([]byte, NonceSize)
    copy(nonce, ciphertext)
    return gcm.Open(nil, nonce[:], ciphertext[NonceSize:], nil)
}

// DecryptString is a convenience function for working with strings
func DecryptString(key, ciphertext string) (string, error) {
    data, err := Decrypt([]byte(key), []byte(ciphertext))
    return string(data), err
}

//---------------------------------------------
// For use with more complex encryption schemes
//---------------------------------------------

// EncryptWithID secures a message and prepends a 4-byte sender ID
// to the message. The end bit is tricky, because gcm.Seal modifies buf, and this is necessary
func EncryptWithID(key, message []byte, sender uint32) ([]byte, error) {
    buf := make([]byte, 4)
    binary.BigEndian.PutUint32(buf, sender)

    c, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }

    gcm, err := cipher.NewGCMWithNonceSize(c, NonceSize)
    if err != nil {
        return nil, err
    }

    nonce, err := generate.Nonce()
    if err != nil {
        return nil, err
    }

    buf = append(buf, nonce[:]...)
    return gcm.Seal(buf, nonce[:], message, buf[:4]), nil
}

// EncryptStringWithID is a helper function to work with strings instead of bytes
func EncryptStringWithID(key, message string, sender uint32) (string, error) {
    data, err := EncryptWithID([]byte(key), []byte(message), sender)
    return string(data), err
}

// DecryptWithID takes an encrypted message and a KeyForID function (to get a key from a cache or a database perhaps)
// It checks the first 4 bytes for prepended header data, in this case, a sender ID
func DecryptWithID(message []byte, k KeyRetriever) ([]byte, error) {

    if len(message) <= NonceSize+4 {
        return nil, encerrors.ErrInvalidMessageLength
    }

    id := binary.BigEndian.Uint32(message[:4])
    key, err := k.KeyForID(id)
    if err != nil {
        return nil, err
    }
    c, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }

    gcm, err := cipher.NewGCMWithNonceSize(c, NonceSize)
    if err != nil {
        return nil, err
    }

    nonce := make([]byte, NonceSize)
    copy(nonce, message[4:])

    ciphertext := message[4+NonceSize:]

    // Decrypt the message, using the sender ID as the additional
    // data requiring authentication.
    out, err := gcm.Open(nil, nonce, ciphertext, message[:4])
    if err != nil {
        return nil, err
    }

    return out, nil
}

// DecryptStringWithID is a helper function to work with strings instead of bytes
func DecryptStringWithID(message string, k KeyRetriever) (string, error) {
    data, err := DecryptWithID([]byte(message), k)
    return string(data), err
}

// KeyRetriever represents a type that should be used in order to retrieve a key from a datastore
type KeyRetriever interface {
    KeyForID(uint32) ([]byte, error)
}

// GCMHelper is designed to make it easy to call EncryptWithID and DecryptWithID by assigning the KeyForIDFunc
// it implements KeyRetriever and provides convenience functions
// It also serves as an example for how to use KeyRetriever
type GCMHelper struct {
    KeyForIDFunc func(uint32) ([]byte, error)
}

// NewGCMHelper returns a new helper
func NewGCMHelper(f func(uint32) ([]byte, error)) *GCMHelper {
    return &GCMHelper{f}
}

// KeyForID implements the KeyRetriever interface, it should be used to get a Key for the given ID
func (h *GCMHelper) KeyForID(u uint32) ([]byte, error) {
    return h.KeyForIDFunc(u)
}

// Package nacl provides encryption by salting a key with a pad
package nacl

// https://en.wikipedia.org/wiki/NaCl_(software)
// work is derived from:
//
// https://github.com/andmarios/golang-nacl-secretbox

import (
    "crypto/rand"
    "errors"
    "fmt"

    "github.com/alistanis/goenc/generate"
    "golang.org/x/crypto/nacl/secretbox"
)

const (
    keySize   = 32
    nonceSize = 24
)

// Cipher to implmement the BlockCipher interface
type Cipher struct {
    Pad []byte
}

// Encrypt implements the BlockCipher interface
func (c *Cipher) Encrypt(key, plaintext []byte) ([]byte, error) {
    return Encrypt(c.Pad, key, plaintext)
}

// Decrypt implements the BlockCipher interface
func (c *Cipher) Decrypt(key, ciphertext []byte) ([]byte, error) {
    return Decrypt(c.Pad, key, ciphertext)
}

// KeySize returns the NaCL keysize
func (c *Cipher) KeySize() int {
    return keySize
}

// Encrypt salts a key using pad and encrypts a message
func Encrypt(pad, key, message []byte) (out []byte, err error) {
    if len(pad) < 32 {
        return nil, fmt.Errorf("pad had a length of %d, it must be at least 32 bytes", len(pad))
    }
    // NaCl's key has a constant size of 32 bytes.
    // The user provided key probably is less than that. We pad it with
    // a long enough string and truncate anything we don't need later on.
    key = append(key, pad...)

    // NaCl's key should be of type [32]byte.
    // Here we create it and truncate key bytes beyond 32
    naclKey := new([keySize]byte)
    copy(naclKey[:], key[:keySize])

    nonce, err := generate.Nonce()
    if err != nil {
        return nil, err
    }
    // out will hold the nonce and the encrypted message (ciphertext)
    out = make([]byte, nonceSize)
    // Copy the nonce to the start of out
    copy(out, nonce[:])
    // SimpleEncrypt the message and append it to out, assign the result to out
    out = secretbox.Seal(out, message, nonce, naclKey)
    return out, err
}

// Decrypt salts a key using pad and decrypts a message
func Decrypt(pad, key, data []byte) (out []byte, err error) {
    key = append(key, pad...)

    // NaCl's key should be of type [32]byte.
    // Here we create it and truncate key bytes beyond 32
    naclKey := new([keySize]byte)
    copy(naclKey[:], key[:keySize])

    // The nonce is of type [24]byte and part of the data we will receive
    nonce := new([nonceSize]byte)

    // Read the nonce from in, it is in the first 24 bytes
    copy(nonce[:], data[:nonceSize])

    // SimpleDecrypt the output of secretbox.Seal which contains the nonce and
    // the encrypted message
    message, ok := secretbox.Open(nil, data[nonceSize:], nonce, naclKey)
    if ok {
        return message, nil
    }
    return nil, errors.New("Decryption failed")
}

// RandomPadEncrypt generates a random pad and returns the encrypted data, the pad, and an error if any
func RandomPadEncrypt(key, message []byte) (pad, out []byte, err error) {
    pad = make([]byte, 32)
    _, err = rand.Read(pad)
    if err != nil {
        return
    }
    out, err = Encrypt(pad, key, message)
    return
}