// 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 awskms_test

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"os"
	"path/filepath"
	"testing"

	// Placeholder for internal flag import.
	// context is used to cancel outstanding requests
	"github.com/tink-crypto/tink-go/v2/aead"
	"github.com/tink-crypto/tink-go/v2/tink"

	"github.com/tink-crypto/tink-go-awskms/v3/integration/awskms"
)

const (
	keyPrefix   = "aws-kms://arn:aws:kms:us-east-2:235739564943:"
	keyAliasURI = "aws-kms://arn:aws:kms:us-east-2:235739564943:alias/unit-and-integration-testing"
	keyID       = "arn:aws:kms:us-east-2:235739564943:key/3ee50705-5a82-4f5b-9753-05c4f473922f"
	keyURI      = "aws-kms://arn:aws:kms:us-east-2:235739564943:key/3ee50705-5a82-4f5b-9753-05c4f473922f"
	keyURI2     = "aws-kms://arn:aws:kms:us-east-2:235739564943:key/b3ca2efd-a8fb-47f2-b541-7e20f8c5cd11"
)

const (
	credCSVFile = "testdata/aws/credentials.csv"
	credINIFile = "testdata/aws/credentials.ini"
)

// Placeholder for internal initialization.

func testFilePath(t *testing.T, filename string) string {
	t.Helper()
	srcDir, ok := os.LookupEnv("TEST_SRCDIR")
	if ok {
		workspaceDir, ok := os.LookupEnv("TEST_WORKSPACE")
		if !ok {
			t.Fatal("TEST_WORKSPACE not found")
		}
		return filepath.Join(srcDir, workspaceDir, filename)
	}
	return filepath.Join("../../..", filename)
}

