종목 최대 40개

# -*- coding: utf-8 -*-
### 모듈 임포트 ###
import os
import sys
import json
import time
import requests
import asyncio
import traceback
import websockets

import Crypto

from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
from base64 import b64decode

# 로깅 설정
import logging
logging.basicConfig(level=logging.INFO)

# AES256 DECODE
def aes_cbc_base64_dec(key, iv, cipher_text):
    cipher = AES.new(key.encode('utf-8'), AES.MODE_CBC, iv.encode('utf-8'))
    return bytes.decode(unpad(cipher.decrypt(b64decode(cipher_text)), AES.block_size))

# 웹소켓 접속키 발급
def get_approval(key, secret):
    url = '<https://openapivts.koreainvestment.com:29443>'
    headers = {"content-type": "application/json"}
    body = {"grant_type": "client_credentials",
            "appkey": key,
            "secretkey": secret}
    PATH = "oauth2/Approval"
    URL = f"{url}/{PATH}"
    res = requests.post(URL, headers=headers, data=json.dumps(body))
    
    if res.status_code == 200:
        approval_key = res.json().get("approval_key")
        return approval_key
    else:
        logging.error("Failed to get approval key: %s" % res.text)
        return None

# 주식 체결 데이터 처리
def stockspurchase_domestic(data_cnt, data):
    print("============================================")
    menulist = "유가증권단축종목코드|주식체결시간|주식현재가|전일대비부호|전일대비|전일대비율|가중평균주식가격|주식시가|주식최고가|주식최저가|매도호가1|매수호가1|체결거래량|누적거래량|누적거래대금|매도체결건수|매수체결건수|순매수체결건수|체결강도|총매도수량|총매수수량|체결구분|매수비율|전일거래량대비등락율|시가시간|시가대비구분|시가대비|최고가시간|고가대비구분|고가대비|최저가시간|저가대비구분|저가대비|영업일자|신장운영구분코드|거래정지여부|매도호가잔량|매수호가잔량|총매도호가잔량|총매수호가잔량|거래량회전율|전일동시간누적거래량|전일동시간누적거래량비율|시간구분코드|임의종료구분코드|정적VI발동기준가"
    menustr = menulist.split('|')
    pValue = data.split('^')
    
    # 데이터 포맷팅을 위한 인덱스
    ticker_index = menustr.index("유가증권단축종목코드")
    current_price_index = menustr.index("주식현재가")

    for cnt in range(data_cnt):  # 넘겨받은 체결데이터 개수만큼 출력
        print("### [%d / %d]" % (cnt + 1, data_cnt))
        
        # 종목 코드와 현재가만 출력
        ticker_code = pValue[ticker_index]
        current_price = pValue[current_price_index]
        
        print(f"종목 코드: {ticker_code}, 현재가: {current_price}")

# 승인 키를 재사용하기 위한 전역 변수
g_approval_key = None

