diff --git a/examples/stream.py b/examples/stream.py new file mode 100644 index 0000000..05c7231 --- /dev/null +++ b/examples/stream.py @@ -0,0 +1,24 @@ +import asyncio +import os + +from fishfish.ws import FishWebsocket, models, WebsocketEvents + + +async def on_domain_create(data: models.WSDomainCreate): + print(f"{data.domain} was just created") + + +async def on_error(error: Exception): + raise error + + +async def main(): + ws: FishWebsocket = FishWebsocket(os.environ["API_KEY"]) + ws.register_error_handler(on_error) + ws.register_listener(WebsocketEvents.DOMAIN_CREATE, on_domain_create) + + task = await ws.start() + await task # Run until complete + + +asyncio.run(main()) diff --git a/fishfish/models/domain.py b/fishfish/models/domain.py index bb72d87..effe3d6 100644 --- a/fishfish/models/domain.py +++ b/fishfish/models/domain.py @@ -32,7 +32,7 @@ class Domain: category: Category added: datetime.datetime checked: datetime.datetime - target: Optional[str] + target: Optional[str] = None @classmethod def from_dict(cls, data) -> Domain: diff --git a/fishfish/models/url.py b/fishfish/models/url.py index 037ac15..87e1f79 100644 --- a/fishfish/models/url.py +++ b/fishfish/models/url.py @@ -32,7 +32,7 @@ class URL: category: Category added: datetime.datetime checked: datetime.datetime - target: Optional[str] + target: Optional[str] = None @classmethod def from_dict(cls, data) -> URL: diff --git a/fishfish/ws/__init__.py b/fishfish/ws/__init__.py new file mode 100644 index 0000000..34ec46e --- /dev/null +++ b/fishfish/ws/__init__.py @@ -0,0 +1,2 @@ +from .events import WebsocketEvents +from .fish_websocket import FishWebsocket diff --git a/fishfish/ws/events.py b/fishfish/ws/events.py new file mode 100644 index 0000000..e89b80a --- /dev/null +++ b/fishfish/ws/events.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from enum import Enum + +from fishfish.ws.models import * + + +class WebsocketEvents(str, Enum): + URL_CREATE = "url_create" + URL_UPDATE = "url_update" + URL_DELETE = "url_delete" + DOMAIN_CREATE = "domain_create" + DOMAIN_UPDATE = "domain_update" + DOMAIN_DELETE = "domain_delete" + + @classmethod + def from_string(cls, event_type: str) -> WebsocketEvents: + """Create the correct event based on the event type. + + Raises + ------ + ValueError + Unknown event type + """ + ctx = getattr(cls, event_type) + if ctx is None: + raise ValueError(f"Unknown event type {event_type}") + + return ctx + + def create_model(self, data: dict): + """Create the correct data model for this event type.""" + factories = { + WebsocketEvents.URL_CREATE: WSUrlCreate, + WebsocketEvents.URL_UPDATE: WSUrlUpdate, + WebsocketEvents.URL_DELETE: WSUrlDelete, + WebsocketEvents.DOMAIN_CREATE: WSDomainCreate, + WebsocketEvents.DOMAIN_UPDATE: WSDomainUpdate, + WebsocketEvents.DOMAIN_DELETE: WSDomainDelete, + } + return factories[self].from_dict(data) diff --git a/fishfish/ws/fish_websocket.py b/fishfish/ws/fish_websocket.py new file mode 100644 index 0000000..8c0ea50 --- /dev/null +++ b/fishfish/ws/fish_websocket.py @@ -0,0 +1,113 @@ +import asyncio +import datetime +import json +import logging +import traceback +from typing import Optional + +import httpx +import websockets as websockets + +from fishfish import Token, Unauthorized +from fishfish.ws import WebsocketEvents + +log = logging.getLogger(__name__) + + +class FishWebsocket: + def __init__(self, api_key: str): + self.__refresh_token: Optional[str] = api_key + self.__current_session_token: Optional[Token] = None + self._listeners: dict[WebsocketEvents, list] = { + WebsocketEvents.URL_CREATE: [], + WebsocketEvents.URL_UPDATE: [], + WebsocketEvents.URL_DELETE: [], + WebsocketEvents.DOMAIN_CREATE: [], + WebsocketEvents.DOMAIN_UPDATE: [], + WebsocketEvents.DOMAIN_DELETE: [], + } + self._error_handler = None + + def _set_session_key(self): + if ( + self.__current_session_token + and not self.__current_session_token.has_expired + ): + return + + r = httpx.post( + "https://api.fishfish.gg/v1/users/@me/tokens", + headers={"Authorization": self.__refresh_token}, + ) + if r.status_code == 401: + raise Unauthorized("Your provided FishFish token is invalid.") + + data = r.json() + token = data["token"] + expires = datetime.datetime.fromtimestamp(data["expires"]) + self.__current_session_token = Token(token, expires) + + async def _process_event(self, data): + try: + event_type: WebsocketEvents = WebsocketEvents.from_string(data["type"]) + model = event_type.create_model(data["data"]) + iters = [coro(model) for coro in self._listeners[event_type]] + if iters: + await asyncio.gather(*iters) + except Exception as e: + if not self._error_handler: + log.error("Attempting to process %s threw %s", str(e), str(e)) + log.error("%s", "".join(traceback.format_exception(e))) + return + + asyncio.create_task(self._error_handler(e)) + + async def _ws_loop(self): + self._set_session_key() + async with websockets.connect( + "wss://api.fishfish.gg/v1/stream/", + extra_headers={"Authorization": self.__current_session_token.token}, + ) as websocket: + async for message in websocket: + asyncio.create_task(self._process_event(json.loads(message))) + + async def start(self) -> asyncio.Task: + return asyncio.create_task(self._ws_loop()) + + def register_listener(self, event: WebsocketEvents, coro): + """Register a listener for a websocket event. + + Parameters + ---------- + event: WebsocketEvents + The event to listen for. + coro + An async function or method to call + when this event is fired. + + The first parameter will be the event data. + + Raises + ------ + ValueError + The provided function or method was not async. + """ + if not asyncio.iscoroutinefunction(coro): + raise ValueError( + f"The provided value for 'coro' must be a coroutine function." + ) + + self._listeners[event].append(coro) + + def register_error_handler(self, coro): + """Register a function to handle errors. + + Parameters + ---------- + coro + An async function or method to call + when an error occurs. + + The first parameter will be the exception. + """ + self._error_handler = coro diff --git a/fishfish/ws/models/__init__.py b/fishfish/ws/models/__init__.py new file mode 100644 index 0000000..ce6bef3 --- /dev/null +++ b/fishfish/ws/models/__init__.py @@ -0,0 +1,2 @@ +from .urls import WSUrlDelete, WSUrlCreate, WSUrlUpdate +from .domains import WSDomainCreate, WSDomainUpdate, WSDomainDelete diff --git a/fishfish/ws/models/domains.py b/fishfish/ws/models/domains.py new file mode 100644 index 0000000..36f0a70 --- /dev/null +++ b/fishfish/ws/models/domains.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class WSDomainDelete: + """A websocket domain delete event model. + + Attributes + ---------- + domain: str + The domain which was deleted. + """ + + domain: str + + @classmethod + def from_dict(cls, data: dict) -> WSDomainDelete: + return cls(**data) + + +@dataclass +class WSDomainCreate(WSDomainDelete): + """A websocket domain create event model. + + Attributes + ---------- + domain: str + The domain which was deleted. + description: Optional[str] + The description for this domain. + category: Optional[str] + The category for this domain. + target: Optional[str] + The target for this domain. + """ + + description: Optional[str] = None + category: Optional[str] = None + target: Optional[str] = None + + +@dataclass +class WSDomainUpdate(WSDomainCreate): + """A websocket domain update event model. + + Attributes + ---------- + domain: str + The domain which was deleted. + description: Optional[str] + The description for this domain. + category: Optional[str] + The category for this domain. + target: Optional[str] + The target for this domain. + checked: Optional[int] + When this domain was checked into the db? + """ + + checked: Optional[int] = None diff --git a/fishfish/ws/models/urls.py b/fishfish/ws/models/urls.py new file mode 100644 index 0000000..896a020 --- /dev/null +++ b/fishfish/ws/models/urls.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class WSUrlDelete: + """A websocket url delete event model. + + Attributes + ---------- + url: str + The url which was deleted. + """ + + url: str + + @classmethod + def from_dict(cls, data: dict) -> WSUrlDelete: + return cls(**data) + + +@dataclass +class WSUrlCreate(WSUrlDelete): + """A websocket url create event model. + + Attributes + ---------- + url: str + The url which was deleted. + description: Optional[str] + The description for this url. + category: Optional[str] + The category for this url. + target: Optional[str] + The target for this url. + """ + + description: Optional[str] = None + category: Optional[str] = None + target: Optional[str] = None + + +@dataclass +class WSUrlUpdate(WSUrlCreate): + """A websocket url update event model. + + Attributes + ---------- + url: str + The url which was deleted. + description: Optional[str] + The description for this url. + category: Optional[str] + The category for this url. + target: Optional[str] + The target for this url. + checked: Optional[int] + When this url was checked into the db? + """ + + checked: Optional[int] = None diff --git a/poetry.lock b/poetry.lock index 4609a46..96ff765 100644 --- a/poetry.lock +++ b/poetry.lock @@ -11,8 +11,8 @@ idna = ">=2.8" sniffio = ">=1.1" [package.extras] -doc = ["packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"] -test = ["contextlib2", "coverage[toml] (>=4.5)", "hypothesis (>=4.0)", "mock (>=4)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (<0.15)", "uvloop (>=0.15)"] +doc = ["packaging", "sphinx-rtd-theme", "sphinx-autodoc-typehints (>=1.2.0)"] +test = ["coverage[toml] (>=4.5)", "hypothesis (>=4.0)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "contextlib2", "uvloop (<0.15)", "mock (>=4)", "uvloop (>=0.15)"] trio = ["trio (>=0.16,<0.22)"] [[package]] @@ -105,8 +105,8 @@ rfc3986 = {version = ">=1.3,<2", extras = ["idna2008"]} sniffio = "*" [package.extras] -brotli = ["brotli", "brotlicffi"] -cli = ["click (>=8.0.0,<9.0.0)", "pygments (>=2.0.0,<3.0.0)", "rich (>=10,<13)"] +brotli = ["brotlicffi", "brotli"] +cli = ["click (>=8.0.0,<9.0.0)", "rich (>=10,<13)", "pygments (>=2.0.0,<3.0.0)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (>=1.0.0,<2.0.0)"] @@ -127,10 +127,10 @@ optional = false python-versions = ">=3.6.1,<4.0" [package.extras] -colors = ["colorama (>=0.4.3,<0.5.0)"] pipfile_deprecated_finder = ["pipreqs", "requirementslib"] +requirements_deprecated_finder = ["pipreqs", "pip-api"] +colors = ["colorama (>=0.4.3,<0.5.0)"] plugins = ["setuptools"] -requirements_deprecated_finder = ["pip-api", "pipreqs"] [[package]] name = "mypy-extensions" @@ -157,8 +157,8 @@ optional = false python-versions = ">=3.7" [package.extras] -docs = ["furo (>=2021.7.5b38)", "proselint (>=0.10.2)", "sphinx (>=4)", "sphinx-autodoc-typehints (>=1.12)"] -test = ["appdirs (==1.4.4)", "pytest (>=6)", "pytest-cov (>=2.7)", "pytest-mock (>=3.6)"] +docs = ["furo (>=2021.7.5b38)", "proselint (>=0.10.2)", "sphinx-autodoc-typehints (>=1.12)", "sphinx (>=4)"] +test = ["appdirs (==1.4.4)", "pytest-cov (>=2.7)", "pytest-mock (>=3.6)", "pytest (>=6)"] [[package]] name = "rfc3986" @@ -198,96 +198,35 @@ category = "dev" optional = false python-versions = ">=3.7" +[[package]] +name = "websockets" +version = "10.4" +description = "An implementation of the WebSocket Protocol (RFC 6455 & 7692)" +category = "main" +optional = false +python-versions = ">=3.7" + [metadata] lock-version = "1.1" python-versions = "^3.8" -content-hash = "2699f562f45d80451148f9fe072f4996c9495b3ae649f9052b883b578d7e8bf9" +content-hash = "5ea5fd3d0dad87e23e1221410f95ba3ecdcec9779135b6b2914e9f1f90466874" [metadata.files] -anyio = [ - {file = "anyio-3.6.2-py3-none-any.whl", hash = "sha256:fbbe32bd270d2a2ef3ed1c5d45041250284e31fc0a4df4a5a6071842051a51e3"}, - {file = "anyio-3.6.2.tar.gz", hash = "sha256:25ea0d673ae30af41a0c442f81cf3b38c7e79fdc7b60335a4c14e05eb0947421"}, -] -black = [ - {file = "black-22.10.0-1fixedarch-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:5cc42ca67989e9c3cf859e84c2bf014f6633db63d1cbdf8fdb666dcd9e77e3fa"}, - {file = "black-22.10.0-1fixedarch-cp311-cp311-macosx_11_0_x86_64.whl", hash = "sha256:5d8f74030e67087b219b032aa33a919fae8806d49c867846bfacde57f43972ef"}, - {file = "black-22.10.0-1fixedarch-cp37-cp37m-macosx_10_16_x86_64.whl", hash = "sha256:197df8509263b0b8614e1df1756b1dd41be6738eed2ba9e9769f3880c2b9d7b6"}, - {file = "black-22.10.0-1fixedarch-cp38-cp38-macosx_10_16_x86_64.whl", hash = "sha256:2644b5d63633702bc2c5f3754b1b475378fbbfb481f62319388235d0cd104c2d"}, - {file = "black-22.10.0-1fixedarch-cp39-cp39-macosx_11_0_x86_64.whl", hash = "sha256:e41a86c6c650bcecc6633ee3180d80a025db041a8e2398dcc059b3afa8382cd4"}, - {file = "black-22.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2039230db3c6c639bd84efe3292ec7b06e9214a2992cd9beb293d639c6402edb"}, - {file = "black-22.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:14ff67aec0a47c424bc99b71005202045dc09270da44a27848d534600ac64fc7"}, - {file = "black-22.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:819dc789f4498ecc91438a7de64427c73b45035e2e3680c92e18795a839ebb66"}, - {file = "black-22.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5b9b29da4f564ba8787c119f37d174f2b69cdfdf9015b7d8c5c16121ddc054ae"}, - {file = "black-22.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8b49776299fece66bffaafe357d929ca9451450f5466e997a7285ab0fe28e3b"}, - {file = "black-22.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:21199526696b8f09c3997e2b4db8d0b108d801a348414264d2eb8eb2532e540d"}, - {file = "black-22.10.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1e464456d24e23d11fced2bc8c47ef66d471f845c7b7a42f3bd77bf3d1789650"}, - {file = "black-22.10.0-cp37-cp37m-win_amd64.whl", hash = "sha256:9311e99228ae10023300ecac05be5a296f60d2fd10fff31cf5c1fa4ca4b1988d"}, - {file = "black-22.10.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:fba8a281e570adafb79f7755ac8721b6cf1bbf691186a287e990c7929c7692ff"}, - {file = "black-22.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:915ace4ff03fdfff953962fa672d44be269deb2eaf88499a0f8805221bc68c87"}, - {file = "black-22.10.0-cp38-cp38-win_amd64.whl", hash = "sha256:444ebfb4e441254e87bad00c661fe32df9969b2bf224373a448d8aca2132b395"}, - {file = "black-22.10.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:974308c58d057a651d182208a484ce80a26dac0caef2895836a92dd6ebd725e0"}, - {file = "black-22.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:72ef3925f30e12a184889aac03d77d031056860ccae8a1e519f6cbb742736383"}, - {file = "black-22.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:432247333090c8c5366e69627ccb363bc58514ae3e63f7fc75c54b1ea80fa7de"}, - {file = "black-22.10.0-py3-none-any.whl", hash = "sha256:c957b2b4ea88587b46cf49d1dc17681c1e672864fd7af32fc1e9664d572b3458"}, - {file = "black-22.10.0.tar.gz", hash = "sha256:f513588da599943e0cde4e32cc9879e825d58720d6557062d1098c5ad80080e1"}, -] -certifi = [ - {file = "certifi-2022.9.24-py3-none-any.whl", hash = "sha256:90c1a32f1d68f940488354e36370f6cca89f0f106db09518524c88d6ed83f382"}, - {file = "certifi-2022.9.24.tar.gz", hash = "sha256:0d9c601124e5a6ba9712dbc60d9c53c21e34f5f641fe83002317394311bdce14"}, -] -click = [ - {file = "click-8.1.3-py3-none-any.whl", hash = "sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48"}, - {file = "click-8.1.3.tar.gz", hash = "sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e"}, -] -colorama = [ - {file = "colorama-0.4.5-py2.py3-none-any.whl", hash = "sha256:854bf444933e37f5824ae7bfc1e98d5bce2ebe4160d46b5edf346a89358e99da"}, - {file = "colorama-0.4.5.tar.gz", hash = "sha256:e6c6b4334fc50988a639d9b98aa429a0b57da6e17b9a44f0451f930b6967b7a4"}, -] -h11 = [ - {file = "h11-0.12.0-py3-none-any.whl", hash = "sha256:36a3cb8c0a032f56e2da7084577878a035d3b61d104230d4bd49c0c6b555a9c6"}, - {file = "h11-0.12.0.tar.gz", hash = "sha256:47222cb6067e4a307d535814917cd98fd0a57b6788ce715755fa2b6c28b56042"}, -] -httpcore = [ - {file = "httpcore-0.15.0-py3-none-any.whl", hash = "sha256:1105b8b73c025f23ff7c36468e4432226cbb959176eab66864b8e31c4ee27fa6"}, - {file = "httpcore-0.15.0.tar.gz", hash = "sha256:18b68ab86a3ccf3e7dc0f43598eaddcf472b602aba29f9aa6ab85fe2ada3980b"}, -] -httpx = [ - {file = "httpx-0.23.0-py3-none-any.whl", hash = "sha256:42974f577483e1e932c3cdc3cd2303e883cbfba17fe228b0f63589764d7b9c4b"}, - {file = "httpx-0.23.0.tar.gz", hash = "sha256:f28eac771ec9eb4866d3fb4ab65abd42d38c424739e80c08d8d20570de60b0ef"}, -] -idna = [ - {file = "idna-3.4-py3-none-any.whl", hash = "sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2"}, - {file = "idna-3.4.tar.gz", hash = "sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4"}, -] -isort = [ - {file = "isort-5.10.1-py3-none-any.whl", hash = "sha256:6f62d78e2f89b4500b080fe3a81690850cd254227f27f75c3a0c491a1f351ba7"}, - {file = "isort-5.10.1.tar.gz", hash = "sha256:e8443a5e7a020e9d7f97f1d7d9cd17c88bcb3bc7e218bf9cf5095fe550be2951"}, -] -mypy-extensions = [ - {file = "mypy_extensions-0.4.3-py2.py3-none-any.whl", hash = "sha256:090fedd75945a69ae91ce1303b5824f428daf5a028d2f6ab8a299250a846f15d"}, - {file = "mypy_extensions-0.4.3.tar.gz", hash = "sha256:2d82818f5bb3e369420cb3c4060a7970edba416647068eb4c5343488a6c604a8"}, -] -pathspec = [ - {file = "pathspec-0.10.1-py3-none-any.whl", hash = "sha256:46846318467efc4556ccfd27816e004270a9eeeeb4d062ce5e6fc7a87c573f93"}, - {file = "pathspec-0.10.1.tar.gz", hash = "sha256:7ace6161b621d31e7902eb6b5ae148d12cfd23f4a249b9ffb6b9fee12084323d"}, -] -platformdirs = [ - {file = "platformdirs-2.5.2-py3-none-any.whl", hash = "sha256:027d8e83a2d7de06bbac4e5ef7e023c02b863d7ea5d079477e722bb41ab25788"}, - {file = "platformdirs-2.5.2.tar.gz", hash = "sha256:58c8abb07dcb441e6ee4b11d8df0ac856038f944ab98b7be6b27b2a3c7feef19"}, -] -rfc3986 = [ - {file = "rfc3986-1.5.0-py2.py3-none-any.whl", hash = "sha256:a86d6e1f5b1dc238b218b012df0aa79409667bb209e58da56d0b94704e712a97"}, - {file = "rfc3986-1.5.0.tar.gz", hash = "sha256:270aaf10d87d0d4e095063c65bf3ddbc6ee3d0b226328ce21e036f946e421835"}, -] -sniffio = [ - {file = "sniffio-1.3.0-py3-none-any.whl", hash = "sha256:eecefdce1e5bbfb7ad2eeaabf7c1eeb404d7757c379bd1f7e5cce9d8bf425384"}, - {file = "sniffio-1.3.0.tar.gz", hash = "sha256:e60305c5e5d314f5389259b7f22aaa33d8f7dee49763119234af3755c55b9101"}, -] -tomli = [ - {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"}, - {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, -] -typing-extensions = [ - {file = "typing_extensions-4.4.0-py3-none-any.whl", hash = "sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e"}, - {file = "typing_extensions-4.4.0.tar.gz", hash = "sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa"}, -] +anyio = [] +black = [] +certifi = [] +click = [] +colorama = [] +h11 = [] +httpcore = [] +httpx = [] +idna = [] +isort = [] +mypy-extensions = [] +pathspec = [] +platformdirs = [] +rfc3986 = [] +sniffio = [] +tomli = [] +typing-extensions = [] +websockets = [] diff --git a/pyproject.toml b/pyproject.toml index 7a50918..aeaa02c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ packages = [{include = "fishfish"}] [tool.poetry.dependencies] python = "^3.8" httpx = "^0.23.0" +websockets = "^10.4" [tool.poetry.dev-dependencies] black = "^22.10.0"