如何使用Okta SAML 2FA通过Python SQLAlchemy连接Redshift

前置依赖安装 #

先安装必要的Python包：

pip install redshift_connector sqlalchemy-redshift pandas beautifulsoup4 pyotp

完整代码实现 #

import redshift_connector
from sqlalchemy import create_engine
import pandas as pd
import requests
from bs4 import BeautifulSoup
import pyotp

# -------------------------- 配置参数 --------------------------
OKTA_DOMAIN = "your-okta-domain.okta.com"
OKTA_APP_ID = "your-okta-app-id"  # Redshift在Okta中的应用ID，可从Okta应用页面URL提取
USERNAME = "your-okta-username"
PASSWORD = "your-okta-password"
TOTP_SECRET = "your-totp-secret"  # Okta账号MFA设置中获取的TOTP密钥
REDSHIFT_HOST = "your-redshift-cluster-url.redshift.amazonaws.com"
REDSHIFT_PORT = 5439
REDSHIFT_DB = "your-database-name"
REDSHIFT_SCHEMA = "your-schema-name"
# -------------------------------------------------------------

def get_okta_saml_assertion():
    # 初始化会话，访问Okta应用登录页
    session = requests.Session()
    login_url = f"https://{OKTA_DOMAIN}/app/{OKTA_APP_ID}/sso/saml"
    response = session.get(login_url)
    soup = BeautifulSoup(response.content, "html.parser")
    
    # 提取登录表单的提交地址和隐藏参数
    form = soup.find("form")
    form_action = form.get("action")
    payload = {input_tag.get("name"): input_tag.get("value") for input_tag in form.find_all("input")}
    
    # 提交账号密码
    payload["username"] = USERNAME
    payload["password"] = PASSWORD
    response = session.post(form_action, data=payload)
    soup = BeautifulSoup(response.content, "html.parser")
    
    # 处理TOTP类型的MFA验证
    mfa_form = soup.find("form", id="mfa-form")
    if mfa_form:
        totp = pyotp.TOTP(TOTP_SECRET)
        current_code = totp.now()
        mfa_payload = {
            input_tag.get("name"): input_tag.get("value") 
            for input_tag in mfa_form.find_all("input") 
            if input_tag.get("name") is not None
        }
        mfa_payload["answer"] = current_code
        mfa_action = mfa_form.get("action")
        response = session.post(mfa_action, data=mfa_payload)
        soup = BeautifulSoup(response.content, "html.parser")
    
    # 提取SAML断言
    saml_response = soup.find("input", {"name": "SAMLResponse"})
    if not saml_response:
        raise Exception("从Okta获取SAML响应失败")
    return saml_response.get("value")

# 获取SAML断言
saml_assertion = get_okta_saml_assertion()

# 建立Redshift连接
conn = redshift_connector.connect(
    host=REDSHIFT_HOST,
    port=REDSHIFT_PORT,
    database=REDSHIFT_DB,
    user=USERNAME,
    password=saml_assertion,
    auth="saml",
    ssl=True
)

# 适配SQLAlchemy引擎
engine = create_engine(
    f'redshift+redshift_connector://{USERNAME}:{saml_assertion}@{REDSHIFT_HOST}:{REDSHIFT_PORT}/{REDSHIFT_DB}?auth=saml'
)

# 加载数据到Pandas DataFrame
query = f"SELECT * FROM {REDSHIFT_SCHEMA}.your_table LIMIT 100"
df = pd.read_sql(query, engine)

# 验证结果
print(df.head())

# 清理连接资源
conn.close()
engine.dispose()

关键说明 #

• Okta参数获取：OKTA_APP_ID可从Okta应用的URL中提取（例如URL为https://xxx.okta.com/app/redshift_xxxx/xxxx，中间的redshift_xxxx即为APP_ID）。
• MFA适配：示例采用TOTP验证，若使用短信等其他MFA方式，需修改MFA部分的表单处理逻辑，对应适配Okta的验证页面结构。
• 安全提示：禁止硬编码敏感信息，建议通过环境变量或密钥管理工具（如AWS Secrets Manager）存储账号、密码及TOTP密钥。

内容的提问来源于stack exchange，提问作者Gabe