Amplify UI + Cognitoで独自仕様のMFA認証をする

# この記事がすること

- Amplify UI(React)のAuthenticatorを使います
- CognitoのMFAを使って、ユーザーにMFA認証をさせます
- MFAの仕様は独自仕様にします

https://github.com/aws-amplify/amplify-ui/issues/664

# 具体的にどうなるのか

完成する画面は以下の通りです。

**1.Amplify UIのログイン画面です**

![1-login.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/233180/360a90fc-a1f3-9fd1-4ac9-f7cd5d23ed60.png)

**2.Sign inを押下すると、MFAの入力画面に遷移します**

![2-mfa-input.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/233180/1f446558-9459-049b-fd15-299e7b92cf87.png)

**3.遷移したタイミングで、メールアドレスにワンタイムパスワードが送られてきます。**
※独自実装だと分かりやすいように、日本語のワンタイムパスワードを発行する仕組みにしました

![S__171999234.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/233180/6a550133-b3f1-ee91-78cf-bf614b351454.jpeg)

**4.Amplify UIの入力欄に、メールで受け取ったワンタイムパスワードを入力します。**  
※ワンタイムパスワードは1時間有効です

![2-mfa-input-ok.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/233180/bff72ef8-670d-740f-f9b4-dfd09b67b48f.png)

**5.TOTPコードを送信をクリックすると、認証が通ります。**

![3-success.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/233180/f228e5aa-57a2-f81b-9d84-502cbf27c9f3.png)

ちなみに、Authenticatorの内側の実装はほぼ公式ドキュメントそのままです。
Cognitoに認証も任せていますので、権限やロールの扱いはCognitoの通常のMFAと同じです。取り出した認証コードもドキュメント通りの方法で使うことができます。

```react: Authenticatorの実装部分
import { Authenticator, Button } from "@aws-amplify/ui-react"; 
<Authenticator className={isMfaView(route) ? "authClassName" : ""}>
    {({ signOut }) => (
      <>
        認証が通りました！！
        <Button variation="primary" onClick={() => signOut && signOut()}>
          ログアウト
        </Button>
      </>
    )}
</Authenticator>
```

# 実装していく

実装の説明をしていきます。
記事では実装するために必要な要所要所の説明だけをしていきますので、実際に動くソースの全体図はGithubで参照してください。

https://github.com/ShotaOki/totp-client-cognito-view

## 画面側の実装

前準備として、Authenticatorよりも上の階層（main.tsx）で、Authenticator.Providerを囲んでおきます。

```diff_jsx: main.tsx
import { Authenticator } from "@aws-amplify/ui-react";

createRoot(document.getElementById("root")!).render(
  <StrictMode>
+    <Authenticator.Provider>
    <App />
+    </Authenticator.Provider>
  </StrictMode>
);
```

Authenticator.Providerで囲むと、Appの中で`useAuthenticator`が使えるようになります。
このフックを使うと、Authenticatorをプログラムで操作したり、入力した情報を拾い上げたり、状態の変更を拾ったりすることができるようになります。

### Authenticatorで遷移をキャプチャする

まず、Authenticatorの状態遷移を拾います。
App.tsxの中で、useAuthenticatorを実行します。

```react: App.tsx
import { useAuthenticator } from "@aws-amplify/ui-react";

function App() {
  // routeが状態遷移の情報
  const { route } =
    useAuthenticator((context) => [
      context.route
    ]);
```

Authenticatorが状態遷移をするたびに、routeが更新されます。
ですので、それをuseEffectや画面で拾い上げます。

```react: App.tsx
import { useAuthenticator } from "@aws-amplify/ui-react";
import { useEffect } from "react";

function App() {
  // routeが状態遷移の情報
  const { route } =
    useAuthenticator((context) => [
      context.route
    ]);

  // 画面が切り替わったことの通知を受ける
  useEffect(() => {
    if (route === "setupTotp") {
       // MFAのコードの初回入力画面が表示されたタイミングで実行される
    }
  }, [route])

  // 画面でrouteを条件文に使えば、画面部品を特定の状態の時だけ表示することができる
  return (
    {route === "setupTotp" && <div>MFAコードの入力中...</div>}
    {route === "signIn" && <div>ユーザー名とパスワードの入力中...</div>}
```

このように、routeを使えば、Authenticatorがどの状態に遷移したのかを受け取ることができます。
また、逆に、visibility: hiddenを設定したCSSのクラスを作っておけば、特定の遷移に来た時にAuthenticatorを隠すことができます。

```css: main.css
.authClassName {
  visibility: hidden;
}
```

