package com.quorum.tessera.enclave;

import static java.util.Collections.*;
import static org.assertj.core.api.Assertions.*;
import static org.mockito.Mockito.*;

import com.quorum.tessera.encryption.*;
import com.quorum.tessera.service.Service;
import java.nio.ByteBuffer;
import java.util.*;
import org.bouncycastle.jcajce.provider.digest.SHA3;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

public class EnclaveTest {

  private Enclave enclave;

  private Encryptor nacl;

  private KeyManager keyManager;

  @Before
  public void onSetUp() {
    this.nacl = mock(Encryptor.class);
    this.keyManager = mock(KeyManager.class);

    this.enclave = new EnclaveImpl(nacl, keyManager);
    enclave.start();
    assertThat(enclave.status()).isEqualTo(Service.Status.STARTED);
  }

  @After
  public void onTearDown() {
    enclave.stop();
    assertThat(enclave.status()).isEqualTo(Service.Status.STARTED);
    verifyNoMoreInteractions(nacl, keyManager);
  }

  @Test
  public void defaultPublicKey() {
    enclave.defaultPublicKey();
    verify(keyManager).defaultPublicKey();
  }

  @Test
  public void getForwardingKeys() {
    enclave.getForwardingKeys();
    verify(keyManager).getForwardingKeys();
  }

  @Test
  public void getPublicKeys() {
    enclave.getPublicKeys();
    verify(keyManager).getPublicKeys();
  }

