Files
stockautomtion/app/core/auth.py

126 lines
4.2 KiB
Python

from __future__ import annotations
import asyncio
import hashlib
import hmac
import time
from datetime import datetime, timedelta
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"
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=10.0)
self._authenticated: bool = False
@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,
}
resp = await self._client.post(url, json=payload)
data = resp.json()
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
logger.info("REST access_token 발급 완료 (만료: %s)", self._token_expires_at)
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,
}
resp = await self._client.post(url, json=payload)
data = resp.json()
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)
logger.info("WebSocket approval_key 발급 완료 (만료: %s)", self._approval_expires_at)
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()