package jose

import (
	"errors"
	"github.com/dvsekhvalnov/jose2go/arrays"	
	"github.com/dvsekhvalnov/jose2go/padding"
	"crypto/hmac"
	"crypto/cipher"
	"crypto/aes"
	"fmt"
)

// AES CBC with HMAC authenticated encryption algorithm implementation
type AesCbcHmac struct{
	keySizeBits int
}

func init() {
	RegisterJwe(&AesCbcHmac{keySizeBits:256})
	RegisterJwe(&AesCbcHmac{keySizeBits:384})
	RegisterJwe(&AesCbcHmac{keySizeBits:512})		
}

func (alg *AesCbcHmac) Name() string {
	switch alg.keySizeBits {
		case 256: return A128CBC_HS256
		case 384: return A192CBC_HS384
		default: return  A256CBC_HS512
	}	
}

func (alg *AesCbcHmac) KeySizeBits() int {
	return alg.keySizeBits
}

func (alg *AesCbcHmac) Encrypt(aad, plainText, cek []byte) (iv, cipherText, authTag []byte, err error) {
	
	cekSizeBits := len(cek)<<3	
	if cekSizeBits != alg.keySizeBits {
		return nil,nil,nil, errors.New(fmt.Sprintf("AesCbcHmac.Encrypt(): expected key of size %v bits, but was given %v bits.",alg.keySizeBits, cekSizeBits))
	}	
	
    hmacKey := cek[0:len(cek)/2] 
 	aesKey := cek[len(cek)/2:]   	
	
	if iv,err = arrays.Random(16);err!=nil {
		return nil,nil,nil,err
	}
		
	var block cipher.Block

	if block, err = aes.NewCipher(aesKey);err!=nil {
		return nil,nil,nil,err
	}
	

	padded := padding.AddPkcs7(plainText,16)

	cipherText = make([]byte,len(padded),cap(padded))
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(cipherText,padded)
	
	authTag = alg.computeAuthTag(aad, iv, cipherText, hmacKey)
	
	return iv,cipherText,authTag,nil		
}


func (alg *AesCbcHmac) Decrypt(aad, cek, iv, cipherText, authTag []byte) (plainText []byte, err error) {
	
	cekSizeBits := len(cek)<<3
	
	if  cekSizeBits != alg.keySizeBits {
		return nil, errors.New(fmt.Sprintf("AesCbcHmac.Decrypt(): expected key of size %v bits, but was given %v bits.",alg.keySizeBits, cekSizeBits))
	}	
	
    hmacKey := cek[0:len(cek)/2] 
 	aesKey := cek[len(cek)/2:]   
		
	// Check MAC
    expectedAuthTag := alg.computeAuthTag(aad, iv, cipherText, hmacKey);

	if !hmac.Equal(expectedAuthTag, authTag) { 
		return nil,errors.New("AesCbcHmac.Decrypt(): Authentication tag do not match.")
	}

	var block cipher.Block

	if block, err = aes.NewCipher(aesKey);err==nil {
		mode := cipher.NewCBCDecrypter(block, iv)
		
		var padded []byte=make([]byte, len(cipherText), cap(cipherText))		
		mode.CryptBlocks(padded, cipherText)
		
		return padding.RemovePkcs7(padded,16),nil
	}
	
	return nil,err
}

func (alg *AesCbcHmac) computeAuthTag(aad []byte, iv []byte, cipherText []byte, hmacKey []byte) (signature []byte) {
	al := arrays.UInt64ToBytes(uint64(len(aad) << 3));	
	hmacInput := arrays.Concat(aad, iv, cipherText, al)
	hmac :=calculateHmac(alg.keySizeBits, hmacInput, hmacKey);

	return hmac[0:len(hmac)/2];
}

