【Snowflake】ユーザのパスワードやKey-Pairの生成と設定を行えるStreamlit in Snowflakeを作ろう

本記事で参考になるケース

• Streamlit in Snowflakeから、自身 or 他のSnowflakeアカウントのパスワード/Public Keyの設定をStreamlit in Snowflakeアプリから実施したい
• kay-pair認証のkeyは自動で生成して欲しい

前提

• Stremalit in Snowflake もしくは Streamlit を利用できること
• Streamlit in Snowflake利用時は、作成するStreamlitのOwnerがユーザの更新権限を保有していること（alter user権限）
• Streamlit in Snowflake利用時は、上記に合わせStreamlitの作成権限を保有していること

さっそく例を!

ユーザアカウントのパスワード及びKey-pair認証の設定を行えるStreamlit

作成する Streamlit in Snowflake のイメージ

---

作ってみよう！

1. Snowsight上（左下から）でStreamlit作成権限及び全ユーザの更新権限を保有するRoleを選択する

ACCOUNTADMINを利用しても構いませんが、くれぐれもセキュリティには注意しましょう

2. Streamlit が立ち上がったら左上の package から cryptography をインストールする
   
3. Streamlitのデフォルトコードを下記に置き換える

import streamlit as st
import datetime
import pandas as pd
from snowflake.snowpark import Session
from snowflake.connector.pandas_tools import write_pandas
from snowflake.snowpark.context import get_active_session



def get_snowflake_connection():
        return get_active_session()
st.session_state.snowflake_connection = get_active_session()

@st.cache_data(experimental_allow_widgets=True)
def get_account_list():
    session = get_snowflake_connection()
    query = f"""
    select name from snowflake.account_usage.users
    """
    re = session.sql(query).collect();
    re_dict = [row.asDict() for row in re]
    result = []
    try:
        result = [name['NAME'].upper() for name in re_dict];
    except:
        pass
    result = list(set(result))
    result.insert(0, None)
    return result


@st.experimental_fragment
def get_file(message,gen_public_key,file_name):
    st.download_button(
        message,
        gen_public_key.decode(encoding="utf-8"), 
        file_name=file_name
    )


def exe_set_user_key(selected_user, public_key):
    try:
        query =f"""
            alter user {selected_user} set RSA_PUBLIC_KEY= '{public_key}'
        """
        st.session_state.snowflake_connection.sql(query).collect();
        st.info('ユーザ '+selected_user+'のPublic Keyをリセットしました', icon="🔑")
    except:
        st.warning('なんらかのエラーが発生しました. Public Keyは更新されていません.', icon="🚨")


@st.experimental_fragment
def set_user_key_pair(selected_user, public_key):
    if st.button(f'''{selected_user} のPublic Keyを上記で生成したKeyに変更''', type="primary"):
        exe_set_user_key(selected_user, public_key)
        url = "https://docs.snowflake.com/ja/user-guide/key-pair-auth"
        st.info(f'''
            参考：{url}\n
            本変更により下記のコマンドが内部で実行されます 
            ```
            ALTER USER {selected_user} SET RSA_PUBLIC_KEY='MIIBIjANBgkqh...';
            ```
        ''')  

@st.experimental_fragment
def set_user_gen_key_pair(selected_user, public_key):
    if st.button(f'''{selected_user} のPublic Keyを更新''', type="primary"):
        exe_set_user_key(selected_user, public_key)
        url = "https://docs.snowflake.com/ja/user-guide/key-pair-auth"
        st.info(f'''
            参考：{url}\n
            本変更により下記のコマンドが内部で実行されます 
            ```
            ALTER USER {selected_user} SET RSA_PUBLIC_KEY='MIIBIjANBgkqh...';
            ```
        ''')  


st.set_page_config(
    page_title="ユーザ情報",
    layout="wide",
    initial_sidebar_state="expanded",
)
st.title("👤 ユーザ情報の更新")
st.divider()
st.header("🔐 パスワードの設定/再設定", divider="rainbow")

if 'selected_column' not in st.session_state:
    st.session_state.selected_column = None

#accout_list=get_account_list(number)
accout_list=get_account_list()
selected_user = st.selectbox("どのアカウントを対象にしますか？",
        accout_list,
    )

