👀

python-joseを使用してjwtの作成と検証を行う

2022/07/20に公開

概要

WebAPIにjwtを使った認証機能を追加しようと考えています。
そのため、下記の要件を満たす関数を作成したいです。

  • ユーザidを保持しているjwtを作成する関数
  • jwtの検証を行い、改ざんされていなければユーザidを含むpayloadを返却する関数
  • jwtの検証を行い、改ざんされていればエラーを出力する関数

この記事のゴール

python-joseを使用して、jwt作成とjwtの検証を行うことを目指します。

python-jose ドキュメント
https://python-jose.readthedocs.io/en/latest/index.html

ディレクトリ構成

tree
.
├── main.py
├── test_main.py
└── requirements.txt

ソースコード

main.py
from jose import jwt

# どのアルゴリズムを使用して電子署名を行うか
# 使用できるアルゴリズムはこちら(https://python-jose.readthedocs.io/en/latest/jws/index.html)
ALGORITHM = "HS256"

# 暗号化に使用する鍵情報
SECRET_KEY = "SECRET_KEY"

# jwtを作成する関数
def create_jwt(payload: dict) -> str:
    return jwt.encode(
        payload,
        SECRET_KEY,
        algorithm=ALGORITHM)

# jwtを検証する関数
def verify_jwt(encoded_jwt: str) -> dict:
    return jwt.decode(
        encoded_jwt, SECRET_KEY, algorithms=[ALGORITHM]
    )
test_main.py
import uuid
import pytest

from base64 import b64encode
from jose.exceptions import JWTError

from main import create_jwt, verify_jwt


def test_successful_validation():

    # ユーザidの作成
    user_id = uuid.uuid4().hex

    # payloadの作成
    payload = {"sub": user_id}

    # jwtの作成
    my_jwt = create_jwt(payload)

    # jwtの検証
    # エラーが出力されなければテストは成功
    assert payload == verify_jwt(my_jwt)


def test_verification_fails():

    # ユーザidの作成
    user_id = uuid.uuid4().hex

    # payloadの作成
    payload = {"sub": user_id}

    # jwtの作成
    my_jwt = create_jwt(payload)

    # jwtからheader, payload, signatureを取得
    [jwt_header, jwt_payload, jwt_signature] = \
        [i for i in my_jwt.split(".")]

    # idを書き換え他のユーザの情報を抜き出そうとするパターンを想定する
    jwt_payload = '{"sub":"aaaaaaaaaaaaaaaaa"}'

    # 書き換えたpayloadをbase64エンコードしjwtを作成する
    attacked_jwt = f"{jwt_header}.{b64encode(jwt_payload.encode()).decode()}.{jwt_signature}"

    # 検証が失敗すればテストは成功とする
    with pytest.raises(JWTError):
        verify_jwt(attacked_jwt)
requirements.txt
python-jose[cryptography]
pytest

テストの作成

まずテストを作成します。
要件は下記の通りのため、検証が成功するパターンと失敗するパターンを作成します。

  • ユーザidを保持しているjwtを作成する関数
  • jwtの検証を行い、改ざんされていなければユーザidを含むpayloadを返却する
  • jwtの検証を行い、改ざんされていればエラーを出力する関数

ユーザidを含む情報を受け取りjwtを返却する関数の箱を作成します。

main.py
# jwtを作成する関数
def create_jwt(payload: dict) -> str:
    pass

次にjwtを受け取り検証を行う関数の箱を作成します。

main.py
# jwtを作成する関数
def create_jwt(payload: dict) -> str:
    pass

# jwtを検証する関数
+ def verify_jwt(encoded_jwt: str) -> dict:
+     pass

jwtを作成し、そのjwtの検証を行うパターンのテストを作成します。

test_main.py
import uuid

from base64 import b64encode

from main import create_jwt, verify_jwt


def test_successful_validation():

    # ユーザidの作成
    user_id = uuid.uuid4().hex

    # payloadの作成
    payload = {"sub": user_id}

    # jwtの作成
    my_jwt = create_jwt(payload)

    # jwtの検証
    # エラーが出力されなければテストは成功
    assert payload == verify_jwt(my_jwt)

次に、jwtを作成し、改ざんされたjwtの検証を行うパターンのテストを作成します。

test_main.py
import uuid
+ import pytest

from base64 import b64encode
+ from jose.exceptions import JWTError

from main import create_jwt, verify_jwt


def test_successful_validation():

    # ユーザidの作成
    user_id = uuid.uuid4().hex

    # payloadの作成
    payload = {"sub": user_id}

    # jwtの作成
    my_jwt = create_jwt(payload)

    # jwtの検証
    # エラーが出力されなければテストは成功
    assert payload == verify_jwt(my_jwt)


+ def test_verification_fails():
+ 
+     # ユーザidの作成
+     user_id = uuid.uuid4().hex
+ 
+     # payloadの作成
+     payload = {"sub": user_id}
+ 
+     # jwtの作成
+     my_jwt = create_jwt(payload)
+ 
+     # jwtからheader, payload, signatureを取得
+     [jwt_header, jwt_payload, jwt_signature] = \
+         [i for i in my_jwt.split(".")]
+ 
+     # idを書き換え他のユーザの情報を抜き出そうとするパターンを想定する
+     jwt_payload = '{"sub":"aaaaaaaaaaaaaaaaa"}'
+ 
+     # 書き換えたpayloadをbase64エンコードしjwtを作成する
+     attacked_jwt = f"{jwt_header}.{b64encode(jwt_payload.encode()).decode()}.{jwt_signature}"
+ 
+     # 検証が失敗すればテストは成功
+     with pytest.raises(JWTError):
+         verify_jwt(attacked_jwt)

テストを実行します。

pytest

~~~省略~~~
FAILED test_main.py::test_successful_validation - AssertionError: assert {'sub': 'userid'} == None
FAILED test_main.py::test_verification_fails - AttributeError: 'NoneType' object has no attribute 'split'

関数の中身がないためテストは失敗します。

機能の実装

下記の機能を関数に実装していきます。

  • ユーザidを保持しているjwtを作成する
  • jwtの検証を行い、改ざんされていなければユーザidを含むpayloadを返却する
  • jwtの検証を行い、改ざんされていればエラーを出力する
main.py
+ from jose import jwt

# どのアルゴリズムを使用して電子署名を行うか
# 使用できるアルゴリズムはこちら(https://python-jose.readthedocs.io/en/latest/jws/index.html)
+ ALGORITHM = "HS256"

# 暗号化に使用する鍵情報
+ SECRET_KEY = "SECRET_KEY"

# jwtを作成する関数
def create_jwt(payload: dict) -> str:
-     pass
+     return jwt.encode(
+         payload,
+         SECRET_KEY,
+         algorithm=ALGORITHM)

# jwtを検証する関数
def verify_jwt(encoded_jwt: str) -> dict:
-     pass
+     return jwt.decode(
+         encoded_jwt, SECRET_KEY, algorithms=[ALGORITHM]
+     )

テストを再度実行します。

pytest

~~~省略~~~
test_main.py ..

作成したテストが全て成功したため、実装は完了です。

Discussion