package jose

import (
	"errors"
	"github.com/dvsekhvalnov/jose2go/arrays"
	"github.com/dvsekhvalnov/jose2go/base64url"
	"github.com/dvsekhvalnov/jose2go/kdf"
	"crypto/sha256"
	"crypto/sha512"	
	"hash"
)

func init() {
	RegisterJwa(&Pbse2HmacAesKW{keySizeBits: 128,aesKW: &AesKW{ keySizeBits: 128}})
	RegisterJwa(&Pbse2HmacAesKW{keySizeBits: 192,aesKW: &AesKW{ keySizeBits: 192}})
	RegisterJwa(&Pbse2HmacAesKW{keySizeBits: 256,aesKW: &AesKW{ keySizeBits: 256}})
}

// PBSE2 with HMAC key management algorithm implementation
type Pbse2HmacAesKW struct{
	keySizeBits int
	aesKW JwaAlgorithm
}

func (alg *Pbse2HmacAesKW) Name() string {
	switch alg.keySizeBits {
		case 128: return PBES2_HS256_A128KW
		case 192: return PBES2_HS384_A192KW
		default: return  PBES2_HS512_A256KW
	}
}

func (alg *Pbse2HmacAesKW) WrapNewKey(cekSizeBits int, key interface{}, header map[string]interface{}) (cek []byte, encryptedCek []byte, err error) {	
	if passphrase,ok:=key.(string); ok {

		algId := []byte(header["alg"].(string))
		
        iterationCount := 8192;
        var saltInput []byte
		
		if saltInput,err = arrays.Random(12);err!=nil {
			return nil,nil,err
		}

        header["p2c"] = iterationCount;
        header["p2s"] = base64url.Encode(saltInput)
		
		salt := arrays.Concat(algId, []byte{0}, saltInput);
		
		kek := kdf.DerivePBKDF2([]byte(passphrase), salt, iterationCount, alg.keySizeBits, alg.prf())
		return alg.aesKW.WrapNewKey(cekSizeBits, kek, header)
	}

	return nil,nil,errors.New("Pbse2HmacAesKW.WrapNewKey(): expected key to be 'string' array")	
}

func (alg *Pbse2HmacAesKW) Unwrap(encryptedCek []byte, key interface{}, cekSizeBits int, header map[string]interface{}) (cek []byte, err error) {		
	
	if passphrase,ok:=key.(string); ok {
		
		var p2s string
		var p2c float64
	
		if p2c,ok = header["p2c"].(float64);!ok {
			return nil,errors.New("Pbse2HmacAesKW.Unwrap(): expected 'p2c' param in JWT header, but was not found.")
		}
	
		if p2s,ok = header["p2s"].(string);!ok {
			return nil,errors.New("Pbse2HmacAesKW.Unwrap(): expected 'p2s' param in JWT header, but was not found.")
		}
		
		var saltInput []byte
		
		algId := []byte(header["alg"].(string))
		
		if saltInput,err = base64url.Decode(p2s);err!=nil {
			return nil,err
		}
		
		salt := arrays.Concat(algId,[]byte{0},saltInput)
		
		kek := kdf.DerivePBKDF2([]byte(passphrase), salt, int(p2c), alg.keySizeBits, alg.prf())
		
		return alg.aesKW.Unwrap(encryptedCek, kek, cekSizeBits, header)
	}
		
	return nil,errors.New("Pbse2HmacAesKW.Unwrap(): expected key to be 'string' array")		
}

func (alg *Pbse2HmacAesKW) prf() hash.Hash {
	switch alg.keySizeBits {
		case 128: return sha256.New()
		case 192: return sha512.New384()
		default: return  sha512.New()
	}
}