// Test vectors for PBKDF2 taken from
// http://tc26.ru/methods/containers_v1/Addition_to_PKCS5_v1_0.pdf test vectors
func TestPBKDF2Vectors(t *testing.T) {
	if bytes.Compare(pbkdf2.Key(
		[]byte("password"),
		[]byte("salt"),
		1,
		32,
		PBKDF2Hash,
	), []byte{0x73, 0x14, 0xe7, 0xc0, 0x4f, 0xb2, 0xe6, 0x62, 0xc5, 0x43, 0x67, 0x42, 0x53, 0xf6, 0x8b, 0xd0, 0xb7, 0x34, 0x45, 0xd0, 0x7f, 0x24, 0x1b, 0xed, 0x87, 0x28, 0x82, 0xda, 0x21, 0x66, 0x2d, 0x58}) != 0 {
		t.Fail()
	}
	if bytes.Compare(pbkdf2.Key(
		[]byte("password"),
		[]byte("salt"),
		2,
		32,
		PBKDF2Hash,
	), []byte{0x99, 0x0d, 0xfa, 0x2b, 0xd9, 0x65, 0x63, 0x9b, 0xa4, 0x8b, 0x07, 0xb7, 0x92, 0x77, 0x5d, 0xf7, 0x9f, 0x2d, 0xb3, 0x4f, 0xef, 0x25, 0xf2, 0x74, 0x37, 0x88, 0x72, 0xfe, 0xd7, 0xed, 0x1b, 0xb3}) != 0 {
		t.Fail()
	}
	if bytes.Compare(pbkdf2.Key(
		[]byte("password"),
		[]byte("salt"),
		4096,
		32,
		PBKDF2Hash,
	), []byte{0x1f, 0x18, 0x29, 0xa9, 0x4b, 0xdf, 0xf5, 0xbe, 0x10, 0xd0, 0xae, 0xb3, 0x6a, 0xf4, 0x98, 0xe7, 0xa9, 0x74, 0x67, 0xf3, 0xb3, 0x11, 0x16, 0xa5, 0xa7, 0xc1, 0xaf, 0xff, 0x9d, 0xea, 0xda, 0xfe}) != 0 {
		t.Fail()
	}
	/*
		// It takes too long
		if bytes.Compare(pbkdf2.Key(
			[]byte("password"),
			[]byte("salt"),
			16777216,
			32,
			PBKDF2Hash,
		), []byte{0xa5, 0x7a, 0xe5, 0xa6, 0x08, 0x83, 0x96, 0xd1, 0x20, 0x85, 0x0c, 0x5c, 0x09, 0xde, 0x0a, 0x52, 0x51, 0x00, 0x93, 0x8a, 0x59, 0xb1, 0xb5, 0xc3, 0xf7, 0x81, 0x09, 0x10, 0xd0, 0x5f, 0xcd, 0x97}) != 0 {
			t.Fail()
		}
	*/
	if bytes.Compare(pbkdf2.Key(
		[]byte("passwordPASSWORDpassword"),
		[]byte("saltSALTsaltSALTsaltSALTsaltSALTsalt"),
		4096,
		40,
		PBKDF2Hash,
	), []byte{0x78, 0x83, 0x58, 0xc6, 0x9c, 0xb2, 0xdb, 0xe2, 0x51, 0xa7, 0xbb, 0x17, 0xd5, 0xf4, 0x24, 0x1f, 0x26, 0x5a, 0x79, 0x2a, 0x35, 0xbe, 0xcd, 0xe8, 0xd5, 0x6f, 0x32, 0x6b, 0x49, 0xc8, 0x50, 0x47, 0xb7, 0x63, 0x8a, 0xcb, 0x47, 0x64, 0xb1, 0xfd}) != 0 {
		t.Fail()
	}
	if bytes.Compare(pbkdf2.Key(
		[]byte("pass\x00word"),
		[]byte("sa\x00lt"),
		4096,
		20,
		PBKDF2Hash,
	), []byte{0x43, 0xe0, 0x6c, 0x55, 0x90, 0xb0, 0x8c, 0x02, 0x25, 0x24, 0x23, 0x73, 0x12, 0x7e, 0xdf, 0x9c, 0x8e, 0x9c, 0x32, 0x91}) != 0 {
		t.Fail()
	}
}

