// Copyright 2020 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 aead_test

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

	internalaead "github.com/tink-crypto/tink-go/v2/internal/aead"
	"github.com/tink-crypto/tink-go/v2/internal/testing/aead"
	"github.com/tink-crypto/tink-go/v2/internal/testing/wycheproof"
	"github.com/tink-crypto/tink-go/v2/subtle/random"
)

func TestAESGCMSIVRejectsInvalidKeyLength(t *testing.T) {
	invalidKeySizes := []uint32{4, 8, 12, 15, 17, 24, 30, 31, 33, 64, 128}

	for _, keySize := range invalidKeySizes {
		key := random.GetRandomBytes(keySize)

		if _, err := internalaead.NewAESGCMSIV(key); err == nil {
			t.Errorf("expected error with invalid key-size %d", keySize)
		}
	}
}

func TestAESGCMSIVRandomNonceProducesDifferentCiphertexts(t *testing.T) {
	nSample := 1 << 17
	key := random.GetRandomBytes(16)
	pt := []byte{}
	ad := []byte{}
	a, err := internalaead.NewAESGCMSIV(key)
	if err != nil {
		t.Fatalf("internalaead.NewAESGCMSIV() err = %v, want nil", err)
	}
	ctSet := make(map[string]bool)

	for i := 0; i < nSample; i++ {
		dst := make([]byte, 0, internalaead.AESGCMSIVNonceSize+internalaead.AESGCMSIVTagSize+len(pt))
		ct, err := a.Encrypt(dst, pt, ad)
		if err != nil {
			t.Fatalf("a.Encrypt() err = %v, want nil", err)
		}
		ctHex := hex.EncodeToString(ct)

		if _, existed := ctSet[ctHex]; existed {
			t.Errorf("nonce is repeated after %d samples", i)
		}
		ctSet[ctHex] = true
	}
}

func TestAESGCMSIVModifyCiphertext(t *testing.T) {
	ad := random.GetRandomBytes(33)
	key := random.GetRandomBytes(16)
	pt := random.GetRandomBytes(32)
	a, err := internalaead.NewAESGCMSIV(key)
	if err != nil {
		t.Fatalf("internalaead.NewAESGCMSIV() err = %v, want nil", err)
	}
	dst := make([]byte, 0, internalaead.AESGCMSIVNonceSize+internalaead.AESGCMSIVTagSize+len(pt))
	ct, err := a.Encrypt(dst, pt, ad)
	if err != nil {
		t.Fatalf("a.Encrypt() err = %v, want nil", err)
	}
	// flipping bits
	for i := 0; i < len(ct); i++ {
		tmp := ct[i]
		for j := 0; j < 8; j++ {
			ct[i] ^= 1 << uint8(j)
			if _, err := a.Decrypt(ct, ad); err == nil {
				t.Errorf("expect an error when flipping bit of ciphertext: byte %d, bit %d", i, j)
			}
			ct[i] = tmp
		}
	}
	// truncated ciphertext
	for i := 1; i < len(ct); i++ {
		if _, err := a.Decrypt(ct[:i], ad); err == nil {
			t.Errorf("expect an error ciphertext is truncated until byte %d", i)
		}
	}
	// modify associated data
	for i := 0; i < len(ad); i++ {
		tmp := ad[i]
		for j := 0; j < 8; j++ {
			ad[i] ^= 1 << uint8(j)
			if _, err := a.Decrypt(ct, ad); err == nil {
				t.Errorf("expect an error when flipping bit of ad: byte %d, bit %d", i, j)
			}
			ad[i] = tmp
		}
	}
}

func TestAESGCMSIVWycheproofCases(t *testing.T) {
	suite := new(aead.WycheproofSuiteV1)
	wycheproof.PopulateSuiteV1(t, suite, "aes_gcm_siv_test.json")
	for _, group := range suite.TestGroups {
		for _, test := range group.Tests {
			caseName := fmt.Sprintf("%s-%s(%d):Case-%d", suite.Algorithm, group.Type, group.KeySize, test.CaseID)
			t.Run("DecryptOnly/"+caseName, func(t *testing.T) { runWycheproofDecryptOnly(t, test) })
			t.Run("EncryptDecrypt/"+caseName, func(t *testing.T) { runWycheproofEncryptDecrypt(t, test) })
		}
	}
}

func runWycheproofDecryptOnly(t *testing.T, testCase *aead.WycheproofCase) {
	aead, err := internalaead.NewAESGCMSIV(testCase.Key)
	if err != nil {
		t.Fatalf("cannot create aead, error: %v", err)
	}

	var combinedCt []byte
	combinedCt = append(combinedCt, testCase.Iv...)
	combinedCt = append(combinedCt, testCase.Ct...)
	combinedCt = append(combinedCt, testCase.Tag...)
	decrypted, err := aead.Decrypt(combinedCt, testCase.Aad)
	switch testCase.Result {
	case "valid":
		if err != nil {
			t.Errorf("unexpected error in test-case: %v", err)
		} else if !bytes.Equal(decrypted, testCase.Msg) {
			t.Errorf(
				"incorrect decryption: actual: %s; expected %s",
				hex.EncodeToString(decrypted), hex.EncodeToString(testCase.Msg))
		}
	case "invalid":
		if err == nil && bytes.Equal(testCase.Ct, decrypted) {
			t.Error("successfully decrypted invalid test-case")
		}
	default:
		t.Errorf("unknown test-case result: %s", testCase.Result)
	}
}

func runWycheproofEncryptDecrypt(t *testing.T, testCase *aead.WycheproofCase) {
	aead, err := internalaead.NewAESGCMSIV(testCase.Key)
	if err != nil {
		t.Fatalf("cannot create aead, error: %v", err)
	}

	dst := make([]byte, 0, internalaead.AESGCMSIVNonceSize+internalaead.AESGCMSIVTagSize+len(testCase.Msg))
	ct, err := aead.Encrypt(dst, testCase.Msg, testCase.Aad)
	if err != nil {
		if testCase.Result != "invalid" {
			t.Errorf("unexpected error in test-case: %v", err)
		}
		return
	}

	decrypted, err := aead.Decrypt(ct, testCase.Aad)
	switch testCase.Result {
	case "valid":
		if err != nil {
			t.Errorf("unexpected error in test-case: %v", err)
		} else if !bytes.Equal(decrypted, testCase.Msg) {
			t.Errorf(
				"incorrect decryption: actual: %s; expected %s",
				hex.EncodeToString(decrypted), hex.EncodeToString(testCase.Msg))
		}
	case "invalid":
		if err == nil && bytes.Equal(ct, decrypted) {
			t.Error("successfully decrypted invalid test-case")
		}
	default:
		t.Errorf("unknown test-case result: %s", testCase.Result)
	}
}