add websocket impl

Author: IAmTomahawkx

twitchio/ext/eventsub/__init__.py

@@ -23,4 +23,5 @@
 """
 
 from .server import EventSubClient
+from .websocket import EventSubWSClient, Websocket
 from .models import *

twitchio/ext/eventsub/http.py

@@ -6,18 +6,19 @@
 
 if TYPE_CHECKING:
     from .server import EventSubClient
+    from .websocket import EventSubWSClient
     from .models import EventData, Subscription
 
 __all__ = ("EventSubHTTP",)
 
 
 class EventSubHTTP:
-    def __init__(self, client: EventSubClient, token: Optional[str]):
+    def __init__(self, client: Union[EventSubClient, EventSubWSClient], token: Optional[str]):
         self._client = client
         self._http = client.client._http
         self._token = token
 
-    async def create_subscription(self, event_type: Tuple[str, int, Type[EventData]], condition: Dict[str, str]):
+    async def create_webhook_subscription(self, event_type: Tuple[str, int, Type[EventData]], condition: Dict[str, str]):
         payload = {
             "type": event_type[0],
             "version": str(event_type[1]),
@@ -27,6 +28,16 @@ async def create_subscription(self, event_type: Tuple[str, int, Type[EventData]]
         route = Route("POST", "eventsub/subscriptions", body=payload, token=self._token)
         return await self._http.request(route, paginate=False, force_app_token=True)
 
+    async def create_websocket_subscription(self, event_type: Tuple[str, int, Type[EventData]], condition: Dict[str, str], session_id: str, token: str):
+        payload = {
+            "type": event_type[0],
+            "version": str(event_type[1]),
+            "condition": condition,
+            "transport": {"method": "websocket", "session_id": session_id},
+        }
+        route = Route("POST", "eventsub/subscriptions", body=payload, token=token)
+        return await self._http.request(route, paginate=False, full_body=True)
+
     async def delete_subscription(self, subscription: Union[str, Subscription]):
         if isinstance(subscription, models.Subscription):
             return await self._http.request(

twitchio/ext/eventsub/models.py

@@ -4,7 +4,7 @@
 import hashlib
 import logging
 from enum import Enum
-from typing import Dict, TYPE_CHECKING, Optional, Type, Union, Tuple, List
+from typing import Dict, TYPE_CHECKING, Optional, Type, Union, Tuple, List, overload
 from typing_extensions import Literal
 
 from aiohttp import web
@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     from .server import EventSubClient
+    from .websocket import EventSubWSClient
 
 try:
     import ujson as json
@@ -36,7 +37,7 @@ def __init__(self, **kwargs):
 
 
 class Subscription:
-    __slots__ = "id", "status", "type", "version", "cost", "condition", "transport", "created_at"
+    __slots__ = "id", "status", "type", "version", "cost", "condition", "transport", "transport_method", "created_at"
 
     def __init__(self, data: dict):
         self.id: str = data["id"]
@@ -47,8 +48,14 @@ def __init__(self, data: dict):
         self.condition: Dict[str, str] = data["condition"]
         self.created_at = _parse_datetime(data["created_at"])
         self.transport = EmptyObject()
-        self.transport.method: str = data["transport"]["method"]  # noqa
-        self.transport.callback: str = data["transport"]["callback"]  # noqa
+        self.transport_method: TransportType = getattr(TransportType, data["transport"]["method"])
+        self.transport.method: str = data["transport"]["method"]  # type: ignore
+
+        if self.transport_method is TransportType.webhook:
+            self.transport.callback: str = data["transport"]["callback"]  # type: ignore
+        else:
+            self.transport.callback: str = "" # type: ignore # compatibility
+            self.transport.session_id: str = data["transport"]["session_id"]  # type: ignore
 
 
 class Headers:
@@ -82,6 +89,42 @@ def __init__(self, request: web.Request):
         self._raw_timestamp = request.headers["Twitch-Eventsub-Message-Timestamp"]
 
 
+class WebsocketHeaders:
+    """
+    The headers of the inbound Websocket EventSub message
+
+    ..warning:
+
+        This is a BETA feature
+
+    Attributes
+    -----------
+    message_id: :class:`str`
+        The unique ID of the message
+    message_type: :class:`str`
+        The type of the message coming through
+    message_retry: :class:`int`
+        Kept for compatibility with :class:`Headers`
+    signature: :class:`str`
+        Kept for compatibility with :class:`Headers`
+    subscription_type: :class:`str`
+        The type of the subscription on the inbound message
+    subscription_version: :class:`str`
+        The version of the subscription.
+    timestamp: :class:`datetime.datetime`
+        The timestamp the message was sent at
+    """
+    def __init__(self, frame: dict):
+        meta = frame["metadata"]
+        self.message_id: str = meta["message_id"]
+        self.timestamp = _parse_datetime(meta["message_timestamp"])
+        self.message_type: Literal["notification", "revocation", "reconnect", "session_keepalive"] = meta["message_type"]
+        self.message_retry: int = 0 # don't make breaking changes with the Header class
+        self.signature: str = ""
+        self.subscription_type: str = frame["payload"]["subscription"]["type"]
+        self.subscription_version: str = frame["payload"]["subscription"]["version"]
+
+
 class BaseEvent:
     """
     The base of all the event classes