// ReadKeyfile reads a randomly generated and encrypted AES-256 key from the
// file with the given filename and returns it in unencrypted form.
// The key is protected by a passphrase, which is processed by PBKDF2 to
// derive the AES-256 key to decrypt the generated key.
func ReadKeyfile(filename string, passphrase []byte) (key []byte, err error) {
	// open keyfile
	keyfile, err := os.Open(filename)
	if err != nil {
		return nil, log.Error(err)
	}
	defer keyfile.Close()
	// read iter and convert to int
	var biter = make([]byte, 8)
	if _, err := keyfile.Read(biter); err != nil {
		return nil, log.Error(err)
	}
	uiter := encode.ToUint64(biter)
	if uiter > 2147483647 {
		return nil, log.Errorf("encdb: ReadKeyfile: invalid iter value")
	}
	iter := int(uiter)
	// read salt
	var salt = make([]byte, 32)
	if _, err := keyfile.Read(salt); err != nil {
		return nil, log.Error(err)
	}
	// read encrypted key
	var encKey = make([]byte, 16+32)
	if _, err := keyfile.Read(encKey); err != nil {
		return nil, log.Error(err)
	}
	// compute derived key from passphrase
	dk := pbkdf2.Key([]byte(passphrase), salt, iter, 32, sha256.New)
	// decrypt key
	return cipher.AES256CBCDecrypt([]byte(dk), encKey), nil
}

func getKDFKey(cryptoJSON cryptoJSON, auth string) ([]byte, error) {
	authArray := []byte(auth)
	salt, err := hex.DecodeString(cryptoJSON.KDFParams["salt"].(string))
	if err != nil {
		return nil, err
	}
	dkLen := ensureInt(cryptoJSON.KDFParams["dklen"])

	if cryptoJSON.KDF == "scrypt" {
		n := ensureInt(cryptoJSON.KDFParams["n"])
		r := ensureInt(cryptoJSON.KDFParams["r"])
		p := ensureInt(cryptoJSON.KDFParams["p"])
		return scrypt.Key(authArray, salt, n, r, p, dkLen)

	} else if cryptoJSON.KDF == "pbkdf2" {
		c := ensureInt(cryptoJSON.KDFParams["c"])
		prf := cryptoJSON.KDFParams["prf"].(string)
		if prf != "hmac-sha256" {
			return nil, fmt.Errorf("Unsupported PBKDF2 PRF: ", prf)
		}
		key := pbkdf2.Key(authArray, salt, c, dkLen, sha256.New)
		return key, nil
	}

	return nil, fmt.Errorf("Unsupported KDF: ", cryptoJSON.KDF)
}

func (kd *KeyDetails) key(password string) ([]byte, error) {
	if kd.Digest != "sha-512" {
		return nil, ErrInvalidKeyConfig
	}

	return pbkdf2.Key([]byte(password), kd.Salt, kd.Iterations, 32, sha512.New), nil
}

// Generates a derived key based on a salt. rails default key size is 64.
func (g *KeyGenerator) Generate(salt []byte, keySize int) []byte {
	// set a default
	if g.Iterations == 0 {
		g.Iterations = 1000 // rails 4 default when setting the session.
	}
	return pbkdf2.Key([]byte(g.Secret), salt, g.Iterations, keySize, sha1.New)
}

// Create a new encryption key and encrypt it using the user-provided
// passphrase. Prints output to stdout that gives text to add to the
// ~/.skicka.config file to store the encryption key.
func generateKey() {
	passphrase := os.Getenv(passphraseEnvironmentVariable)
	if passphrase == "" {
		printErrorAndExit(fmt.Errorf(passphraseEnvironmentVariable +
			" environment variable not set."))
	}

	// Derive a 64-byte hash from the passphrase using PBKDF2 with 65536
	// rounds of SHA256.
	salt := getRandomBytes(32)
	hash := pbkdf2.Key([]byte(passphrase), salt, 65536, 64, sha256.New)
	if len(hash) != 64 {
		printErrorAndExit(fmt.Errorf("incorrect key size returned by pbkdf2 %d", len(hash)))
	}

	// We'll store the first 32 bytes of the hash to use to confirm the
	// correct passphrase is given on subsequent runs.
	passHash := hash[:32]
	// And we'll use the remaining 32 bytes as a key to encrypt the actual
	// encryption key. (These bytes are *not* stored).
	keyEncryptKey := hash[32:]

	// Generate a random encryption key and encrypt it using the key
	// derived from the passphrase.
	key := getRandomBytes(32)
	iv := getRandomBytes(16)
	encryptedKey := encryptBytes(keyEncryptKey, iv, key)

	fmt.Printf("; Add the following lines to the [encryption] section\n")
	fmt.Printf("; of your ~/.skicka.config file.\n")
	fmt.Printf("\tsalt=%s\n", hex.EncodeToString(salt))
	fmt.Printf("\tpassphrase-hash=%s\n", hex.EncodeToString(passHash))
	fmt.Printf("\tencrypted-key=%s\n", hex.EncodeToString(encryptedKey))
	fmt.Printf("\tencrypted-key-iv=%s\n", hex.EncodeToString(iv))
}