if selected_user  is None: st.write( "アカウントを選択してください" );
else: 
    st.header("更新するアカウント【 "+selected_user+" 】",divider="green")
    
    password = st.text_input("新しいパスワードを入力", type="password").replace("\'","")
    st.info('シングルクウォート「\'」はパスワードに含んではいけません')
    if st.button("パスワードを変更", type="primary"):
        try:
            query =f"""
            alter user {selected_user} set password = '{password}'
            """
            st.session_state.snowflake_connection.sql(query).collect();
            st.info('ユーザ '+selected_user+' のパスワードをリセットしました', icon="🔑")
        except:
            st.warning('エラーが発生しました. パスワードは更新されていません.', icon="🚨")

    with st.expander('key-pair認証を行う場合'):
        st.header("key-pair認証のkey生成とユーザの認証設定を実施します",divider="blue")
        if st.button(f'''key-pairを生成する''', type="primary"):
            from cryptography.hazmat.primitives import serialization
            from cryptography.hazmat.backends import default_backend
            from cryptography.hazmat.primitives.asymmetric import rsa
            import base64
            key = rsa.generate_private_key(
                key_size=2048,
                public_exponent=65537,
                backend=default_backend()            
            )
            gen_private_key = key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
            gen_public_key = key.public_key().public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            )
            st.text_area("private_key", gen_private_key.decode(encoding="utf-8"), height=200)
            get_file("Download private key",gen_private_key,"private.key")

            gen_public_key_der = base64.b64encode( 
                key.public_key().public_bytes(
                    serialization.Encoding.DER,
                    serialization.PublicFormat.SubjectPublicKeyInfo
                )).decode('utf-8')
            
            st.text_area("public_key", gen_public_key.decode(encoding="utf-8"), height=250)
            get_file("Download public key",gen_public_key,"public.pub")
            set_user_key_pair(selected_user,gen_public_key_der)
    
    with st.expander('【Option】key-pair認証　:　保有するkey-pairを認証に利用したい場合'):
        st.header("作成済みのPublic Keyを設定してください",divider="blue")
        public_key_der = st.text_input("Public Key", type="password")
        st.warning(f'''
            🔸間違った入力
            ```
            -----BEGIN ENCRYPTED PRIVATE KEY-----
            MIIE6T...
            -----END ENCRYPTED PRIVATE KEY-----
            ```
            🔸正しい入力
            ```
            MIIE6T...
            ```
        ''')  
        set_user_gen_key_pair(selected_user,public_key_der)

できること
パスワード/Kay-Pair認証設定対象ユーザアカウントの選択
Kay-Pairの生成とPrivate Key / Public Keyのダウンロード
保有しているKey-pairのPublic Keyをユーザアカウントに設定

---

発展

上記の例では、基本的にStreamlitのOwner Roleが権限を保有しいれば全てのユーザのパスワードやkey-pairの設定更新が可能である.

更新対象のアカウントを制限する①

• 本Streamlitアプリで更新できるユーザを制限する場合

def get_account_list():
    session = get_snowflake_connection()
    query = f"""
    select name from snowflake.account_usage.users
    """

上記で実行されているクエリはユーザを選択する際に表示されるユーザのリストである.

したがって、ここでユーザの制限を行うことが可能である.

更新対象のアカウントを制限する②

• 本Streamlitアプリを開いたユーザの情報を取得し、それにより更新ができるユーザを絞り込みたい場合

current_role = st.session_state.snowflake_connection.get_current_role()

上記を用いることで、本Streamlitアプリを開いたユーザのRoleを取得することが可能である

user_name = st.experimental_user["user_name"].

上記を用いることで、本Streamlitアプリを開いたユーザネームを取得することが可能である

これらの情報を用いて、必要に応じて制限をかけてきます.

まとめ

Streamlit in Snowflakeを活用してできることは多くあります. 特にSQL操作が苦手な人への情報伝達や短時間での情報理解にはうってつけのツールです.
適宜活用して行きましょう！

参考

https://docs.snowflake.com/ja/user-guide/key-pair-auth

捕捉

.replace("\'","")がある理由

password = st.text_input("新しいパスワードを入力", type="password").replace("\'","")

Streamlitのst.text_inputにおいて入力された更新用のパスワードは下記でユーザに設定を行っている

    if st.button("パスワードを変更", type="primary"):
        try:
            query =f"""
            alter user {selected_user} set password = '{password}'
            """

上記 alter userのパスワードはログ上マスキングされて、ACCOUNTADMINであっても見ることはできない.しかし、ある条件下においてログに表示されてしまう場合があります.
例えば、シングルクウォートがパスワードに含まれ更新に失敗したときなどです.
今回、ユーザが入力する文字にシングルクウォート「'」が混ざっていた場合、パスワード更新は失敗するどころかログにパスワードが一部表示されてしまうことになってしまった.
したがって、これを排除する設定にしました.