async def connect_single_socket(url, tr_keys):
    global g_approval_key

    try:
        if g_approval_key is None:
            g_appkey = "PSKJKSd75HbuPliK70C0cUdl4OYAvyIGmROi"
            g_appsecret = "ZtH0c5eLI1BJnBZ1f9LVn9ggNT7Af2NHVTOu7dMBdvxCy7OhDxnuBjKU8YkpiYXgxxHJgvuuOSbqKV/HeuQ+smMz0K88mNkiRtXGZASlotr+nMaOPCerfs/fduLBhbimdiimL4IobO6Ne0VPZEPanlFzz0x91nbDi/W8wmLOa5ZkkmRrDus="
            g_approval_key = get_approval(g_appkey, g_appsecret)
            if not g_approval_key:
                logging.error("Failed to get approval key. Exiting.")
                return

        logging.info("WebSocket connection established for %s" % ', '.join(tr_keys))
        async with websockets.connect(url, ping_interval=None) as websocket:
            for tr_key in tr_keys:
                senddata = '{"header":{"approval_key": "%s","custtype":"P","tr_type":"1","content-type":"utf-8"},"body":{"input":{"tr_id":"H0STCNT0","tr_key":"%s"}}}' % (g_approval_key, tr_key)
                await websocket.send(senddata)
                logging.info(f"Input Command is: {senddata}")

            while True:
                try:
                    data = await websocket.recv()
                    # logging.info(f"Received Command is: {data}")

                    if data[0] == '0':
                        recvstr = data.split('|')
                        trid0 = recvstr[1]

                        if trid0 == "H0STCNT0":  # 주식체결 데이터 처리
                            logging.info("#### 국내주식 체결 ####")
                            data_cnt = int(recvstr[2])  # 체결데이터 개수
                            stockspurchase_domestic(data_cnt, recvstr[3])
                
                    else:
                        jsonObject = json.loads(data)
                        trid = jsonObject["header"]["tr_id"]

                        if trid != "PINGPONG":
                            rt_cd = jsonObject["body"]["rt_cd"]

                            if rt_cd == '1':  # 에러일 경우 처리
                                logging.error("### ERROR RETURN CODE [ %s ][ %s ] MSG [ %s ]" % (jsonObject["header"]["tr_key"], rt_cd, jsonObject["body"]["msg1"]))
                                if jsonObject["body"]["msg1"] == 'invalid approval : NOT FOUND':
                                    g_approval_key = None  # 승인 키를 초기화하여 다음 연결 시 새로 요청
                                break

                            elif rt_cd == '0':  # 정상일 경우 처리
                                logging.info("### RETURN CODE [ %s ][ %s ] MSG [ %s ]" % (jsonObject["header"]["tr_key"], rt_cd, jsonObject["body"]["msg1"]))
                                if trid in ["H0STCNI0", "H0STCNI9"]:
                                    aes_key = jsonObject["body"]["output"]["key"]
                                    aes_iv = jsonObject["body"]["output"]["iv"]
                                    logging.info("### TRID [%s] KEY[%s] IV[%s]" % (trid, aes_key, aes_iv))

                        elif trid == "PINGPONG":
                            logging.info("### RECV [PINGPONG] [%s]" % (data))
                            await websocket.pong(data)
                            logging.info("### SEND [PINGPONG] [%s]" % (data))

                except websockets.ConnectionClosed:
                    logging.error("WebSocket connection closed.")
                    break

    except Exception as e:
        logging.error('Exception Raised!')
        logging.error(e)
        await asyncio.sleep(1)
        await connect_single_socket(url, tr_keys)  # 재연결 시도 

async def connect_multiple_sockets(code_list):
    url = 'ws://ops.koreainvestment.com:31000'
    tasks = []

    # 종목 코드를 40개씩 묶어서 소켓 연결
    for i in range(0, len(code_list), 40):
        batch = code_list[i:i + 40]
        tr_keys = [k[2] for k in batch]  # 종목 코드만 추출
        tasks.append(connect_single_socket(url, tr_keys))

    await asyncio.gather(*tasks)