func generatePbkdf2KeyCombo(pw string, salt []byte) *gocrypt.KeyCombo {
	key := pbkdf2.Key([]byte(pw), salt, pbkdf2Iterations, aes.KeyLength*2, sha256.New)
	return &gocrypt.KeyCombo{
		CryptoKey: key[:aes.KeyLength],
		AuthKey:   key[aes.KeyLength:],
	}
}

// SetPassword is a function that allows a password to be hashed and added to
// an InMemPwdStore instance.
func GetPbkdf2Hash(
	password string,
	iterations uint16,
) (*Pbkdf2Hash, error) {
	if iterations < Pbkdf2MinIterations {
		return nil, InsufficientIterationsError
	}

	var hashStruct Pbkdf2Hash
	randCount, err := rand.Read(hashStruct.Salt[:])
	if err != nil {
		return nil, err
	} else if randCount != Pbkdf2KeyLength {
		return nil, InsufficientEntropyError
	}

	hashStruct.Iterations = iterations

	subtle.ConstantTimeCopy(1, hashStruct.Hash[:], pbkdf2.Key(
		[]byte(password),
		hashStruct.Salt[:],
		int(hashStruct.Iterations),
		Pbkdf2KeyLength,
		sha256.New,
	))

	return &hashStruct, nil
}

func decryptKey(masterPwd []byte, encryptedKey []byte, salt []byte, iterCount int, validation []byte) ([]byte, error) {
	const keyLen = 32
	derivedKey := pbkdf2.Key(masterPwd, salt, iterCount, keyLen, sha1.New)

	aesKey := derivedKey[0:16]
	iv := derivedKey[16:32]
	decryptedKey, err := aesCbcDecrypt(aesKey, encryptedKey, iv)
	if err != nil {
		return nil, err
	}

	validationSalt, validationCipherText, err := extractSaltAndCipherText(validation)
	if err != nil {
		return nil, fmt.Errorf("Invalid validation: %v", err)
	}

	validationAesKey, validationIv := openSslKey(decryptedKey, validationSalt)
	decryptedValidation, err := aesCbcDecrypt(validationAesKey, validationCipherText, validationIv)
	if err != nil {
		return nil, fmt.Errorf("Failed to decrypt validation: %v", err)
	}

	if string(decryptedValidation) != string(decryptedKey) {
		return nil, errors.New("Validation decryption failed")
	}

	return decryptedKey, nil
}

func authMakePassword(password string) string {
	salt := make([]byte, 16)
	_, err := rand.Read(salt)
	checkErr(err)
	passwordHash := pbkdf2.Key([]byte(password), salt, 8192, 64, sha512.New)
	return hex.EncodeToString(salt) + ":" + hex.EncodeToString(passwordHash)
}

func (p *Profile) setMasterPassword(pwd string) error {

	var (
		dk            = pbkdf2.Key([]byte(pwd), p.Salt, p.Iterations, 64, sha512.New)
		derivedEncKey = dk[:32]
		derivedMacKey = dk[32:]
	)

	masterKey, err := decrypt(nil, p.MasterKey, derivedEncKey, derivedMacKey)
	if err != nil {
		return err
	}

	overviewKey, err := decrypt(nil, p.OverviewKey, derivedEncKey, derivedMacKey)
	if err != nil {
		return err
	}

	mac := sha512.New()
	mac.Write(masterKey)
	macData := mac.Sum(nil)
	p.masterEncKey = macData[:32]
	p.masterMacKey = macData[32:]

	mac.Reset()
	mac.Write(overviewKey)
	macData = mac.Sum(nil)
	p.overviewEncKey = macData[:32]
	p.overviewMacKey = macData[32:]

	return nil
}