func TestNewClientWithCredentialsGetAEADEncryptDecrypt(t *testing.T) {
	credFilePath := testFilePath(t, credCSVFile)
	client, err := awskms.NewClientWithOptions(t.Context(), keyURI, awskms.WithCredentialPath(credFilePath))
	if err != nil {
		t.Fatalf("error setting up AWS client: %v", err)
	}
	a, err := client.GetAEAD(keyURI)
	if err != nil {
		t.Fatalf("client.GetAEAD(keyURI) err = %v, want nil", err)
	}
	plaintext := []byte("plaintext")
	associatedData := []byte("associatedData")
	ciphertext, err := a.Encrypt(plaintext, associatedData)
	if err != nil {
		t.Fatalf("a.Encrypt(plaintext, associatedData) err = %v, want nil", err)
	}

	t.Logf("ciphertext: %s", hex.EncodeToString(ciphertext))

	gotPlaintext, err := a.Decrypt(ciphertext, associatedData)
	if err != nil {
		t.Fatalf("a.Decrypt(ciphertext, associatedData) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintext, plaintext) {
		t.Errorf("a.Decrypt() = %q, want %q", gotPlaintext, plaintext)
	}

	invalidAssociatedData := []byte("invalidAssociatedData")
	_, err = a.Decrypt(ciphertext, invalidAssociatedData)
	if err == nil {
		t.Error("a.Decrypt(ciphertext, invalidAssociatedData) err = nil, want error")
	}
}

func TestNewAEADWithContextEncryptDecrypt(t *testing.T) {
	credFilePath := testFilePath(t, credCSVFile)
	a, err := awskms.NewAEADWithContext(t.Context(), keyID, awskms.WithCredentialPath(credFilePath))
	if err != nil {
		t.Fatalf("awskms.NewAEADWithContext() err = %v", err)
	}
	plaintext := []byte("plaintext")
	associatedData := []byte("associatedData")
	ciphertext, err := a.EncryptWithContext(t.Context(), plaintext, associatedData)
	if err != nil {
		t.Fatalf("a.Encrypt(plaintext, associatedData) err = %v, want nil", err)
	}

	gotPlaintext, err := a.DecryptWithContext(t.Context(), ciphertext, associatedData)
	if err != nil {
		t.Fatalf("a.Decrypt(ciphertext, associatedData) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintext, plaintext) {
		t.Errorf("a.Decrypt() = %q, want %q", gotPlaintext, plaintext)
	}

	invalidAssociatedData := []byte("invalidAssociatedData")
	_, err = a.DecryptWithContext(t.Context(), ciphertext, invalidAssociatedData)
	if err == nil {
		t.Error("a.Decrypt(ciphertext, invalidAssociatedData) err = nil, want error")
	}

	// Check that the context is not ignored.
	cancelledCtx, cancel := context.WithCancelCause(t.Context())
	causeErr := errors.New("cause error message")
	cancel(causeErr)
	_, err = a.EncryptWithContext(cancelledCtx, plaintext, associatedData)
	if err == nil {
		t.Error("a.EncryptWithContext(cancelledCtx, plaintext, associatedData) err = nil, want error")
	}
	_, err = a.DecryptWithContext(cancelledCtx, ciphertext, associatedData)
	if err == nil {
		t.Error("a.DecryptWithContext(cancelledCtx, ciphertext, associatedData) err = nil, want error")
	}

	// Check that it is compatible with the AEAD returned by NewClientWithOptions.
	client, err := awskms.NewClientWithOptions(t.Context(), keyURI, awskms.WithCredentialPath(credFilePath))
	if err != nil {
		t.Fatalf("error setting up AWS client: %v", err)
	}
	clientAEAD, err := client.GetAEAD(keyURI)
	if err != nil {
		t.Fatalf("client.GetAEAD(keyURI) err = %v, want nil", err)
	}
	gotPlaintextFromClient, err := clientAEAD.Decrypt(ciphertext, associatedData)
	if err != nil {
		t.Fatalf("a.Decrypt(ciphertext, associatedData) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintextFromClient, plaintext) {
		t.Errorf("a.Decrypt() = %q, want %q", gotPlaintextFromClient, plaintext)
	}
}

func TestEmptyAssociatedDataEncryptDecrypt(t *testing.T) {
	credFilePath := testFilePath(t, credCSVFile)
	client, err := awskms.NewClientWithOptions(t.Context(), keyURI, awskms.WithCredentialPath(credFilePath))
	if err != nil {
		t.Fatalf("error setting up AWS client: %v", err)
	}
	a, err := client.GetAEAD(keyURI)
	if err != nil {
		t.Fatalf("client.GetAEAD(keyURI) err = %v, want nil", err)
	}
	plaintext := []byte("plaintext")
	emptyAssociatedData := []byte{}
	ciphertext, err := a.Encrypt(plaintext, emptyAssociatedData)
	if err != nil {
		t.Fatalf("a.Encrypt(plaintext, emptyAssociatedData) err = %v, want nil", err)
	}
	gotPlaintext, err := a.Decrypt(ciphertext, emptyAssociatedData)
	if err != nil {
		t.Fatalf("a.Decrypt(ciphertext, associatedData) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintext, plaintext) {
		t.Errorf("a.Decrypt() = %q, want %q", gotPlaintext, plaintext)
	}

	gotPlaintext2, err := a.Decrypt(ciphertext, nil)
	if err != nil {
		t.Fatalf("a.Decrypt(ciphertext, nil) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintext2, plaintext) {
		t.Errorf("a.Decrypt() = %q, want %q", gotPlaintext, plaintext)
	}
}

func TestKeyCommitment(t *testing.T) {
	credFilePath := testFilePath(t, credCSVFile)
	client, err := awskms.NewClientWithOptions(t.Context(), keyPrefix, awskms.WithCredentialPath(credFilePath))
	if err != nil {
		t.Fatalf("error setting up AWS client: %v", err)
	}

	// Create AEAD primitives for two keys.
	keys := []string{keyURI, keyURI2}
	aeads := make([]tink.AEAD, 0, len(keys))
	for _, k := range keys {
		a, err := client.GetAEAD(k)
		if err != nil {
			t.Fatalf("client.GetAEAD(keyURI) err = %v, want nil", err)
		}
		aeads = append(aeads, a)
	}

	// Create a ciphertext using the first primitive.
	plaintext := []byte("plaintext")
	associatedData := []byte("associated data")
	ciphertext, err := aeads[0].Encrypt(plaintext, associatedData)
	if err != nil {
		t.Fatalf("aeads[0].Encrypt(plaintext, associatedData) err = %v, want nil", err)
	}
	gotPlaintext, err := aeads[0].Decrypt(ciphertext, associatedData)
	if err != nil {
		t.Fatalf("aeads[0].Decrypt(ciphertext, associatedData) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintext, plaintext) {
		t.Errorf("aeads[0].Decrypt() = %q, want %q", gotPlaintext, plaintext)
	}

	// Attempt to decrypt using the other primitive.
	_, err = aeads[1].Decrypt(ciphertext, associatedData)
	if err == nil {
		t.Fatalf("aeads[1].Decrypt(ciphertext, associatedData) err = nil, want non-nil")
	}
}

func TestKMSEnvelopeAEADEncryptAndDecrypt(t *testing.T) {
	for _, credFile := range []string{credCSVFile, credINIFile} {
		credFilePath := testFilePath(t, credFile)
		client, err := awskms.NewClientWithOptions(t.Context(), keyURI, awskms.WithCredentialPath(credFilePath))
		if err != nil {
			t.Fatalf("awskms.NewClientWithOptions() err = %q, want nil", err)
		}

		kekAEAD, err := client.GetAEAD(keyURI)
		if err != nil {
			t.Fatalf("client.GetAEAD(keyURI) err = %q, want nil", err)
		}

		dekTemplate := aead.AES128CTRHMACSHA256KeyTemplate()
		a := aead.NewKMSEnvelopeAEAD2(dekTemplate, kekAEAD)
		if err != nil {
			t.Fatalf("aead.NewKMSEnvelopeAEAD2(dekTemplate, kekAEAD) err = %q, want nil", err)
		}
		plaintext := []byte("plaintext")
		for _, associatedData := range [][]byte{nil, []byte("associated data")} {
			ciphertext, err := a.Encrypt(plaintext, associatedData)
			if err != nil {
				t.Fatalf("a.Encrypt(plaintext, associatedData) err = %q, want nil", err)
			}
			gotPlaintext, err := a.Decrypt(ciphertext, associatedData)
			if err != nil {
				t.Fatalf("a.Decrypt(ciphertext, associatedData) err = %q, want nil", err)
			}
			if !bytes.Equal(gotPlaintext, plaintext) {
				t.Errorf("a.Decrypt() = %q, want %q", gotPlaintext, plaintext)
			}
		}
	}
}

func TestDecryptValidCiphertextForKeyURI(t *testing.T) {
	// This is a valid ciphertext for keyURI that was generated with
	// TestNewClientWithCredentialsGetAEADEncryptDecrypt in aws_kms_integration_test.go of the v2 package.
	// If the key changes, this ciphertext needs to be updated.
	ciphertext, err := hex.DecodeString("01020200787a5511910e44ea9520df697fb6150f6261b59cee44cc846915252c083f03aad00168fe1a9ca1a6af5d571c1aa6e0afcba000000067306506092a864886f70d010706a0583056020100305106092a864886f70d010701301e060960864801650304012e3011040cb085893bc4757850d73a68ca02011080241fcff6b61f510d20c236ee7292b626e35a22b615e57326911daafaa399525431748d3e24")
	if err != nil {
		t.Fatalf("hex.DecodeString() err = %v, want nil", err)
	}
	wantPlaintext := []byte("plaintext")
	associatedData := []byte("associatedData")
	credFilePath := testFilePath(t, credCSVFile)

	client, err := awskms.NewClientWithOptions(t.Context(), keyURI, awskms.WithCredentialPath(credFilePath))
	if err != nil {
		t.Fatalf("error setting up AWS client: %v", err)
	}
	a, err := client.GetAEAD(keyURI)
	if err != nil {
		t.Fatalf("client.GetAEAD(keyURI) err = %v, want nil", err)
	}
	gotPlaintext, err := a.Decrypt(ciphertext, associatedData)
	if err != nil {
		t.Fatalf("a.Decrypt(ciphertext, associatedData) err = %v, want nil", err)
	}
	if !bytes.Equal(gotPlaintext, wantPlaintext) {
		t.Errorf("a.Decrypt() = %q, want %q", gotPlaintext, wantPlaintext)
	}
}