```react:App.tsx
import {
  Authenticator,
  Card,
  Flex,
  Grid
} from "@aws-amplify/ui-react";
      {/* MFA入力中は独自の画面を代わりに出す */}
      {route === "setupTotp" && (
        <Grid>
          <Card
            variation="outlined"
            style={{
              width:
                "var(--amplify-components-authenticator-container-width-max)",
              placeSelf: "center",
              padding: "32px",
            }}
          >
            <Flex direction="column">
              オリジナルな入力画面
            </Flex>
          </Card>
        </Grid>
      )}
      {/* MFA入力中はもともとのAuthenticatorを隠す */}
      <Authenticator className={route === "setupTotp" ? "authClassName" : ""}>
      </Authenticator>
```

### Authenticatorの入力情報を拾い上げる

useAuthenticatorを使うと、Authenticatorに入力された情報を参照することもできます。

```react
import { useAuthenticator } from "@aws-amplify/ui-react";

function AppInternal() {
  // useAuthenticatorから認証情報を受け取る
  const { totpSecretCode, username } =
    useAuthenticator((context) => [
      context.totpSecretCode,
      context.username,
    ]);
```

totpSecretCodeはMFAのシークレット情報、usernameはユーザーが画面で入力したログインユーザーのIDです。

情報の更新があるたびにリアルタイムで反映されます。

### プログラムからAuthenticatorを操作する

useAuthenticatorを使うと、プログラムでAuthenticatorを操作することもできます。
もちろんvisibility: hiddenで非表示にしたAuthenticatorも動かせます。

```react
import { useAuthenticator } from "@aws-amplify/ui-react";

function AppInternal() {
  // useAuthenticatorから操作関数を受け取る
  const { updateForm, submitForm } =
    useAuthenticator((context) => [
      context.updateForm,
      context.submitForm,
    ]);

  // ユーザー名をmail@address.com、パスワードをpasswordにして、ログイン操作を実行する
  const executeLogin = () => {
    updateForm({name: "username", value: "mail@address.com"})
    updateForm({name: "password", value: "password"})
    submitForm();
  }

  // MFAの認証コードを123456にして、MFAの認証操作をする
  const executeMFA = () => {
    updateForm({name: "confirmation_code", value: "123456"})
    submitForm();
  }
```

この関数を実行するだけで、プログラムからAuthenticatorを動かすことができます。

:::note info
2024年9月1日時点で、ドキュメントにupdateFormとsubmitFormの説明はありません。ただ、関数そのものは昔からあって、2021年のIssueでも「ドキュメントに書かないとね」と話題に上がっているのですが、まだIssueの状態です。
:::

https://github.com/aws-amplify/amplify-ui/issues/664

## CognitoのMFAのコード発行

AuthenticatorのMFAのコードは、PyOTPを使ってLambdaで発行することができます。

```python: Lambdaの処理
import pyotp

# 画面側で取ったtotpSecretCodeを、pyotpに渡す
otp = pyotp.TOTP(totpSecretCode).now()
```

https://pyauth.github.io/pyotp/

画面側は、Lambdaが発行したotpを受け取って、先ほど紹介したsubmitFormを実行します。
たったこれだけでCognitoのMFAの認証を通すことができます。簡単です。

```react: 画面側
  const executeMFA = () => {
    fetch("APIGatewayのURL", {
      method: "POST",
      body: JSON.stringify({
        "totpSecretCode": totpSecretCode
      })
    }).then((res) => res.json()).then(({ otp }) => {
      // OTPのコードを使って、認証画面をプログラムからクリアする
      updateForm({name: "confirmation_code", value: otp})
      submitForm();
    })
  }
```

ただ、いくつかの課題があります。

- OTPのコードは30秒で失効する（＝LambdaからSNSで送ると、開く頃に失効する）
- totpSecretCodeは初回の認証以外では画面から参照できない

ですので、DynamoDBを使ってtotpSecretCodeを保管する仕組みと、メールで渡した認証情報からOTPのコードを発行する仕組みが必要になります。

## これらの機能を使って、独自MFAを実装する

あらためて情報をまとめます。

useAuthenticatorを使うと、以下のことができます。
- 遷移の特定の状態になったとき、元のAuthenticatorを隠す
- 遷移の特定の状態になったとき、独自の画面に差し替える
- 独自の画面から非表示のAuthenticatorを操作して、認証を実行する

OTPのメール送信に必要な仕組みは以下の通りです。
- メールで送った独自の認証情報から、OTPを発行する
- DynamoDBでtotpSecretCodeを保管する

ですので、次のような仕組みを作ります。

![proc.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/233180/6f69765e-f3a4-85f9-9de0-6feff7cb6ae1.png)


Lambdaを使ってMFAのワンタイムパスワードを発行、発行したものを非表示のAmplifyUIに渡して認証します。

## バックエンド側の実装

### 独自パスワードの発行