// Decrypts the encrypted encryption key using values from the config file
// and the user's passphrase.
func decryptEncryptionKey() []byte {
	if key != nil {
		panic("key aready decrypted!")
	}

	salt := decodeHexString(config.Encryption.Salt)
	passphraseHash := decodeHexString(config.Encryption.Passphrase_hash)
	encryptedKey := decodeHexString(config.Encryption.Encrypted_key)
	encryptedKeyIv := decodeHexString(config.Encryption.Encrypted_key_iv)

	passphrase := os.Getenv(passphraseEnvironmentVariable)
	if passphrase == "" {
		fmt.Fprintf(os.Stderr, "skicka: "+passphraseEnvironmentVariable+
			" environment variable not set")
		os.Exit(1)
	}

	derivedKey := pbkdf2.Key([]byte(passphrase), salt, 65536, 64, sha256.New)
	// Make sure the first 32 bytes of the derived key match the bytes stored
	// when we first generated the key; if they don't, the user gave us
	// the wrong passphrase.
	if !bytes.Equal(derivedKey[:32], passphraseHash) {
		fmt.Fprintf(os.Stderr, "skicka: incorrect passphrase")
		os.Exit(1)
	}

	// Use the last 32 bytes of the derived key to decrypt the actual
	// encryption key.
	keyEncryptKey := derivedKey[32:]
	return decryptBytes(keyEncryptKey, encryptedKeyIv, encryptedKey)
}

func blockCrypt(key, crypt, salt string) (block kcp.BlockCrypt) {
	pass := pbkdf2.Key([]byte(key), []byte(salt), 4096, 32, sha1.New)

	switch crypt {
	case "tea":
		block, _ = kcp.NewTEABlockCrypt(pass[:16])
	case "xor":
		block, _ = kcp.NewSimpleXORBlockCrypt(pass)
	case "none":
		block, _ = kcp.NewNoneBlockCrypt(pass)
	case "aes-128":
		block, _ = kcp.NewAESBlockCrypt(pass[:16])
	case "aes-192":
		block, _ = kcp.NewAESBlockCrypt(pass[:24])
	case "blowfish":
		block, _ = kcp.NewBlowfishBlockCrypt(pass)
	case "twofish":
		block, _ = kcp.NewTwofishBlockCrypt(pass)
	case "cast5":
		block, _ = kcp.NewCast5BlockCrypt(pass[:16])
	case "3des":
		block, _ = kcp.NewTripleDESBlockCrypt(pass[:24])
	case "xtea":
		block, _ = kcp.NewXTEABlockCrypt(pass[:16])
	case "salsa20":
		block, _ = kcp.NewSalsa20BlockCrypt(pass)
	case "aes":
		fallthrough
	default: // aes
		block, _ = kcp.NewAESBlockCrypt(pass)
	}
	return
}

func main() {
	var rs [128][3][]byte
	for i := range rs {
		p := make([]byte, i)
		if _, err := io.ReadFull(rand.Reader, p[:]); err != nil {
			panic(err)
		}
		s := make([]byte, i)
		if _, err := io.ReadFull(rand.Reader, s[:]); err != nil {
			panic(err)
		}
		c := 128 - i + 2
		dk := pbkdf2.Key(p, s, c, i+8, sha256.New)
		rs[i][0] = p
		rs[i][1] = s
		rs[i][2] = dk
	}
	out, err := json.MarshalIndent(rs, "", "")
	if err != nil {
		panic(err)
	}
	fmt.Print("module.exports = ")
	fmt.Print(string(out))
	fmt.Println(";")
}

