from __future__ import annotations import asyncio import json from datetime import datetime, timedelta from pathlib import Path import httpx from app.core.config import get_base_url, settings from app.core.logger import setup_logger from app.core.rate_limiter import RateLimiter logger = setup_logger("auth") _TOKEN_URL = "/oauth2/tokenP" _APPROVAL_URL = "/oauth2/Approval" _HASHKEY_URL = "/uapi/hashkey" _TOKEN_CACHE_PATH = Path(settings.app.db_path).parent / ".auth_cache.json" class TokenManager: def __init__(self) -> None: self._access_token: str = "" self._token_expires_at: datetime = datetime.min self._approval_key: str = "" self._approval_expires_at: datetime = datetime.min self._rate_limiter = RateLimiter(requests_per_second=1.0) self._lock = asyncio.Lock() self._client = httpx.AsyncClient(timeout=15.0) self._authenticated: bool = False self._load_cache() def _load_cache(self) -> None: if not _TOKEN_CACHE_PATH.exists(): return try: data = json.loads(_TOKEN_CACHE_PATH.read_text()) now = datetime.now() token_exp = datetime.fromisoformat(data.get("token_expires_at", "")) if data.get("access_token") and now < token_exp: self._access_token = data["access_token"] self._token_expires_at = token_exp self._authenticated = True logger.info("캐시에서 토큰 복원 완료 (만료: %s)", token_exp) approval_exp = datetime.fromisoformat(data.get("approval_expires_at", "")) if data.get("approval_key") and now < approval_exp: self._approval_key = data["approval_key"] self._approval_expires_at = approval_exp logger.info("캐시에서 approval_key 복원 완료 (만료: %s)", approval_exp) except Exception as e: logger.debug("토큰 캐시 로드 실패: %s", e) def _save_cache(self) -> None: try: _TOKEN_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True) _TOKEN_CACHE_PATH.write_text(json.dumps({ "access_token": self._access_token, "token_expires_at": self._token_expires_at.isoformat(), "approval_key": self._approval_key, "approval_expires_at": self._approval_expires_at.isoformat(), })) except Exception as e: logger.debug("토큰 캐시 저장 실패: %s", e) @property def is_authenticated(self) -> bool: return self._authenticated @property def is_vps(self) -> bool: return settings.kis.server_mode == "vps" async def _request_token(self) -> None: if not settings.kis.app_key or not settings.kis.app_secret: logger.debug("API 키가 설정되지 않음") return url = f"{get_base_url()}{_TOKEN_URL}" payload = { "grant_type": "client_credentials", "appkey": settings.kis.app_key, "appsecret": settings.kis.app_secret, } for attempt in range(3): resp = await self._client.post(url, json=payload) data = resp.json() if data.get("error_code") == "EGW00133": wait = 65 * (attempt + 1) logger.warning("토큰 발급 1분 제한 - %d초 후 재시도 (%d/3)", wait, attempt + 1) await asyncio.sleep(wait) continue if resp.status_code != 200 or "access_token" not in data: logger.warning("토큰 발급 실패: %s", data.get("message", resp.status_code)) return self._access_token = data["access_token"] expires_in = int(data.get("expires_in", 7776000)) self._token_expires_at = datetime.now() + timedelta(seconds=expires_in) self._authenticated = True self._save_cache() logger.info("REST access_token 발급 완료 (만료: %s)", self._token_expires_at) return logger.error("토큰 발급 3회 모두 실패 (1분 제한)") async def _request_approval_key(self) -> None: if not settings.kis.app_key or not settings.kis.app_secret: return url = f"{get_base_url()}{_APPROVAL_URL}" payload = { "grant_type": "client_credentials", "appkey": settings.kis.app_key, "secretkey": settings.kis.app_secret, } for attempt in range(3): resp = await self._client.post(url, json=payload) data = resp.json() if data.get("error_code") == "EGW00133": wait = 65 * (attempt + 1) logger.warning( "approval_key 발급 1분 제한 - %d초 후 재시도 (%d/3)", wait, attempt + 1 ) await asyncio.sleep(wait) continue if resp.status_code != 200 or "approval_key" not in data: logger.warning("approval_key 발급 실패: %s", data.get("message", resp.status_code)) return self._approval_key = data["approval_key"] self._approval_expires_at = datetime.now() + timedelta(hours=23, minutes=50) self._save_cache() logger.info("WebSocket approval_key 발급 완료 (만료: %s)", self._approval_expires_at) return logger.error("approval_key 발급 3회 모두 실패 (1분 제한)") async def get_access_token(self) -> str: async with self._lock: if datetime.now() >= self._token_expires_at: await self._request_token() return self._access_token async def get_approval_key(self) -> str: async with self._lock: if datetime.now() >= self._approval_expires_at: await self._request_approval_key() return self._approval_key async def generate_hashkey(self, data: dict) -> str: await self._rate_limiter.acquire() url = f"{get_base_url()}{_HASHKEY_URL}" headers = { "content-type": "application/json", "appkey": settings.kis.app_key, "appsecret": settings.kis.app_secret, } resp = await self._client.post(url, json=data, headers=headers) if resp.status_code != 200: raise RuntimeError(f"hashkey 생성 실패: {resp.status_code}") return resp.json()["HASH"] def get_auth_headers(self, tr_id: str) -> dict[str, str]: return { "Content-Type": "application/json; charset=utf-8", "authorization": f"Bearer {self._access_token}", "appKey": settings.kis.app_key, "appSecret": settings.kis.app_secret, "tr_id": tr_id, "custtype": "P", } async def close(self) -> None: await self._client.aclose() token_manager = TokenManager()