GetKeys.dfy

// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

include "../Model/AwsCryptographyKeyStoreTypes.dfy"
include "Structure.dfy"
include "KMSKeystoreOperations.dfy"
include "ErrorMessages.dfy"
include "KmsArn.dfy"
include "../../AwsCryptographicMaterialProviders/src/CanonicalEncryptionContext.dfy"

module GetKeys {
  import opened StandardLibrary
  import opened Wrappers
  import opened Seq

  import Structure
  import CanonicalEncryptionContext
  import DefaultKeyStorageInterface
  import KMSKeystoreOperations
  import ErrorMessages = KeyStoreErrorMessages
  import AtomicPrimitives

  import Types = AwsCryptographyKeyStoreTypes
  import DDB = ComAmazonawsDynamodbTypes
  import KMS = ComAmazonawsKmsTypes
  import UTF8
  import KmsArn

  method GetActiveKeyAndUnwrap(
    input: Types.GetActiveBranchKeyInput,
    logicalKeyStoreName: string,
    kmsConfiguration: Types.KMSConfiguration,
    grantTokens: KMS.GrantTokenList,
    kmsClient: KMS.IKMSClient,
    storage: Types.IKeyStorageInterface
  )
    returns (output: Result<Types.GetActiveBranchKeyOutput, Types.Error>)
    requires storage.Modifies !! kmsClient.Modifies

    requires storage is DefaultKeyStorageInterface.DynamoDBKeyStorageInterface
             ==>
               logicalKeyStoreName == (storage as DefaultKeyStorageInterface.DynamoDBKeyStorageInterface).logicalKeyStoreName

    requires kmsClient.ValidState() && storage.ValidState()
    modifies storage.Modifies, kmsClient.Modifies
    ensures storage.ValidState() && kmsClient.ValidState()

    //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
    //= type=implication
    //# GetActiveBranchKey MUST get the active version for the branch key id from the keystore
    //# by calling the configured [KeyStorage interface's](./key-store/key-storage.md#interface)
    //# [GetEncryptedActiveBranchKey](./key-store/key-storage.md#getencryptedactivebranchkey)
    //# using the supplied `branch-key-id`.
    ensures
      && |storage.History.GetEncryptedActiveBranchKey| == |old(storage.History.GetEncryptedActiveBranchKey)| + 1
      && Seq.Last(storage.History.GetEncryptedActiveBranchKey).input.Identifier == input.branchKeyIdentifier

    ensures output.Success?
            ==>
              && Seq.Last(storage.History.GetEncryptedActiveBranchKey).output.Success?
              && var activeItem := Seq.Last(storage.History.GetEncryptedActiveBranchKey).output.value.Item;

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
              //= type=implication
              //# GetActiveBranchKey MUST verify that the returned EncryptedHierarchicalKey MUST have the requested `branch-key-id`.
              && activeItem.Identifier == input.branchKeyIdentifier

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
              //= type=implication
              //# GetActiveBranchKey MUST verify that the returned EncryptedHierarchicalKey is an ActiveHierarchicalSymmetricVersion.
              && Structure.ActiveHierarchicalSymmetricKey?(activeItem)

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
              //= type=implication
              //# GetActiveBranchKey MUST verify that the returned EncryptedHierarchicalKey MUST have a logical table name equal to the configured logical table name.
              && activeItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#discovery
              //= type=implication
              //# The Keystore MAY use the KMS Key ARNs already
              //# persisted to the backing DynamoDB table,
              //# provided they are in records created
              //# with an identical Logical Keystore Name.
              && (kmsConfiguration.kmsKeyArn? ==> activeItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName)

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#mrdiscovery
              //= type=implication
              //# The Keystore MAY use the KMS Key ARNs already
              //# persisted to the backing DynamoDB table,
              //# provided they are in records created
              //# with an identical Logical Keystore Name.
              && (kmsConfiguration.kmsMRKeyArn? ==> activeItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName)

