from pyqqq import get_api_key
from pyqqq.utils.kvstore import KVStore
from tinydb import TinyDB, Query
from typing import Optional, Tuple
import datetime as dtm
import pyqqq.config as c
import requests
[docs]
class TossAuth:
"""
토스증권 인증 정보를 담고 있는 객체
Args:
client_id (str): 클라이언트 ID
client_secret (str): 클라이언트 시크릿
"""
[docs]
def __init__(self, client_id: str, client_secret: str):
assert client_id is not None, "Client id must be set"
assert client_secret is not None, "Client secret must be set"
self.client_id = client_id
self.client_secret = client_secret
self.repo = TossTokenRepository()
@property
def host_url(self) -> str:
"""
API 호스트 URL을 반환합니다.
Returns:
str: API 호스트 URL
"""
return "https://openapi.tossinvest.com"
[docs]
def get_token(self, refresh=False) -> str:
"""
액세스 토큰을 반환합니다.
Args:
refresh (bool): 강제 재발급
Returns:
str: 액세스 토큰
"""
token = self.repo.find(TossKeyTypes.ACCESS_TOKEN, self.client_id, self.client_secret)
if token is None or refresh is True:
token, expiry = self.issue_access_token()
self.repo.save(TossKeyTypes.ACCESS_TOKEN, self.client_id, self.client_secret, token, expiry)
return token
[docs]
def issue_access_token(self) -> Tuple[str, int]:
"""
액세스 토큰을 발급합니다.
OAuth 2.0 Client Credentials Grant 방식으로 토큰을 발급받습니다.
Returns:
tuple: 액세스 토큰과 만료 시간(UNIX 타임스탬프)을 담은 튜플
"""
if self.client_id is None or self.client_secret is None:
raise ValueError("Client id and client secret must be set")
headers = {"content-type": "application/x-www-form-urlencoded"}
body = {
"grant_type": "client_credentials",
"client_id": self.client_id,
"client_secret": self.client_secret,
}
r = requests.post(f"{self.host_url}/oauth2/token", headers=headers, data=body)
r.raise_for_status()
data = r.json()
access_token = data["access_token"]
expires_in = data["expires_in"]
# 만료 시간 직전 갱신을 피하기 위해 60초의 안전 여유를 둔다.
expiry = int(dtm.datetime.now().timestamp() + expires_in - 60)
return access_token, expiry
class TossTokenRepository:
def __init__(self):
if get_api_key():
self.repo = _RemoteTokenRepository()
else:
self.repo = _LocalTokenRepository()
def find(self, type, client_id, client_secret) -> Optional[str]:
"""
클라이언트 ID와 시크릿에 해당하는 액세스 토큰을 반환합니다.
Args:
type (str): 토큰 타입 (access_token)
client_id (str): 클라이언트 ID
client_secret (str): 클라이언트 시크릿
Returns:
str|None: 액세스 토큰
"""
return self.repo.find(type, client_id, client_secret)
def save(self, type, client_id, client_secret, token, expiry):
"""
클라이언트 ID와 시크릿에 해당하는 액세스 토큰을 저장합니다.
Args:
type (str): 토큰 타입 (access_token)
client_id (str): 클라이언트 ID
client_secret (str): 클라이언트 시크릿
token (str): 액세스 토큰
expiry (int): 만료 시간(UNIX 타임스탬프)
"""
return self.repo.save(type, client_id, client_secret, token, expiry)
class TossKeyTypes:
ACCESS_TOKEN = "access_token"
class _LocalTokenRepository:
def __init__(self):
self.db = TinyDB(c.get_tiny_db_path())
def find(self, type, client_id, client_secret) -> Optional[str]:
TokenCache = Query()
result = self.db.search(
TokenCache.type == type
and TokenCache.client_id == client_id
and TokenCache.client_secret == client_secret
and TokenCache.expiry > dtm.datetime.now().timestamp()
)
if len(result) > 0:
return result[0]["token"]
else:
return None
def save(self, type, client_id, client_secret, token, expiry):
TokenCache = Query()
self.db.upsert(
{
"type": type,
"client_id": client_id,
"client_secret": client_secret,
"token": token,
"expiry": expiry,
},
TokenCache.client_id == client_id and TokenCache.client_secret == client_secret,
)
class _RemoteTokenRepository:
def __init__(self):
self.kvstore = KVStore("toss_token_cache")
self.cache = {}
def _cache_key(self, type, client_id, client_secret):
return f"{type}_{client_id}_{client_secret}"
def find(self, type, client_id, client_secret) -> Optional[str]:
key = self._cache_key(type, client_id, client_secret)
# 1st. check local data
local_data = self.cache.get(key, None)
if local_data:
if local_data["expiry"] > dtm.datetime.now().timestamp():
return local_data["token"]
else:
self.cache.pop(key)
# 2nd. check remote data
if self.kvstore.get(key):
remote_data = self.kvstore.get(key)
if remote_data["expiry"] > dtm.datetime.now().timestamp():
self.cache[key] = remote_data
return remote_data["token"]
else:
self.kvstore.delete(key)
return None
def save(self, type, client_id, client_secret, token, expiry):
key = self._cache_key(type, client_id, client_secret)
data = {"token": token, "expiry": expiry}
self.cache[key] = data
self.kvstore.set(key, data)