// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

package hybrid

import (
	"bytes"
	"encoding/hex"
	"testing"

	"github.com/tink-crypto/tink-go/v2/aead"
	"github.com/tink-crypto/tink-go/v2/daead"
	"github.com/tink-crypto/tink-go/v2/hybrid/subtle"
	"github.com/tink-crypto/tink-go/v2/subtle/random"
	"github.com/tink-crypto/tink-go/v2/testutil"
	tinkpb "github.com/tink-crypto/tink-go/v2/proto/tink_go_proto"
)

func modifyDecrypt(t *testing.T, c string, k *tinkpb.KeyTemplate) {
	t.Helper()
	curve, err := subtle.GetCurve(c)
	if err != nil {
		t.Fatalf("error getting %s curve: %s ", c, err)
	}
	pvt, err := subtle.GenerateECDHKeyPair(curve)
	if err != nil {
		t.Fatalf("error generating ECDH key pair: %s", err)
	}
	salt := random.GetRandomBytes(8)
	pt := random.GetRandomBytes(4)
	context := random.GetRandomBytes(4)
	rDem, err := newRegisterECIESAEADHKDFDemHelper(k)
	if err != nil {
		t.Fatalf("error generating a DEM helper :%s", err)
	}
	e, err := subtle.NewECIESAEADHKDFHybridEncrypt(&pvt.PublicKey, salt, "SHA256", "UNCOMPRESSED", rDem)
	if err != nil {
		t.Fatalf("error generating an encryption construct :%s", err)
	}
	d, err := subtle.NewECIESAEADHKDFHybridDecrypt(pvt, salt, "SHA256", "UNCOMPRESSED", rDem)
	if err != nil {
		t.Fatalf("error generating an decryption construct :%s", err)
	}
	ct, err := e.Encrypt(pt, context)
	if err != nil {
		t.Fatalf("encryption error :%s", err)
	}
	dt, err := d.Decrypt(ct, context)
	if err != nil {
		t.Fatalf("decryption error :%s", err)
	}
	if !bytes.Equal(dt, pt) {
		t.Fatalf("decryption not inverse of encryption")
	}

	for _, g := range testutil.GenerateMutations(ct) {
		if _, err := d.Decrypt(g, context); err == nil {
			t.Fatalf("invalid cipher text should throw exception")
		}
	}
	for _, g := range testutil.GenerateMutations(context) {
		if _, err := d.Decrypt(ct, g); err == nil {
			t.Fatalf("invalid context should throw exception")
		}
	}
	mSalt := make([]byte, len(salt))

	for i := 0; i < len(salt); i++ {
		for j := 0; j < 8; j++ {
			copy(mSalt, salt)
			mSalt[i] ^= (1 << uint8(j))
			d, err = subtle.NewECIESAEADHKDFHybridDecrypt(pvt, mSalt, "SHA256", "UNCOMPRESSED", rDem)
			if err != nil {
				t.Fatalf("subtle.NewECIESAEADHKDFHybridDecrypt:%v", err)
			}
			if _, err := d.Decrypt(ct, context); err == nil {
				t.Fatalf("invalid salt should throw exception")
			}
		}
	}
}

func TestECAESCTRHMACSHA256Decrypt(t *testing.T) {
	modifyDecrypt(t, "NIST_P256", aead.AES256CTRHMACSHA256KeyTemplate())
	modifyDecrypt(t, "NIST_P384", aead.AES256CTRHMACSHA256KeyTemplate())
	modifyDecrypt(t, "NIST_P521", aead.AES256CTRHMACSHA256KeyTemplate())
	modifyDecrypt(t, "NIST_P224", aead.AES256CTRHMACSHA256KeyTemplate())

	modifyDecrypt(t, "NIST_P256", aead.AES128CTRHMACSHA256KeyTemplate())
	modifyDecrypt(t, "NIST_P384", aead.AES128CTRHMACSHA256KeyTemplate())
	modifyDecrypt(t, "NIST_P521", aead.AES128CTRHMACSHA256KeyTemplate())
	modifyDecrypt(t, "NIST_P224", aead.AES128CTRHMACSHA256KeyTemplate())
}

