diff --git a/inventory.py b/inventory.py index 29cd941..007e76e 100644 --- a/inventory.py +++ b/inventory.py @@ -127,5 +127,17 @@ class DropsCampaign: def active(self): return self.status == "ACTIVE" + @property + def total_drops(self) -> int: + return len(self.timed_drops) + + @property + def claimed_drops(self) -> int: + return sum(d.is_claimed for d in self.timed_drops.values()) + + @property + def remaining_drops(self) -> int: + return sum(not d.is_claimed for d in self.timed_drops.values()) + def get_drop(self, drop_id: str) -> Optional[TimedDrop]: return self.timed_drops.get(drop_id) diff --git a/main.py b/main.py index 03cc651..bce8e52 100644 --- a/main.py +++ b/main.py @@ -7,6 +7,7 @@ import json import ctypes import logging import asyncio +import warnings import traceback import threading from typing import Any, Dict, NoReturn @@ -20,7 +21,9 @@ try: except (ImportError, ModuleNotFoundError): raise ImportError("You have to run 'python -m pip install pywin32' first") - +# disable some warnings +warnings.simplefilter("ignore", RuntimeWarning) +warnings.simplefilter("ignore", DeprecationWarning) # nice console title try: ctypes.windll.kernel32.SetConsoleTitleW(f"Twitch Drops Miner v{__version__} (by DevilXD)") diff --git a/twitch.py b/twitch.py index 40fbbeb..e5b2493 100644 --- a/twitch.py +++ b/twitch.py @@ -63,10 +63,14 @@ class Twitch: self._channel_change = asyncio.Event() # Inventory self.inventory: List[DropsCampaign] = [] + self._campaign_change = asyncio.Event() def wait_until_login(self): return self._is_logged_in.wait() + def reevaluate_campaigns(self): + self._campaign_change.set() + async def close(self): print("Exiting...") self._session.cookie_jar.save(COOKIES_PATH) # type: ignore @@ -87,75 +91,96 @@ class Twitch: • Selecting a stream to watch, and watching it • Changing the stream that's being watched if necessary """ - # Start our websocket connection - shouldn't require task tracking - asyncio.create_task(self.websocket.connect()) - # Claim the drops we can - self.inventory = await self.get_inventory() - games = set() - for campaign in self.inventory: - if campaign.status == "UPCOMING": - # we have no use in processing upcoming campaigns here - continue - for drop in campaign.timed_drops.values(): - if drop.can_earn: - games.add(campaign.game) - if drop.can_claim: - await drop.claim() - # games now has all games we want to farm drops for - if not channel_names: - # get a list of all channels with drops enabled - print("Fetching suitable live channels to watch...") - live_streams: Dict[Game, List[Channel]] = await self.get_live_streams( - games, [DROPS_ENABLED_TAG] - ) - for game, channels in live_streams.items(): - for channel in channels: - self.channels[channel.id] = channel - print(f"Added channel: {channel.name} for game: {game.name}") - else: - # Fetch information about all channels we're supposed to handle - for channel_name in channel_names: - channel: Channel = await Channel(self, channel_name) # type: ignore - self.channels[channel.id] = channel - # Sub to these channel updates - topics: List[WebsocketTopic] = [] - for channel_id in self.channels: - topics.append( - get_topic( - "VideoPlayback", channel_id, partial(self.process_stream_state, channel_id) - ) - ) - await self.websocket.add_topics(topics) - - # Repeat: Change into a channel we can watch, then reset the flag - self._channel_change.set() - refresh_channels = False # we're entering having fresh channel data already + # Start our websocket connection + self.websocket.start() while True: - # wait for the change channel signal - await self._channel_change.wait() - for channel in self.channels.values(): - if ( - channel.stream is not None # steam online - and channel.stream.drops_enabled # drops are enabled - and channel.stream.game in games # streams a game we can earn drops in - ): - self.watch(channel) - refresh_channels = True - self._channel_change.clear() - break - else: - # there's no available channel to watch - if refresh_channels: - # refresh the status of all channels, - # to make sure that our websocket didn't miss anything til this point - print("No suitable channel to watch, refreshing...") - for channel in self.channels.values(): - await channel.get_stream() - await asyncio.sleep(0.5) - refresh_channels = False + # Claim the drops we can + self.inventory = await self.get_inventory() + games = set() + for campaign in self.inventory: + if campaign.status == "UPCOMING": + # we have no use in processing upcoming campaigns here continue - print("No suitable channel to watch, retrying in 120 seconds") - await asyncio.sleep(120) + for drop in campaign.timed_drops.values(): + if drop.can_earn: + games.add(campaign.game) + if drop.can_claim: + await drop.claim() + # games now has all games we want to farm drops for + # if it's empty, there's no point in continuing + if not games: + print( + "No active campaigns to farm drops for.\n\n" + "Application Terminated.\nClose the console window to exit the application.", + ) + await asyncio.Future() + if not channel_names: + # get a list of all channels with drops enabled + print("Fetching suitable live channels to watch...") + live_streams: Dict[Game, List[Channel]] = await self.get_live_streams( + games, [DROPS_ENABLED_TAG] + ) + for game, channels in live_streams.items(): + for channel in channels: + self.channels[channel.id] = channel + print(f"Added channel: {channel.name} for game: {game.name}") + else: + # Fetch information about all channels we're supposed to handle + for channel_name in channel_names: + channel: Channel = await Channel(self, channel_name) # type: ignore + self.channels[channel.id] = channel + # Sub to these channel updates + topics: List[WebsocketTopic] = [] + for channel_id in self.channels: + topics.append( + get_topic( + "VideoPlayback", channel_id, partial(self.process_stream_state, channel_id) + ) + ) + await self.websocket.add_topics(topics) + + # Repeat: Change into a channel we can watch, then reset the flag + self._channel_change.set() + refresh_channels = False # we're entering having fresh channel data already + while True: + # wait for either the change channel signal, or campaign change signal + await asyncio.wait( + ( + self._channel_change.wait(), + self._campaign_change.wait(), + ), + return_when=asyncio.FIRST_COMPLETED, + ) + if self._campaign_change.is_set(): + # we need to reevaluate all campaigns + # reconnect the websocket + await self.websocket.close() + self.websocket.start() + break # cycles the outer loop + # otherwise, it was the channel change one + for channel in self.channels.values(): + if ( + channel.stream is not None # steam online + and channel.stream.drops_enabled # drops are enabled + and channel.stream.game in games # streams a game we can earn drops in + ): + self.watch(channel) + refresh_channels = True + self._channel_change.clear() + break + else: + # there's no available channel to watch + if refresh_channels: + # refresh the status of all channels, + # to make sure that our websocket didn't miss anything til this point + print("No suitable channel to watch, refreshing...") + for channel in self.channels.values(): + await channel.get_stream() + await asyncio.sleep(0.5) + refresh_channels = False + continue + print("No suitable channel to watch, retrying in 120 seconds") + await asyncio.sleep(120) def watch(self, channel: Channel): if self._watching_task is not None: @@ -217,7 +242,7 @@ class Twitch: assert channel.stream is not None viewers = message["viewers"] channel.stream.viewer_count = viewers - logger.info(f"{channel.name} viewers: {viewers}") + logger.debug(f"{channel.name} viewers: {viewers}") async def _validate_password(self, password: str) -> bool: """ diff --git a/websocket.py b/websocket.py index 67232d2..b82306a 100644 --- a/websocket.py +++ b/websocket.py @@ -5,13 +5,15 @@ import random import string import asyncio import logging -from typing import Any, Optional, Dict, Tuple, Set, Iterable, TYPE_CHECKING +from functools import wraps +from collections import deque +from typing import Any, Optional, Union, Dict, Tuple, Set, Iterable, cast, TYPE_CHECKING from websockets.exceptions import ConnectionClosed, ConnectionClosedOK from websockets.client import WebSocketClientProtocol, connect as websocket_connect -from exceptions import MinerException from inventory import TimedDrop +from exceptions import MinerException from constants import WEBSOCKET_URL, PING_INTERVAL, WebsocketTopic, get_topic if TYPE_CHECKING: @@ -21,118 +23,27 @@ if TYPE_CHECKING: logger = logging.getLogger("TwitchDrops") +def task_wrapper(afunc): + @wraps(afunc) + async def wrapper(self: Websocket, *args, **kwargs): + try: + await afunc(self, *args, **kwargs) + except Exception: + logger.exception("Exception in websocket task") + return wrapper + + class Websocket: def __init__(self, twitch: Twitch): self._twitch = twitch self._ws: Optional[WebSocketClientProtocol] = None self.connected = asyncio.Event() # set when there's an active websocket connection self.reconnect = asyncio.Event() # set when the websocket needs to reconnect - self._send_queue: asyncio.Queue[Tuple[str, Dict[str, Any]]] = asyncio.Queue() + self._send_queue: deque[Tuple[str, Dict[str, Any]]] = deque() self._recv_dict: Dict[str, asyncio.Future[Any]] = {} self._topics: Set[WebsocketTopic] = set() self._ping_task: Optional[asyncio.Task[Any]] = None - - async def connect(self): - # ensure we're logged in before connecting - await self._twitch.wait_until_login() - assert self._twitch._user_id is not None - logger.info("Connecting to Websocket") - # Listen to our events of choice - user_id = self._twitch._user_id - # Defer topics via a task so we can proceed to the actual websocket connection - asyncio.create_task( - self.add_topics([ - get_topic("UserDrops", user_id, self.process_drops), - get_topic("UserCommunityPoints", user_id, self.process_points), - ]) - ) - # Connect/Reconnect loop - async for websocket in websocket_connect(WEBSOCKET_URL, ssl=True, ping_interval=None): - websocket.BACKOFF_MAX = 3 * 60 # type: ignore # 3 minutes - self._ws = websocket - self.reconnect.clear() - self.connected.set() - logger.info("Websocket Connected") - self._ping_task = asyncio.create_task(self._ping_loop()) - try: - while not self.reconnect.is_set(): - # Process receive - try: - # Wait up to 0.5s for a message we're supposed to receive - raw_message = await asyncio.wait_for(websocket.recv(), timeout=0.5) - except asyncio.TimeoutError: - # nothing - skip handling - pass - else: - # we've got something to process - message = json.loads(raw_message) - logger.debug(f"Websocket received: {message}") - msg_type = message["type"] - # handle the simple PING case - if msg_type == "PONG": - ping_future = self._recv_dict.pop("PING", None) - if ping_future is not None and not ping_future.done(): - ping_future.set_result(message) - elif msg_type == "RESPONSE": - try: - self._recv_dict.pop(message["nonce"]).set_result(message) - except KeyError: - logger.exception("Received response for a request we didn't send") - elif msg_type == "RECONNECT": - # We've received a reconnect request - logger.warning("Received a Websocket Reconnect Request") - self.reconnect.set() - elif msg_type == "MESSAGE": - # request the assigned topic to process the response - target_topic = message["data"]["topic"] - for topic in self._topics: - if target_topic == topic: - try: - await topic.process(json.loads(message["data"]["message"])) - except Exception: - logger.exception("Exception in websocket topic process") - break - else: - logger.error(f"Received unknown websocket payload: {message}") - - # Early exit if needed - if self.reconnect.is_set(): - break - - # Process send - while not self._send_queue.empty(): - nonce, message = self._send_queue.get_nowait() - if nonce != "PING": - message["nonce"] = nonce - await websocket.send(json.dumps(message, separators=(',', ':'))) - logger.debug(f"Websocket sent: {message}") - # A reconnect was requested - continue - except ConnectionClosed as exc: - self.connected.clear() - if self._ping_task is not None: - self._ping_task.cancel() - self._ping_task = None - if isinstance(exc, ConnectionClosedOK): - if exc.rcvd_then_sent: - # server closed the connection, not us - reconnect - logger.warning("Websocket Disconnected - Reconnecting") - continue - # we closed it - exit - break - # otherwise, reconnect - logger.warning("Websocket Closed - Reconnecting") - continue - except Exception: - logger.exception("Exception in Websocket - Reconnecting") - continue - - async def close(self): - self.connected.clear() - if self._ping_task is not None: - self._ping_task.cancel() - if self._ws is not None: - await self._ws.close() + self._connect_task: Optional[asyncio.Task[Any]] = None async def _ping_loop(self): await self.connected.wait() @@ -147,6 +58,136 @@ class Websocket: break await asyncio.sleep(ping_every) + def change_connection_state(self, state: bool): + if state: + # websocket is considered connected + logger.info("Websocket Connected") + self.connected.set() + self._ping_task = asyncio.create_task(self._ping_loop()) + else: + # websocket is considered disconnected + self.connected.clear() + if self._ping_task is not None: + self._ping_task.cancel() + self._ping_task = None + + def start(self): + self._connect_task = asyncio.create_task(self.connect()) + + def stop(self): + self.change_connection_state(False) + if self._connect_task is not None: + self._connect_task.cancel() + self._connect_task = None + + @task_wrapper + async def connect(self): + # ensure we're logged in before connecting + await self._twitch.wait_until_login() + logger.info("Connecting to Websocket") + # Listen to our events of choice + user_id = cast(int, self._twitch._user_id) + # Add default topics + self.add_topics([ + get_topic("UserDrops", user_id, self.process_drops), + get_topic("UserCommunityPoints", user_id, self.process_points), + ]) + # Connect/Reconnect loop + async for websocket in websocket_connect(WEBSOCKET_URL, ssl=True, ping_interval=None): + websocket.BACKOFF_MAX = 3 * 60 # type: ignore # 3 minutes + self._ws = websocket + self.reconnect.clear() + self.change_connection_state(True) + # Send all our chosen topics + topics_list = list(map(str, self._topics)) + logger.info(f"Listening for: {', '.join(topics_list)}") + self.send( + { + "type": "LISTEN", + "data": { + "topics": topics_list, + "auth_token": self._twitch._access_token, + } + } + ) + try: + while not self.reconnect.is_set(): + # Process receive + try: + # Wait up to 0.5s for a message we're supposed to receive + raw_message = await asyncio.wait_for(websocket.recv(), timeout=0.5) + except asyncio.TimeoutError: + # nothing - skip handling + pass + else: + # we've got something to process + # separate method solely because the indent was getting rather ridiculus + await self.process_message(raw_message) + + # Early exit if needed + if self.reconnect.is_set(): + break + + # Process send + while self._send_queue: + nonce, message = self._send_queue.popleft() + if nonce != "PING": + message["nonce"] = nonce + await websocket.send(json.dumps(message, separators=(',', ':'))) + logger.debug(f"Websocket sent: {message}") + # A reconnect was requested + self.change_connection_state(False) + continue + except ConnectionClosed as exc: + self.change_connection_state(False) + if isinstance(exc, ConnectionClosedOK): + if exc.rcvd_then_sent: + # server closed the connection, not us - reconnect + logger.warning("Server Disconnected - Reconnecting") + continue + # we closed it - exit + return + # otherwise, reconnect + logger.warning("Websocket Closed - Reconnecting") + continue + except Exception: + logger.exception("Exception in Websocket - Reconnecting") + continue + + async def process_message(self, raw_message: Union[bytes, str]): + message = json.loads(raw_message) + logger.debug(f"Websocket received: {message}") + msg_type = message["type"] + # handle the simple PING case + if msg_type == "PONG": + ping_future = self._recv_dict.pop("PING", None) + if ping_future is not None and not ping_future.done(): + ping_future.set_result(message) + elif msg_type == "RESPONSE": + try: + self._recv_dict.pop(message["nonce"]).set_result(message) + except KeyError: + logger.exception("Received response for a request we didn't send") + elif msg_type == "RECONNECT": + # We've received a reconnect request + logger.warning("Received a Websocket Reconnect Request") + self.reconnect.set() + elif msg_type == "MESSAGE": + # request the assigned topic to process the response + target_topic = message["data"]["topic"] + for topic in self._topics: + if target_topic == topic: + # use a task to not block the websocket + asyncio.create_task(topic.process(json.loads(message["data"]["message"]))) + break + else: + logger.error(f"Received unknown websocket payload: {message}") + + async def close(self): + self.stop() + if self._ws is not None: + await self._ws.close() + def create_nonce(self, length: int = 30) -> str: available_chars = string.ascii_letters + string.digits return ''.join(random.choices(available_chars, k=length)) @@ -158,12 +199,12 @@ class Websocket: nonce = "PING" else: nonce = self.create_nonce() - self._send_queue.put_nowait((nonce, message)) + self._send_queue.append((nonce, message)) future: asyncio.Future[Dict[str, Any]] = asyncio.get_running_loop().create_future() self._recv_dict[nonce] = future return future - async def add_topics(self, topics: Iterable[WebsocketTopic]): + def add_topics(self, topics: Iterable[WebsocketTopic]): # ensure no topics end up duplicated topics = set(topics) topics.difference_update(self._topics) @@ -174,21 +215,24 @@ class Websocket: if len(self._topics) >= 50: # TODO: Handle multiple connections (up to 10) since one allows only up to 50 topics raise MinerException("Too many topics") - await self.connected.wait() - assert self._twitch._access_token is not None - auth_token = self._twitch._access_token - topics_list = list(map(str, topics)) - logger.info(f"Listening for: {', '.join(topics_list)}") - await self.send( - { - "type": "LISTEN", - "data": { - "topics": topics_list, - "auth_token": auth_token, + if self.connected.is_set(): + # we're already connected, so we have to send the topics list ourselves + topics_list = list(map(str, topics)) + logger.info(f"Listening for: {', '.join(topics_list)}") + return self.send( + { + "type": "LISTEN", + "data": { + "topics": topics_list, + "auth_token": self._twitch._access_token, + } } - } - ) + ) + else: + # no connection is made, so let it wait until there is one + return self.connected.wait() + @task_wrapper async def process_drops(self, message: Dict[str, Any]): drop_id = message["data"]["drop_id"] drop: Optional[TimedDrop] = None @@ -203,13 +247,20 @@ class Websocket: msg_type = message["type"] if msg_type == "drop-progress": print( - f"Drop: {drop.rewards_text()}: {drop.progress:5.1%} " + f"Drop: {drop.rewards_text()}: {drop.progress:6.1%} " f"({drop.remaining_minutes} minutes remaining)" ) elif msg_type == "drop-claim": + campaign = drop.campaign await drop.claim() - print(f"Claimed drop: {drop.rewards_text()}") + print( + f"Claimed drop: {drop.rewards_text()} " + f"({campaign.claimed_drops}/{campaign.total_drops})" + ) + if campaign.remaining_drops == 0: + self._twitch.reevaluate_campaigns() + @task_wrapper async def process_points(self, message: Dict[str, Any]): msg_type = message["type"] if msg_type == "points-earned":