실습 5 — API 키로 접근 막기

import hashlib, hmac, json, os, secrets
import joblib
import pandas as pd

# 실제로는 관리자가 키를 만들어 사용자에게 한 번만 알려 주고, 서버의 환경 변수(또는 비밀 저장소)에 넣는다
issued = {"물류팀": secrets.token_urlsafe(24), "고객센터": secrets.token_urlsafe(24)}
os.environ["DELAY_API_KEYS"] = json.dumps({k: v for k, v in issued.items()})
print({team: key[:4] + "…" for team, key in issued.items()})    # 키는 앞 네 글자만 보인다

from fastapi import Depends, FastAPI, HTTPException, Security
from fastapi.security import APIKeyHeader
from fastapi.testclient import TestClient
from pydantic import BaseModel, Field

# 서버가 뜰 때 환경 변수의 키를 읽어 해시만 남긴다(키 원문은 들고 있지 않는다)
digest = lambda k: hashlib.sha256(k.encode()).hexdigest()
KEY_HASHES = {digest(k): team for team, k in json.loads(os.environ["DELAY_API_KEYS"]).items()}
header = APIKeyHeader(name="X-API-Key", auto_error=False)      # X-API-Key 머리글을 읽는다. 없으면 None


def require_key(key: str | None = Security(header)):
    """키가 없거나 맞지 않으면 401. 맞으면 키 주인(팀 이름)을 돌려준다"""
    if not key:
        raise HTTPException(status_code=401, detail="API 키가 없습니다")
    for h, team in KEY_HASHES.items():
        if hmac.compare_digest(h, digest(key)):
            return team
    raise HTTPException(status_code=401, detail="API 키가 맞지 않습니다")


print("등록된 팀:", sorted(KEY_HASHES.values()))

BUNDLE = joblib.load("out_delay_model.joblib")
app = FastAPI(title="배송 지연 예측 API(키 필요)")


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


@app.get("/health")                                 # 감시용 주소는 키 없이 열어 둔다
def health():
    return {"상태": "ok"}


@app.post("/predict")
def predict(order: Order, team: str = Depends(require_key)):     # Depends: 함수 전에 require_key 를 먼저 부른다
    p = BUNDLE["파이프라인"].predict_proba(pd.DataFrame([order.model_dump()]))[0, 1]
    return {"지연확률": round(float(p), 4), "부른 팀": team}


client = TestClient(app)
order = {"주문시각": 22, "거리km": 197.6, "무게kg": 4.6, "지역": "지방", "상품분류": "도서", "날씨": "눈", "택배사": "빠른별택배"}
print(client.get("/health").status_code)

cases = {"키 없음": {}, "틀린 키": {"X-API-Key": "wrong-key-0000"}, "물류팀 키": {"X-API-Key": issued["물류팀"]},
         "고객센터 키": {"X-API-Key": issued["고객센터"]}}
for name, headers in cases.items():
    r = client.post("/predict", json=order, headers=headers)
    print(f"{name}: {r.status_code} {r.json()}")
# 문서(OpenAPI)에도 키가 필요하다는 것이 적힌다
print(client.get("/openapi.json").json()["components"]["securitySchemes"])