From a77d1b947bea719f4bc1383985e55af291b3f06e Mon Sep 17 00:00:00 2001
From: uzlonewolf <github_com@hacker-nin.com>
Date: Sat, 31 Aug 2024 11:55:35 -0700
Subject: [PATCH] Add Ping/Pong functions

---
 README.md                                     | 19 ++++++++++++
 .../SimpleWebSocketServer.py                  | 31 +++++++++++++++++--
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 7d699f0..61a056f 100644
--- a/README.md
+++ b/README.md
@@ -114,10 +114,29 @@ handleMessage: gets called when there is an incoming message from the client end
  - self.data: bytearray (BINARY frame) or unicode string payload (TEXT frame)  
  - self.request: HTTP details from the WebSocket handshake (refer to BaseHTTPRequestHandler)
 
+handlePing: gets called when there is an incoming PING request from the client endpoint
+ - self.address: TCP address port tuple of the endpoint
+ - self.opcode: always PING
+ - self.data: bytearray payload
+ - self.request: HTTP details from the WebSocket handshake (refer to BaseHTTPRequestHandler)
+ - **You must send the PONG reply yourself with sendPong(self.data) [or other data] if you define this function**
+
+handlePong: gets called when there is an incoming PONG reply from the client endpoint
+ - self.address: TCP address port tuple of the endpoint
+ - self.opcode:	always PONG
+ - self.data: bytearray payload
+ - self.request: HTTP details from the WebSocket handshake (refer to BaseHTTPRequestHandler)
+
 sendMessage: send some text or binary data to the client endpoint
  - sending data as a unicode object will send a TEXT frame
  - sending data as a bytearray object will send a BINARY frame
 
+sendPing: send a PING request to the client endpoint
+ - send data as a bytearray object
+
+sendPong: send a PONG reply to the client endpoint
+ - send data as a bytearray object
+
 sendClose: send close frame to endpoint
 
 
diff --git a/SimpleWebSocketServer/SimpleWebSocketServer.py b/SimpleWebSocketServer/SimpleWebSocketServer.py
index ccb9949..e2ea3a1 100644
--- a/SimpleWebSocketServer/SimpleWebSocketServer.py
+++ b/SimpleWebSocketServer/SimpleWebSocketServer.py
@@ -138,6 +138,19 @@ def handleClose(self):
       """
       pass
 
+   def handlePing(self):
+      """
+          Called when a websocket server gets a Ping frame from a client.
+          If overridden, make sure to send the PONG reply
+      """
+      self.sendPong(self.data)
+
+   def handlePong(self):
+     """
+         Called when a websocket server gets a Pong frame from a client.
+     """
+     pass
+
    def _handlePacket(self):
       if self.opcode == CLOSE:
          pass
@@ -229,10 +242,10 @@ def _handlePacket(self):
               self.frag_buffer = None
 
           elif self.opcode == PING:
-              self._sendMessage(False, PONG, self.data)
+              self.handlePing()
 
           elif self.opcode == PONG:
-              pass
+              self.handlePong()
 
           else:
               if self.frag_start is True:
@@ -401,6 +414,20 @@ def sendMessage(self, data):
          opcode = TEXT
       self._sendMessage(False, opcode, data)
 
+   def sendPing(self, data=None):
+      """
+          Send websocket Ping
+      """
+      if not data:
+          data = bytearray()
+      self._sendMessage(False, PING, data)
+
+   def sendPong(self, data):
+      """
+           Send websocket Pong
+      """
+      self._sendMessage(False, PONG, data)
+
 
    def _sendMessage(self, fin, opcode, data):