【立花証券e支店】PKI認証（簡易版）

import base64
from datetime import datetime
import json
import logging
import threading
import time
from typing import Any

import requests

from Cryptodome.PublicKey import RSA
from Cryptodome.Cipher import PKCS1_OAEP
from Cryptodome.Hash import SHA256


class ApiClient:
    def __init__(self):
        self.logger = logging.getLogger(__class__.__name__)

        self.session = requests.Session()
        self.request_counter: int = 0
        self.session_info: dict = {}

        self._lock = threading.Lock()
        self.timeout = 60        


    def _get_next_p_no(self) -> str:
        """
        APIリクエストごとにインクリメントされるp_no（リクエスト番号）を生成します。
        """
        self.request_counter += 1
        return str(self.request_counter)

    def _get_current_datetime(self) -> str:
        """
        p_sd_date用の日時文字列（YYYY.MM.DD-HH:MI:SS.ms）を生成します。
        """
        # APIの仕様に合わせ、マイクロ秒を示す末尾3桁を削除しミリ秒までの精度にします。
        return datetime.now().strftime("%Y.%m.%d-%H:%M:%S.%f")[:-3]

    def _call_api(
        self,
        url: str,
        payload: dict,
        method: str = "POST",
        headers: dict[str, str] = {},
        timeout: int | None = None,
        stream: bool = False,
    ) -> Any:
        """
        APIサーバーへリクエストを送信し、結果を取得する共通メソッド。

        APIをコールし、取得したJSONの辞書オブジェクトを返します。
        このメソッド内で、リクエスト共通のパラメータが付与されます。

        Args:
            url (str): エンドポイントURL。
            payload (dict): 送信パラメータ。
            method (str): "GET" または "POST"。
            headers (dict[str, str]): HTTPヘッダー。
            timeout (int): タイムアウト秒数。
            stream (bool): ストリーミングモード。デフォルトは False です。

        Raises:
            requests.exceptions.RequestException: APIサーバーへ接続出来なかった場合。
            Exception: その他の例外発生時。

        Returns:
            dict: APIから取得したJSONを辞書へ変換したオブジェクト。
        """
        self.logger.debug("URL: %s", url)
        self.logger.debug("METHOD: %s", method)
        self.logger.debug("PAYLOAD: %s", json.dumps(payload))
        self.logger.debug("HEADERS: %s", json.dumps(headers))

        dict_result: dict = {}
        with self._lock:
            try:
                # 共通パラメーターの付与
                payload["p_no"] = self._get_next_p_no()
                payload["p_sd_date"] = self._get_current_datetime()
                payload["sJsonOfmt"] = "4"

                if timeout is None:
                    timeout = self.timeout

                response = self.session.request(
                    method=method,
                    url=url,
                    headers=headers,
                    json=payload,
                    timeout=timeout,
                )

                content_type = response.headers.get("Content-Type", "")
                self.logger.debug("CONTENT-TYPE: %s", content_type)

                response.raise_for_status()
                response.encoding = "cp932"
                self.logger.debug(
                    "STATUS: %s %s", response.status_code, response.reason
                )
                if stream is True:
                    return response

                if "application/json" in content_type:
                    dict_result = response.json()
                    return dict_result
                else:
                    self.logger.error("JSON形式の応答がありませんでした。")
                    self.logger.error(f"受信したテキスト: {response.text}")
                    return {"p_errno": "9999", "p_err": "No JSON"}

            except requests.exceptions.RequestException as exp:
                self.logger.error("URL: %s", exp.request.url if exp.request else url)
                self.logger.error(exp)
                if exp.response is not None:
                    self.logger.error(
                        "CODE: %s, REASON: %s",
                        exp.response.status_code,
                        exp.response.reason,
                    )
                    try:
                        self.logger.error("Message: %s", exp.response.json())
                    except json.JSONDecodeError:
                        self.logger.error("Message: %s", exp.response.text)
                raise exp

        return dict_result

    def login_pki(self, url: str, auth_id: str, private_pem: str) -> bool:
        """
        APIサーバーへログインします。

        成功した場合、セッション情報を`self.session_info`に格納します。

        Args:
            auth_id (str): ログインに使用する認証ID。
            private_pem (str): PEM形式の秘密鍵。

        Returns:
            bool: ログインに成功した場合はTrue、失敗した場合はFalse。
        """

        payload = {
            "sCLMID": "CLMAuthLoginRequest",
            "sAuthId": auth_id,
        }

        results = self._call_api(url, payload)

        self.session_info = results

        # 暗号化されている仮想URLを復号
        private_key = RSA.import_key(private_pem.encode())
        decryptor = PKCS1_OAEP.new(private_key, hashAlgo=SHA256)

        target_url_keys = ['sUrlRequest', 'sUrlMaster', 'sUrlPrice', 'sUrlEvent', 'sUrlEventWebSocket']
        for key in target_url_keys:
            url = self.session_info.get(key, "")

            self.logger.debug("復号前: %s", url)

            url = self.decrypt(url, decryptor)

            self.logger.debug("復号後: %s", url)

            self.session_info[key] = url.strip()

        return True

    def decrypt(self, b64_data: str, decryptor: PKCS1_OAEP.PKCS1OAEP_Cipher):
        """
        公開鍵で暗号化されたデータを復号します。

        Args:
            b64_data (str): 公開鍵で暗号化されたデータ（Base64形式でエンコードされた文字列）。
            decryptor (PKCS1_OAEP.PKCS1OAEP_Cipher): 復号器。

        Returns:
            str: 復号されたデータ。
        """
        encrypted_data = base64.b64decode(b64_data)
        return decryptor.decrypt(encrypted_data).decode()


    def logout(self) -> bool:
        """
        サーバーからログアウトし、セッション情報をクリアします。

        Returns:
            bool: ログアウトに成功した場合はTrue、失敗した場合はFalse。
        """
        payload = {
            "sCLMID": "CLMAuthLogoutRequest",
        }
        url = self.session_info.get("sUrlRequest")
        if url is None:
            return False
        
        results = self._call_api(url, payload)

        self.session_info = {}
        return True




if __name__ == "__main__":

    # logger の初期化
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    api_client = ApiClient()

    with open(".pki/e_api_authid.txt", "rb") as f:
        auth_id = f.read().decode()

    with open(".pki/e_api_private_key.pem", "rb") as f:
        private_pem = f.read().decode()

    url = "https://kabuka.e-shiten.jp/e_api_v4r9/auth/"

    is_login = api_client.login_pki(url, auth_id, private_pem)

    if is_login:
        time.sleep(10.0)
        api_client.logout()