func Decrypt(password string, data []byte) ([]byte, error) {
	version := data[:1]
	options := data[1:2]
	encSalt := data[2:10]
	hmacSalt := data[10:18]
	iv := data[18:34]
	cipherText := data[34:(len(data) - 66 + 34)]
	expectedHmac := data[len(data)-32 : len(data)]

	msg := make([]byte, 0)
	msg = append(msg, version...)
	msg = append(msg, options...)
	msg = append(msg, encSalt...)
	msg = append(msg, hmacSalt...)
	msg = append(msg, iv...)
	msg = append(msg, cipherText...)

	hmacKey := pbkdf2.Key([]byte(password), hmacSalt, 10000, 32, sha1.New)
	testHmac := hmac.New(sha256.New, hmacKey)
	testHmac.Write(msg)
	testHmacVal := testHmac.Sum(nil)

	// its important to use hmac.Equal to not leak time
	// information. See https://github.com/RNCryptor/RNCryptor-Spec
	verified := hmac.Equal(testHmacVal, expectedHmac)

	if !verified {
		return nil, errors.New("Password may be incorrect, or the data has been corrupted. (HMAC could not be verified)")
	}

	cipherKey := pbkdf2.Key([]byte(password), encSalt, 10000, 32, sha1.New)
	cipherBlock, err := aes.NewCipher(cipherKey)
	if err != nil {
		return nil, err
	}

	decrypted := make([]byte, len(cipherText))
	copy(decrypted, cipherText)
	decrypter := cipher.NewCBCDecrypter(cipherBlock, iv)
	decrypter.CryptBlocks(decrypted, decrypted)

	// un-padd decrypted data
	length := len(decrypted)
	unpadding := int(decrypted[length-1])

	return decrypted[:(length - unpadding)], nil
}

// New creates a new Password from a plain text password.
func New(password string) Password {
	salt := make([]byte, 8, 28)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		panic(err.Error())
	}
	gen := pbkdf2.Key([]byte(password), salt, 4096, 20, sha1.New)
	return Password(base64.StdEncoding.EncodeToString(append(salt, gen...)))
}

func blobKey(username string, secret []byte) []byte {
	data := pbkdf2.Key(secret, []byte(username), 0x100, 20, sha1.New)[0:20]

	hash := sha1.Sum(data)
	length := make([]byte, 4)
	binary.BigEndian.PutUint32(length, 20)
	return append(hash[:], length...)
}

func generateKeys(password, salt []byte, keySize int) (encKey, authKey, pwv []byte) {
	totalSize := (keySize * 2) + 2 // enc + auth + pv sizes
	key := pbkdf2.Key(password, salt, 1000, totalSize, sha1.New)
	encKey = key[:keySize]
	authKey = key[keySize : keySize*2]
	pwv = key[keySize*2:]
	return
}

// FIXME: If the rand stuff in here fails, it
// will probably crash the app, there is no
// error handling here
func hashPass(password string) ([]byte, int, []byte) {
	salt := make([]byte, 32)
	rand.Read(salt)
	ii, _ := rand.Int(rand.Reader, big.NewInt(16000))
	iterations := int(ii.Int64()) + 64000
	hash := pbkdf2.Key([]byte(password), salt, iterations, 32, sha256.New)
	return salt, iterations, hash
}

// Function to decrypt a file
func decryptFile(file, key string) {

	fmt.Println("Decrypting...")

	//Generate a key of required length using the pbkd2 lib and the input
	cipherKey := pbkdf2.Key([]byte(key), []byte(salt), 4096, 32, sha1.New)
	// Define a new AES cipher with our generated key
	block, err := aes.NewCipher(cipherKey)
	HandleError(err, "cipher")

	// Open input file to be encrypted
	fin, err := os.Open(file)
	HandleError(err, "open input file")
	defer fin.Close()
	//Get input file size
	size := FileSize(file)
	// Open ouput file
	fout, err := os.OpenFile(file+".dec", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0600)
	HandleError(err, "open output file")
	defer fout.Close()

	iv := make([]byte, aes.BlockSize)
	_, err = fin.Read(iv)
	HandleError(err, "reading iv")

	// If file size is greater than 32KB, make a byte buffer of 32KB
	// Otherwise, create a buffer of file size
	var buf []byte
	if size > (4 << 20) {
		buf = make([]byte, 32768)
	} else {
		buf = make([]byte, size)
	}
	// Loop until we reach end of file
	for {
		// Read data
		res, err := fin.Read(buf)
		// If there is any error, exit
		if err != nil && err != io.EOF {
			panic(err)
		}
		// If end of file is reached or there is no data to be read, break
		if res == 0 || err == io.EOF {
			break
		}
		// Create a byte array for decrypted data
		cipherText := make([]byte, len(buf))
		// Decrypt the input data
		stream := cipher.NewCFBDecrypter(block, iv)
		stream.XORKeyStream(cipherText, buf[:res])
		//Write the decrypted data to output file
		_, err = fout.Write(cipherText)
		HandleError(err, "writing cipher block")

	}

	fmt.Println("Done.")
}