From ff664fa9d3345881040de015ba33d34fd3afcd2b Mon Sep 17 00:00:00 2001
From: Jaime Idolpx <jaime@idolpx.com>
Date: Mon, 8 Jun 2026 14:08:37 -0400
Subject: [PATCH] feat(WebSocket): implement WebSocket handling and upgrade
 support

---
 webdav3.py | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 57 insertions(+), 1 deletion(-)

diff --git a/webdav3.py b/webdav3.py
index 3839087..80603dc 100644
--- a/webdav3.py
+++ b/webdav3.py
@@ -43,7 +43,7 @@ from http.server import HTTPServer, BaseHTTPRequestHandler
 from socketserver import ThreadingMixIn
 import urllib.request, urllib.parse, urllib.error, urllib.parse
 from time import timezone, strftime, localtime, gmtime
-import os, sys, re, shutil, uuid, hashlib, mimetypes, base64, socket
+import os, sys, re, shutil, struct, uuid, hashlib, mimetypes, base64, socket
 
 # Debug message ( True / False )
 sys_debug = False
@@ -687,7 +687,63 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
     self.end_headers()
     w.flush()
 
+  def _ws_recv(self, sock, n):
+    data = b''
+    while len(data) < n:
+      chunk = sock.recv(n - len(data))
+      if not chunk:
+        raise ConnectionError('closed')
+      data += chunk
+    return data
+
+  def _handle_websocket(self):
+    key = self.headers.get('Sec-WebSocket-Key', '')
+    accept = base64.b64encode(
+      hashlib.sha1((key + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11').encode()).digest()
+    ).decode()
+    self.send_response(101, 'Switching Protocols')
+    self.send_header('Upgrade', 'websocket')
+    self.send_header('Connection', 'Upgrade')
+    self.send_header('Sec-WebSocket-Accept', accept)
+    self.end_headers()
+    sock = self.connection
+    sock.settimeout(None)
+    try:
+      while True:
+        b0, b1 = self._ws_recv(sock, 2)
+        opcode = b0 & 0x0F
+        masked = bool(b1 & 0x80)
+        length = b1 & 0x7F
+        if length == 126:
+          length = struct.unpack('>H', self._ws_recv(sock, 2))[0]
+        elif length == 127:
+          length = struct.unpack('>Q', self._ws_recv(sock, 8))[0]
+        mask = self._ws_recv(sock, 4) if masked else b''
+        payload = self._ws_recv(sock, length)
+        if masked:
+          payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
+        if opcode == 0x8:  # close
+          sock.send(b'\x88\x00')
+          break
+        elif opcode == 0x9:  # ping → pong
+          frame = bytes([0x8A, len(payload)]) + payload
+          sock.send(frame)
+        elif opcode in (0x1, 0x2):  # text or binary — echo back as text
+          if length < 126:
+            header = bytes([0x81, length])
+          elif length < 65536:
+            header = bytes([0x81, 126]) + struct.pack('>H', length)
+          else:
+            header = bytes([0x81, 127]) + struct.pack('>Q', length)
+          sock.send(header + payload)
+    except Exception:
+      pass
+
   def do_GET(self, onlyhead=False):
+    if (self.path == '/ws'
+        and self.headers.get('Upgrade', '').lower() == 'websocket'):
+      self._handle_websocket()
+      return
     if self.WebAuth():
       return
     path, elem = self.path_elem()