              && KMSKeystoreOperations.AttemptKmsOperation?(kmsConfiguration, activeItem.EncryptionContext)
              && |kmsClient.History.Decrypt| == |old(kmsClient.History.Decrypt)| + 1

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
              //= type=implication
              //# The operation MUST decrypt the EncryptedHierarchicalKey according to the [AWS KMS Branch Key Decryption](#aws-kms-branch-key-decryption) section.
              && KMSKeystoreOperations.AwsKmsBranchKeyDecryption?(
                   activeItem,
                   kmsConfiguration,
                   grantTokens,
                   kmsClient,
                   Seq.Last(kmsClient.History.Decrypt)
                 )

              && var decryptResponse := Seq.Last(kmsClient.History.Decrypt).output.value;

              && Structure.ToBranchKeyMaterials(activeItem, decryptResponse.Plaintext.value).Success?

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
              //= type=implication
              //# This GetActiveBranchKey MUST construct [branch key materials](./structures.md#branch-key-materials)
              //# according to [Branch Key Materials From Authenticated Encryption Context](#branch-key-materials-from-authenticated-encryption-context).
              && var branchKeyMaterials :=  Structure.ToBranchKeyMaterials(
                                              activeItem,
                                              decryptResponse.Plaintext.value
                                            ).value;

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
              //= type=implication
              //# This operation MUST return the constructed [branch key materials](./structures.md#branch-key-materials).
              && output.value.branchKeyMaterials == branchKeyMaterials

              && output.value.branchKeyMaterials.branchKeyIdentifier == input.branchKeyIdentifier

    ensures
      || (&& |storage.History.GetEncryptedActiveBranchKey| == |old(storage.History.GetEncryptedActiveBranchKey)| + 1
          && Seq.Last(storage.History.GetEncryptedActiveBranchKey).output.Failure?
          ==> output.Failure?)

      || (&& |storage.History.GetEncryptedActiveBranchKey| == |old(storage.History.GetEncryptedActiveBranchKey)| + 1
          && Seq.Last(storage.History.GetEncryptedActiveBranchKey).output.Success?
          && !Structure.ActiveHierarchicalSymmetricKey?(Seq.Last(storage.History.GetEncryptedActiveBranchKey).output.value.Item)
          ==> output.Failure?)