@@ -94,21 +137,47 @@ class BaseEvent:
         The headers received with the message
     """
 
-    __slots__ = "_client", "_raw_data", "subscription", "headers"
+    __slots__ = ("_client", "_raw_data", "subscription", "headers")
+
+    @overload
+    def __init__(self, client: EventSubClient, _data: str, request: web.Request):
+        ...
 
-    def __init__(self, client: EventSubClient, data: str, request: web.Request):
+    @overload
+    def __init__(self, client: EventSubWSClient, _data: dict, request: None):
+        ...
+
+    def __init__(self, client: Union[EventSubClient, EventSubWSClient], _data: Union[str, dict], request: Optional[web.Request]):
         self._client = client
-        self._raw_data = data
-        _data: dict = _loads(data)
-        self.subscription = Subscription(_data["subscription"])
-        self.headers = Headers(request)
-        self.setup(_data)
+        self._raw_data = _data
+
+        if isinstance(_data, str):
+            data: dict = _loads(_data)
+        else:
+            data = _data
+
+        self.headers: Union[Headers, WebsocketHeaders]
+        self.subscription: Subscription
+
+        if request:
+            data: dict = _loads(_data)
+            self.headers = Headers(request)
+            self.subscription = Subscription(data["subscription"])
+            self.setup(data)
+        else:
+            self.headers = WebsocketHeaders(data)
+            self.subscription = Subscription(data["payload"]["subscription"])
+            self.setup(data["payload"])
+
 
     def setup(self, data: dict):
         pass
 
     def verify(self):
-        hmac_message = (self.headers.message_id + self.headers._raw_timestamp + self._raw_data).encode("utf-8")
+        """
+        Only used in webhook transport types. Verifies the message is valid
+        """
+        hmac_message = (self.headers.message_id + self.headers._raw_timestamp + self._raw_data).encode("utf-8") # type: ignore
         secret = self._client.secret.encode("utf-8")
         digest = hmac.new(secret, msg=hmac_message, digestmod=hashlib.sha256).hexdigest()
 
@@ -127,6 +196,9 @@ class ChallengeEvent(BaseEvent):
     """
     A challenge event.
 
+    .. note::
+        These are only dispatched when using :class:`~twitchio.ext.eventsub.EventSubClient`
+
     Attributes
     -----------
     challenge: :class`str`
@@ -139,7 +211,7 @@ def setup(self, data: dict):
         self.challenge: str = data["challenge"]
 
     def verify(self):
-        hmac_message = (self.headers.message_id + self.headers._raw_timestamp + self._raw_data).encode("utf-8")
+        hmac_message = (self.headers.message_id + self.headers._raw_timestamp + self._raw_data).encode("utf-8") # type: ignore
         secret = self._client.secret.encode("utf-8")
         digest = hmac.new(secret, msg=hmac_message, digestmod=hashlib.sha256).hexdigest()
 
@@ -150,6 +222,39 @@ def verify(self):
         return web.Response(status=200, text=self.challenge)
 
 
