kis_bot/kis_approval_manager.py

#!/usr/bin/env python3
"""
kis_approval_manager.py — KIS WebSocket approval_key 파일 캐시 (국내·해외 WS 공유)
==================================================================================
REST access_token(.kis_token_cache_*.json) 과 동일 패턴:
  - 프로세스·재시작·국내/해외 WS 가 **하나의 approval_key** 를 공유
  - KIS 정책: 24h 유효, 6h 이내 REST 재발급 금지
  - 파일 잠금으로 동시 발급(국내+해외·재시작) 경합 방지

문제(기존):
  - kis_ws / kis_ws_overseas 가 각각 /oauth2/Approval 호출 → 나중 키가 앞 키 무효화
  - 메모리만 사용 → 재시작마다 새 발급 → invalid approval 루프
"""
from __future__ import annotations

import json
import logging
import os
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

import requests

logger = logging.getLogger(__name__)

ROOT = Path(__file__).parent
CACHE_MOCK = ROOT / ".kis_approval_cache_mock.json"
CACHE_REAL = ROOT / ".kis_approval_cache_real.json"
LOCK_FILE = ROOT / ".kis_approval_manager.lock"
LOCK_TIMEOUT_S = 60

DEFAULT_MIN_REISSUE_SEC = 21600.0
DEFAULT_MAX_VALID_SEC = 82800.0

_am_instances: dict = {}
_am_instances_lock = threading.Lock()


def _min_reissue_sec() -> float:
    try:
        raw = os.environ.get("KIS_WS_APPROVAL_MIN_REISSUE_SEC", "").strip()
        return max(0.0, float(raw)) if raw else DEFAULT_MIN_REISSUE_SEC
    except (ValueError, TypeError):
        return DEFAULT_MIN_REISSUE_SEC


def _max_valid_sec() -> float:
    try:
        raw = os.environ.get("KIS_WS_APPROVAL_MAX_VALID_SEC", "").strip()
        return max(3600.0, float(raw)) if raw else DEFAULT_MAX_VALID_SEC
    except (ValueError, TypeError):
        return DEFAULT_MAX_VALID_SEC