  // Case 1
  @Test
  public void unencryptPSVTransaction() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("RecipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientNonce(recipientNonce)
            .withRecipientKeys(List.of(recipientKey, senderKey))
            .withPrivacyMode(PrivacyMode.PRIVATE_STATE_VALIDATION)
            .withExecHash("EXEC_HASH".getBytes())
            .build();

    final PrivateKey recipientPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(recipientKey)).thenReturn(recipientPrivateKey);

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(senderKey, recipientPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, recipientKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(senderKey, recipientPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(recipientKey);
  }

  // Case 1 - error with unknown key
  @Test
  public void unencryptPSVTransactionWithUnknownKey() {
    final PublicKey nonRecipientKey = PublicKey.from("unknown".getBytes());

    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("RecipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientNonce(recipientNonce)
            .withRecipientKeys(List.of(recipientKey, senderKey))
            .withPrivacyMode(PrivacyMode.PRIVATE_STATE_VALIDATION)
            .withExecHash("EXEC_HASH".getBytes())
            .build();

    final Throwable throwable =
        catchThrowable(() -> enclave.unencryptTransaction(payload, nonRecipientKey));

    assertThat(throwable)
        .isInstanceOf(EnclaveException.class)
        .hasMessage("recipient not found in listed keys");
  }

  // Case 2
  @Test
  public void unencryptNoRecipientsSingleBoxTransaction() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("RecipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final PrivateKey recipientPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(recipientKey)).thenReturn(recipientPrivateKey);

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(senderKey, recipientPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, recipientKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(senderKey, recipientPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(recipientKey);
  }

  // Case 3 - decrypt using the recipient key
  @Test
  public void unencryptTransactionWeSentUsingRecipientKey() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final RecipientBox senderBox = RecipientBox.from("senderBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBoxes(List.of(senderBox.getData(), recipientBox.getData()))
            .withRecipientKeys(List.of(senderKey, recipientKey))
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final PrivateKey recipientPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(recipientKey)).thenReturn(recipientPrivateKey);

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(senderKey, recipientPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, recipientKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(senderKey, recipientPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(recipientKey);
  }

  // Case 3 - decrypt using the sender key
  @Test
  public void unencryptTransactionWeSentUsingSenderKey() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final RecipientBox senderBox = RecipientBox.from("senderBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBoxes(List.of(senderBox.getData(), recipientBox.getData()))
            .withRecipientKeys(List.of(senderKey, recipientKey))
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final PrivateKey senderPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(senderKey)).thenReturn(senderPrivateKey);

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(senderKey, senderPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(senderBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, senderKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(senderBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(senderKey, senderPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(senderKey);
  }

  // Case 3 - error when decrypt using the unknown key
  @Test
  public void unencryptTransactionWeSentUsingUnknownKey() {
    final PublicKey nonRecipientKey = PublicKey.from("unknown".getBytes());

    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final RecipientBox senderBox = RecipientBox.from("senderBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBoxes(List.of(senderBox.getData(), recipientBox.getData()))
            .withRecipientKeys(List.of(senderKey, recipientKey))
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final Throwable throwable =
        catchThrowable(() -> enclave.unencryptTransaction(payload, nonRecipientKey));

    assertThat(throwable)
        .isInstanceOf(EnclaveException.class)
        .hasMessage("recipient not found in listed keys");
  }

  // Case 5 - decrypt normal transaction which we didn't send
  @Test
  public void unencryptTransactionFromOtherNode() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientKey(recipientKey)
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final PrivateKey recipientPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(recipientKey)).thenReturn(recipientPrivateKey);
    when(keyManager.getPublicKeys()).thenReturn(Set.of(recipientKey));

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(senderKey, recipientPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, recipientKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(senderKey, recipientPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(recipientKey);
    verify(keyManager).getPublicKeys();
  }

  // Case 5 - error case where wrong key is used
  @Test
  public void unencryptTransactionFromOtherNodeUsingWrongKey() {
    final PublicKey nonRecipientKey = PublicKey.from("unknown".getBytes());

    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientKey(recipientKey)
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    when(keyManager.getPublicKeys()).thenReturn(Set.of(recipientKey));

    final Throwable throwable =
        catchThrowable(() -> enclave.unencryptTransaction(payload, nonRecipientKey));

    assertThat(throwable)
        .isInstanceOf(EnclaveException.class)
        .hasMessage("recipient not found in listed keys");

    verify(keyManager).getPublicKeys();
  }

  // Case 4.1 - decrypt transaction we sent from a very early version
  // where the sender didn't get an entry
  @Test
  public void unencryptTransactionFromEarlyVersionUsingSender() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientKey(recipientKey)
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final PrivateKey senderPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(senderKey)).thenReturn(senderPrivateKey);
    when(keyManager.getPublicKeys()).thenReturn(Set.of(senderKey));

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(recipientKey, senderPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, senderKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(recipientKey, senderPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(senderKey);
    verify(keyManager).getPublicKeys();
  }

  // Case 4.2 - decrypt transaction we sent from a very early version
  // where the sender didn't get an entry
  @Test
  public void unencryptTransactionFromEarlyVersionUsingRecipient() {
    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientKey(recipientKey)
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    final PrivateKey recipientPrivateKey = PrivateKey.from("private-key".getBytes());
    when(keyManager.getPrivateKeyForPublicKey(recipientKey)).thenReturn(recipientPrivateKey);
    when(keyManager.getPublicKeys()).thenReturn(Set.of(senderKey, recipientKey));

    final SharedKey sharedKey = SharedKey.from("shared-key".getBytes());
    when(nacl.computeSharedKey(senderKey, recipientPrivateKey)).thenReturn(sharedKey);

    final byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());
    when(nacl.openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    final byte[] result = enclave.unencryptTransaction(payload, recipientKey);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(recipientBox.getData(), recipientNonce, sharedKey);
    verify(nacl)
        .openAfterPrecomputation(any(byte[].class), eq(cipherTextNonce), any(MasterKey.class));
    verify(nacl).computeSharedKey(senderKey, recipientPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(recipientKey);
    verify(keyManager).getPublicKeys();
  }

  // Case 4.2 - error where given key wasn't recipient
  @Test
  public void unencryptTransactionFromEarlyVersionUsingUnknownRecipient() {
    final PublicKey nonRecipientKey = PublicKey.from("unknown".getBytes());

    final PublicKey senderKey = PublicKey.from("senderKey".getBytes());
    final PublicKey recipientKey = PublicKey.from("recipientKey".getBytes());
    final byte[] cipherText = "cipherText".getBytes();
    final Nonce cipherTextNonce = mock(Nonce.class);
    final RecipientBox recipientBox = RecipientBox.from("recipientBox".getBytes());
    final Nonce recipientNonce = mock(Nonce.class);

    final EncodedPayload payload =
        EncodedPayload.Builder.create()
            .withSenderKey(senderKey)
            .withCipherText(cipherText)
            .withCipherTextNonce(cipherTextNonce)
            .withRecipientBox(recipientBox.getData())
            .withRecipientKey(recipientKey)
            .withRecipientNonce(recipientNonce)
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .build();

    when(keyManager.getPublicKeys()).thenReturn(Set.of(senderKey, recipientKey));

    final Throwable throwable =
        catchThrowable(() -> enclave.unencryptTransaction(payload, nonRecipientKey));

    assertThat(throwable)
        .isInstanceOf(EnclaveException.class)
        .hasMessage("recipient not found in listed keys");

    verify(keyManager).getPublicKeys();
  }

  @Test
  public void unencryptRawPayload() {

    PublicKey senderKey = mock(PublicKey.class);

    byte[] cipherText = "cipherText".getBytes();

    byte[] recipientBox = "RecipientBox".getBytes();

    Nonce nonce = mock(Nonce.class);

    RawTransaction rawTransaction = new RawTransaction(cipherText, recipientBox, nonce, senderKey);

    PrivateKey senderPrivateKey = mock(PrivateKey.class);

    when(keyManager.getPrivateKeyForPublicKey(senderKey)).thenReturn(senderPrivateKey);

    SharedKey sharedKey = mock(SharedKey.class);
    when(nacl.computeSharedKey(senderKey, senderPrivateKey)).thenReturn(sharedKey);

    byte[] expectedOutcome = "SUCCESS".getBytes();

    when(nacl.openAfterPrecomputation(any(byte[].class), any(Nonce.class), any(SharedKey.class)))
        .thenReturn("sharedOrMasterKeyBytes".getBytes());

    when(nacl.openAfterPrecomputation(any(byte[].class), any(Nonce.class), any(MasterKey.class)))
        .thenReturn(expectedOutcome);

    byte[] result = enclave.unencryptRawPayload(rawTransaction);

    assertThat(result).isNotNull().isSameAs(expectedOutcome);

    verify(nacl).openAfterPrecomputation(any(byte[].class), any(Nonce.class), any(SharedKey.class));
    verify(nacl).openAfterPrecomputation(any(byte[].class), any(Nonce.class), any(MasterKey.class));
    verify(keyManager).getPrivateKeyForPublicKey(senderKey);
    verify(nacl).computeSharedKey(senderKey, senderPrivateKey);
  }

  @Test
  public void encryptPayload() {

    byte[] message = "MESSAGE".getBytes();

    PublicKey senderPublicKey = mock(PublicKey.class);
    PublicKey recipientPublicKey = mock(PublicKey.class);

    byte[] masterKeyBytes = "masterKeyBytes".getBytes();
    MasterKey masterKey = MasterKey.from(masterKeyBytes);
    Nonce cipherNonce = mock(Nonce.class);
    Nonce recipientNonce = mock(Nonce.class);

    byte[] cipherText = "cipherText".getBytes();

    when(nacl.createMasterKey()).thenReturn(masterKey);
    when(nacl.randomNonce()).thenReturn(cipherNonce, recipientNonce);

    when(nacl.sealAfterPrecomputation(message, cipherNonce, masterKey)).thenReturn(cipherText);

    PrivateKey senderPrivateKey = mock(PrivateKey.class);
    when(keyManager.getPrivateKeyForPublicKey(senderPublicKey)).thenReturn(senderPrivateKey);

    SharedKey sharedKey = mock(SharedKey.class);
    when(nacl.computeSharedKey(recipientPublicKey, senderPrivateKey)).thenReturn(sharedKey);

    byte[] encryptedMasterKeys = "encryptedMasterKeys".getBytes();
    when(nacl.sealAfterPrecomputation(masterKeyBytes, recipientNonce, sharedKey))
        .thenReturn(encryptedMasterKeys);

    final PrivacyMetadata metaData =
        PrivacyMetadata.Builder.create().withPrivacyMode(PrivacyMode.STANDARD_PRIVATE).build();

    EncodedPayload result =
        enclave.encryptPayload(
            message, senderPublicKey, Arrays.asList(recipientPublicKey), metaData);

    assertThat(result).isNotNull();
    assertThat(result.getRecipientKeys()).containsExactly(recipientPublicKey);
    assertThat(result.getCipherText()).isEqualTo(cipherText);
    assertThat(result.getCipherTextNonce()).isEqualTo(cipherNonce);
    assertThat(result.getSenderKey()).isEqualTo(senderPublicKey);
    assertThat(result.getRecipientBoxes()).containsExactly(RecipientBox.from(encryptedMasterKeys));
    assertThat(result.getPrivacyGroupId()).isNotPresent();

    verify(nacl).createMasterKey();
    verify(nacl, times(2)).randomNonce();
    verify(nacl).sealAfterPrecomputation(message, cipherNonce, masterKey);
    verify(nacl).sealAfterPrecomputation(masterKeyBytes, recipientNonce, sharedKey);
    verify(nacl).computeSharedKey(recipientPublicKey, senderPrivateKey);
    verify(keyManager).getPrivateKeyForPublicKey(senderPublicKey);
  }

  @Test
  public void encryptPayloadWithAffectedTransactions() {

    byte[] message = "MESSAGE".getBytes();

    PublicKey senderPublicKey = mock(PublicKey.class);
    PublicKey recipientPublicKey = mock(PublicKey.class);

    byte[] masterKeyBytes = "masterKeyBytes".getBytes();
    MasterKey masterKey = MasterKey.from(masterKeyBytes);
    Nonce cipherNonce = mock(Nonce.class);
    Nonce recipientNonce = mock(Nonce.class);
    final RecipientBox closedbox = RecipientBox.from("closed".getBytes());
    final byte[] openbox = "open".getBytes();
    byte[] cipherText = "cipherText".getBytes();

    when(nacl.createMasterKey()).thenReturn(masterKey);
    when(nacl.randomNonce()).thenReturn(cipherNonce, recipientNonce);

    when(nacl.sealAfterPrecomputation(message, cipherNonce, masterKey)).thenReturn(cipherText);
    when(keyManager.getPublicKeys()).thenReturn(Collections.singleton(senderPublicKey));

    PrivateKey senderPrivateKey = mock(PrivateKey.class);
    when(keyManager.getPrivateKeyForPublicKey(senderPublicKey)).thenReturn(senderPrivateKey);

    SharedKey sharedKey = mock(SharedKey.class);
    when(nacl.computeSharedKey(recipientPublicKey, senderPrivateKey)).thenReturn(sharedKey);
    when(nacl.openAfterPrecomputation(closedbox.getData(), recipientNonce, sharedKey))
        .thenReturn(openbox);
    byte[] encryptedMasterKeys = "encryptedMasterKeys".getBytes();
    when(nacl.sealAfterPrecomputation(masterKeyBytes, recipientNonce, sharedKey))
        .thenReturn(encryptedMasterKeys);

    EncodedPayload affectedTxPayload = mock(EncodedPayload.class);
    when(affectedTxPayload.getSenderKey()).thenReturn(senderPublicKey);
    when(affectedTxPayload.getCipherText()).thenReturn(cipherText);
    when(affectedTxPayload.getCipherTextNonce()).thenReturn(cipherNonce);
    when(affectedTxPayload.getRecipientBoxes()).thenReturn(singletonList(closedbox));
    when(affectedTxPayload.getRecipientNonce()).thenReturn(recipientNonce);
    when(affectedTxPayload.getRecipientKeys()).thenReturn(singletonList(recipientPublicKey));
    when(affectedTxPayload.getPrivacyMode()).thenReturn(PrivacyMode.STANDARD_PRIVATE);
    when(affectedTxPayload.getAffectedContractTransactions()).thenReturn(emptyMap());
    when(affectedTxPayload.getExecHash()).thenReturn(new byte[0]);

    TxHash txnHash =
        new TxHash(
            "bfMIqWJ/QGQhkK4USxMBxduzfgo/SIGoCros5bWYfPKUBinlAUCqLVOUAP9q+BgLlsWni1M6rnzfmaqSw2J5hQ==");

    AffectedTransaction affectedTransaction = mock(AffectedTransaction.class);
    when(affectedTransaction.getHash()).thenReturn(txnHash);
    when(affectedTransaction.getPayload()).thenReturn(affectedTxPayload);

    List<AffectedTransaction> affectedContractTransactions = List.of(affectedTransaction);

    final PrivacyMetadata metaData =
        PrivacyMetadata.Builder.create()
            .withPrivacyMode(PrivacyMode.STANDARD_PRIVATE)
            .withAffectedTransactions(affectedContractTransactions)
            .build();

    final EncodedPayload result =
        enclave.encryptPayload(
            message, senderPublicKey, Arrays.asList(recipientPublicKey), metaData);

    assertThat(result).isNotNull();
    assertThat(result.getRecipientKeys()).containsExactly(recipientPublicKey);
    assertThat(result.getCipherText()).isEqualTo(cipherText);
    assertThat(result.getCipherTextNonce()).isEqualTo(cipherNonce);
    assertThat(result.getSenderKey()).isEqualTo(senderPublicKey);
    assertThat(result.getRecipientBoxes()).containsExactly(RecipientBox.from(encryptedMasterKeys));
    assertThat(result.getAffectedContractTransactions()).containsOnlyKeys(txnHash);

    verify(nacl).createMasterKey();
    verify(nacl, times(2)).randomNonce();
    verify(nacl).sealAfterPrecomputation(message, cipherNonce, masterKey);
    verify(nacl).sealAfterPrecomputation(masterKeyBytes, recipientNonce, sharedKey);
    verify(nacl).openAfterPrecomputation(closedbox.getData(), recipientNonce, sharedKey);
    verify(nacl, times(2)).computeSharedKey(recipientPublicKey, senderPrivateKey);
    verify(keyManager, times(2)).getPrivateKeyForPublicKey(senderPublicKey);
    verify(keyManager).getPublicKeys();
  }

  @Test
  public void encryptPayloadWithMandatoryRecipients() {

    byte[] message = "MESSAGE".getBytes();

    PublicKey senderPublicKey = mock(PublicKey.class);
    PublicKey recipientPublicKey = mock(PublicKey.class);

    byte[] masterKeyBytes = "masterKeyBytes".getBytes();
    MasterKey masterKey = MasterKey.from(masterKeyBytes);
    Nonce cipherNonce = mock(Nonce.class);
    Nonce recipientNonce = mock(Nonce.class);

    byte[] cipherText = "cipherText".getBytes();

    when(nacl.createMasterKey()).thenReturn(masterKey);
    when(nacl.randomNonce()).thenReturn(cipherNonce, recipientNonce);

    when(nacl.sealAfterPrecomputation(message, cipherNonce, masterKey)).thenReturn(cipherText);

    PrivateKey senderPrivateKey = mock(PrivateKey.class);
    when(keyManager.getPrivateKeyForPublicKey(senderPublicKey)).thenReturn(senderPrivateKey);

    SharedKey sharedKey = mock(SharedKey.class);
    when(nacl.computeSharedKey(recipientPublicKey, senderPrivateKey)).thenReturn(sharedKey);

    byte[] encryptedMasterKeys = "encryptedMasterKeys".getBytes();
    when(nacl.sealAfterPrecomputation(masterKeyBytes, recipientNonce, sharedKey))
        .thenReturn(encryptedMasterKeys);

    final PrivacyMetadata metaData =
        PrivacyMetadata.Builder.create()
            .withPrivacyMode(PrivacyMode.MANDATORY_RECIPIENTS)
            .withMandatoryRecipients(Set.of(recipientPublicKey))
            .build();

    EncodedPayload result =
        enclave.encryptPayload(
            message, senderPublicKey, Arrays.asList(recipientPublicKey), metaData);

    assertThat(result).isNotNull();
    assertThat(result.getRecipientKeys()).containsExactly(recipientPublicKey);
    assertThat(result.getCipherText()).isEqualTo(cipherText);
    assertThat(result.getCipherTextNonce()).isEqualTo(cipherNonce);
    assertThat(result.getSenderKey()).isEqualTo(senderPublicKey);
    assertThat(result.getRecipientBoxes()).containsExactly(RecipientBox.from(encryptedMasterKeys));
    assertThat(result.getPrivacyGroupId()).isNotPresent();
    assertThat(result.getMandatoryRecipients()).hasSize(1);
    assertThat(result.getMandatoryRecipients()).containsExactly(recipientPublicKey);

    verify(nacl).createMasterKey();
    verify(nacl, times(2)).randomNonce();
    verify(nacl).sealAfterPrecomputation(message, cipherNonce, masterKey);
    verify(nacl).sealAfterPrecomputation(masterKeyBytes, recipientNonce, sharedKey);
    verify(nacl).computeSharedKey(recipientPublicKey, senderPrivateKey);