+class ReconnectEvent(BaseEvent):
+    """
+    A reconnect event. Called by twitch when the websocket needs to be disconnected for maintenance or other reasons
+
+    .. note::
+        These are only dispatched when using :class:`~twitchio.ext.eventsub.EventSubWSClient
+
+    Attributes
+    -----------
+    reconnect_url: :class:`str`
+        The URL to reconnect to
+    connected_at: :class:`~datetime.datetime`
+        When the original websocket connected
+    """
+
+    __slots__ = ("reconnect_url", "connected_at")
+
+    def setup(self, data: dict):
+        self.reconnect_url: str = data["session"]["reconnect_url"]
+        self.connected_at: datetime.datetime = _parse_datetime(data["session"]["connected_at"])
+
+
+class KeepAliveEvent(BaseEvent):
+    """
+    A keep-alive event. Called by twitch when no message has been sent for more than ``keepalive_timeout``
+
+    .. note::
+        These are only dispatched when using :class:`~twitchio.ext.eventsub.EventSubWSClient
+
+    """
+    pass
+
+
 class NotificationEvent(BaseEvent):
     """
     A notification event
@@ -1410,3 +1515,7 @@ class _SubscriptionTypes(metaclass=_SubTypesMeta):
 
 
 SubscriptionTypes = _SubscriptionTypes()
+
+class TransportType(Enum):
+    webhook = "webhook"
+    websocket = "websocket"

twitchio/ext/eventsub/server.py

@@ -68,7 +68,7 @@ async def subscribe_user_updated(self, user: Union[PartialUser, str, int]):
             user = user.id
 
         user = str(user)
-        return await self._http.create_subscription(models.SubscriptionTypes.user_update, {"user_id": user})
+        return await self._http.create_webhook_subscription(models.SubscriptionTypes.user_update, {"user_id": user})
 
     async def subscribe_channel_raid(
         self, from_broadcaster: Union[PartialUser, str, int] = None, to_broadcaster: Union[PartialUser, str, int] = None
@@ -87,7 +87,7 @@ async def subscribe_channel_raid(
             broadcaster = broadcaster.id
 
         broadcaster = str(broadcaster)
-        return await self._http.create_subscription(models.SubscriptionTypes.raid, {who: broadcaster})
+        return await self._http.create_webhook_subscription(models.SubscriptionTypes.raid, {who: broadcaster})
 
     async def _subscribe_channel_points_reward(
         self, event, broadcaster: Union[PartialUser, str, int], reward_id: str = None
@@ -100,7 +100,7 @@ async def _subscribe_channel_points_reward(
         if reward_id:
             data["reward_id"] = reward_id
 
-        return await self._http.create_subscription(event, data)
+        return await self._http.create_webhook_subscription(event, data)
 
     async def _subscribe_with_broadcaster(
         self, event: Tuple[str, int, Type[models._DataType]], broadcaster: Union[PartialUser, str, int]
@@ -109,7 +109,7 @@ async def _subscribe_with_broadcaster(
             broadcaster = broadcaster.id
 
         broadcaster = str(broadcaster)
-        return await self._http.create_subscription(event, {"broadcaster_user_id": broadcaster})
+        return await self._http.create_webhook_subscription(event, {"broadcaster_user_id": broadcaster})
 
     def subscribe_channel_bans(self, broadcaster: Union[PartialUser, str, int]):
         return self._subscribe_with_broadcaster(models.SubscriptionTypes.ban, broadcaster)
@@ -215,12 +215,12 @@ def subscribe_channel_prediction_end(self, broadcaster: Union[PartialUser, str,
         return self._subscribe_with_broadcaster(models.SubscriptionTypes.prediction_end, broadcaster)
 
     async def subscribe_user_authorization_granted(self):
-        return await self._http.create_subscription(
+        return await self._http.create_webhook_subscription(
             models.SubscriptionTypes.user_authorization_grant, {"client_id": self.client._http.client_id}
         )
 
     async def subscribe_user_authorization_revoked(self):
-        return await self._http.create_subscription(
+        return await self._http.create_webhook_subscription(
             models.SubscriptionTypes.user_authorization_revoke, {"client_id": self.client._http.client_id}
         )