common.py

# common.py
import extra_streamlit_components as stx
import streamlit as st
import logging
import os

from time import time
from requests_oauthlib import OAuth2Session
from streamlit import runtime
from streamlit.runtime.scriptrunner import get_script_run_ctx
import ipaddress

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("__name__")
logger.debug("調査用ログ")

# 接続元制御
ALLOW_IP_ADDRESS = os.environ["ALLOW_IP_ADDRESS"]

# Azure AD app registration details
CLIENT_ID = os.environ["CLIENT_ID"]
TENANT_ID = os.environ["TENANT_ID"]

# Azure API
AUTHORITY = f"https://login.microsoftonline.com/{TENANT_ID}"
REDIRECT_PATH = os.environ["REDIRECT_PATH"]
AUTHORIZATION_URL = f"{AUTHORITY}/oauth2/v2.0/authorize"
SCOPES = ["openid", "profile", "User.Read"]

# 認証用URL取得
def authorization_request():
    oauth = OAuth2Session(CLIENT_ID, redirect_uri=REDIRECT_PATH, scope=SCOPES)
    authorization_url, state = oauth.authorization_url(AUTHORIZATION_URL)
    return authorization_url, state

# 接続元IP取得
def get_remote_ip():
    ctx = get_script_run_ctx()
    session_info = runtime.get_instance().get_client(ctx.session_id)
    return session_info.request.remote_ip

# 接続元IP許可判定
def is_allow_ip_address():
    remote_ip = get_remote_ip()
    logger.info("remote_ip")
    logger.info(remote_ip)
    # localhost
    if remote_ip == "::1":
        return True

    # プライベートIP
    ipaddr = ipaddress.IPv4Address(remote_ip)
    logger.info("ipaddr")
    logger.info(ipaddr)
    if ipaddr.is_private:
        return True

    # その他（許可リスト判定）
    return remote_ip in ALLOW_IP_ADDRESS

#ログインの確認
def check_login():
    # 接続元IP許可判定
    if not is_allow_ip_address():
        st.title("HTTP 403 Forbidden")
        return

    if "token" not in st.session_state or st.session_state["token"] is None or float(st.session_state["token_expires"]) <= time():
        # 認証用リンク表示
        authorization_url, st.session_state["authorization_state"] = authorization_request()
        st.markdown(f'[Click here to log in]({authorization_url})', unsafe_allow_html=True)
        st.stop()