      //= aws-encryption-sdk-specification/framework/branch-key-store.md#getactivebranchkey
      //= type=implication
      //# If the branch key fails to decrypt, GetActiveBranchKey MUST fail.
      || (&& |kmsClient.History.Decrypt| == |old(kmsClient.History.Decrypt)| + 1
          && Seq.Last(kmsClient.History.Decrypt).output.Failure?
          ==> output.Failure?)
  {

    var ActiveOutput :- storage.GetEncryptedActiveBranchKey(
      Types.GetEncryptedActiveBranchKeyInput(
        Identifier := input.branchKeyIdentifier
      )
    );

    var branchKeyItem := ActiveOutput.Item;

    :- Need(
      || storage is DefaultKeyStorageInterface.DynamoDBKeyStorageInterface
      || (
           && Structure.ActiveHierarchicalSymmetricKey?(branchKeyItem)
           && branchKeyItem.Identifier == input.branchKeyIdentifier
           && branchKeyItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName
         ),
      Types.KeyStoreException(
        message := ErrorMessages.INVALID_ACTIVE_BRANCH_KEY_FROM_STORAGE)
    );
    :- Need(
      branchKeyItem.EncryptionContext[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_1 ||
      branchKeyItem.EncryptionContext[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_2,
      Types.KeyStoreException(
        message := ErrorMessages.INVALID_HIERARCHY_VERSION
      )
    );
    if (branchKeyItem.EncryptionContext[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_1) {
      var branchKey: KMS.DecryptResponse :- KMSKeystoreOperations.DecryptKey(
        branchKeyItem,
        kmsConfiguration,
        grantTokens,
        kmsClient
      );

      var branchKeyMaterials :- Structure.ToBranchKeyMaterials(
        branchKeyItem,
        branchKey.Plaintext.value
      );
      return Success(
        Types.GetActiveBranchKeyOutput(
          branchKeyMaterials := branchKeyMaterials
        ));
    } else if (branchKeyItem.EncryptionContext[Structure.HIERARCHY_VERSION] == Structure.HIERARCHY_VERSION_2) {
      var hv2EC := getHV2EC(branchKeyItem.EncryptionContext);
      var hv2BranchKey := Types.EncryptedHierarchicalKey(
        Identifier := branchKeyItem.Identifier,
        Type := branchKeyItem.Type,
        CreateTime := branchKeyItem.CreateTime,
        KmsArn := branchKeyItem.KmsArn,
        EncryptionContext := hv2EC,
        CiphertextBlob := branchKeyItem.CiphertextBlob
      );
      var branchKey: KMS.DecryptResponse :- KMSKeystoreOperations.DecryptKey(
        hv2BranchKey,
        kmsConfiguration,
        grantTokens,
        kmsClient
      );
      mdDigestFromTable := getMdDigestFromEC(branchKeyItem.EncryptionContext)
      var utf8MdDigest :- UnstringifyEncryptionContext(mdDigestFromTable);
      var crypto := ProvideCryptoClient();
      if (crypto.Failure?) {
        var e := Types.KeyStoreException(
          message :=
            "Local Cryptography error: " + AtomicPrimitives.ErrorUtils.MessageOrUnknown(crypto.error));
        return Failure(e);
      }
      var mdDigestShaFromTable := CanonicalEncryptionContext.EncryptionContextDigest(crypto.value, utf8MdDigest);
      if (ecDigest.Failure?) {
        var e := Types.KeyStoreException(
          message :=
            "Failed to create mdDigest");
        return Failure(e);
      }
      var plaintextBranchKeyWithMdDigest := branchKey.Plaintext.value;
      var plaintextBranchKey := plaintextBranchKeyWithMdDigest[0..|plaintextBranchKeyWithMdDigest|-48];
      var decryptedMdDigest := plaintextBranchKeyWithMdDigest[|plaintextBranchKeyWithMdDigest|-48..];
      if (decryptedMdDigest != mdDigestShaFromTable) {
        var e := Types.KeyStoreException(
          message :=
            ErrorMessages.MD_DIGEST_SHA_NOT_MATCHED);
        return Failure(e);
      }
      var branchKeyMaterials :- Structure.ToBranchKeyMaterials(
        branchKeyItem,
        plaintextBranchKey
      );
      return Success(
        Types.GetActiveBranchKeyOutput(
          branchKeyMaterials := branchKeyMaterials
        ));
    }

  }

  method getMdDigestFromEC(
    item: Types.EncryptionContextString
  ) returns (output: Types.EncryptionContextString)
  {
    mdDigest := map k | k in Structure.BRANCH_KEY_RESTRICTED_FIELD_NAMES && k in item :: item[k];
  }

  method getHV2EC(
    item: Types.EncryptionContextString
  ) returns (output: Types.EncryptionContextString)
  {
    var withoutReserved := set k | 
      k in item && 
      k !in Structure.BRANCH_KEY_RESTRICTED_FIELD_NAMES :: k;
    var newMap := map[];
    var remaining := withoutReserved;
    while (remaining != {})
        decreases remaining
    {
        var key :| key in remaining;
        var value := item[key];
        remaining := remaining - {key};
        if (|key| >= |Structure.ENCRYPTION_CONTEXT_PREFIX| && key[..|Structure.ENCRYPTION_CONTEXT_PREFIX|] == Structure.ENCRYPTION_CONTEXT_PREFIX) {
            key := key[|Structure.ENCRYPTION_CONTEXT_PREFIX|..];
        }
        newMap := newMap + map[key := value];
    }
    return newMap;
  }

  function method UnstringifyEncryptionContext(stringEncCtx: Types.EncryptionContextString) : (res: Result<Types.EncryptionContext, Types.Error>)
  {
      if |stringEncCtx| == 0 then
        Success(map[])
      else
        var parseResults: map<string, Result<(UTF8.ValidUTF8Bytes, UTF8.ValidUTF8Bytes), Types.Error>> :=
          map strKey | strKey in stringEncCtx.Keys :: strKey := UnstringifyEncryptionContextPair(strKey, stringEncCtx[strKey]);
        if exists r | r in parseResults.Values :: r.Failure?
        then Failure(
              Types.KeyStoreException(message := "Encryption context contains invalid UTF8")
            )
        else
          assert forall r | r in parseResults.Values :: r.Success?;
          var utf8KeysUnique := forall k, k' | k in parseResults && k' in parseResults
                                  :: k != k' ==> parseResults[k].value.0 != parseResults[k'].value.0;
          if !utf8KeysUnique then Failure(Types.KeyStoreException(
                                            message := "Encryption context keys are not unique"))  // this should never happen...
          else Success(map r | r in parseResults.Values :: r.value.0 := r.value.1)
  }

  function method UnstringifyEncryptionContextPair(strKey: string, strValue: string) : (res: Result<(UTF8.ValidUTF8Bytes, UTF8.ValidUTF8Bytes), Types.Error>)
      ensures (UTF8.Encode(strKey).Success? && UTF8.Encode(strValue).Success?) <==> res.Success?
  {
      var key :- UTF8
                .Encode(strKey)
                .MapFailure(WrapStringToError);
      var value :- UTF8
                  .Encode(strValue)
                  .MapFailure(WrapStringToError);

      Success((key, value))
  }

  function method WrapStringToError(e: string)
    :(ret: Types.Error)
  {
    Types.KeyStoreException( message := e )
  }

  method ProvideCryptoClient(
    // Crypto?: Option<AtomicPrimitives.Types.IAwsCryptographicPrimitivesClient> := None
    Crypto?: Option<AtomicPrimitives.AtomicPrimitivesClient> := None
  )
    returns (output: Result<AtomicPrimitives.AtomicPrimitivesClient, AtomicPrimitives.Types.Error>)
    requires Crypto?.Some? ==> Crypto?.value.ValidState()
    modifies (if Crypto?.Some? then Crypto?.value.Modifies else {})
    ensures output.Success?
            ==>
              && output.value.ValidState()
              && fresh(output.value)
              && fresh(output.value.Modifies)
  {
    var Crypto: AtomicPrimitives.AtomicPrimitivesClient; //AtomicPrimitives.Types.IAwsCryptographicPrimitivesClient;
    if (Crypto?.None?) {
      Crypto :- AtomicPrimitives.AtomicPrimitives();
    } else {
      Crypto := Crypto?.value;
    }
    // If the customer gave us the Crypto Client, it is fresh
    // If we create the Crypto Client, it is fresh
    assume {:axiom} fresh(Crypto) && fresh(Crypto.Modifies);
    return Success(Crypto);
  }

  method GetBranchKeyVersion(
    input: Types.GetBranchKeyVersionInput,
    logicalKeyStoreName: string,
    kmsConfiguration: Types.KMSConfiguration,
    grantTokens: KMS.GrantTokenList,
    kmsClient: KMS.IKMSClient,
    storage: Types.IKeyStorageInterface
  )
    returns (output: Result<Types.GetBranchKeyVersionOutput, Types.Error>)
    requires storage.Modifies !! kmsClient.Modifies

    requires kmsClient.ValidState() && storage.ValidState()
    modifies storage.Modifies, kmsClient.Modifies
    ensures storage.ValidState() && kmsClient.ValidState()

    requires storage is DefaultKeyStorageInterface.DynamoDBKeyStorageInterface
             ==>
               logicalKeyStoreName == (storage as DefaultKeyStorageInterface.DynamoDBKeyStorageInterface).logicalKeyStoreName

    ensures
      //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
      //= type=implication
      //# GetBranchKeyVersion MUST get the requested version for the branch key id from the keystore
      //# by calling the configured [KeyStorage interface's](./key-store/key-storage.md#interface)
      //# [GetEncryptedActiveBranchKey](./key-store/key-storage.md#getencryptedbranchkeyversion)
      //# using the supplied `branch-key-id`.
      && |storage.History.GetEncryptedBranchKeyVersion| == |old(storage.History.GetEncryptedBranchKeyVersion)| + 1
      && Seq.Last(storage.History.GetEncryptedBranchKeyVersion).input
         == Types.GetEncryptedBranchKeyVersionInput(
              Identifier := input.branchKeyIdentifier,
              Version := input.branchKeyVersion
            )

    ensures output.Success?
            ==>
              && Seq.Last(storage.History.GetEncryptedBranchKeyVersion).output.Success?
              && var versionItem := Seq.Last(storage.History.GetEncryptedBranchKeyVersion).output.value.Item;

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# GetBranchKeyVersion MUST verify that the returned EncryptedHierarchicalKey MUST have the requested `branch-key-id`.
              && versionItem.Identifier == input.branchKeyIdentifier

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# GetActiveBranchKey MUST verify that the returned EncryptedHierarchicalKey is an HierarchicalSymmetricVersion.
              && Structure.DecryptOnlyHierarchicalSymmetricKey?(versionItem)

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# GetBranchKeyVersion MUST verify that the returned EncryptedHierarchicalKey MUST have the requested `branchKeyVersion`.
              && versionItem.Type == Types.HierarchicalSymmetricVersion(
                                       Types.HierarchicalSymmetric(
                                         Version := input.branchKeyVersion
                                       ))

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# GetBranchKeyVersion MUST verify that the returned EncryptedHierarchicalKey MUST have a logical table name equal to the configured logical table name.
              && versionItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#discovery
              //= type=implication
              //# The Keystore MAY use the KMS Key ARNs already
              //# persisted to the backing DynamoDB table,
              //# provided they are in records created
              //# with an identical Logical Keystore Name.
              && (kmsConfiguration.kmsKeyArn? ==> versionItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName)

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#mrdiscovery
              //= type=implication
              //# The Keystore MAY use the KMS Key ARNs already
              //# persisted to the backing DynamoDB table,
              //# provided they are in records created
              //# with an identical Logical Keystore Name.
              && (kmsConfiguration.kmsMRKeyArn? ==> versionItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName)

              && KMSKeystoreOperations.AttemptKmsOperation?(kmsConfiguration, versionItem.EncryptionContext)
              && |kmsClient.History.Decrypt| == |old(kmsClient.History.Decrypt)| + 1

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# The operation MUST decrypt the branch key according to the [AWS KMS Branch Key Decryption](#aws-kms-branch-key-decryption) section.
              && KMSKeystoreOperations.AwsKmsBranchKeyDecryption?(
                   versionItem,
                   kmsConfiguration,
                   grantTokens,
                   kmsClient,
                   Seq.Last(kmsClient.History.Decrypt)
                 )

              // && var versionEncryptionContext := Structure.ToBranchKeyContext(versionItem, logicalKeyStoreName);
              && var decryptResponse := Seq.Last(kmsClient.History.Decrypt).output.value;

              && Structure.ToBranchKeyMaterials(versionItem, decryptResponse.Plaintext.value).Success?

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# This GetBranchKeyVersion MUST construct [branch key materials](./structures.md#branch-key-materials)
              //# according to [Branch Key Materials From Authenticated Encryption Context](#branch-key-materials-from-authenticated-encryption-context).
              && var branchKeyMaterials := Structure
                                           .ToBranchKeyMaterials(
                                             versionItem,
                                             decryptResponse.Plaintext.value
                                           )
                                           .value;

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
              //= type=implication
              //# This operation MUST return the constructed [branch key materials](./structures.md#branch-key-materials).
              && output.value.branchKeyMaterials == branchKeyMaterials

              && output.value.branchKeyMaterials.branchKeyIdentifier == input.branchKeyIdentifier
              && UTF8.Encode(input.branchKeyVersion).Success?
              && output.value.branchKeyMaterials.branchKeyVersion == UTF8.Encode(input.branchKeyVersion).value

    ensures
      || (&& |storage.History.GetEncryptedBranchKeyVersion| == |old(storage.History.GetEncryptedBranchKeyVersion)| + 1
          && Seq.Last(storage.History.GetEncryptedBranchKeyVersion).output.Failure?
          ==> output.Failure?)

      || (&& |storage.History.GetEncryptedBranchKeyVersion| == |old(storage.History.GetEncryptedBranchKeyVersion)| + 1
          && Seq.Last(storage.History.GetEncryptedBranchKeyVersion).output.Success?
          && !Structure.ActiveHierarchicalSymmetricKey?(Seq.Last(storage.History.GetEncryptedBranchKeyVersion).output.value.Item)
          ==> output.Failure?)

      //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbranchkeyversion
      //= type=implication
      //# If the branch key fails to decrypt, this operation MUST fail.
      || (&& |kmsClient.History.Decrypt| == |old(kmsClient.History.Decrypt)| + 1
          && Seq.Last(kmsClient.History.Decrypt).output.Failure?
          ==> output.Failure?)
  {
    var VersionItem :- storage.GetEncryptedBranchKeyVersion(
      Types.GetEncryptedBranchKeyVersionInput(
        Identifier := input.branchKeyIdentifier,
        Version := input.branchKeyVersion
      )
    );

    var branchKeyItem := VersionItem.Item;

    :- Need(
      || storage is DefaultKeyStorageInterface.DynamoDBKeyStorageInterface
      || (
           && Structure.DecryptOnlyHierarchicalSymmetricKey?(branchKeyItem)
           && branchKeyItem.Identifier == input.branchKeyIdentifier
           && branchKeyItem.Type == Types.HierarchicalSymmetricVersion(
                                      Types.HierarchicalSymmetric(
                                        Version := input.branchKeyVersion
                                      ))
           && branchKeyItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName
         ),
      Types.KeyStoreException(
        message := ErrorMessages.INVALID_BRANCH_KEY_VERSION_FROM_STORAGE)
    );

    var branchKey: KMS.DecryptResponse :- KMSKeystoreOperations.DecryptKey(
      branchKeyItem,
      kmsConfiguration,
      grantTokens,
      kmsClient
    );

    var branchKeyMaterials :- Structure.ToBranchKeyMaterials(
      branchKeyItem,
      branchKey.Plaintext.value
    );

    return Success(
        Types.GetBranchKeyVersionOutput(
          branchKeyMaterials := branchKeyMaterials
        ));
  }

  method {:vcs_split_on_every_assert} GetBeaconKeyAndUnwrap(
    input: Types.GetBeaconKeyInput,
    logicalKeyStoreName: string,
    kmsConfiguration: Types.KMSConfiguration,
    grantTokens: KMS.GrantTokenList,
    kmsClient: KMS.IKMSClient,
    storage: Types.IKeyStorageInterface
  )
    returns (output: Result<Types.GetBeaconKeyOutput, Types.Error>)
    requires storage.Modifies !! kmsClient.Modifies
    requires storage is DefaultKeyStorageInterface.DynamoDBKeyStorageInterface
             ==>
               logicalKeyStoreName == (storage as DefaultKeyStorageInterface.DynamoDBKeyStorageInterface).logicalKeyStoreName

    requires kmsClient.ValidState() && storage.ValidState()
    modifies storage.Modifies, kmsClient.Modifies
    ensures storage.ValidState() && kmsClient.ValidState()

    ensures
      //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
      //= type=implication
      //# GetBeaconKey MUST get the requested beacon key from the keystore
      //# by calling the configured [KeyStorage interface's](./key-store/key-storage.md#interface)
      //# [GetEncryptedBeaconKey](./key-store/key-storage.md#getencryptedbeaconkey)
      //# using the supplied `branch-key-id`.
      && |storage.History.GetEncryptedBeaconKey| == |old(storage.History.GetEncryptedBeaconKey)| + 1
      && Seq.Last(storage.History.GetEncryptedBeaconKey).input.Identifier == input.branchKeyIdentifier

    ensures output.Success? ==>
              && Seq.Last(storage.History.GetEncryptedBeaconKey).output.Success?
              && var beaconItem := Seq.Last(storage.History.GetEncryptedBeaconKey).output.value.Item;

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
              //= type=implication
              //# GetBeaconKey MUST verify that the returned EncryptedHierarchicalKey MUST have the requested `branch-key-id`.
              && beaconItem.Identifier == input.branchKeyIdentifier

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
              //= type=implication
              //# GetBeaconKey MUST verify that the returned EncryptedHierarchicalKey is an ActiveHierarchicalSymmetricBeacon.
              && Structure.ActiveHierarchicalSymmetricBeaconKey?(beaconItem)

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
              //= type=implication
              //# GetBeaconKey MUST verify that the returned EncryptedHierarchicalKey MUST have a logical table name equal to the configured logical table name.
              && beaconItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#discovery
              //= type=implication
              //# The Keystore MAY use the KMS Key ARNs already
              //# persisted to the backing DynamoDB table,
              //# provided they are in records created
              //# with an identical Logical Keystore Name.
              && (kmsConfiguration.kmsKeyArn? ==> beaconItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName)

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#mrdiscovery
              //= type=implication
              //# The Keystore MAY use the KMS Key ARNs already
              //# persisted to the backing DynamoDB table,
              //# provided they are in records created
              //# with an identical Logical Keystore Name.
              && (kmsConfiguration.kmsMRKeyArn? ==> beaconItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName)

              && KMSKeystoreOperations.AttemptKmsOperation?(kmsConfiguration, beaconItem.EncryptionContext)
              && |kmsClient.History.Decrypt| == |old(kmsClient.History.Decrypt)| + 1

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
              //= type=implication
              //# The operation MUST decrypt the beacon key according to the [AWS KMS Branch Key Decryption](#aws-kms-branch-key-decryption) section.
              && KMSKeystoreOperations.AwsKmsBranchKeyDecryption?(
                   beaconItem,
                   kmsConfiguration,
                   grantTokens,
                   kmsClient,
                   Seq.Last(kmsClient.History.Decrypt)
                 )

              && var decryptResponse := Seq.Last(kmsClient.History.Decrypt).output.value;

              && Structure.ToBeaconKeyMaterials(beaconItem, decryptResponse.Plaintext.value).Success?

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
              //= type=implication
              //# This GetBeaconKey MUST construct [beacon key materials](./structures.md#beacon-key-materials) from the decrypted branch key material
              //# and the `branchKeyId` from the returned `branch-key-id` field.
              && var beaconKeyMaterials := Structure.ToBeaconKeyMaterials(
                                             beaconItem,
                                             decryptResponse.Plaintext.value
                                           ).value;

              //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
              //= type=implication
              //# This operation MUST return the constructed [beacon key materials](./structures.md#beacon-key-materials).
              && output.value.beaconKeyMaterials == beaconKeyMaterials

              && output.value.beaconKeyMaterials.beaconKeyIdentifier == input.branchKeyIdentifier

    ensures
      || (&& |storage.History.GetEncryptedBeaconKey| == |old(storage.History.GetEncryptedBeaconKey)| + 1
          && Seq.Last(storage.History.GetEncryptedBeaconKey).output.Failure?
          ==> output.Failure?)

      || (&& |storage.History.GetEncryptedBeaconKey| == |old(storage.History.GetEncryptedBeaconKey)| + 1
          && Seq.Last(storage.History.GetEncryptedBeaconKey).output.Success?
          && !Structure.ActiveHierarchicalSymmetricKey?(Seq.Last(storage.History.GetEncryptedBeaconKey).output.value.Item)
          ==> output.Failure?)

      //= aws-encryption-sdk-specification/framework/branch-key-store.md#getbeaconkey
      //= type=implication
      //# If the beacon key fails to decrypt, this operation MUST fail.
      || (&& |kmsClient.History.Decrypt| == |old(kmsClient.History.Decrypt)| + 1
          && Seq.Last(kmsClient.History.Decrypt).output.Failure?
          ==> output.Failure?)
  {

    var BeaconOutput :- storage.GetEncryptedBeaconKey(
      Types.GetEncryptedBeaconKeyInput(
        Identifier := input.branchKeyIdentifier
      )
    );

    var branchKeyItem := BeaconOutput.Item;

    :- Need(
      || storage is DefaultKeyStorageInterface.DynamoDBKeyStorageInterface
      || (
           && branchKeyItem.Identifier == input.branchKeyIdentifier
           && Structure.ActiveHierarchicalSymmetricBeaconKey?(branchKeyItem)
           && branchKeyItem.EncryptionContext[Structure.TABLE_FIELD] == logicalKeyStoreName
         ),
      Types.KeyStoreException(
        message := ErrorMessages.INVALID_BEACON_KEY_FROM_STORAGE)
    );

    var branchKey: KMS.DecryptResponse :- KMSKeystoreOperations.DecryptKey(
      branchKeyItem,
      kmsConfiguration,
      grantTokens,
      kmsClient
    );

    var branchKeyMaterials :- Structure.ToBeaconKeyMaterials(
      branchKeyItem,
      branchKey.Plaintext.value
    );

    return Success(
        Types.GetBeaconKeyOutput(
          beaconKeyMaterials := branchKeyMaterials
        ));
  }

}