def _acquire_lock() -> bool:
    deadline = time.time() + LOCK_TIMEOUT_S
    while time.time() < deadline:
        try:
            fd = os.open(str(LOCK_FILE), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            os.write(fd, str(os.getpid()).encode())
            os.close(fd)
            return True
        except FileExistsError:
            try:
                if time.time() - LOCK_FILE.stat().st_mtime > 300:
                    LOCK_FILE.unlink(missing_ok=True)
                    continue
            except Exception:
                pass
            time.sleep(1)
        except Exception as exc:
            logger.warning("approval 잠금 획득 실패: %s", exc)
            return False
    logger.warning("approval 잠금 타임아웃 → 파일 캐시 재사용 시도")
    return False


def _release_lock() -> None:
    try:
        LOCK_FILE.unlink(missing_ok=True)
    except Exception:
        pass


def _issue_rest(app_key: str, app_secret: str, base_url: str) -> Optional[str]:
    url = f"{base_url.rstrip('/')}/oauth2/Approval"
    body = {
        "grant_type": "client_credentials",
        "appkey": app_key,
        "secretkey": app_secret,
    }
    try:
        resp = requests.post(url, json=body, timeout=10)
        data = resp.json()
        key = str(data.get("approval_key") or "").strip()
        if key:
            return key
        logger.error("❌ approval_key REST 발급 실패: %s", data)
    except Exception as exc:
        logger.error("❌ approval_key REST 요청 예외: %s", exc)
    return None


class KISApprovalManager:
    """KIS WebSocket approval_key 싱글톤 — 실전/모의 각 1개."""

    @classmethod
    def instance(cls, is_mock: bool) -> "KISApprovalManager":
        with _am_instances_lock:
            if is_mock not in _am_instances:
                _am_instances[is_mock] = cls(is_mock)
            return _am_instances[is_mock]

    def __init__(self, is_mock: bool) -> None:
        self._is_mock = bool(is_mock)
        self._mode_str = "모의" if is_mock else "실전"
        self._cache_path = CACHE_MOCK if is_mock else CACHE_REAL
        self._lock = threading.Lock()
        self._approval_key: Optional[str] = None
        self._issued_ts: float = 0.0
        self._app_key_prefix: str = ""
        self._load_from_file()

    def _load_from_file(self) -> None:
        if not self._cache_path.exists():
            return
        try:
            data = json.loads(self._cache_path.read_text(encoding="utf-8"))
            key = str(data.get("approval_key") or "").strip()
            ts = float(data.get("issued_at_epoch") or 0.0)
            prefix = str(data.get("app_key_prefix") or "")
            if not key or ts <= 0:
                return
            age = time.time() - ts
            if age >= _max_valid_sec():
                return
            self._approval_key = key
            self._issued_ts = ts
            self._app_key_prefix = prefix
        except Exception as exc:
            logger.debug("approval 파일 로드 실패 [%s]: %s", self._mode_str, exc)

    def _save_to_file(self) -> None:
        if not self._approval_key or self._issued_ts <= 0:
            return
        try:
            self._cache_path.write_text(
                json.dumps(
                    {
                        "approval_key": self._approval_key,
                        "issued_at": datetime.fromtimestamp(self._issued_ts).strftime(
                            "%Y-%m-%d %H:%M:%S"
                        ),
                        "issued_at_epoch": self._issued_ts,
                        "app_key_prefix": self._app_key_prefix,
                        "mock": self._is_mock,
                    },
                    ensure_ascii=False,
                    indent=2,
                ),
                encoding="utf-8",
            )
        except Exception as exc:
            logger.warning("approval 파일 저장 실패 (%s): %s", self._cache_path.name, exc)

    def age_sec(self) -> float:
        if not self._issued_ts:
            return 999999.0
        return max(0.0, time.time() - self._issued_ts)

    def reload_from_file(self) -> Optional[str]:
        """REST 없이 파일 캐시만 다시 읽기 (다른 프로세스 갱신 반영)."""
        with self._lock:
            self._load_from_file()
            return self._approval_key

    @property
    def issued_ts(self) -> float:
        return self._issued_ts

    def get_approval_key(
        self,
        app_key: str,
        app_secret: str,
        base_url: str,
        *,
        force_refresh: bool = False,
    ) -> Optional[str]:
        """
        유효한 approval_key 반환.
        - 24h 이내: 파일/메모리 캐시 재사용 (재연결 포함)
        - 6h 미경과: force_refresh=True 여도 REST 재발급 금지
        - 6h 경과 + 24h 미만 + force_refresh: REST 재발급 허용
        - 24h 경과: REST 재발급
        """
        app_key = (app_key or "").strip()
        app_secret = (app_secret or "").strip()
        if not app_key or not app_secret:
            return None

        min_reissue = _min_reissue_sec()
        max_valid = _max_valid_sec()

        with self._lock:
            self._load_from_file()
            age = self.age_sec()
            prefix_ok = (
                not self._app_key_prefix
                or app_key.startswith(self._app_key_prefix)
                or self._app_key_prefix.startswith(app_key[:8])
            )
            if self._approval_key and age < max_valid and prefix_ok:
                if force_refresh and age >= min_reissue:
                    pass
                elif force_refresh and age < min_reissue:
                    logger.info(
                        "🔑 approval_key REST 재발급 생략 (KIS 6h — %.0f분/%.0fh, %s 재사용)",
                        age / 60,
                        min_reissue / 3600,
                        self._cache_path.name,
                    )
                    return self._approval_key
                else:
                    return self._approval_key

        if not _acquire_lock():
            with self._lock:
                self._load_from_file()
                return self._approval_key

        try:
            with self._lock:
                self._load_from_file()
                age = self.age_sec()
                if self._approval_key and age < max_valid:
                    if not force_refresh or age < min_reissue:
                        return self._approval_key

            key = _issue_rest(app_key, app_secret, base_url)
            if not key:
                with self._lock:
                    return self._approval_key

            now = time.time()
            with self._lock:
                self._approval_key = key
                self._issued_ts = now
                self._app_key_prefix = app_key[:8]
                self._save_to_file()
            logger.info(
                "✅ WebSocket approval_key 발급·저장 (%s, 앞8자: %s…, "
                "다음 REST 가능: %.0fh 후 → %s)",
                self._mode_str,
                key[:8],
                min_reissue / 3600,
                self._cache_path.name,
            )
            return key
        finally:
            _release_lock()


def get_approval_status(is_mock: bool) -> dict:
    mgr = KISApprovalManager.instance(is_mock)
    with mgr._lock:
        mgr._load_from_file()
        age_h = round(mgr.age_sec() / 3600.0, 2) if mgr._approval_key else -1.0
        valid = bool(mgr._approval_key) and mgr.age_sec() < _max_valid_sec()
        return {
            "valid": valid,
            "key_prefix": (mgr._approval_key[:8] + "…") if mgr._approval_key else "",
            "age_h": age_h,
            "file": str(mgr._cache_path.name),
        }