// Copyright (C) MongoDB, Inc. 2017-present.
//
// 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

package mongo

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"log"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo/options"
)

func Example_clientSideEncryption() {
	// This would have to be the same master key that was used to create the
	// encryption key.
	localKey := make([]byte, 96)
	if _, err := rand.Read(localKey); err != nil {
		log.Fatal(err)
	}
	kmsProviders := map[string]map[string]interface{}{
		"local": {
			"key": localKey,
		},
	}
	keyVaultNamespace := "encryption.__keyVault"

	uri := "mongodb://localhost:27017"
	autoEncryptionOpts := options.AutoEncryption().
		SetKeyVaultNamespace(keyVaultNamespace).
		SetKmsProviders(kmsProviders)
	clientOpts := options.Client().
		ApplyURI(uri).
		SetAutoEncryptionOptions(autoEncryptionOpts)
	client, err := Connect(context.TODO(), clientOpts)
	if err != nil {
		log.Fatalf("Connect error: %v", err)
	}
	defer func() {
		if err = client.Disconnect(context.TODO()); err != nil {
			log.Fatalf("Disconnect error: %v", err)
		}
	}()

	collection := client.Database("test").Collection("coll")
	if err := collection.Drop(context.TODO()); err != nil {
		log.Fatalf("Collection.Drop error: %v", err)
	}

	_, err = collection.InsertOne(
		context.TODO(),
		bson.D{{"encryptedField", "123456789"}})
	if err != nil {
		log.Fatalf("InsertOne error: %v", err)
	}
	res, err := collection.FindOne(context.TODO(), bson.D{}).Raw()
	if err != nil {
		log.Fatalf("FindOne error: %v", err)
	}
	fmt.Println(res)
}

func Example_clientSideEncryptionCreateKey() {
	keyVaultNamespace := "encryption.__keyVault"
	uri := "mongodb://localhost:27017"
	// kmsProviders would have to be populated with the correct KMS provider
	// information before it's used.
	var kmsProviders map[string]map[string]interface{}

	// Create Client and ClientEncryption
	clientEncryptionOpts := options.ClientEncryption().
		SetKeyVaultNamespace(keyVaultNamespace).
		SetKmsProviders(kmsProviders)
	keyVaultClient, err := Connect(
		context.TODO(),
		options.Client().ApplyURI(uri))
	if err != nil {
		log.Fatalf("Connect error for keyVaultClient: %v", err)
	}
	clientEnc, err := NewClientEncryption(keyVaultClient, clientEncryptionOpts)
	if err != nil {
		log.Fatalf("NewClientEncryption error: %v", err)
	}
	defer func() {
		// this will disconnect the keyVaultClient as well
		if err = clientEnc.Close(context.TODO()); err != nil {
			log.Fatalf("Close error: %v", err)
		}
	}()

	// Create a new data key and encode it as base64
	dataKeyID, err := clientEnc.CreateDataKey(context.TODO(), "local")
	if err != nil {
		log.Fatalf("CreateDataKey error: %v", err)
	}
	dataKeyBase64 := base64.StdEncoding.EncodeToString(dataKeyID.Data)

	// Create a JSON schema using the new data key. This schema could also be
	// written in a separate file and read in using I/O functions.
	schema := `{
		"properties": {
			"encryptedField": {
				"encrypt": {
					"keyId": [{
						"$binary": {
							"base64": "%s",
							"subType": "04"
						}
					}],
					"bsonType": "string",
					"algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic"
				}
			}
		},
		"bsonType": "object"
	}`
	schema = fmt.Sprintf(schema, dataKeyBase64)
	var schemaDoc bson.Raw
	err = bson.UnmarshalExtJSON([]byte(schema), true, &schemaDoc)
	if err != nil {
		log.Fatalf("UnmarshalExtJSON error: %v", err)
	}

	// Configure a Client with auto encryption using the new schema
	dbName := "test"
	collName := "coll"
	schemaMap := map[string]interface{}{
		dbName + "." + collName: schemaDoc,
	}
	autoEncryptionOpts := options.AutoEncryption().
		SetKmsProviders(kmsProviders).
		SetKeyVaultNamespace(keyVaultNamespace).
		SetSchemaMap(schemaMap)

	clientOptions := options.Client().
		ApplyURI(uri).
		SetAutoEncryptionOptions(autoEncryptionOpts)
	client, err := Connect(context.TODO(), clientOptions)
	if err != nil {
		log.Fatalf("Connect error for encrypted client: %v", err)
	}
	defer func() {
		_ = client.Disconnect(context.TODO())
	}()

	// Use client for operations.
}

