import os import time import hmac import hashlib import base64 import json import pandas as pd import threading import requests import websocket import logging class OKXClient: PUBLIC_WS_URL = "wss://ws.okx.com:8443/ws/v5/public" PRIVATE_WS_URL = "wss://ws.okx.com:8443/ws/v5/private" REST_URL = "https://www.okx.com" def __init__(self, authenticate: bool = True): self.authenticated = False self.api_key = None self.api_secret = None self.api_passphrase = None self.ws = None self.ws_private = None self._lock = threading.Lock() self._private_lock = threading.Lock() if authenticate: config_path = os.path.join(os.path.dirname(__file__), '../credentials/okx_creds.json') try: with open(config_path, 'r') as f: config = json.load(f) except FileNotFoundError: raise FileNotFoundError(f"Credentials file not found at {config_path}. Please create it with the required keys.") except json.JSONDecodeError: raise ValueError(f"Credentials file at {config_path} is not valid JSON.") self.api_key = config.get("OKX_API_KEY") self.api_secret = config.get("OKX_API_SECRET") self.api_passphrase = config.get("OKX_API_PASSPHRASE") if not self.api_key or not self.api_secret or not self.api_passphrase: raise ValueError("API key, secret, and passphrase must be set in the credentials JSON file.") self._authenticate() self._connect_ws() def _connect_ws(self): if self.ws is None: self.ws = websocket.create_connection(self.PUBLIC_WS_URL, timeout=10) if self.authenticated and self.api_key and self.api_secret and self.api_passphrase and self.ws_private is None: self.ws_private = websocket.create_connection(self.PRIVATE_WS_URL, timeout=10) def _get_timestamp(self): return str(round(time.time(), 3)) def _sign(self, timestamp, method, request_path, body): if not body: body = '' message = f'{timestamp}{method}{request_path}{body}' mac = hmac.new(self.api_secret.encode('utf-8'), message.encode('utf-8'), hashlib.sha256) return base64.b64encode(mac.digest()).decode() def _authenticate(self): import websocket timestamp = self._get_timestamp() sign = self._sign(timestamp, 'GET', '/users/self/verify', '') login_params = { "op": "login", "args": [{ "apiKey": self.api_key, "passphrase": self.api_passphrase, "timestamp": timestamp, "sign": sign }] } self.ws_private.send(json.dumps(login_params)) logging.info("Waiting for login response from OKX...") while True: try: resp = self.ws_private.recv() logging.debug(f"Received from OKX private WS: {resp}") if not resp: continue try: msg = json.loads(resp) except Exception: logging.warning(f"Non-JSON message received: {resp}") continue if msg.get("event") == "login": if msg.get("code") == "0": logging.info("OKX WebSocket login successful.") self.authenticated = True break else: raise Exception(f"WebSocket authentication failed: {msg}") except websocket._exceptions.WebSocketConnectionClosedException as e: logging.error(f"WebSocket connection closed during authentication: {e}") raise except Exception as e: logging.error(f"Exception during authentication: {e}") raise def subscribe_candlesticks(self, instrument="BTC-USDT", timeframe="1m"): # OKX uses candle1m, candle5m, etc. tf_map = {"1m": "candle1m", "5m": "candle5m", "15m": "candle15m", "1h": "candle1H"} channel = tf_map.get(timeframe, f"candle{timeframe}") params = { "op": "subscribe", "args": [{"channel": channel, "instId": instrument}] } logging.debug(f"[PUBLIC WS] Sending subscription: {json.dumps(params)}") self.ws.send(json.dumps(params)) def subscribe_trades(self, instrument="BTC-USDT"): params = { "op": "subscribe", "args": [{"channel": "trades", "instId": instrument}] } logging.debug(f"[PUBLIC WS] Sending subscription: {json.dumps(params)}") self.ws.send(json.dumps(params)) def subscribe_ticker(self, instrument="BTC-USDT"): params = { "op": "subscribe", "args": [{"channel": "tickers", "instId": instrument}] } logging.debug(f"[PUBLIC WS] Sending subscription: {json.dumps(params)}") self.ws.send(json.dumps(params)) def subscribe_book(self, instrument="BTC-USDT", depth=5, channel="books5"): # OKX supports books5, books50, books-l2-tbt # channel = "books5" if depth <= 5 else "books50" params = { "op": "subscribe", "args": [{"channel": channel, "instId": instrument}] } logging.debug(f"[PUBLIC WS] Sending subscription: {json.dumps(params)}") self.ws.send(json.dumps(params)) def subscribe_user_order(self): if not self.authenticated: logging.warning("Attempted to subscribe to user order channel without authentication.") return params = { "op": "subscribe", "args": [{"channel": "orders", "instType": "SPOT"}] } logging.debug(f"[PRIVATE WS] Sending subscription: {json.dumps(params)}") self.ws_private.send(json.dumps(params)) def subscribe_user_trade(self): if not self.authenticated: logging.warning("Attempted to subscribe to user trade channel without authentication.") return params = { "op": "subscribe", "args": [{"channel": "trades", "instType": "SPOT"}] } logging.debug(f"[PRIVATE WS] Sending subscription: {json.dumps(params)}") self.ws_private.send(json.dumps(params)) def subscribe_user_balance(self): if not self.authenticated: logging.warning("Attempted to subscribe to user balance channel without authentication.") return params = { "op": "subscribe", "args": [{"channel": "balance_and_position"}] } logging.debug(f"[PRIVATE WS] Sending subscription: {json.dumps(params)}") self.ws_private.send(json.dumps(params)) def get_balance(self, currency=None): url = f"{self.REST_URL}/api/v5/account/balance" timestamp = self._get_timestamp() method = "GET" request_path = "/api/v5/account/balance" body = '' sign = self._sign(timestamp, method, request_path, body) headers = { "OK-ACCESS-KEY": self.api_key, "OK-ACCESS-SIGN": sign, "OK-ACCESS-TIMESTAMP": timestamp, "OK-ACCESS-PASSPHRASE": self.api_passphrase, "Content-Type": "application/json" } resp = requests.get(url, headers=headers) if resp.status_code == 200: data = resp.json() balances = data.get("data", [{}])[0].get("details", []) if currency: return [b for b in balances if b.get("ccy") == currency] return balances return [] def place_order(self, side, amount, instrument="BTC-USDT"): url = f"{self.REST_URL}/api/v5/trade/order" timestamp = self._get_timestamp() method = "POST" request_path = "/api/v5/trade/order" body_dict = { "instId": instrument, "tdMode": "cash", "side": side.lower(), "ordType": "market", "sz": str(amount) } body = json.dumps(body_dict) sign = self._sign(timestamp, method, request_path, body) headers = { "OK-ACCESS-KEY": self.api_key, "OK-ACCESS-SIGN": sign, "OK-ACCESS-TIMESTAMP": timestamp, "OK-ACCESS-PASSPHRASE": self.api_passphrase, "Content-Type": "application/json" } resp = requests.post(url, headers=headers, data=body) return resp.json() def buy_btc(self, amount, instrument="BTC-USDT"): return self.place_order("buy", amount, instrument) def sell_btc(self, amount, instrument="BTC-USDT"): return self.place_order("sell", amount, instrument) def get_instruments(self): url = f"{self.REST_URL}/api/v5/public/instruments?instType=SPOT" resp = requests.get(url) if resp.status_code == 200: data = resp.json() return data.get("data", []) return []