실습 6 — 토큰과 역할로 권한 나누기

import hashlib, hmac, os, secrets, time
from datetime import datetime, timedelta, timezone

def hash_password(password, salt=None):
    """소금(salt)을 붙여 20만 번 해시한다. 같은 비밀번호도 사용자마다 다른 해시가 된다"""
    salt = salt or secrets.token_bytes(16)
    return salt, hashlib.pbkdf2_hmac("sha256", password.encode(), salt, 200_000)

# 시험용 사용자. 실제로는 사용자가 정한 비밀번호의 해시만 DB 에 있다
TEST_PASSWORDS = {"상담원1": secrets.token_urlsafe(9), "분석가1": secrets.token_urlsafe(9), "관리자1": secrets.token_urlsafe(9)}
USERS = {}
for name, role in [("상담원1", "상담원"), ("분석가1", "분석가"), ("관리자1", "관리자")]:
    salt, h = hash_password(TEST_PASSWORDS[name])
    USERS[name] = {"역할": role, "salt": salt, "hash": h}
t = time.perf_counter()
hash_password("시험", b"0" * 16)
print("사용자:", {n: u["역할"] for n, u in USERS.items()}, f"· 해시 한 번 {(time.perf_counter() - t) * 1000:.0f}ms")

import jwt
import joblib
import pandas as pd
from fastapi import Depends, FastAPI, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from fastapi.testclient import TestClient
from pydantic import BaseModel, Field

SECRET = os.environ.get("DELAY_JWT_SECRET") or secrets.token_hex(32)      # 서명 비밀 키. 없으면 시험용으로 만든다
BUNDLE = joblib.load("out_delay_model.joblib")
app = FastAPI(title="배송 지연 예측 API(역할별 권한)")


class Login(BaseModel):
    아이디: str
    비밀번호: str


@app.post("/token")
def login(body: Login):
    user = USERS.get(body.아이디)
    ok = user is not None and hmac.compare_digest(hash_password(body.비밀번호, user["salt"])[1], user["hash"])
    if not ok:                                         # 아이디가 없는지 비밀번호가 틀렸는지 알려 주지 않는다
        raise HTTPException(status_code=401, detail="아이디나 비밀번호가 맞지 않습니다")
    exp = datetime.now(timezone.utc) + timedelta(minutes=30)               # 30분 뒤 만료
    token = jwt.encode({"sub": body.아이디, "role": user["역할"], "exp": exp}, SECRET, algorithm="HS256")
    return {"access_token": token, "token_type": "bearer", "역할": user["역할"]}


client = TestClient(app)
r = client.post("/token", json={"아이디": "분석가1", "비밀번호": TEST_PASSWORDS["분석가1"]})
token = r.json()["access_token"]
print(r.status_code, r.json()["역할"], token[:20] + "…")
print("토큰 속:", jwt.decode(token, SECRET, algorithms=["HS256"]))
print("틀린 비밀번호:", client.post("/token", json={"아이디": "분석가1", "비밀번호": "틀림"}).status_code)

bearer = HTTPBearer(auto_error=False)                 # Authorization: Bearer 토큰 머리글을 읽는다


def current_user(cred: HTTPAuthorizationCredentials | None = Depends(bearer)):
    """토큰이 없거나, 서명이 틀리거나, 만료됐으면 401"""
    if cred is None:
        raise HTTPException(status_code=401, detail="토큰이 없습니다")
    try:
        return jwt.decode(cred.credentials, SECRET, algorithms=["HS256"])
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="토큰이 만료됐습니다")
    except jwt.InvalidTokenError:
        raise HTTPException(status_code=401, detail="토큰이 올바르지 않습니다")


def allow(*roles):
    """허락할 역할을 받아, 그 역할이 아니면 403 을 내는 의존성을 만든다"""
    def check(user=Depends(current_user)):
        if user["role"] not in roles:
            raise HTTPException(status_code=403, detail=f"{user['role']} 은(는) 쓸 수 없습니다")
        return user
    return check


class Order(BaseModel):
    주문시각: int = Field(ge=0, le=23)
    거리km: float = Field(gt=0, le=1000)
    무게kg: float = Field(ge=0, le=100)
    지역: str
    상품분류: str
    날씨: str = "모름"
    택배사: str


def probs(orders):
    return BUNDLE["파이프라인"].predict_proba(pd.DataFrame([o.model_dump() for o in orders]))[:, 1].round(4).tolist()


@app.post("/predict")
def predict(order: Order, user=Depends(allow("상담원", "분석가", "관리자"))):
    return {"지연확률": probs([order])[0], "부른 사람": user["sub"]}


@app.post("/predict/batch")
def predict_batch(orders: list[Order], user=Depends(allow("분석가", "관리자"))):
    return {"건수": len(orders), "지연확률": probs(orders)}


@app.get("/model")
def model_info(user=Depends(allow("관리자"))):
    return {"판": BUNDLE["판"], "시험성능": BUNDLE["시험성능"]}


print([r.path for r in app.routes if r.path.startswith(("/predict", "/model", "/token"))])

order = {"주문시각": 22, "거리km": 197.6, "무게kg": 4.6, "지역": "지방", "상품분류": "도서", "날씨": "눈", "택배사": "빠른별택배"}
calls = [("POST", "/predict", order), ("POST", "/predict/batch", [order, order]), ("GET", "/model", None)]
rows = {}
for name in USERS:
    tok = client.post("/token", json={"아이디": name, "비밀번호": TEST_PASSWORDS[name]}).json()["access_token"]
    head = {"Authorization": f"Bearer {tok}"}
    rows[USERS[name]["역할"]] = [client.request(m, p, json=b, headers=head).status_code for m, p, b in calls]
# 토큰 없이, 위조한 토큰, 만료된 토큰
expired = jwt.encode({"sub": "관리자1", "role": "관리자", "exp": datetime.now(timezone.utc) - timedelta(minutes=1)}, SECRET, algorithm="HS256")
forged = jwt.encode({"sub": "누군가", "role": "관리자", "exp": datetime.now(timezone.utc) + timedelta(minutes=5)}, secrets.token_hex(32), algorithm="HS256")
for name, head in [("토큰 없음", {}), ("위조 토큰", {"Authorization": f"Bearer {forged}"}), ("만료 토큰", {"Authorization": f"Bearer {expired}"})]:
    rows[name] = [client.request(m, p, json=b, headers=head).status_code for m, p, b in calls]
print(pd.DataFrame(rows, index=[f"{m} {p}" for m, p, _ in calls]).T.to_string())