func Example_explictEncryption() {
	// localMasterKey must be the same master key that was used to create the
	// encryption key.
	var localMasterKey []byte
	kmsProviders := map[string]map[string]interface{}{
		"local": {
			"key": localMasterKey,
		},
	}

	// The MongoDB namespace (db.collection) used to store the encryption data
	// keys.
	keyVaultDBName, keyVaultCollName := "encryption", "testKeyVault"
	keyVaultNamespace := keyVaultDBName + "." + keyVaultCollName

	// The Client used to read/write application data.
	client, err := Connect(
		context.TODO(),
		options.Client().ApplyURI("mongodb://localhost:27017"))
	if err != nil {
		panic(err)
	}
	defer func() { _ = client.Disconnect(context.TODO()) }()

	// Get a handle to the application collection and clear existing data.
	coll := client.Database("test").Collection("coll")
	_ = coll.Drop(context.TODO())

	// Set up the key vault for this example.
	keyVaultColl := client.Database(keyVaultDBName).Collection(keyVaultCollName)
	_ = keyVaultColl.Drop(context.TODO())
	// Ensure that two data keys cannot share the same keyAltName.
	keyVaultIndex := IndexModel{
		Keys: bson.D{{"keyAltNames", 1}},
		Options: options.Index().
			SetUnique(true).
			SetPartialFilterExpression(bson.D{
				{"keyAltNames", bson.D{
					{"$exists", true},
				}},
			}),
	}
	_, err = keyVaultColl.Indexes().CreateOne(context.TODO(), keyVaultIndex)
	if err != nil {
		panic(err)
	}

	// Create the ClientEncryption object to use for explicit
	// encryption/decryption. The Client passed to NewClientEncryption is used
	// to read/write to the key vault. This can be the same Client used by the
	// main application.
	clientEncryptionOpts := options.ClientEncryption().
		SetKmsProviders(kmsProviders).
		SetKeyVaultNamespace(keyVaultNamespace)
	clientEncryption, err := NewClientEncryption(client, clientEncryptionOpts)
	if err != nil {
		panic(err)
	}
	defer func() { _ = clientEncryption.Close(context.TODO()) }()

	// Create a new data key for the encrypted field.
	dataKeyOpts := options.DataKey().
		SetKeyAltNames([]string{"go_encryption_example"})
	dataKeyID, err := clientEncryption.CreateDataKey(
		context.TODO(),
		"local",
		dataKeyOpts)
	if err != nil {
		panic(err)
	}

	// Create a bson.RawValue to encrypt and encrypt it using the key that was
	// just created.
	rawValueType, rawValueData, err := bson.MarshalValue("123456789")
	if err != nil {
		panic(err)
	}
	rawValue := bson.RawValue{Type: rawValueType, Value: rawValueData}
	encryptionOpts := options.Encrypt().
		SetAlgorithm("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic").
		SetKeyID(dataKeyID)
	encryptedField, err := clientEncryption.Encrypt(
		context.TODO(),
		rawValue,
		encryptionOpts)
	if err != nil {
		panic(err)
	}

	// Insert a document with the encrypted field and then find it.
	_, err = coll.InsertOne(
		context.TODO(),
		bson.D{{"encryptedField", encryptedField}})
	if err != nil {
		panic(err)
	}
	var foundDoc bson.M
	err = coll.FindOne(context.TODO(), bson.D{}).Decode(&foundDoc)
	if err != nil {
		panic(err)
	}

	// Decrypt the encrypted field in the found document.
	decrypted, err := clientEncryption.Decrypt(
		context.TODO(),
		foundDoc["encryptedField"].(primitive.Binary))
	if err != nil {
		panic(err)
	}
	fmt.Printf("Decrypted value: %s\n", decrypted)
}

