106 lines
3.5 KiB
Python
106 lines
3.5 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)
|
|
|
|
@property
|
|
def is_vps(self) -> bool:
|
|
return settings.kis.server_mode == "vps"
|
|
|
|
async def _request_token(self) -> None:
|
|
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)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
self._access_token = data["access_token"]
|
|
expires_in = int(data.get("expires_in", 7776000))
|
|
self._token_expires_at = datetime.now() + timedelta(seconds=expires_in)
|
|
logger.info("REST access_token 발급 완료 (만료: %s)", self._token_expires_at)
|
|
|
|
async def _request_approval_key(self) -> None:
|
|
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)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
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)
|
|
resp.raise_for_status()
|
|
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()
|