smartWebSocketOrderUpdate.py

import ssl
import websocket
import time
import logging 
from logzero import logger
import logzero
import os

class SmartWebSocketOrderUpdate(object):
    WEBSOCKET_URI = "wss://tns.angelone.in/smart-order-update"
    HEARTBEAT_MESSAGE = "ping"  # Heartbeat message to maintain Socket connection.
    HEARTBEAT_INTERVAL_SECONDS = 10  # Interval for sending heartbeat messages to keep the connection alive.
    MAX_CONNECTION_RETRY_ATTEMPTS = 2  # Max retry attempts to establish Socket connection in case of failure.
    RETRY_DELAY_SECONDS = 10  # Delay between retry attempts when reconnecting to Socket in case of failure.
    wsapp = None  #Socket connection instance
    last_pong_timestamp = None #Timestamp of the last received pong message
    current_retry_attempt = 0  #Current retry attempt count

    def __init__(self, auth_token, api_key, client_code, feed_token):
        self.auth_token = auth_token
        self.api_key = api_key
        self.client_code = client_code
        self.feed_token = feed_token
        # Create a log folder based on the current date
        log_folder = time.strftime("%Y-%m-%d", time.localtime())
        log_folder_path = os.path.join("logs", log_folder)  # Construct the full path to the log folder
        os.makedirs(log_folder_path, exist_ok=True) # Create the log folder if it doesn't exist
        log_path = os.path.join(log_folder_path, "app.log") # Construct the full path to the log file
        logzero.logfile(log_path, loglevel=logging.INFO)  # Output logs to a date-wise log file

    def on_message(self, wsapp, message):
        logger.info("Received message: %s", message)

    def on_data(self, wsapp, message, data_type, continue_flag):
        self.on_message(wsapp, message)

    def on_open(self, wsapp):
        logger.info("Connection opened")

    def on_error(self, wsapp, error):
        logger.error("Error: %s", error)

    def on_close(self, wsapp, close_status_code, close_msg):
        logger.info("Connection closed")
        self.retry_connect()

    def on_ping(self, wsapp, data):
        timestamp = time.time()
        formatted_timestamp = time.strftime("%d-%m-%y %H:%M:%S", time.localtime(timestamp))
        logger.info("In on ping function ==> %s, Timestamp: %s", data, formatted_timestamp)

    def on_pong(self, wsapp, data):
        if data == self.HEARTBEAT_MESSAGE:
            timestamp = time.time()
            formatted_timestamp = time.strftime("%d-%m-%y %H:%M:%S", time.localtime(timestamp))
            logger.info("In on pong function ==> %s, Timestamp: %s", data, formatted_timestamp)
            self.last_pong_timestamp = timestamp
        else:
            self.on_data(wsapp, data, websocket.ABNF.OPCODE_BINARY, False)

    def check_connection_status(self):
        current_time = time.time()
        if self.last_pong_timestamp is not None and current_time - self.last_pong_timestamp > 2 * self.HEARTBEAT_INTERVAL_SECONDS:
            self.close_connection()

    def connect(self):
        headers = {
            "Authorization": self.auth_token,
            "x-api-key": self.api_key,
            "x-client-code": self.client_code,
            "x-feed-token": self.feed_token
        }
        try:
            self.wsapp = websocket.WebSocketApp(self.WEBSOCKET_URI, header=headers, on_open=self.on_open,
                                                on_error=self.on_error, on_close=self.on_close,
                                                on_data=self.on_data, on_ping=self.on_ping, on_pong=self.on_pong)
            self.wsapp.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE}, ping_interval=self.HEARTBEAT_INTERVAL_SECONDS)
        except Exception as e:
            logger.error("Error connecting to WebSocket: %s", e)
            self.retry_connect()

    def retry_connect(self):
        if self.current_retry_attempt < self.MAX_CONNECTION_RETRY_ATTEMPTS:
            logger.info("Retrying connection (Attempt %s)...", self.current_retry_attempt + 1)
            time.sleep(self.RETRY_DELAY_SECONDS)
            self.current_retry_attempt += 1
            self.connect()
        else:
            logger.warning("Max retry attempts reached.")

    def close_connection(self):
        if self.wsapp:
            self.wsapp.close()