func Example_explictEncryptionWithAutomaticDecryption() {
	// Automatic encryption requires MongoDB 4.2 enterprise, but automatic
	// decryption is supported for all users.

	// localMasterKey must be the same master key that was used to create the
	// encryption key.
	var localMasterKey []byte
	kmsProviders := map[string]map[string]interface{}{
		"local": {
			"key": localMasterKey,
		},
	}

	// The MongoDB namespace (db.collection) used to store the encryption data
	// keys.
	keyVaultDBName, keyVaultCollName := "encryption", "testKeyVault"
	keyVaultNamespace := keyVaultDBName + "." + keyVaultCollName

	// Create the Client for reading/writing application data. Configure it with
	// BypassAutoEncryption=true to disable automatic encryption but keep
	// automatic decryption. Setting BypassAutoEncryption will also bypass
	// spawning mongocryptd in the driver.
	autoEncryptionOpts := options.AutoEncryption().
		SetKmsProviders(kmsProviders).
		SetKeyVaultNamespace(keyVaultNamespace).
		SetBypassAutoEncryption(true)
	clientOpts := options.Client().
		ApplyURI("mongodb://localhost:27017").
		SetAutoEncryptionOptions(autoEncryptionOpts)
	client, err := Connect(context.TODO(), clientOpts)
	if err != nil {
		panic(err)
	}
	defer func() { _ = client.Disconnect(context.TODO()) }()

	// Get a handle to the application collection and clear existing data.
	coll := client.Database("test").Collection("coll")
	_ = coll.Drop(context.TODO())

	// Set up the key vault for this example.
	keyVaultColl := client.Database(keyVaultDBName).Collection(keyVaultCollName)
	_ = keyVaultColl.Drop(context.TODO())
	// Ensure that two data keys cannot share the same keyAltName.
	keyVaultIndex := IndexModel{
		Keys: bson.D{{"keyAltNames", 1}},
		Options: options.Index().
			SetUnique(true).
			SetPartialFilterExpression(bson.D{
				{"keyAltNames", bson.D{
					{"$exists", true},
				}},
			}),
	}

	_, err = keyVaultColl.Indexes().CreateOne(context.TODO(), keyVaultIndex)
	if err != nil {
		panic(err)
	}

	// Create the ClientEncryption object to use for explicit
	// encryption/decryption. The Client passed to NewClientEncryption is used
	// to read/write to the key vault. This can be the same Client used by the
	// main application.
	clientEncryptionOpts := options.ClientEncryption().
		SetKmsProviders(kmsProviders).
		SetKeyVaultNamespace(keyVaultNamespace)
	clientEncryption, err := NewClientEncryption(client, clientEncryptionOpts)
	if err != nil {
		panic(err)
	}
	defer func() { _ = clientEncryption.Close(context.TODO()) }()

	// Create a new data key for the encrypted field.
	dataKeyOpts := options.DataKey().
		SetKeyAltNames([]string{"go_encryption_example"})
	dataKeyID, err := clientEncryption.CreateDataKey(
		context.TODO(),
		"local",
		dataKeyOpts)
	if err != nil {
		panic(err)
	}

	// Create a bson.RawValue to encrypt and encrypt it using the key that was
	// just created.
	rawValueType, rawValueData, err := bson.MarshalValue("123456789")
	if err != nil {
		panic(err)
	}
	rawValue := bson.RawValue{Type: rawValueType, Value: rawValueData}
	encryptionOpts := options.Encrypt().
		SetAlgorithm("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic").
		SetKeyID(dataKeyID)
	encryptedField, err := clientEncryption.Encrypt(
		context.TODO(),
		rawValue,
		encryptionOpts)
	if err != nil {
		panic(err)
	}

	// Insert a document with the encrypted field and then find it. The FindOne
	// call will automatically decrypt the field in the document.
	_, err = coll.InsertOne(
		context.TODO(),
		bson.D{{"encryptedField", encryptedField}})
	if err != nil {
		panic(err)
	}
	var foundDoc bson.M
	err = coll.FindOne(context.TODO(), bson.D{}).Decode(&foundDoc)
	if err != nil {
		panic(err)
	}
	fmt.Printf("Decrypted document: %v\n", foundDoc)
}