独自パスワードを発行するタイミングで、画面側で取得できる`totpSecretCode`をDynamoDBで永続化しておきます。  
初回のMFAのsetupTotpではこの変数に値が入るのですが、2回目以降のMFAでは画面が変わる（※confirmSignInになる）ため、totpSecretCodeの変数に値が入らなくなります。
ですので、ユーザー名に紐づけてtotpSecretCodeを保管しておきます。

```python: 独自パスワードの発行(Python)
dynamodb_client = boto3.client("dynamodb")

try:
    # シークレットキーをDynamoDBに登録する
    dynamodb_client.put_item(
        TableName=table_name,
        Item={
            "user_id": {"S": input.user_id},
            "secret_key": {"S": input.secret_key}, # totpSecretCode
        },
        # 初回更新が登録済みなら、登録リクエストを無視する
        ConditionExpression="attribute_not_exists(protected)",
    )
except ClientError as e:
    print(e)
```

一時的に有効な独自パスワードを発行して、SNSでメール送信します。
独自パスワードは数値でもJWTでもいいのですが、今回は日本語で発行したかったため、What3Wordsを利用しました。

```python: 独自パスワードの発行(Python)
# ワンタイムパスワードを生成する
# 今回はwhat3words APIを利用して、日本語のワンタイムパスワードを生成する
geocoder = what3words.Geocoder(api_key=api_key, language="ja")
words = geocoder.convert_to_3wa(
    what3words.Coordinates(
        lat=random.random() * 180 - 90.0,  # 緯度をランダムに生成
        lng=random.random() * 360 - 180.0,  # 経度をランダムに生成
    )
)
```

What3Wordsを実行すると、数値がひらがなの3単語に変換されます。
「えいしょう。いのり。ねんじろ」みたいなキーワードになります。

このキーワードと有効期限もDynamoDBに保管します。
また、キーワードはSNSを使ってメール送信します。

```python: 独自パスワードの発行(Python)

# 独自パスワードの有効期限を設定する
expired = datetime.datetime.now() + datetime.timedelta(minutes=60)
expired_timestamp = str(expired.timestamp())

# 独自パスワードをDynamoDBに登録する
dynamodb_client.update_item(
    TableName=table_name,
    Key={"user_id": {"S": input.user_id}},
    AttributeUpdates={
        "words": {"Value": {"S": words["words"]}},
        "expired_time": {"Value": {"N": expired_timestamp}},
    },
)

# 独自パスワードをSMSで送信する
sns_client.publish(
    TopicArn=topic_arn, # SNSの送信先
    Message=f"ワンタイムパスワードは「{words['words']}」です",
    Subject="OTP",
)
```

### 独自パスワードの検証

メールでキーワードを受け取ったユーザーは、画面にキーワードを入力して、Lambdaを実行します。

```python: 独自パスワードの検証(Python)
# DynamoDBからユーザー情報を取得する
item = dynamodb_client.get_item(
    TableName=table_name,
    Key={"user_id": {"S": input.user_id}},
)

# 独自パスワードの期限切れ、または誤りを検証する
now = datetime.datetime.now().timestamp()
expired_time = item.get("Item", {}).get("expired_time", {}).get("N", "0")
words = item.get("Item", {}).get("words", {}).get("S", "")

if now >= float(expired_time) or words != input.totp_key:
    # 有効期限切れ、または独自パスワードが誤っている
```

誤っていないのなら、pyotpを使って、Cognito向けのワンタイムパスワードを発行します。
ここでpyotp.TOTPが発行した6桁の数値を画面に渡して、updateFormとsubmitFormを実行すれば認証を通すことができます。

ちなみにCognitoのMFAのワンタイムパスワードは30秒で次のパスワードに切り替わるため、直接ワンタイムパスワードをメールで送ってしまうと、ちょうど開く頃に失効します。
独自パスワードを噛ませることで余裕を持って開くことができるようになります。

```python: 独自パスワードの検証(Python)
# シークレットキーを元に、認証用のワンタイムパスワードを発行する
secret_key = item.get("Item", {}).get("secret_key", {}).get("S", "")
otp = pyotp.TOTP(secret_key).now()
```

また、認証用のワンタイムパスワードを返すときに、DynamoDBのシークレットキーに保護をかけておきます。

```python: 独自パスワードの検証(Python)
# 初回の認証が通ったのなら、ユーザー情報を保護する
dynamodb_client.update_item(
    TableName=table_name,
    Key={"user_id": {"S": input.user_id}},
    AttributeUpdates={
        "protected": {"Value": {"BOOL": True}},
    },
)
```

以上で実装できると思います。
あらためての紹介になりますが、実際に実装したものはこちらです。

https://github.com/ShotaOki/totp-client-cognito-view