155 lines
5.7 KiB
Python
155 lines
5.7 KiB
Python
import asyncio
|
|
import json
|
|
import websockets
|
|
import logging
|
|
from typing import Dict, Set, Callable, Optional
|
|
|
|
from app.services.kis_auth import kis_auth
|
|
from app.core.crypto import aes_cbc_base64_dec
|
|
from app.db.database import SessionLocal
|
|
# from app.db.crud import update_stock_price # TODO: Implement CRUD
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class RealtimeManager:
|
|
"""
|
|
Manages KIS WebSocket Connection.
|
|
Handles: Connection, Subscription, Decryption, PINGPONG.
|
|
"""
|
|
WS_URL_REAL = "ws://ops.koreainvestment.com:21000"
|
|
|
|
def __init__(self):
|
|
self.ws: Optional[websockets.WebSocketClientProtocol] = None
|
|
self.approval_key: Optional[str] = None
|
|
self.subscribed_codes: Set[str] = set()
|
|
self.running = False
|
|
self.data_map: Dict[str, Dict] = {} # Store IV/Key for encrypted TRs
|
|
|
|
async def start(self):
|
|
"""
|
|
Main loop: Connect -> Authenticate -> Listen
|
|
"""
|
|
self.running = True
|
|
while self.running:
|
|
try:
|
|
# 1. Get Approval Key
|
|
self.approval_key = await kis_auth.get_approval_key()
|
|
|
|
logger.info(f"Connecting to KIS WS: {self.WS_URL_REAL}")
|
|
async with websockets.connect(self.WS_URL_REAL, ping_interval=None) as websocket:
|
|
self.ws = websocket
|
|
logger.info("Connected.")
|
|
|
|
# 2. Resubscribe if recovering connection
|
|
if self.subscribed_codes:
|
|
await self._resubscribe_all()
|
|
|
|
# 3. Listen Loop
|
|
await self._listen()
|
|
|
|
except Exception as e:
|
|
logger.error(f"WS Connection Error: {e}. Retrying in 5s...")
|
|
await asyncio.sleep(5)
|
|
|
|
async def stop(self):
|
|
self.running = False
|
|
if self.ws:
|
|
await self.ws.close()
|
|
|
|
async def subscribe(self, stock_code: str, type="price"):
|
|
"""
|
|
Subscribe to a stock.
|
|
type: 'price' (H0STCNT0 - 체결가)
|
|
"""
|
|
if not self.ws or not self.approval_key:
|
|
logger.warning("WS not ready. Adding to pending list.")
|
|
self.subscribed_codes.add(stock_code)
|
|
return
|
|
|
|
# Domestic Realtime Price TR ID: H0STCNT0
|
|
tr_id = "H0STCNT0"
|
|
tr_key = stock_code
|
|
|
|
payload = {
|
|
"header": {
|
|
"approval_key": self.approval_key,
|
|
"custtype": "P",
|
|
"tr_type": "1", # 1=Register, 2=Unregister
|
|
"content-type": "utf-8"
|
|
},
|
|
"body": {
|
|
"input": {
|
|
"tr_id": tr_id,
|
|
"tr_key": tr_key
|
|
}
|
|
}
|
|
}
|
|
|
|
await self.ws.send(json.dumps(payload))
|
|
self.subscribed_codes.add(stock_code)
|
|
logger.info(f"Subscribed to {stock_code}")
|
|
|
|
async def _resubscribe_all(self):
|
|
for code in self.subscribed_codes:
|
|
await self.subscribe(code)
|
|
|
|
async def _listen(self):
|
|
async for message in self.ws:
|
|
try:
|
|
# Message can be plain text provided by library, or bytes
|
|
if isinstance(message, bytes):
|
|
message = message.decode('utf-8')
|
|
|
|
# KIS sends data in specific formats.
|
|
# 1. JSON (Control Messages, PINGPONG, Subscription Ack)
|
|
# 2. Text/Pipe separated (Real Data) - Usually starts with 0 or 1
|
|
|
|
first_char = message[0]
|
|
|
|
if first_char in ['{', '[']:
|
|
# JSON Message
|
|
data = json.loads(message)
|
|
header = data.get('header', {})
|
|
tr_id = header.get('tr_id')
|
|
|
|
if tr_id == "PINGPONG":
|
|
await self.ws.send(message) # Echo back
|
|
logger.debug("PINGPONG handled")
|
|
|
|
elif 'body' in data:
|
|
# Subscription Ack
|
|
# Store IV/Key if encryption is enabled (msg1 often contains 'ENCRYPT')
|
|
# But for Brokerage API, H0STCNT0 is usually plaintext unless configured otherwise.
|
|
# If encrypted, 'iv' and 'key' are in body['output']
|
|
pass
|
|
|
|
elif first_char in ['0', '1']:
|
|
# Real Data: 0|TR_ID|DATA_CNT|DATA...
|
|
parts = message.split('|')
|
|
if len(parts) < 4:
|
|
continue
|
|
|
|
tr_id = parts[1]
|
|
raw_data = parts[3]
|
|
|
|
# Decryption Check
|
|
# If this tr_id was registered as encrypted, decrypt it.
|
|
# For now assuming Plaintext for H0STCNT0 as per standard Personal API.
|
|
|
|
# Parse Data
|
|
if tr_id == "H0STCNT0": # Domestic Price
|
|
# Data format: TIME^PRICE^...
|
|
# We need to look up format spec.
|
|
# Simple implementation: just log or split
|
|
fields = raw_data.split('^')
|
|
if len(fields) > 2:
|
|
current_price = fields[2] # Example index
|
|
# TODO: Update DB
|
|
# print(f"Price Update: {current_price}")
|
|
pass
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing WS message: {e}")
|
|
|
|
realtime_manager = RealtimeManager()
|