func TestECAES256GCMDecrypt(t *testing.T) {
	modifyDecrypt(t, "NIST_P256", aead.AES256GCMKeyTemplate())
	modifyDecrypt(t, "NIST_P384", aead.AES256GCMKeyTemplate())
	modifyDecrypt(t, "NIST_P521", aead.AES256GCMKeyTemplate())
	modifyDecrypt(t, "NIST_P224", aead.AES256GCMKeyTemplate())

	modifyDecrypt(t, "NIST_P256", aead.AES128GCMKeyTemplate())
	modifyDecrypt(t, "NIST_P384", aead.AES128GCMKeyTemplate())
	modifyDecrypt(t, "NIST_P521", aead.AES128GCMKeyTemplate())
	modifyDecrypt(t, "NIST_P224", aead.AES128GCMKeyTemplate())
}

func TestECAESSIVDecrypt(t *testing.T) {
	modifyDecrypt(t, "NIST_P256", daead.AESSIVKeyTemplate())
	modifyDecrypt(t, "NIST_P384", daead.AESSIVKeyTemplate())
	modifyDecrypt(t, "NIST_P521", daead.AESSIVKeyTemplate())
	modifyDecrypt(t, "NIST_P224", daead.AESSIVKeyTemplate())
}

func TestECAESSIVTestVectors(t *testing.T) {
	// These are the same test vectors used to test the c++ implementation in
	// //third_party/tink/cc/hybrid/ecies_aead_hkdf_hybrid_decrypt_test.cc.
	testvectors := []struct {
		name, key, ciphertext, context, plaintext string
	}{
		{
			name:       "Empty plaintext",
			key:        "32588172ed65830571bb83748f7fddd383323208a7825c80a71bef846333eb02",
			ciphertext: "0401b11f8c9bafe30ae13f8bd15528714e752631a4328bf146009068e99489c8e9fae1ec39e3fe9994723711417fcab2af4b3c9b60117d47d33d35175c87b483b8935a73312940d1fbf8da3944a89b5e8b",
			context:    "some context info",
			plaintext:  "",
		},
		{
			name:       "Small plaintext",
			key:        "32588172ed65830571bb83748f7fddd383323208a7825c80a71bef846333eb02",
			ciphertext: "040230023d1547b55af5a735a7f460722612126d7539d7cd0f677d308b29c6f52a964e66e7b0cb44cff1673df9e2c793f1477ca755807bfbeadcae1ab20b45ecb1501ca5e3f5b0626d3ca40aa5d010443d506e4df90b",
			context:    "some context info",
			plaintext:  "hello",
		},
		{
			name:       "Large plaintext",
			key:        "32588172ed65830571bb83748f7fddd383323208a7825c80a71bef846333eb02",
			ciphertext: " ",
			context:    "some context info",
			plaintext:  "08b8b2b733424243760fe426a4b54908632110a66c2f6591eabd3345e3e4eb98fa6e264bf09efe12ee50f8f54e9f77b1e355f6c50544e23fb1433ddf73be84d879de7c0046dc4996d9e773f4bc9efe5738829adb26c81b37c93a1b270b20329d658675fc6ea534e0810a4432826bf58c941efb65d57a338bbd2e26640f89ffbc1a858efcb8550ee3a5e1998bd177e93a7363c344fe6b199ee5d02e82d522c4feba15452f80288a821a579116ec6dad2b3b310da903401aa62100ab5d1a36553e06203b33890cc9b832f79ef80560ccb9a39ce767967ed628c6ad573cb116dbefefd75499da96bd68a8a97b928a8bbc103b6621fcde2beca1231d206be6cd9ec7aff6f6c94fcd7204ed3455c68c83f4a41da4af2b74ef5c53f1d8ac70bdcb7ed185ce81bd84359d44254d95629e9855a94a7c1958d1f8ada5d0532ed8a5aa3fb2d17ba70eb6248e594e1a2297acbbb39d502f1a8c6eb6f1ce22b3de1a1f40cc24554119a831a9aad6079cad88425de6bde1a9187ebb6092cf67bf2b13fd65f27088d78b7e883c8759d2c4f5c65adb7553878ad575f9fad878e80a0c9ba63bcbcc2732e69485bbc9c90bfbd62481d9089beccf80cfe2df16a2cf65bd92dd597b0707e0917af48bbb75fed413d238f5555a7a569d80c3414a8d0859dc65a46128bab27af87a71314f318c782b23ebfe808b82b0ce26401d2e22f04d83d1255dc51addd3b75a2b1ae0784504df543af8969be3ea7082ff7fc9888c144da2af58429ec96031dbcad3dad9af0dcbaaaf268cb8fcffead94f3c7ca495e056a9b47acdb751fb73e666c6c655ade8297297d07ad1ba5e43f1bca32301651339e22904cc8c42f58c30c04aafdb038dda0847dd988dcda6f3bfd15c4b4c4525004aa06eeff8ca61783aacec57fb3d1f92b0fe2fd1a85f6724517b65e614ad6808d6f6ee34dff7310fdc82aebfd904b01e1dc54b2927094b2db68d6f903b68401adebf5a7e08d78ff4ef5d63653a65040cf9bfd4aca7984a74d37145986780fc0b16ac451649de6188a7dbdf191f64b5fc5e2ab47b57f7f7276cd419c17a3ca8e1b939ae49e488acba6b965610b5480109c8b17b80e1b7b750dfc7598d5d5011fd2dcc5600a32ef5b52a1ecc820e308aa342721aac0943bf6686b64b2579376504ccc493d97e6aed3fb0f9cd71a43dd497f01f17c0e2cb3797aa2a2f256656168e6c496afc5fb93246f6b1116398a346f1a641f3b041e989f7914f90cc2c7fff357876e506b50d334ba77c225bc307ba537152f3f1610e4eafe595f6d9d90d11faa933a15ef1369546868a7f3a45a96768d40fd9d03412c091c6315cf4fde7cb68606937380db2eaaa707b4c4185c32eddcdd306705e4dc1ffc872eeee475a64dfac86aba41c0618983f8741c5ef68d3a101e8a3b8cac60c905c15fc910840b94c00a0b9d0",
		},
	}
	for _, tv := range testvectors {
		t.Run(tv.name, func(t *testing.T) {
			key, err := hex.DecodeString(tv.key)
			if err != nil {
				t.Fatalf("error decoding test vector key :%s", err)
			}
			ct, err := hex.DecodeString(tv.ciphertext)
			if err != nil {
				t.Fatalf("error decoding test vector ciphertext :%s", err)
			}
			ctx := []byte(tv.context)
			pt := []byte(tv.plaintext)

			curve, _ := subtle.GetCurve("NIST_P256")
			pvt := subtle.GetECPrivateKey(curve, key)
			var salt []byte
			helper, err := newRegisterECIESAEADHKDFDemHelper(daead.AESSIVKeyTemplate())
			if err != nil {
				t.Fatalf("error generating a DEM helper :%s", err)
			}
			d, err := subtle.NewECIESAEADHKDFHybridDecrypt(pvt, salt, "SHA256", "UNCOMPRESSED", helper)
			if err != nil {
				t.Fatalf("error generating an decryption construct :%s", err)
			}
			dt, err := d.Decrypt(ct, ctx)
			if err != nil {
				t.Fatalf("decryption error :%s", err)
			}

			if !bytes.Equal(dt, pt) {
				t.Errorf("Decryption did not match test vector's plaintext")
			}
		})
	}

}
