From 486087cdb349dbc07b479d2286a02bdca310ea38 Mon Sep 17 00:00:00 2001
From: Tom van der Lee <tom@vanderlee.io>
Date: Wed, 20 Mar 2024 21:48:45 +0100
Subject: Added websocket support

---
 ttun_server/types.py | 51 +++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 45 insertions(+), 6 deletions(-)

(limited to 'ttun_server/types.py')

diff --git a/ttun_server/types.py b/ttun_server/types.py
index 8a4d929..8591e7d 100644
--- a/ttun_server/types.py
+++ b/ttun_server/types.py
@@ -3,7 +3,7 @@ from enum import Enum
 from typing import TypedDict, Optional
 
 
-class MessageType(Enum):
+class HttpMessageType(Enum):
     request = 'request'
     response = 'response'
 
@@ -13,23 +13,62 @@ class Config(TypedDict):
     client_version: str
 
 
-class RequestData(TypedDict):
+class HttpRequestData(TypedDict):
     method: str
     path: str
     headers: list[tuple[str, str]]
     body: Optional[str]
 
 
-class ResponseData(TypedDict):
+class HttpResponseData(TypedDict):
     status: int
     headers: list[tuple[str, str]]
     body: Optional[str]
 
 
-class Message(TypedDict):
-    type: MessageType
+class HttpMessage(TypedDict):
+    type: HttpMessageType
     identifier: str
-    payload: Config | RequestData | ResponseData
+    payload: Config | HttpRequestData | HttpResponseData
+
+
+class WebsocketMessageType(Enum):
+    connect = 'connect'
+    disconnect = 'disconnect'
+    message = 'message'
+    ack = 'ack'
+
+
+class WebsocketConnectData(TypedDict):
+    path: str
+    headers: list[tuple[str, str]]
+
+
+class WebsocketDisconnectData(TypedDict):
+    close_code: int
+
+
+class WebsocketMessageData(TypedDict):
+    body: Optional[str]
+
+
+class WebsocketMessage(TypedDict):
+    type: WebsocketMessageType
+    identifier: str
+    payload: WebsocketConnectData | WebsocketDisconnectData | WebsocketMessageData
+
+
+class MessageType(Enum):
+    request = 'request'
+    response = 'response'
+
+    ws_connect = 'connect'
+    ws_disconnect = 'disconnect'
+    ws_message = 'message'
+    ws_ack = 'ack'
+
+
+Message = HttpMessage | WebsocketMessage
 
 
 class MemoryConnection(TypedDict):
-- 
cgit v1.2.3