# 메인 실행 부분
if __name__ == "__main__":
    code_list = [
                        ['1', 'H0STCNT0', '005930'],  # 삼성전자
                        ['1', 'H0STCNT0', '066570'],  # LG전자
                        ['1', 'H0STCNT0', '000660'],  # SK하이닉스
                        ['1', 'H0STCNT0', '035420'],  # NAVER
                        ['1', 'H0STCNT0', '051910'],  # 삼성바이오로직스
                        ['1', 'H0STCNT0', '005380'],  # 현대차
                        ['1', 'H0STCNT0', '017670'],  # SK텔레콤
                        ['1', 'H0STCNT0', '032830'],  # 삼성SDI
                        ['1', 'H0STCNT0', '096770'],  # 카카오
                        ['1', 'H0STCNT0', '003550'],  # LG화학

                        ['1', 'H0STCNT0', '012330'],  # 현대모비스
                        ['1', 'H0STCNT0', '028260'],  # 삼성물산
                        ['1', 'H0STCNT0', '035720'],  # 카카오게임즈
                        ['1', 'H0STCNT0', '006400'],  # 현대제철
                        ['1', 'H0STCNT0', '009150'],  # 삼성전기
                        ['1', 'H0STCNT0', '018880'],  # S-Oil
                        ['1', 'H0STCNT0', '034730'],  # LG전자
                        ['1', 'H0STCNT0', '017390'],  # 포스코
                        ['1', 'H0STCNT0', '019170'],  # KB금융
                        ['1', 'H0STCNT0', '105560'],  # 삼성생명
                        
                        ['1', 'H0STCNT0', '086790'],  # 신한지주
                        ['1', 'H0STCNT0', '032640'],  # 하나금융지주
                        ['1', 'H0STCNT0', '003670'],  # LG디스플레이
                        ['1', 'H0STCNT0', '028050'],  # 아모레퍼시픽
                        ['1', 'H0STCNT0', '030200'],  # 아시아나항공
                        ['1', 'H0STCNT0', '002790'],  # SK증권
                        ['1', 'H0STCNT0', '004020'],  # 한국전력
                        ['1', 'H0STCNT0', '006800'],  # 한화솔루션
                        ['1', 'H0STCNT0', '001040'],  # 삼성엔지니어링
                        ['1', 'H0STCNT0', '007570'],  # 에스원
                        
                        ['1', 'H0STCNT0', '010950'],  # 삼성카드
                        ['1', 'H0STCNT0', '001680'],  # LG상사
                        ['1', 'H0STCNT0', '017550'],  # KT&G
                        ['1', 'H0STCNT0', '006260'],  # 삼성중공업
                        ['1', 'H0STCNT0', '009830'],  # 한화생명
                        ['1', 'H0STCNT0', '004990'],  # 두산중공업
                        ['1', 'H0STCNT0', '035590'],  # 카카오뱅크
                        ['1', 'H0STCNT0', '008930'],  # 한국가스공사
                        ['1', 'H0STCNT0', '010140'],  # 현대건설
                        ['1', 'H0STCNT0', '005250'],  # SK이노베이션
                        
                        ['1', 'H0STCNT0', '004560'],  # LG유플러스
                        ['1', 'H0STCNT0', '027410'],  # 한국조선해양
                        ['1', 'H0STCNT0', '052690'],  # GS건설
                        ['1', 'H0STCNT0', '025540'],  # CJ대한통운
                        ['1', 'H0STCNT0', '028670'],  # LG화학
                        ['1', 'H0STCNT0', '036570'],  # 카카오게임즈
                        ['1', 'H0STCNT0', '036810'],  # SK온
                        ['1', 'H0STCNT0', '006740'],  # 한화
                        ['1', 'H0STCNT0', '008770'],  # 한미약품
                        ['1', 'H0STCNT0', '003830'],  # 삼성제약
                        
                        ['1', 'H0STCNT0', '018260'],  # 삼성에스디아이
                        ['1', 'H0STCNT0', '007680'],  # 한국타이어
                        ['1', 'H0STCNT0', '006400'],  # 현대자동차
                        ['1', 'H0STCNT0', '014820'],  # 삼성전기
                        ['1', 'H0STCNT0', '005940'],  # SK텔레콤
                        ['1', 'H0STCNT0', '002350'],  # 삼성중공업
                        ['1', 'H0STCNT0', '000270'],  # 기아
                        ['1', 'H0STCNT0', '004690'],  # 한미사이언스
                        ['1', 'H0STCNT0', '010600'],  # 신풍제약
                        ['1', 'H0STCNT0', '000100'],  # 유한양행
                        
                        ['1', 'H0STCNT0', '036810'],  # SK바이오사이언스
                        ['1', 'H0STCNT0', '007700'],  # 동양생명
                        ['1', 'H0STCNT0', '002960'],  # 하나금융지주
                        ['1', 'H0STCNT0', '000660'],  # SK하이닉스
                        ['1', 'H0STCNT0', '004370'],  # 삼양사
                        ['1', 'H0STCNT0', '003830'],  # 삼성제약
                        ['1', 'H0STCNT0', '002120'],  # 한국타이어
                        ['1', 'H0STCNT0', '005490'],  # 삼성정밀화학
                        ['1', 'H0STCNT0', '009140'],  # 대우건설
                    ]
    asyncio.run(connect_multiple_sockets(code_list))