AWS SDKを使ってKMSを操作する

## はじめに

最近現場でAWS Key Management Service (KMS)を作成・管理する経験をしました。ただ、実際にKMSキーを利用してデータを暗号化・復号することはやらなかったため、SDKを使って操作してみたいと思いました。

本記事では、PythonのSDK（Boto3）を利用してKMSを操作します。KMSを操作する前準備として、AWS Security Token Service（STS）にリクエストを送り、セッショントークンを取得します。

---

## AWS SDKの概要

### AWS SDKとは

**AWS SDK** はAWSサービスをプログラムから操作するためのライブラリです。今回は **Python** 用の **Boto3** を使用します。  
公式ドキュメントが非常に充実しているため、実装時に大きなつまずきはありませんでした。

#### 公式リソース:
- **サンプルコード**: [AWS SDK コード例](https://docs.aws.amazon.com/ja_jp/code-library/latest/ug/python_3_kms_code_examples.html)  
- **詳細リファレンス**: [Boto3 ドキュメント](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/index.html)  

### SDKのインストール方法

Python用のBoto3をインストールするには、以下のコマンドを実行します:

```bash
pip install boto3
```

pipが入っていれば簡単です。

---

## 検証

### Terraformによるキー作成

キーの作成・エイリアスの作成にはTerraformを利用します。SDKでも作成はできます。

**`main.tf`**:
```hcl
provider "aws" {
  region = var.region
}

resource "aws_kms_key" "this" {
  description             = "Example KMS Key"
  deletion_window_in_days = 7
  enable_key_rotation     = true
  rotation_period_in_days = 365

  key_usage                = "ENCRYPT_DECRYPT"
  customer_master_key_spec = "SYMMETRIC_DEFAULT"
}

resource "aws_kms_alias" "this" {
  name          = "alias/${var.kms_alias}"
  target_key_id = aws_kms_key.this.key_id
}
```

- **description**: キーの説明を設定する
- **deletion_window_in_days**: 削除待機期間 (例: `7` → 削除を予約してから7日後に削除)  
- **enable_key_rotation**: キーの自動ローテーションを有効化。今回の検証では関係なし
- **rotation_period_in_days**: キーのローテーション間隔。今回の検証では関係なし
- **key_usage**: キーの用途を指定。暗号化と復号に使用
- **customer_master_key_spec**: キーの暗号方式。対称暗号方式を使用。
- **name**: この例ではエイリアスの名前

リソースを作成します。
```
terraform apply
```

---

### セッショントークン取得

検証環境の都合上、MFAを使用してSTSから一時的なセキュリティ認証情報を取得します。

```
# access sts to get temporary security token.
def get_security_token(access_key, secret_key, mfa_name, mfa_code, duration=3600):
    sts_client = boto3.client(
        'sts',
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
    )

    # get temporary credentials
    response = sts_client.get_session_token(DurationSeconds=duration, SerialNumber=mfa_name, TokenCode=mfa_code)

    credentials = response['Credentials']
    logging.info("Temporary Security Credentials Retrieved:")
    logging.info(f"Access Key: {credentials['AccessKeyId']}")
    logging.info(f"Secret Key: {credentials['SecretAccessKey']}")
    logging.info(f"Session Token: {credentials['SessionToken']}")

    return credentials
```

### KMS client作成

ほとんどAWSのドキュメント通りです。認証情報はSTSで取得したクレデンシャルを指定。

encrypt, decryptメソッドは後ほど作成します。

```
class KeyManagementService:
    def __init__(self, kms_client):
        self.kms_client = kms_client

    @classmethod
    def from_client(cls, region_name, credentials) -> "KeyManagementService":
        """
        Creates a KeyManagementService instance with a default KMS client.

        :return: An instance of KeyManagementService initialized with the default KMS client.
        """
        
        kms_client = boto3.client(
            "kms",
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken'],
            region_name=region_name
        )

        return cls(kms_client)

    def encrypt(self, key_id: str, text: str) -> str:
        ...
    def decrypt(self, ciphertext: bytes) -> str:
        ...
```

### データの暗号化

encryptメソッドを実装します。

```python
def encrypt(self, key_id: str, text: str) -> str:
    response = self.kms_client.encrypt(KeyId=key_id, Plaintext=text.encode(CHARACTER_ENCODING))
    logger.info(f"The string was encrypted with algorithm {response['EncryptionAlgorithm']}")
    logger.info(f"Encrypted data: {response['CiphertextBlob']}")

    return response["CiphertextBlob"]
```

**説明**:  
- **`key_id`** はキーのIDまたはエイリアスを指定します。
- **`text`** は暗号化対象のデータを指定します。
- 暗号化後のデータは **CiphertextBlob** に含まれています。

### データの復号

decryptメソッドを実装します。

```python
def decrypt(self, ciphertext: bytes) -> str:
    response = self.kms_client.decrypt(CiphertextBlob=ciphertext)
    plaintext = response["Plaintext"].decode(CHARACTER_ENCODING)
    logger.info(f"Decrypted data: {plaintext}")

    return plaintext
```

**説明**:  
- **`ciphertext`** は暗号化データ（CiphertextBlob）を指定します。
  - 暗号化データ内に暗号化したときのキー情報が含まれているようで、再度キー指定は不要です。

### 動作確認

``` main.py
def load_env() -> dict:
    load_dotenv()

    missing_vars = [var for var in required_env_vars if not os.getenv(var)]
    if missing_vars:
        raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
    
    # Return the environment variables as a dictionary
    return {var: os.getenv(var) for var in required_env_vars}


def main():
    # load environment variables and check if required variables are set
    env_variables = load_env()

    # get mfa things
    mfa_code = sys.argv[1]
    mfa_name = env_variables['MFA_SERIAL']

    # get temporary credentials using mfa
    credentials = get_security_token(env_variables['AWS_ACCESS_KEY_ID'], env_variables['AWS_SECRET_ACCESS_KEY'], mfa_name, mfa_code)
    
    # setup KMS client
    kms = KeyManagementService.from_client(env_variables['AWS_DEFAULT_REGION'], credentials)

    # Test data
    data = "This is a test message."
    logging.info(f"Original data: {data}")

    ### Encrypt and decrypt

    # The key is in a pending deletion state.
    # encrypted_data = kms.encrypt("953a1cc8-e60a-4748-95d6-9b3aa8cdd9e4", data)

    encrypted_data = kms.encrypt(env_variables['KMS_ALIAS'], data)
    decrypted_data = kms.decrypt(encrypted_data)
```

```
python main.py <mfa_code>
2024-12-15 20:38:42,459 - INFO - Temporary Security Credentials Retrieved:
2024-12-15 20:38:42,459 - INFO - Access Key: <access_key>
2024-12-15 20:38:42,460 - INFO - Secret Key: <secret_key>
2024-12-15 20:38:42,460 - INFO - Session Token: <session_token>
2024-12-15 20:38:42,474 - INFO - Original data: This is a test message.
2024-12-15 20:38:42,601 - INFO - The string was encrypted with algorithm SYMMETRIC_DEFAULT
2024-12-15 20:38:42,602 - INFO - Encrypted data: b'\x01\x02\x02\x00x\x1c\xc7\xef\xab\x03S\xe4\xdcV\x12\xaf\x99=\xf3\xdak\x00\x9b\xc2\xe8b\xeb5\xbei\xa4H\xf4&R\xf0F\x01\x90\x02(\xb5\xd1\xb0\xc7k\xddc\xe72"\xcdV5\x00\x00\x00u0s\x06\t*\x86H\x86\xf7\r\x01\x07\x06\xa0f0d\x02\x01\x000_\x06\t*\x86H\x86\xf7\r\x01\x07\x010\x1e\x06\t`\x86H\x01e\x03\x04\x01.0\x11\x04\x0c\xed\xbe\xfbW\xff:$\xc7\xeb\xf5\xb7`\x02\x01\x10\x802\x92\xa8x\xc6\x80qNU\xa1?\x17\xd5\xddmt\xa2\xc8\xac\xdb\xa3\x17\xc3\x17v\x08:Yi\xfe\'\x1b\xd8\xb2u\xef\xe4\xce0\x9a\x9b\xfd\x19\x10[\x95\x04h\x88\xc8\x8f'
2024-12-15 20:38:42,622 - INFO - Decrypted data: This is a test message.
```

細かいところは本記事内で紹介できてません。
コードサンプルをGitHubで公開していますので、そちらでご確認ください:  
[GitHubリポジトリ](https://github.com/hsmto25519/boto3-sample/tree/main)

## おまけ

#### 削除対象のキーに対する暗号化リクエスト
```
botocore.errorfactory.KMSInvalidStateException: An error occurred (KMSInvalidStateException) when calling the Encrypt operation: arn:aws:kms:ap-northeast-1:<account_id>:key/953a1cc8-e60a-4748-95d6-9b3aa8cdd9e4 is pending deletion.
```

キーの状態とそれに対応するAPIリクエストの結果はすべてまとまっているようです。
https://docs.aws.amazon.com/ja_jp/kms/latest/developerguide/key-state.html

#### EncryptionContext検証

暗号化、復号時にEncryptionContextというものが利用できて、これを利用するとデータが改ざんされていないことを検知できるようです。
コードは以下のような感じです。（EncryptionContext=ENCRYPTION_CONTEXTを追加するだけ）

```
ENCRYPTION_CONTEXT = { "System": "Test" }
...

class KeyManagementService:
...
    def encrypt(self, key_id: str, text: str) -> str:
        response = self.kms_client.encrypt(KeyId=key_id, Plaintext=text.encode(CHARACTER_ENCODING), EncryptionContext=ENCRYPTION_CONTEXT)
        logger.info(f"The string was encrypted with algorithm {response['EncryptionAlgorithm']}")
        logging.info(f"Encrypted data: {response["CiphertextBlob"]}")

        return response["CiphertextBlob"]


    def decrypt(self, ciphertext: bytes) -> str:
        response = self.kms_client.decrypt(CiphertextBlob=ciphertext, EncryptionContext=ENCRYPTION_CONTEXT)
        plaintext = response["Plaintext"].decode(CHARACTER_ENCODING)
        logger.info(f"Decrypted data: {plaintext}")

        return plaintext
```

暗号化時、復号時ともに同じkey, valueを指定しないとリクエストが失敗するようになります。
`botocore.errorfactory.InvalidCiphertextException: An error occurred (InvalidCiphertextException) when calling the Decrypt operation:`

Encryption Context自体は平文で登録されます。（ドキュメントに記載ある通り）
（cloudtrailのログ抜粋）
``` 
    "requestParameters": {
        "keyId": "alias/kms-verification-key",
        "encryptionContext": {
            "System": "Test"
        },
        "encryptionAlgorithm": "SYMMETRIC_DEFAULT"
    },
```

boto3のドキュメントでは推奨されていたので、導入を検討してみたいですね。
>On operations with symmetric encryption KMS keys, an encryption context is optional, but it is strongly recommended.
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/kms/client/encrypt.html

以上です!

---

## まとめ

本ブログでは **AWS SDK** を試してみました。ドキュメントはすごく読みやすくてとてもよかったです。本当はキー作成もSDKでやりたかったのですが、間違えていくつもキーを作ってしまうのが怖くて慣れているterraformにしました。