diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml index 42468a5..17062ce 100644 --- a/.github/workflows/format.yml +++ b/.github/workflows/format.yml @@ -9,6 +9,6 @@ jobs: - uses: actions/checkout@v3 - uses: psf/black@stable with: - options: "--check --verbose --exclude 'rustplus/api/remote/camera/camera_constants.py'" + options: "--check --verbose --exclude 'rustplus/remote/camera/camera_constants.py'" src: "./rustplus" version: "~= 24.4.2" \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in index 24ed4e9..7009f2f 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,3 @@ recursive-include rustplus *.png -recursive-include rustplus *.ttf \ No newline at end of file +recursive-include rustplus *.ttf +include requirements.txt diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index e5265e9..b23524b 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -21,10 +21,8 @@ * [Getting the Time](api-methods/getting-the-time.md) * [Getting Entity Information](api-methods/getting-entity-information.md) * [Getting Map Markers](api-methods/getting-map-markers.md) -* [Getting Current Map Events](api-methods/getting-current-map-events.md) * [Getting Contents of Monitors](api-methods/getting-contents-of-monitors.md) * [Promoting Players to Team Leader](api-methods/promoting-players-to-team-leader.md) -* [Toggling Smart Switches](api-methods/toggling-smart-switches.md) ## Command System diff --git a/docs/api-methods/getting-current-map-events.md b/docs/api-methods/getting-current-map-events.md deleted file mode 100644 index c0f1532..0000000 --- a/docs/api-methods/getting-current-map-events.md +++ /dev/null @@ -1,12 +0,0 @@ -# Getting Current Map Events - -The following are all defined as "events": - -* Explosions (Bradley / Attack Helicopter) -* Cargo Ship -* CH47 (Chinook) -* Locked Crates -* Attack Helicopter - -Calling `rust_socket.get_current_events()` returns a list of all current `RustMarker`'s that are the above events. This can be used for working out whether Cargo Ship / Oil Rig etc has been taken / is being taken. See [Here ](getting-map-markers.md)for information on `RustMarker` - diff --git a/docs/api-methods/getting-entity-information.md b/docs/api-methods/getting-entity-information.md index 7b774d5..d74884f 100644 --- a/docs/api-methods/getting-entity-information.md +++ b/docs/api-methods/getting-entity-information.md @@ -27,3 +27,7 @@ Alarm = 2 StorageMonitor = 3 ``` +## Setting Entity Information + +Calling `rust_socket.set_entity_value(entity_id: int, value: bool)` will set the value of the entity with the given ID to the given value. + diff --git a/docs/api-methods/removing-listeners.md b/docs/api-methods/removing-listeners.md index 16f0245..f6e0a7b 100644 --- a/docs/api-methods/removing-listeners.md +++ b/docs/api-methods/removing-listeners.md @@ -2,9 +2,21 @@ ### Registered Listeners -A registered listener is a wrapper object around the coroutine itself that will allow the listener to be removed later on. Should you need the coroutine back, call `RegisteredListener.get_coro()`. +A registered listener is a wrapper object around the coroutine itself that will allow the listener to be removed later +on. Should you need the coroutine back, call `RegisteredListener.get_coro()`. ### Removing The listener -Removing a listener is as simple as calling `RustSocket.remove_listener(RegisteredListener)` and will return a boolean value. True if a listener was removed and false otherwise +Removing a listener is as simple as using an Event's HandlerList. This is one example: +```python +@EntityEvent(server_details, 25743493) +async def on_entity_event(payload: EntityEventPayload): + await rust_socket.set_entity_value(payload.entity_id, not payload.value) + + +EntityEventPayload.HANDLER_LIST.unregister(on_entity_event, server_details) + +# You can also unregister all listeners for a specific event +EntityEventPayload.HANDLER_LIST.unregister_all() +``` diff --git a/docs/api-methods/toggling-smart-switches.md b/docs/api-methods/toggling-smart-switches.md deleted file mode 100644 index 86d4f26..0000000 --- a/docs/api-methods/toggling-smart-switches.md +++ /dev/null @@ -1,8 +0,0 @@ -# Toggling Smart Switches - -Calling `rust_socket.turn_on_smart_switch(eid: int)` and `rust_socket.turn_off_smart_switch(eid: int)` can be used to toggle smart switches on and off from code. To get the entity\_ids look at: - -{% content-ref url="../getting-started/getting-player-details/getting-entity-ids.md" %} -[getting-entity-ids.md](../getting-started/getting-player-details/getting-entity-ids.md) -{% endcontent-ref %} - diff --git a/docs/command-system/command-decorator.md b/docs/command-system/command-decorator.md index ad6f780..1c1e5de 100644 --- a/docs/command-system/command-decorator.md +++ b/docs/command-system/command-decorator.md @@ -1,16 +1,18 @@ +from main import server_details + # Command Decorator The command decorator is used to mark a coroutine as a command listener. Usage: ```python -@rust_socket.command -async def hi(command: Command): +@Command(server_details) +async def hi(command: ChatCommand): print("Command Ran!") ``` The fact that the coroutine's name is `hi` means that the command will be `hi` . -You also get access to this `Command` object which has a slew of useful information about the how the command was called. +You also get access to this `ChatCommand` object which has a slew of useful information about the how the command was called. | Field | Value | | ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | @@ -27,16 +29,16 @@ This decorator returns a [`RegisteredListener`](../api-methods/removing-listener You don't want to have to register 100's of commands for every permutation of phrasing, so why should you! ```python -@rust_socket.command(aliases=["hello", "hey"]) -async def hi(command: Command): +@ChatCommand(server_details, aliases=["hello", "hey"]) +async def hi(command: ChatCommand): print("Command Ran!") ``` This is a simple example of how you could incorporate different function names into one command, but sometimes we need more than that! ```python -@rust_socket.command(alais_func=lambda x: x.lower() == "test") -async def pair(command: Command): +@ChatCommand(server_details, alais_func=lambda x: x.lower() == "test") +async def pair(command: ChatCommand): print("Command Ran!") ``` diff --git a/docs/command-system/command-options.md b/docs/command-system/command-options.md index 5d2848c..b189f87 100644 --- a/docs/command-system/command-options.md +++ b/docs/command-system/command-options.md @@ -1,14 +1,11 @@ # Command Options -Command options are what you use to tell the [`RustSocket`](../getting-started/rustsocket/) what the general structure of your commands will be. These define the prefix for the command, as well as any "overruling commands" which are commands that do not require a prefix. Usage: +Command options are what you use to tell the [`RustSocket`](../getting-started/rustsocket/) what the general structure of your commands will be. +These define the prefix for the command, as well as any "overruling commands" which are commands that do not require a prefix. Usage: ```python -from rustplus import RustSocket +from rustplus import CommandOptions -options = CommandOptions(prefix="!", overruling_commands = ["time"]) - -# Prefix is a string, and the overruling_commands are a list of strings which would be the name of the coroutines +options = CommandOptions(prefix="!") ``` -You can then just pass these into the RustSocket constructor using the `overruling_commands` kwarg. - diff --git a/docs/command-system/commands-overview.md b/docs/command-system/commands-overview.md index 962709d..a167683 100644 --- a/docs/command-system/commands-overview.md +++ b/docs/command-system/commands-overview.md @@ -4,13 +4,14 @@ Commands allow the triggering of custom coroutines when a specific keyword is se {% code title="main.py" %} ```python -from rustplus import RustSocket, CommandOptions, Command +from rustplus import RustSocket, CommandOptions, Command, ServerDetails, Command, ChatCommand options = CommandOptions(prefix="!") # Use whatever prefix you want here -rust_socket = RustSocket("IP", "PORT", STEAMID, PLAYERTOKEN, command_options=options) +server_details = ServerDetails("IP", "PORT", STEAMID, PLAYERTOKEN) +socket = RustSocket(server_details) -@rust_socket.command -async def hi(command : Command): +@Command(server_details) +async def hi(command : ChatCommand): await socket.send_team_message(f"Hi, {command.sender_name}") ``` {% endcode %} diff --git a/docs/community-examples/examples.md b/docs/community-examples/examples.md new file mode 100644 index 0000000..43d1479 --- /dev/null +++ b/docs/community-examples/examples.md @@ -0,0 +1,4 @@ +# Community Examples + +This page contains a list of community-contributed examples that demonstrate how to use the Library. +Contact me on Discord if you would like to add your example to this list! \ No newline at end of file diff --git a/docs/event-system/events-overview.md b/docs/event-system/events-overview.md index 40e8019..0b2e844 100644 --- a/docs/event-system/events-overview.md +++ b/docs/event-system/events-overview.md @@ -10,25 +10,30 @@ Socket Events are called whenever a specific action / event happens. There are c These will be called by the socket when the respective events occur. Here are some example usages: {% code title="listeners.py" %} + ```python -from rustplus import EntityEvent, TeamEvent, ChatEvent - -@rust_socket.entity_event(ENTITYID) -async def alarm(event : EntityEvent): - value = "On" if event.value else "Off" - print(f"{entity_type_to_string(event.type)} has been turned {value}") - -@rust_socket.team_event -async def team(event : TeamEvent): - print(f"The team leader's steamId is: {event.team_info.leader_steam_id}") - -@rust_socket.chat_event -async def chat(event : ChatEvent): - print(f"{event.message.name}: {event.message.message}") - -@rust_socket.protobuf_received +from rustplus import EntityEventPayload, TeamEventPayload, ChatEventPayload, ProtobufEvent, ChatEvent, EntityEvent, TeamEvent + + +@EntityEvent(server_details, 25743493) +async def alarm(event: EntityEventPayload): + value = "On" if event.value else "Off" + print(f"Entity has been turned {value}") + + +@TeamEvent(server_details) +async def team(event: TeamEventPayload): + print(f"The team leader's steamId is: {event.team_info.leader_steam_id}") + + +@ChatEvent(server_details) +async def chat(event: ChatEventPayload): + print(f"{event.message.name}: {event.message.message}") + + +@ProtobufEvent(server_details) async def proto(data: bytes): - print(data) + print(data) ``` {% endcode %} @@ -37,8 +42,7 @@ async def proto(data: bytes): The `entity_event` decorator takes an extra parameter of the [entity id](../getting-started/getting-player-details/getting-entity-ids.md) that you are listening for changes to. The `EntityEvent` object holds information on the entity: | Name | Description | -| ------------------- | ------------------------------------- | -| `type` | The type of entity, as an `int` | +|---------------------|---------------------------------------| | `entity_id` | The Entity Id | | `value` | The value of the entity, `boolean` | | `capacity` | The capacity of the entity | @@ -51,7 +55,7 @@ The `entity_event` decorator takes an extra parameter of the [entity id](../gett This event is typically called when the team changes, e.g. a player leaves or joins. The `team_event` decorator will pass a `TeamEvent` object as a parameter with the following information: | Name | Description | -| ------------- | ------------------------------------------------------------------------------- | +|---------------|---------------------------------------------------------------------------------| | `player_info` | The `player_id` of the changed information | | `team_info` | The [`team info`](../api-methods/getting-team-info.md) on the team that changed | @@ -60,7 +64,7 @@ This event is typically called when the team changes, e.g. a player leaves or jo This event is called when a message is sent to the team chat. It will give you a `ChatEvent` object when called with this information: | Name | Description | -| --------- | ---------------------------------------------------------------- | +|-----------|------------------------------------------------------------------| | `message` | The [message](../api-methods/getting-team-chat.md) that was sent | ### Protobuf Event diff --git a/docs/getting-started/getting-player-details/README.md b/docs/getting-started/getting-player-details/README.md index 2ae2c26..fafbf4c 100644 --- a/docs/getting-started/getting-player-details/README.md +++ b/docs/getting-started/getting-player-details/README.md @@ -1,5 +1,5 @@ --- -description: This will show you how to get your personal player details using the RustCli +description: This will show you how to get your personal player details using the Web tool --- # Getting Player Details diff --git a/docs/getting-started/getting-player-details/fcm-listener.md b/docs/getting-started/getting-player-details/fcm-listener.md index bf13f09..f1b4a35 100644 --- a/docs/getting-started/getting-player-details/fcm-listener.md +++ b/docs/getting-started/getting-player-details/fcm-listener.md @@ -19,7 +19,7 @@ FCM(fcm_details).start() ``` {% endcode %} -The `on_notification` method will be called everytime a message is recieved from the game server. +The `on_notification` method will be called everytime a message is received from the game server. The `rustplus.py.config.json` is the file created by the RustCli, when you register for FCM notifications. See: diff --git a/docs/getting-started/quick-start.md b/docs/getting-started/quick-start.md index 2360b5f..6475069 100644 --- a/docs/getting-started/quick-start.md +++ b/docs/getting-started/quick-start.md @@ -6,7 +6,7 @@ In order to access the API, you must first install the package using pip: pip install rustplus ``` -You must then get your Personal details using the RustCli, as shown here: +You must then get your Personal details using the web tool, as shown here: {% content-ref url="getting-player-details/" %} [getting-player-details](getting-player-details/) @@ -15,10 +15,11 @@ You must then get your Personal details using the RustCli, as shown here: {% code title="main.py" %} ```python import asyncio -from rustplus import RustSocket +from rustplus import RustSocket, ServerDetails async def main(): - socket = RustSocket("IP", "PORT", STEAMID, PLAYERTOKEN) + server_details = ServerDetails("IP", "PORT", STEAMID, PLAYERTOKEN) + socket = RustSocket(server_details) await socket.connect() print(f"It is {(await socket.get_time()).time}") @@ -29,5 +30,6 @@ asyncio.run(main()) ``` {% endcode %} -This will run, and print the time on the Rust Server +This will run, and print the time on the Rust Server. +API methods will return `None` if they are not successful. This is to allow for better error handling, and reconnecting to the server. diff --git a/rustplus/__init__.py b/rustplus/__init__.py index ff6bce3..cbaef7a 100644 --- a/rustplus/__init__.py +++ b/rustplus/__init__.py @@ -2,25 +2,15 @@ RustPlus, An API wrapper for interfacing with the Rust+ App API """ -from .api import RustSocket -from .api.remote.events import ( - EntityEvent, - TeamEvent, - ChatEvent, - MarkerEvent, - ProtobufEvent, - RegisteredListener, -) -from .api.structures import RustMarker, Vector -from .api.remote.fcm_listener import FCMListener -from .api.remote.ratelimiter import RateLimiter -from .api.remote.camera import CameraManager, MovementControls, CameraMovementOptions -from .commands import CommandOptions, Command -from .exceptions import * -from .conversation import ConversationFactory, Conversation, ConversationPrompt -from .utils import * +from .rust_api import RustSocket +from .identification import ServerDetails +from .annotations import Command, ChatEvent, ProtobufEvent, TeamEvent, EntityEvent +from .remote.fcm import FCMListener +from .commands import CommandOptions, ChatCommand +from .events import ChatEventPayload, TeamEventPayload, EntityEventPayload +from .utils import convert_event_type_to_name, Emoji __name__ = "rustplus" __author__ = "olijeffers0n" -__version__ = "5.6.18" +__version__ = "6.0.0" __support__ = "Discord: https://discord.gg/nQqJe8qvP8" diff --git a/rustplus/annotations/__init__.py b/rustplus/annotations/__init__.py new file mode 100644 index 0000000..6d3211f --- /dev/null +++ b/rustplus/annotations/__init__.py @@ -0,0 +1,5 @@ +from .command import Command +from .entity_event import EntityEvent +from .chat_event import ChatEvent +from .team_event import TeamEvent +from .protobuf_event import ProtobufEvent diff --git a/rustplus/annotations/chat_event.py b/rustplus/annotations/chat_event.py new file mode 100644 index 0000000..595398d --- /dev/null +++ b/rustplus/annotations/chat_event.py @@ -0,0 +1,21 @@ +from typing import Callable + +from .. import ServerDetails +from ..identification import RegisteredListener +from ..events import ChatEventPayload as ChatEventManager + + +def ChatEvent(server_details: ServerDetails) -> Callable: + + def wrapper(func) -> RegisteredListener: + + if isinstance(func, RegisteredListener): + func = func.get_coro() + + listener = RegisteredListener(func.__name__, func) + + ChatEventManager.HANDLER_LIST.register(listener, server_details) + + return listener + + return wrapper diff --git a/rustplus/annotations/command.py b/rustplus/annotations/command.py new file mode 100644 index 0000000..24ecb4c --- /dev/null +++ b/rustplus/annotations/command.py @@ -0,0 +1,23 @@ +from typing import Callable + +from ..identification import RegisteredListener, ServerDetails +from ..commands import ChatCommand, ChatCommandData + + +def Command( + server_details: ServerDetails, aliases: list = None, alias_func: Callable = None +) -> Callable: + + def wrapper(func): + + if isinstance(func, RegisteredListener): + func = func.get_coro() + + command_data = ChatCommandData( + coroutine=func, aliases=aliases, callable_func=alias_func + ) + ChatCommand.REGISTERED_COMMANDS[server_details][func.__name__] = command_data + + return RegisteredListener(func.__name__, func) + + return wrapper diff --git a/rustplus/annotations/entity_event.py b/rustplus/annotations/entity_event.py new file mode 100644 index 0000000..7d27de3 --- /dev/null +++ b/rustplus/annotations/entity_event.py @@ -0,0 +1,19 @@ +from typing import Callable + +from .. import ServerDetails +from ..identification import RegisteredListener +from ..events import EntityEventPayload as EntityEventManager + + +def EntityEvent(server_details: ServerDetails, eid: int) -> Callable: + def wrapper(func) -> RegisteredListener: + if isinstance(func, RegisteredListener): + func = func.get_coro() + + listener = RegisteredListener(str(eid), func) + + EntityEventManager.HANDLER_LIST.register(listener, server_details) + + return listener + + return wrapper diff --git a/rustplus/annotations/protobuf_event.py b/rustplus/annotations/protobuf_event.py new file mode 100644 index 0000000..889aaee --- /dev/null +++ b/rustplus/annotations/protobuf_event.py @@ -0,0 +1,19 @@ +from typing import Callable + +from .. import ServerDetails +from ..identification import RegisteredListener +from ..events import ProtobufEventPayload as ProtobufEventManager + + +def ProtobufEvent(server_details: ServerDetails) -> Callable: + def wrapper(func) -> RegisteredListener: + if isinstance(func, RegisteredListener): + func = func.get_coro() + + listener = RegisteredListener(func.__name__, func) + + ProtobufEventManager.HANDLER_LIST.register(listener, server_details) + + return listener + + return wrapper diff --git a/rustplus/annotations/team_event.py b/rustplus/annotations/team_event.py new file mode 100644 index 0000000..c25801b --- /dev/null +++ b/rustplus/annotations/team_event.py @@ -0,0 +1,19 @@ +from typing import Callable + +from .. import ServerDetails +from ..identification import RegisteredListener +from ..events import TeamEventPayload as TeamEventManager + + +def TeamEvent(server_details: ServerDetails) -> Callable: + def wrapper(func) -> RegisteredListener: + if isinstance(func, RegisteredListener): + func = func.get_coro() + + listener = RegisteredListener(func.__name__, func) + + TeamEventManager.HANDLER_LIST.register(listener, server_details) + + return listener + + return wrapper diff --git a/rustplus/api/__init__.py b/rustplus/api/__init__.py deleted file mode 100644 index 2ac9e55..0000000 --- a/rustplus/api/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .rust_api import RustSocket diff --git a/rustplus/api/base_rust_api.py b/rustplus/api/base_rust_api.py deleted file mode 100644 index 17a4622..0000000 --- a/rustplus/api/base_rust_api.py +++ /dev/null @@ -1,637 +0,0 @@ -import asyncio -from typing import List, Union, Coroutine, Callable, Dict, Tuple -from PIL import Image - -from .remote.events.event_loop_manager import EventLoopManager -from .structures import * -from .remote.rustplus_proto import AppEmpty, AppRequest -from .remote import RustRemote, HeartBeat, MapEventListener, ServerChecker, RateLimiter -from .remote.camera import CameraManager -from ..commands import CommandOptions, CommandHandler -from ..commands.command_data import CommandData -from ..exceptions import * -from .remote.events import ( - RegisteredListener, - EntityEvent, - TeamEvent, - ChatEvent, - ProtobufEvent, -) -from ..utils import deprecated -from ..conversation import ConversationFactory -from ..utils import ServerID - - -class BaseRustSocket: - def __init__( - self, - ip: str = None, - port: str = None, - steam_id: int = None, - player_token: int = None, - command_options: CommandOptions = None, - raise_ratelimit_exception: bool = False, - ratelimit_limit: int = 25, - ratelimit_refill: int = 3, - heartbeat: HeartBeat = None, - use_proxy: bool = False, - use_test_server: bool = False, - event_loop: asyncio.AbstractEventLoop = None, - rate_limiter: RateLimiter = None, - debug: bool = False, - ) -> None: - if ip is None: - raise ValueError("Ip cannot be None") - if steam_id is None: - raise ValueError("SteamID cannot be None") - if player_token is None: - raise ValueError("PlayerToken cannot be None") - - try: - steam_id = int(steam_id) - except ValueError: - raise ValueError("SteamID must be an integer") - - try: - player_token = int(player_token) - except ValueError: - raise ValueError("PlayerToken must be an integer") - - self.server_id = ServerID(ip, port, steam_id, player_token) - self.seq = 1 - self.command_options = command_options - self.raise_ratelimit_exception = raise_ratelimit_exception - self.ratelimit_limit = ratelimit_limit - self.ratelimit_refill = ratelimit_refill - self.marker_listener = MapEventListener(self) - self.use_test_server = use_test_server - self.event_loop = event_loop - - self.remote = RustRemote( - server_id=self.server_id, - command_options=command_options, - ratelimit_limit=ratelimit_limit, - ratelimit_refill=ratelimit_refill, - use_proxy=use_proxy, - api=self, - use_test_server=use_test_server, - rate_limiter=rate_limiter, - debug=debug, - ) - - if heartbeat is None: - raise ValueError("Heartbeat cannot be None") - self.heartbeat = heartbeat - - async def _handle_ratelimit(self, amount=1) -> None: - """ - Handles the ratelimit for a specific request. Will sleep if tokens are not currently available and is set to wait - :param amount: The amount to consume - :raises RateLimitError - If the tokens are not available and is not set to wait - :return: None - """ - while True: - if await self.remote.ratelimiter.can_consume(self.server_id, amount): - await self.remote.ratelimiter.consume(self.server_id, amount) - break - - if self.raise_ratelimit_exception: - raise RateLimitError("Out of tokens") - - await asyncio.sleep( - await self.remote.ratelimiter.get_estimated_delay_time( - self.server_id, amount - ) - ) - - self.heartbeat.reset_rhythm() - - def _generate_protobuf(self) -> AppRequest: - """ - Generates the default protobuf for a request - - :return: AppRequest - The default request object - """ - app_request = AppRequest() - app_request.seq = self.seq - app_request.player_id = self.server_id.player_id - app_request.player_token = self.server_id.player_token - - self.seq += 1 - - return app_request - - async def connect( - self, - retries: int = float("inf"), - delay: int = 20, - on_failure: Union[Coroutine, Callable[[], None], None] = None, - on_success: Union[Coroutine, Callable[[], None], None] = None, - on_success_args_kwargs: Tuple[List, Dict] = ([], {}), - on_failure_args_kwargs: Tuple[List, Dict] = ([], {}), - ) -> None: - """ - Attempts to open a connection to the rust game server specified in the constructor - - :param retries: The number of times to attempt reconnecting. Defaults to infinite. - :param delay: The delay (in seconds) between reconnection attempts. - :param on_failure: Optional function to be called when connecting fails. - :param on_success: Optional function to be called when connecting succeeds. - :param on_success_args_kwargs: Optional tuple holding keyword and regular arguments - for on_success in this format (args, kwargs) - :param on_failure_args_kwargs: Optional tuple holding keyword and regular arguments - for on_failure in this format (args, kwargs) - - :return: None - """ - EventLoopManager.set_loop( - ( - self.event_loop - if self.event_loop is not None - else asyncio.get_event_loop() - ), - self.server_id, - ) - - if not self.use_test_server: - ServerChecker(self.server_id.ip, self.server_id.port).run() - - EventLoopManager.set_loop( - ( - self.event_loop - if self.event_loop is not None - else asyncio.get_event_loop() - ), - self.server_id, - ) - - try: - if self.remote.ws is None: - await self.remote.connect( - retries=retries, - delay=delay, - on_failure=on_failure, - on_success=on_success, - on_success_args_kwargs=on_success_args_kwargs, - on_failure_args_kwargs=on_failure_args_kwargs, - ) - await self.heartbeat.start_beat() - except ConnectionRefusedError: - raise ServerNotResponsiveError("Cannot Connect") - - async def close_connection(self) -> None: - """ - Disconnects from the Rust Server - - :return: None - """ - await self.remote.close() - - async def disconnect(self) -> None: - """ - Disconnects from the Rust Server - - :return: None - """ - await self.close_connection() - - async def send_wakeup_request(self) -> None: - """ - Sends a request to the server to wake up broadcast responses - - :return: None - """ - await self._handle_ratelimit() - - app_request = self._generate_protobuf() - app_request.get_time = AppEmpty() - - await self.remote.add_ignored_response(app_request.seq) - - await self.remote.send_message(app_request) - - async def switch_server( - self, - ip: str = None, - port: str = None, - steam_id: int = None, - player_token: int = None, - command_options: CommandOptions = None, - raise_ratelimit_exception: bool = True, - connect: bool = False, - use_proxy: bool = False, - ) -> None: - """ - Disconnects and replaces server params, allowing the socket to connect to a new server. - - :param raise_ratelimit_exception: Whether to raise an exception or wait - :param command_options: The command options - :param ip: IP of the server - :param port: Port of the server - :param player_token: The player Token - :param steam_id: Steam id of the player - :param connect: bool indicating if socket should automatically self.connect() - :param use_proxy: Whether to use the facepunch proxy - :return: None - """ - - if self.use_test_server: - raise ServerSwitchDisallowedError("Cannot switch server") - - if ip is None: - raise ValueError("Ip cannot be None") - if port is None: - raise ValueError("Port cannot be None") - if steam_id is None: - raise ValueError("SteamID cannot be None") - if player_token is None: - raise ValueError("PlayerToken cannot be None") - - # disconnect before redefining - await self.disconnect() - - # Reset basic credentials - self.server_id = ServerID(ip, port, steam_id, player_token) - self.seq = 1 - - # Deal with commands - - if command_options is not None: - self.command_options = command_options - self.remote.command_options = command_options - if self.remote.use_commands: - self.remote.command_handler.command_options = command_options - else: - self.remote.use_commands = True - self.remote.command_handler = CommandHandler(self.command_options, self) - - self.raise_ratelimit_exception = raise_ratelimit_exception - - self.remote.pending_entity_subscriptions = [] - self.remote.server_id = ServerID(ip, port, steam_id, player_token) - - # reset ratelimiter - self.remote.use_proxy = use_proxy - await self.remote.ratelimiter.remove(self.server_id) - self.remote.ratelimiter.add_socket( - self.server_id, - self.ratelimit_limit, - self.ratelimit_limit, - 1, - self.ratelimit_refill, - ) - self.remote.conversation_factory = ConversationFactory(self) - # remove entity events - EntityEvent.handlers.unregister_all() - # reset marker listener - self.marker_listener.persistent_ids.clear() - self.marker_listener.highest_id = 0 - - if connect: - await self.connect() - - def command( - self, - coro: Callable = None, - aliases: List[str] = None, - alias_func: Callable = None, - ) -> Union[Callable, RegisteredListener]: - """ - A coroutine decorator used to register a command executor - - :param alias_func: The function to test the aliases against - :param aliases: The aliases to register the command under - :param coro: The coroutine to call when the command is called - :return: RegisteredListener - The listener object | Callable - The callable func for the decorator - """ - - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - if self.remote.command_handler is None: - raise CommandsNotEnabledError("Not enabled") - - if asyncio.iscoroutinefunction(coro): - cmd_data = CommandData( - coro, - aliases, - alias_func, - ) - self.remote.command_handler.register_command(cmd_data) - return RegisteredListener(coro.__name__, cmd_data.coro) - - def wrap_func(coro): - if self.command_options is None: - raise CommandsNotEnabledError("Not enabled") - - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - cmd_data = CommandData( - coro, - aliases, - alias_func, - ) - self.remote.command_handler.register_command(cmd_data) - return RegisteredListener(coro.__name__, cmd_data.coro) - - return wrap_func - - def team_event(self, coro) -> RegisteredListener: - """ - A Decorator to register an event listener for team changes - - :param coro: The coroutine to call when a change happens - :return: RegisteredListener - The listener object - """ - - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - listener = RegisteredListener("team_changed", coro) - TeamEvent.handlers.register(listener, self.server_id) - return listener - - def chat_event(self, coro) -> RegisteredListener: - """ - A Decorator to register an event listener for chat messages - - :param coro: The coroutine to call when a message is sent - :return: RegisteredListener - The listener object - """ - - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - listener = RegisteredListener("chat_message", coro) - ChatEvent.handlers.register(listener, self.server_id) - return listener - - def entity_event(self, eid): - """ - Decorator to register a smart device listener - - :param eid: The entity id of the entity - :return: RegisteredListener - The listener object - :raises SmartDeviceRegistrationError - """ - - def wrap_func(coro) -> RegisteredListener: - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - self.remote.handle_subscribing_entity(eid, coro) - - return RegisteredListener(eid, coro) - - return wrap_func - - async def start_marker_event_listener(self, delay: int = 5) -> None: - """ - Starts the marker event listener - :param delay: The delay between marker checking - :return: None - """ - self.marker_listener.start(delay) - - def marker_event(self, coro) -> RegisteredListener: - """ - A Decorator to register an event listener for new map markers - - :param coro: The coroutine to call when the command is called - :return: RegisteredListener - The listener object - """ - - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - if not self.marker_listener: - raise ValueError("Marker listener not started") - - listener = RegisteredListener("map_marker", coro) - self.marker_listener.add_listener(listener) - return listener - - def protobuf_received(self, coro) -> RegisteredListener: - """ - A Decorator to register an event listener for protobuf being received on the websocket - - :param coro: The coroutine to call when the command is called - :return: RegisteredListener - The listener object - """ - - if isinstance(coro, RegisteredListener): - coro = coro.get_coro() - - listener = RegisteredListener("protobuf_received", coro) - ProtobufEvent.handlers.register(listener, self.server_id) - return listener - - def remove_listener(self, listener) -> bool: - """ - This will remove a listener, command or event. Takes a RegisteredListener instance - - :return: Success of removal. True = Removed. False = Not Removed - """ - if isinstance(listener, RegisteredListener): - if listener.listener_id == "map_marker": - return self.marker_listener.remove_listener(listener) - - if ChatEvent.handlers.has(listener, self.server_id): - ChatEvent.handlers.unregister(listener, self.server_id) - return True - - if TeamEvent.handlers.has(listener, self.server_id): - TeamEvent.handlers.unregister(listener, self.server_id) - return True - - if EntityEvent.handlers.has(listener, self.server_id): - EntityEvent.handlers.unregister(listener, self.server_id) - return True - - if ProtobufEvent.handlers.has(listener, self.server_id): - ProtobufEvent.handlers.unregister(listener, self.server_id) - return True - - return False - - @staticmethod - async def hang() -> None: - """ - This Will permanently put your script into a state of 'hanging' Cannot be Undone. Only do this in scripts - using commands - - :returns Nothing, This will never return - """ - - while True: - await asyncio.sleep(1) - - @deprecated("Implement this yourself. This will be removed in thed future") - def get_conversation_factory(self) -> ConversationFactory: - """ - Gets the current ConversationFactory object - - :returns ConversationFactory: the factory - """ - return self.remote.conversation_factory - - async def get_time(self) -> RustTime: - """ - Gets the current in-game time from the server. - - :returns RustTime: The Time - """ - raise NotImplementedError("Not Implemented") - - async def send_team_message(self, message: str) -> None: - """ - Sends a message to the in-game team chat - - :param message: The string message to send - """ - raise NotImplementedError("Not Implemented") - - async def get_info(self) -> RustInfo: - """ - Gets information on the Rust Server - :return: RustInfo - The info of the server - """ - raise NotImplementedError("Not Implemented") - - async def get_team_chat(self) -> List[RustChatMessage]: - """ - Gets the team chat from the server - - :return List[RustChatMessage]: The chat messages in the team chat - """ - raise NotImplementedError("Not Implemented") - - async def get_team_info(self) -> RustTeamInfo: - """ - Gets Information on the members of your team - - :return RustTeamInfo: The info of your team - """ - raise NotImplementedError("Not Implemented") - - async def get_markers(self) -> List[RustMarker]: - """ - Gets all the map markers from the server - - :return List[RustMarker]: All the markers on the map - """ - raise NotImplementedError("Not Implemented") - - async def get_map( - self, - add_icons: bool = False, - add_events: bool = False, - add_vending_machines: bool = False, - override_images: dict = None, - add_grid: bool = False, - ) -> Image.Image: - """ - Gets an image of the map from the server with the specified additions - - :param add_icons: To add the monument icons - :param add_events: To add the Event icons - :param add_vending_machines: To add the vending icons - :param override_images: To override the images pre-supplied with RustPlus.py - :param add_grid: To add the grid to the map - :return Image: PIL Image - """ - raise NotImplementedError("Not Implemented") - - async def get_raw_map_data(self) -> RustMap: - """ - Gets the raw map data from the server - - :return RustMap: The raw map of the server - """ - raise NotImplementedError("Not Implemented") - - async def get_entity_info(self, eid: int = None) -> RustEntityInfo: - """ - Gets entity info from the server - - :param eid: The Entities ID - :return RustEntityInfo: The entity Info - """ - raise NotImplementedError("Not Implemented") - - async def turn_on_smart_switch(self, eid: int = None) -> None: - """ - Turns on a given smart switch by entity ID - - :param eid: The Entities ID - :return None: - """ - raise NotImplementedError("Not Implemented") - - async def turn_off_smart_switch(self, eid: int = None) -> None: - """ - Turns off a given smart switch by entity ID - - :param eid: The Entities ID - :return None: - """ - raise NotImplementedError("Not Implemented") - - async def promote_to_team_leader(self, steamid: int = None) -> None: - """ - Promotes a given user to the team leader by their 64-bit Steam ID - - :param steamid: The SteamID of the player to promote - :return None: - """ - raise NotImplementedError("Not Implemented") - - @deprecated("Use RustSocket#get_markers") - async def get_current_events(self) -> List[RustMarker]: - """ - Returns all the map markers that are for events: - Can detect: - - Explosion - - CH47 (Chinook) - - Cargo Ship - - Locked Crate - - Attack Helicopter - - :return List[RustMarker]: All current events - """ - raise NotImplementedError("Not Implemented") - - async def get_contents( - self, eid: int = None, combine_stacks: bool = False - ) -> RustContents: - """ - Gets the contents of a storage monitor-attached container - - :param eid: The EntityID Of the storage Monitor - :param combine_stacks: Whether to combine alike stacks together - :return RustContents: The contents on the monitor - """ - raise NotImplementedError("Not Implemented") - - @deprecated("Use RustSocket#get_contents") - async def get_tc_storage_contents( - self, eid: int = None, combine_stacks: bool = False - ) -> RustContents: - """ - Gets the Information about TC Upkeep and Contents. - Do not use this for any other storage monitor than a TC - """ - raise NotImplementedError("Not Implemented") - - async def get_camera_manager(self, cam_id: str) -> CameraManager: - """ - Gets a camera manager for a given camera ID - - NOTE: This will override the current camera manager if one exists for the given ID so you cannot have multiple - - :param cam_id: The ID of the camera - :return CameraManager: The camera manager - :raises RequestError: If the camera is not found, or you cannot access it. See reason for more info - """ - raise NotImplementedError("Not Implemented") diff --git a/rustplus/api/remote/__init__.py b/rustplus/api/remote/__init__.py deleted file mode 100644 index 3a622b8..0000000 --- a/rustplus/api/remote/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .rustplus_proto import * -from .rustws import RustWebsocket -from .ratelimiter import RateLimiter -from .rust_remote_interface import RustRemote -from .heartbeat import HeartBeat -from .server_checker import ServerChecker -from ..remote.events.event_handler import EventHandler -from ..remote.events.map_event_listener import MapEventListener diff --git a/rustplus/api/remote/events/__init__.py b/rustplus/api/remote/events/__init__.py deleted file mode 100644 index e303ed1..0000000 --- a/rustplus/api/remote/events/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .registered_listener import RegisteredListener -from .events import EntityEvent, TeamEvent, ChatEvent, MarkerEvent, ProtobufEvent -from .event_loop_manager import EventLoopManager -from .event_handler import EventHandler diff --git a/rustplus/api/remote/events/event_handler.py b/rustplus/api/remote/events/event_handler.py deleted file mode 100644 index 6c45f9e..0000000 --- a/rustplus/api/remote/events/event_handler.py +++ /dev/null @@ -1,44 +0,0 @@ -from typing import Set, Union - -from ....utils import ServerID -from .events import EntityEvent, TeamEvent, ChatEvent, ProtobufEvent -from .registered_listener import RegisteredListener -from ..rustplus_proto import AppMessage - - -class EventHandler: - @staticmethod - async def run_entity_event( - name: Union[str, int], app_message: AppMessage, server_id: ServerID - ) -> None: - handlers: Set[RegisteredListener] = EntityEvent.handlers.get_handlers( - server_id - ).get(str(name)) - - if handlers is None: - return - - for handler in handlers.copy(): - await handler.get_coro()( - EntityEvent(app_message, handler.get_entity_type()) - ) - - @staticmethod - async def run_team_event(app_message: AppMessage, server_id: ServerID) -> None: - handlers: Set[RegisteredListener] = TeamEvent.handlers.get_handlers(server_id) - for handler in handlers.copy(): - await handler.get_coro()(TeamEvent(app_message)) - - @staticmethod - async def run_chat_event(app_message: AppMessage, server_id: ServerID) -> None: - handlers: Set[RegisteredListener] = ChatEvent.handlers.get_handlers(server_id) - for handler in handlers.copy(): - await handler.get_coro()(ChatEvent(app_message)) - - @staticmethod - async def run_proto_event(byte_data: bytes, server_id: ServerID) -> None: - handlers: Set[RegisteredListener] = ProtobufEvent.handlers.get_handlers( - server_id - ) - for handler in handlers.copy(): - await handler.get_coro()(ProtobufEvent(byte_data)) diff --git a/rustplus/api/remote/events/event_loop_manager.py b/rustplus/api/remote/events/event_loop_manager.py deleted file mode 100644 index 16392cd..0000000 --- a/rustplus/api/remote/events/event_loop_manager.py +++ /dev/null @@ -1,25 +0,0 @@ -import asyncio -from typing import Dict - -from ....utils import ServerID - - -class EventLoopManager: - _loop: Dict[ServerID, asyncio.AbstractEventLoop] = {} - - @staticmethod - def get_loop(server_id: ServerID) -> asyncio.AbstractEventLoop: - if ( - EventLoopManager._loop is None - or EventLoopManager._loop.get(server_id) is None - ): - raise RuntimeError("Event loop is not set") - - if EventLoopManager._loop.get(server_id).is_closed(): - raise RuntimeError("Event loop is not running") - - return EventLoopManager._loop.get(server_id) - - @staticmethod - def set_loop(loop: asyncio.AbstractEventLoop, server_id: ServerID) -> None: - EventLoopManager._loop[server_id] = loop diff --git a/rustplus/api/remote/events/events.py b/rustplus/api/remote/events/events.py deleted file mode 100644 index 6a6afaa..0000000 --- a/rustplus/api/remote/events/events.py +++ /dev/null @@ -1,126 +0,0 @@ -from typing import List - -from ..rustplus_proto import AppMessage, AppEntityPayloadItem -from ...structures import RustChatMessage -from ...structures.rust_team_info import RustTeamInfo -from ...structures.rust_marker import RustMarker -from .handler_list import HandlerList, EntityHandlerList - - -class Item: - def __init__(self, app_message: AppEntityPayloadItem) -> None: - self._item_id: int = app_message.item_id - self._quantity: int = app_message.quantity - self._item_is_blueprint: bool = app_message.item_is_blueprint - - @property - def item_id(self) -> int: - return self._item_id - - @property - def quantity(self) -> int: - return self._quantity - - @property - def item_is_blueprint(self) -> bool: - return self._item_is_blueprint - - -class TeamEvent: - handlers = HandlerList() - - def __init__(self, app_message: AppMessage) -> None: - self._player_id: int = app_message.broadcast.team_changed.player_id - self._team_info = RustTeamInfo(app_message.broadcast.team_changed.team_info) - - @property - def player_id(self) -> int: - return self._player_id - - @property - def team_info(self) -> RustTeamInfo: - return self._team_info - - -class ChatEvent: - handlers = HandlerList() - - def __init__(self, app_message: AppMessage) -> None: - self._message = RustChatMessage(app_message.broadcast.team_message.message) - - @property - def message(self) -> RustChatMessage: - return self._message - - -class EntityEvent: - handlers = EntityHandlerList() - - def __init__(self, app_message: AppMessage, entity_type) -> None: - self._type = int(entity_type) - self._entity_id: int = app_message.broadcast.entity_changed.entity_id - self._value: bool = app_message.broadcast.entity_changed.payload.value - self._capacity: int = app_message.broadcast.entity_changed.payload.capacity - self._has_protection: bool = ( - app_message.broadcast.entity_changed.payload.has_protection - ) - self._protection_expiry: int = ( - app_message.broadcast.entity_changed.payload.protection_expiry - ) - - self._items: List[Item] = [ - Item(item) for item in app_message.broadcast.entity_changed.payload.items - ] - - @property - def type(self) -> int: - return self._type - - @property - def entity_id(self) -> int: - return self._entity_id - - @property - def value(self) -> bool: - return self._value - - @property - def capacity(self) -> int: - return self._capacity - - @property - def has_protection(self) -> bool: - return self._has_protection - - @property - def protection_expiry(self) -> int: - return self._protection_expiry - - @property - def items(self) -> List[Item]: - return self._items - - -class MarkerEvent: - def __init__(self, marker, is_new) -> None: - self._marker = marker - self._is_new = is_new - - @property - def marker(self) -> RustMarker: - return self._marker - - @property - def is_new(self) -> bool: - return self._is_new - - -class ProtobufEvent: - handlers = HandlerList() - - def __init__(self, byte_data) -> None: - self._byte_data = byte_data - - @property - def byte_data(self) -> bytes: - return self._byte_data diff --git a/rustplus/api/remote/events/handler_list.py b/rustplus/api/remote/events/handler_list.py deleted file mode 100644 index d98c989..0000000 --- a/rustplus/api/remote/events/handler_list.py +++ /dev/null @@ -1,63 +0,0 @@ -from collections import defaultdict -from typing import Set, Dict -from .registered_listener import RegisteredListener -from ....utils import ServerID - - -class HandlerList: - def __init__(self) -> None: - self._handlers: Dict[ServerID, Set[RegisteredListener]] = defaultdict(set) - - def unregister(self, listener: RegisteredListener, server_id: ServerID) -> None: - self._handlers[server_id].remove(listener) - - def register(self, listener: RegisteredListener, server_id: ServerID) -> None: - self._handlers[server_id].add(listener) - - def has(self, listener: RegisteredListener, server_id: ServerID) -> bool: - return listener in self._handlers[server_id] - - def unregister_all(self) -> None: - self._handlers.clear() - - def get_handlers( - self, server_id: ServerID - ) -> Dict[ServerID, Set[RegisteredListener]]: - return self._handlers.get(server_id, set()) - - -class EntityHandlerList(HandlerList): - def __init__(self) -> None: - super().__init__() - self._handlers: Dict[ServerID, Dict[str, Set[RegisteredListener]]] = ( - defaultdict(dict) - ) - - def unregister(self, listener: RegisteredListener, server_id: ServerID) -> None: - if listener.listener_id in self._handlers.get(server_id): - self._handlers.get(server_id).get(listener.listener_id).remove(listener) - - def register(self, listener: RegisteredListener, server_id: ServerID) -> None: - if server_id not in self._handlers: - self._handlers[server_id] = defaultdict(set) - - if listener.listener_id not in self._handlers.get(server_id): - self._handlers.get(server_id)[listener.listener_id] = set() - - self._handlers.get(server_id).get(listener.listener_id).add(listener) - - def has(self, listener: RegisteredListener, server_id: ServerID) -> bool: - if server_id in self._handlers and listener.listener_id in self._handlers.get( - server_id - ): - return listener in self._handlers.get(server_id).get(listener.listener_id) - - return False - - def unregister_all(self) -> None: - self._handlers.clear() - - def get_handlers( - self, server_id: ServerID - ) -> Dict[ServerID, Set[RegisteredListener]]: - return self._handlers.get(server_id, dict()) diff --git a/rustplus/api/remote/events/map_event_listener.py b/rustplus/api/remote/events/map_event_listener.py deleted file mode 100644 index 412ff52..0000000 --- a/rustplus/api/remote/events/map_event_listener.py +++ /dev/null @@ -1,104 +0,0 @@ -import asyncio -import threading -import time - -from .event_loop_manager import EventLoopManager -from .events import MarkerEvent - - -class MapEventListener: - def __init__(self, api) -> None: - self.api = api - self.thread = None - self.gc = None - self._iter_delay = 5 - self.persistent_ids = {} - self.highest_id = 0 - self.listeners = [] - - def add_listener(self, listener) -> None: - self.listeners.append(listener) - - def remove_listener(self, listener) -> bool: - try: - self.listeners.remove(listener) - return True - except ValueError: - return False - - def start(self, delay) -> None: - self._iter_delay = delay - self.thread = threading.Thread( - target=self._run, daemon=True, name="MapEventListener" - ) - self.thread.start() - self.gc = IDGarbageCollector(self.persistent_ids) - self.gc.start() - - def _run(self) -> None: - while True: - try: - future = asyncio.run_coroutine_threadsafe( - self.api.get_markers(), - EventLoopManager.get_loop(self.api.server_id), - ) - new_highest_id = 0 - for marker in future.result(): - new = False - - if marker.id in self.persistent_ids: - self.call_event(marker, new) - continue - - if marker.id > self.highest_id: - new = True - if marker.id > new_highest_id: - new_highest_id = marker.id - - # Removal Times - removal_time = time.time() - - if marker.type == 3 or marker.type == 1: - removal_time = float("inf") - else: - removal_time += 120 * 60 - - self.persistent_ids[marker.id] = removal_time - - self.call_event(marker, new) - - self.highest_id = new_highest_id - - except Exception as e: - print(e) - - time.sleep(self._iter_delay) - - def call_event(self, marker, is_new) -> None: - for listener in self.listeners: - asyncio.run_coroutine_threadsafe( - listener.get_coro()(MarkerEvent(marker, is_new)), - EventLoopManager.get_loop(self.api.server_id), - ).result() - - -class IDGarbageCollector: - def __init__(self, target: dict) -> None: - self.target = target - self.thread = None - - def start(self) -> None: - self.thread = threading.Thread( - target=self._run, daemon=True, name="IDGarbageCollector" - ) - self.thread.start() - - def _run(self) -> None: - while True: - try: - for key in list(self.target.keys()): - if self.target[key] < time.time(): - del self.target[key] - except Exception as e: - print(e) - time.sleep(5) diff --git a/rustplus/api/remote/heartbeat.py b/rustplus/api/remote/heartbeat.py deleted file mode 100644 index bedf410..0000000 --- a/rustplus/api/remote/heartbeat.py +++ /dev/null @@ -1,32 +0,0 @@ -import asyncio -import time - - -class HeartBeat: - def __init__(self, rust_api) -> None: - self.rust_api = rust_api - self.next_run = time.time() - self.running = False - - async def start_beat(self) -> None: - if self.running: - return - - self.running = True - - asyncio.create_task(self._heart_beat()) - - async def _heart_beat(self) -> None: - while True: - if time.time() >= self.next_run: - await self.beat() - - else: - await asyncio.sleep(1) - - async def beat(self) -> None: - if self.rust_api.remote.ws is not None and self.rust_api.remote.is_open(): - await self.rust_api.send_wakeup_request() - - def reset_rhythm(self) -> None: - self.next_run = time.time() + 240 diff --git a/rustplus/api/remote/rplus_version_handler.py b/rustplus/api/remote/rplus_version_handler.py deleted file mode 100644 index ce4a5e8..0000000 --- a/rustplus/api/remote/rplus_version_handler.py +++ /dev/null @@ -1,19 +0,0 @@ -import requests -import logging - - -class MagicValueGrabber: - @staticmethod - def get_magic_value() -> int: - data = requests.get("https://companion-rust.facepunch.com/api/version") - - if data.status_code == 200: - data = data.json() - time = data.get("minPublishedTime", None) - if time is not None: - return time + 1 - - logging.getLogger("rustplus.py").warning( - "[Rustplus.py] Failed to get magic value from RustPlus Server" - ) - return 9999999999999 diff --git a/rustplus/api/remote/rust_remote_interface.py b/rustplus/api/remote/rust_remote_interface.py deleted file mode 100644 index 0686932..0000000 --- a/rustplus/api/remote/rust_remote_interface.py +++ /dev/null @@ -1,256 +0,0 @@ -import asyncio -import logging -from asyncio import Future -from typing import Union, Dict - -import betterproto - -from .camera.camera_manager import CameraManager -from .events import EventLoopManager, EntityEvent, RegisteredListener -from .rustplus_proto import AppRequest, AppMessage, AppEmpty, AppCameraSubscribe -from .rustws import RustWebsocket, CONNECTED, PENDING_CONNECTION -from .ratelimiter import RateLimiter -from .rplus_version_handler import MagicValueGrabber -from ...utils import ServerID, YieldingEvent -from ...conversation import ConversationFactory -from ...commands import CommandHandler -from ...exceptions import ( - ClientNotConnectedError, - RequestError, - SmartDeviceRegistrationError, -) - - -class RustRemote: - def __init__( - self, - server_id: ServerID, - command_options, - ratelimit_limit, - ratelimit_refill, - use_proxy: bool = False, - api=None, - use_test_server: bool = False, - rate_limiter: RateLimiter = None, - debug: bool = False, - ) -> None: - self.server_id = server_id - self.api = api - self.command_options = command_options - self.ratelimit_limit = ratelimit_limit - self.ratelimit_refill = ratelimit_refill - self.use_proxy = use_proxy - if isinstance(rate_limiter, RateLimiter): - self.ratelimiter = rate_limiter - else: - self.ratelimiter = RateLimiter.default() - - self.ratelimiter.add_socket( - self.server_id, ratelimit_limit, ratelimit_limit, 1, ratelimit_refill - ) - self.ws = None - self.logger = logging.getLogger("rustplus.py") - - self.ignored_responses = set() - self.pending_response_events: Dict[int, YieldingEvent] = {} - - self.command_handler = None - if command_options is None: - self.use_commands = False - else: - self.use_commands = True - self.command_handler = CommandHandler(self.command_options, api) - - self.magic_value = MagicValueGrabber.get_magic_value() - self.conversation_factory = ConversationFactory(api) - self.use_test_server = use_test_server - self.pending_entity_subscriptions = [] - self.camera_manager: Union[CameraManager, None] = None - self.debug = debug - - async def connect( - self, - retries, - delay, - on_failure, - on_success, - on_success_args_kwargs, - on_failure_args_kwargs, - ) -> None: - self.ws = RustWebsocket( - server_id=self.server_id, - remote=self, - use_proxy=self.use_proxy, - magic_value=self.magic_value, - use_test_server=self.use_test_server, - on_failure=on_failure, - on_success=on_success, - delay=delay, - on_success_args_kwargs=on_success_args_kwargs, - on_failure_args_kwargs=on_failure_args_kwargs, - debug=self.debug, - ) - await self.ws.connect(retries=retries) - - for entity_id, coroutine in self.pending_entity_subscriptions: - self.handle_subscribing_entity(entity_id, coroutine) - - async def close(self) -> None: - if self.ws is not None: - await self.ws.close() - del self.ws - self.ws = None - - def is_pending(self) -> bool: - if self.ws is not None: - return self.ws.connection_status == PENDING_CONNECTION - return False - - def is_open(self) -> bool: - if self.ws is not None: - return self.ws.connection_status == CONNECTED - return False - - async def send_message(self, request: AppRequest) -> None: - if self.ws is None: - raise ClientNotConnectedError("No Current Websocket Connection") - - if self.debug: - self.logger.info( - f"[RustPlus.py] Sending Message with seq {request.seq}: {request}" - ) - - self.pending_response_events[request.seq] = YieldingEvent() - await self.ws.send_message(request) - - async def get_response( - self, seq: int, app_request: AppRequest, error_check: bool = True - ) -> AppMessage: - """ - Returns a given response from the server. - """ - - attempts = 1 - - while True: - event = self.pending_response_events.get(seq) - if event is None: - raise Exception("Event Doesn't exist") - - response: AppMessage = await event.event_wait_for(4) - if response is not None: - break - - await self.send_message(app_request) - - if attempts % 150 == 0: - self.logger.info( - f"[RustPlus.py] Been waiting 10 minutes for a response for seq {seq}" - ) - - attempts += 1 - - self.pending_response_events.pop(seq) - - if response.response.error.error == "rate_limit": - logging.getLogger("rustplus.py").warning( - "[Rustplus.py] RateLimit Exception Occurred. Retrying after bucket is full" - ) - - # Fully Refill the bucket - bucket = self.ratelimiter.socket_buckets.get(self.server_id) - bucket.current = 0 - - while bucket.current < bucket.max: - await asyncio.sleep(1) - bucket.refresh() - - # Reattempt the sending with a full bucket - cost = self.ws.get_proto_cost(app_request) - - while True: - if await self.ratelimiter.can_consume(self.server_id, cost): - await self.ratelimiter.consume(self.server_id, cost) - break - - await asyncio.sleep( - await self.ratelimiter.get_estimated_delay_time( - self.server_id, cost - ) - ) - - await self.send_message(app_request) - response = await self.get_response(seq, app_request) - - elif self.ws.error_present(response.response.error.error) and error_check: - raise RequestError(response.response.error.error) - - return response - - def handle_subscribing_entity(self, entity_id: int, coroutine) -> None: - if not self.is_open(): - self.pending_entity_subscriptions.append((entity_id, coroutine)) - return - - async def get_entity_info(remote: RustRemote, eid): - await remote.api._handle_ratelimit() - - app_request: AppRequest = remote.api._generate_protobuf() - app_request.entity_id = eid - app_request.get_entity_info = AppEmpty() - - await remote.send_message(app_request) - - return await remote.get_response(app_request.seq, app_request, False) - - def entity_event_callback(future_inner: Future) -> None: - entity_info: AppMessage = future_inner.result() - - if betterproto.serialized_on_wire(entity_info.response.error): - raise SmartDeviceRegistrationError( - f"Entity: '{entity_id}' has not been found" - ) - - EntityEvent.handlers.register( - RegisteredListener( - entity_id, coroutine, entity_info.response.entity_info.type - ), - self.server_id, - ) - - future = asyncio.run_coroutine_threadsafe( - get_entity_info(self, entity_id), EventLoopManager.get_loop(self.server_id) - ) - future.add_done_callback(entity_event_callback) - - async def subscribe_to_camera( - self, entity_id: int, ignore_response: bool = False - ) -> AppRequest: - await self.api._handle_ratelimit() - app_request: AppRequest = self.api._generate_protobuf() - subscribe = AppCameraSubscribe() - subscribe.camera_id = entity_id - app_request.camera_subscribe = subscribe - - await self.send_message(app_request) - - if ignore_response: - await self.add_ignored_response(app_request.seq) - - return app_request - - async def create_camera_manager(self, cam_id) -> CameraManager: - if self.camera_manager is not None: - if self.camera_manager._cam_id == cam_id: - return self.camera_manager - - app_request = await self.subscribe_to_camera(cam_id) - app_message = await self.get_response(app_request.seq, app_request) - - self.camera_manager = CameraManager( - self.api, cam_id, app_message.response.camera_subscribe_info - ) - return self.camera_manager - - async def add_ignored_response(self, seq) -> None: - self.ignored_responses.add(seq) diff --git a/rustplus/api/remote/rustws.py b/rustplus/api/remote/rustws.py deleted file mode 100644 index 1f0bfb6..0000000 --- a/rustplus/api/remote/rustws.py +++ /dev/null @@ -1,394 +0,0 @@ -import asyncio -import base64 -import logging -import time -from datetime import datetime -from typing import Optional, Union, Coroutine -import betterproto -from asyncio import Task, AbstractEventLoop -from websockets.client import connect -from websockets.legacy.client import WebSocketClientProtocol - -from .camera.structures import RayPacket -from .rustplus_proto import AppMessage, AppRequest -from .events import EventHandler -from ..structures import RustChatMessage -from ...exceptions import ClientNotConnectedError -from ...conversation import Conversation -from ...utils import ServerID, YieldingEvent - -CONNECTED = 1 -PENDING_CONNECTION = 2 -CLOSING = 4 -CLOSED = 3 - - -class RustWebsocket: - def __init__( - self, - server_id: ServerID, - remote, - use_proxy, - magic_value, - use_test_server, - on_failure, - on_success, - delay, - on_success_args_kwargs, - on_failure_args_kwargs, - debug: bool = False, - ): - self.connection: Union[WebSocketClientProtocol, None] = None - self.task: Union[Task, None] = None - self.server_id = server_id - self.connection_status = CLOSED - self.use_proxy = use_proxy - self.remote = remote - self.logger = logging.getLogger("rustplus.py") - self.connected_time = time.time() - self.magic_value = magic_value - self.use_test_server = use_test_server - self.outgoing_conversation_messages = [] - self.on_failure = on_failure - self.on_success = on_success - self.delay = delay - self.on_success_args_kwargs = on_success_args_kwargs - self.on_failure_args_kwargs = on_failure_args_kwargs - self.debug = debug - - async def connect( - self, retries=float("inf"), ignore_open_value: bool = False - ) -> None: - if ( - not self.connection_status == CONNECTED or ignore_open_value - ) and not self.remote.is_pending(): - attempts = 0 - - while True: - if attempts >= retries: - raise ConnectionAbortedError("Reached Retry Limit") - - self.connection_status = PENDING_CONNECTION - - try: - address = ( - ( - f"wss://{self.server_id.ip}" - if self.server_id.port is None - else f"ws://{self.server_id.ip}:{self.server_id.port}" - ) - if self.use_test_server - else ( - f"wss://companion-rust.facepunch.com/game/{self.server_id.ip}/{self.server_id.port}" - if self.use_proxy - else f"ws://{self.server_id.ip}:{self.server_id.port}" - ) - ) - address += f"?v={str(self.magic_value)}" - - if self.debug: - self.logger.info(f"[RustPlus.py] Connecting to {address}") - - self.connection = await connect( - address, - close_timeout=0, - ping_interval=None, - max_size=1_000_000_000, - ) - self.connected_time = time.time() - - if self.on_success is not None: - try: - if asyncio.iscoroutinefunction(self.on_success): - await self.on_success( - *self.on_success_args_kwargs[0], - **self.on_success_args_kwargs[1], - ) - else: - self.on_success( - *self.on_success_args_kwargs[0], - **self.on_success_args_kwargs[1], - ) - except Exception as e: - self.logger.warning(e) - break - - except Exception as exception: - self.logger.info(f"[RustPlus.py] {exception}") - - print_error = True - if not isinstance(exception, KeyboardInterrupt): - # Run the failure callback - if self.on_failure is not None: - try: - if asyncio.iscoroutinefunction(self.on_failure): - val = await self.on_failure( - *self.on_failure_args_kwargs[0], - **self.on_failure_args_kwargs[1], - ) - else: - val = self.on_failure( - *self.on_failure_args_kwargs[0], - **self.on_failure_args_kwargs[1], - ) - - if val is not None: - print_error = val - - except Exception as e: - self.logger.warning(e) - - if print_error: - self.logger.warning( - f"{datetime.now().strftime('%d/%m/%Y %H:%M:%S')} " - f"[RustPlus.py] Cannot Connect to server. Retrying in {str(self.delay)} second/s" - ) - attempts += 1 - await asyncio.sleep(self.delay) - - self.connection_status = CONNECTED - - if not ignore_open_value: - self.task = asyncio.create_task( - self.run(), name="[RustPlus.py] Websocket Polling Task" - ) - - async def close(self) -> None: - self.connection_status = CLOSING - await self.connection.close() - self.connection = None - self.task.cancel() - self.task = None - self.connection_status = CLOSED - - if self.debug: - self.logger.info(f"[RustPlus.py] Connection Closed") - - async def send_message(self, message: AppRequest) -> None: - """ - Send the Protobuf to the server - """ - - if self.connection_status == CLOSED: - raise ClientNotConnectedError("Not Connected") - - try: - if self.use_test_server: - await self.connection.send( - base64.b64encode(bytes(message)).decode("utf-8") - ) - else: - await self.connection.send(bytes(message)) - except Exception: - self.logger.exception("An exception occurred whilst sending a message") - - while self.remote.is_pending(): - await asyncio.sleep(0.5) - return await self.send_message(message) - - async def run(self) -> None: - while self.connection_status == CONNECTED: - try: - data = await self.connection.recv() - - # See below for context on why this is needed - await self.run_coroutine_non_blocking( - EventHandler.run_proto_event(data, self.server_id) - ) - - app_message = AppMessage() - app_message.parse( - base64.b64decode(data) if self.use_test_server else data - ) - - except Exception as e: - if self.connection_status == CONNECTED: - print(e) - self.logger.exception( - f"{datetime.now().strftime('%d/%m/%Y %H:%M:%S')} [RustPlus.py] Connection interrupted, Retrying" - ) - await self.connect(ignore_open_value=True) - - continue - return - - try: - # This creates an asyncio task rather than awaiting the coroutine directly. - # This fixes the bug where if you called a BaseRustSocket#get... from within a RegisteredListener or callback, - # It would hang the websocket. This is because the websocket event loop would be stuck on the callback rather than polling the socket. - # This way, we can schedule the execution of all logic for this message, but continue polling the WS - await self.run_coroutine_non_blocking(self.handle_message(app_message)) - except Exception: - self.logger.exception( - "An Error occurred whilst handling the message from the server" - ) - - async def handle_message(self, app_message: AppMessage) -> None: - if self.debug: - self.logger.info( - f"[RustPlus.py] Received Message with seq {app_message.response.seq}: {app_message}" - ) - - if app_message.response.seq in self.remote.ignored_responses: - self.remote.ignored_responses.remove(app_message.response.seq) - return - - prefix = self.get_prefix( - str(app_message.broadcast.team_message.message.message) - ) - - if prefix is not None: - # This means it is a command - - if self.debug: - self.logger.info( - f"[RustPlus.py] Attempting to run Command: {app_message}" - ) - - message = RustChatMessage(app_message.broadcast.team_message.message) - await self.remote.command_handler.run_command(message, prefix) - - if self.is_entity_broadcast(app_message): - # This means that an entity has changed state - - if self.debug: - self.logger.info(f"[RustPlus.py] Running Entity Event: {app_message}") - - await EventHandler.run_entity_event( - app_message.broadcast.entity_changed.entity_id, - app_message, - self.server_id, - ) - - elif self.is_camera_broadcast(app_message): - if self.debug: - self.logger.info(f"[RustPlus.py] Running Camera Event: {app_message}") - - if self.remote.camera_manager is not None: - await self.remote.camera_manager.add_packet( - RayPacket(app_message.broadcast.camera_rays) - ) - - elif self.is_team_broadcast(app_message): - if self.debug: - self.logger.info(f"[RustPlus.py] Running Team Event: {app_message}") - - # This means that the team of the current player has changed - await EventHandler.run_team_event(app_message, self.server_id) - - elif self.is_message(app_message): - # This means that a message has been sent to the team chat - - if self.debug: - self.logger.info(f"[RustPlus.py] Running Chat Event: {app_message}") - - steam_id = int(app_message.broadcast.team_message.message.steam_id) - message = str(app_message.broadcast.team_message.message.message) - - # Conversation API - if self.remote.conversation_factory.has_conversation(steam_id): - if message not in self.outgoing_conversation_messages: - conversation: Conversation = ( - self.remote.conversation_factory.get_conversation(steam_id) - ) - - conversation.get_answers().append(message) - await conversation.get_current_prompt().on_response(message) - - if conversation.has_next(): - conversation.increment_prompt() - prompt = conversation.get_current_prompt() - prompt_string = await prompt.prompt() - await conversation.send_prompt(prompt_string) - - else: - prompt = conversation.get_current_prompt() - prompt_string = await prompt.on_finish() - if prompt_string != "": - await conversation.send_prompt(prompt_string) - self.remote.conversation_factory.abort_conversation(steam_id) - else: - self.outgoing_conversation_messages.remove(message) - - # Conversation API end - - await EventHandler.run_chat_event(app_message, self.server_id) - - else: - # This means that it wasn't sent by the server and is a message from the server in response to an action - event: YieldingEvent = self.remote.pending_response_events.get( - app_message.response.seq, None - ) - if event is not None: - if self.debug: - self.logger.info( - f"[RustPlus.py] Running Response Event: {app_message}" - ) - - event.set_with_value(app_message) - - def get_prefix(self, message: str) -> Optional[str]: - if self.remote.use_commands: - if message.startswith(self.remote.command_options.prefix): - return self.remote.command_options.prefix - else: - return None - - for overrule in self.remote.command_options.overruling_commands: - if message.startswith(overrule): - return overrule - - return None - - @staticmethod - def is_message(app_message: AppMessage) -> bool: - return betterproto.serialized_on_wire( - app_message.broadcast.team_message.message - ) - - @staticmethod - def is_camera_broadcast(app_message: AppMessage) -> bool: - return betterproto.serialized_on_wire(app_message.broadcast.camera_rays) - - @staticmethod - def is_entity_broadcast(app_message: AppMessage) -> bool: - return betterproto.serialized_on_wire(app_message.broadcast.entity_changed) - - @staticmethod - def is_team_broadcast(app_message: AppMessage) -> bool: - return betterproto.serialized_on_wire(app_message.broadcast.team_changed) - - @staticmethod - def get_proto_cost(app_request: AppRequest) -> int: - """ - Gets the cost of an AppRequest - """ - costs = [ - (app_request.get_time, 1), - (app_request.send_team_message, 2), - (app_request.get_info, 1), - (app_request.get_team_chat, 1), - (app_request.get_team_info, 1), - (app_request.get_map_markers, 1), - (app_request.get_map, 5), - (app_request.set_entity_value, 1), - (app_request.get_entity_info, 1), - (app_request.promote_to_leader, 1), - ] - for request, cost in costs: - if betterproto.serialized_on_wire(request): - return cost - - raise ValueError() - - @staticmethod - def error_present(message) -> bool: - """ - Checks message for error - """ - return message != "" - - @staticmethod - async def run_coroutine_non_blocking(coroutine: Coroutine) -> Task: - loop: AbstractEventLoop = asyncio.get_event_loop_policy().get_event_loop() - return loop.create_task(coroutine) diff --git a/rustplus/api/remote/server_checker.py b/rustplus/api/remote/server_checker.py deleted file mode 100644 index c4c4aa0..0000000 --- a/rustplus/api/remote/server_checker.py +++ /dev/null @@ -1,26 +0,0 @@ -import logging -import threading -import requests - - -class ServerChecker: - def __init__(self, ip: str, port: str) -> None: - self.ip = ip - self.port = port - self.logger = logging.getLogger("rustplus.py") - - def run(self) -> None: - threading.Thread(target=self._check_server, daemon=True).start() - - def _check_server(self) -> None: - try: - req = requests.post( - f"https://companion-rust.facepunch.com/api/server/test_connection?address={self.ip}&port={self.port}" - ) - for msg in req.json()["messages"]: - if "does not match your outgoing IP address" not in msg: - self.logger.warning(f"Error from server Checker: {msg}") - except Exception: - self.logger.exception( - f"Unable to test connection to server - {self.ip}:{self.port}" - ) diff --git a/rustplus/api/rust_api.py b/rustplus/api/rust_api.py deleted file mode 100644 index a99cd01..0000000 --- a/rustplus/api/rust_api.py +++ /dev/null @@ -1,428 +0,0 @@ -import asyncio -import requests -from typing import List, Union -from PIL import Image -from io import BytesIO -from datetime import datetime -from collections import defaultdict -from importlib import resources - -from .base_rust_api import BaseRustSocket -from .remote.camera.camera_manager import CameraManager -from .structures import ( - RustInfo, - RustMap, - RustMarker, - RustChatMessage, - RustTeamInfo, - RustEntityInfo, - RustContents, - RustItem, -) -from .remote.rustplus_proto import ( - AppEmpty, - AppSendMessage, - AppSetEntityValue, - AppPromoteToLeader, -) -from .remote import HeartBeat, RateLimiter -from ..commands import CommandOptions -from ..exceptions import * -from ..utils import ( - RustTime, - format_time, - format_coord, - convert_marker, - convert_monument, - translate_id_to_stack, - deprecated, - generate_grid, - avatar_processing, -) - - -class RustSocket(BaseRustSocket): - def __init__( - self, - ip: str = None, - port: str = None, - steam_id: int = None, - player_token: int = None, - command_options: CommandOptions = None, - raise_ratelimit_exception: bool = False, - ratelimit_limit: int = 25, - ratelimit_refill: int = 3, - use_proxy: bool = False, - use_test_server: bool = False, - event_loop: asyncio.AbstractEventLoop = None, - rate_limiter: RateLimiter = None, - debug: bool = False, - ) -> None: - super().__init__( - ip=ip, - port=port, - steam_id=steam_id, - player_token=player_token, - command_options=command_options, - raise_ratelimit_exception=raise_ratelimit_exception, - ratelimit_limit=ratelimit_limit, - ratelimit_refill=ratelimit_refill, - heartbeat=HeartBeat(self), - use_proxy=use_proxy, - use_test_server=use_test_server, - event_loop=event_loop, - rate_limiter=rate_limiter, - debug=debug, - ) - - async def get_time(self) -> RustTime: - await self._handle_ratelimit() - - app_request = self._generate_protobuf() - app_request.get_time = AppEmpty() - - await self.remote.send_message(app_request) - - response = await self.remote.get_response(app_request.seq, app_request) - - return format_time(response) - - async def send_team_message(self, message: Union[str, object]) -> None: - await self._handle_ratelimit(2) - - app_send_message = AppSendMessage() - app_send_message.message = str(message) - - app_request = self._generate_protobuf() - app_request.send_team_message = app_send_message - - await self.remote.add_ignored_response(app_request.seq) - - await self.remote.send_message(app_request) - - async def get_info(self) -> RustInfo: - await self._handle_ratelimit() - - app_request = self._generate_protobuf() - app_request.get_info = AppEmpty() - - await self.remote.send_message(app_request) - - response = await self.remote.get_response(app_request.seq, app_request) - - return RustInfo(response.response.info) - - async def get_team_chat(self) -> List[RustChatMessage]: - await self._handle_ratelimit() - - app_request = self._generate_protobuf() - app_request.get_team_chat = AppEmpty() - - await self.remote.send_message(app_request) - - messages = ( - await self.remote.get_response(app_request.seq, app_request) - ).response.team_chat.messages - - return [RustChatMessage(message) for message in messages] - - async def get_team_info(self) -> RustTeamInfo: - await self._handle_ratelimit() - - app_request = self._generate_protobuf() - app_request.get_team_info = AppEmpty() - - await self.remote.send_message(app_request) - - app_message = await self.remote.get_response(app_request.seq, app_request) - - return RustTeamInfo(app_message.response.team_info) - - async def get_markers(self) -> List[RustMarker]: - await self._handle_ratelimit() - - app_request = self._generate_protobuf() - app_request.get_map_markers = AppEmpty() - - await self.remote.send_message(app_request) - - app_message = await self.remote.get_response(app_request.seq, app_request) - - return [ - RustMarker(marker) for marker in app_message.response.map_markers.markers - ] - - async def get_raw_map_data(self) -> RustMap: - await self._handle_ratelimit(5) - - app_request = self._generate_protobuf() - app_request.get_map = AppEmpty() - - await self.remote.send_message(app_request) - - app_message = await self.remote.get_response(app_request.seq, app_request) - - return RustMap(app_message.response.map) - - async def get_map( - self, - add_icons: bool = False, - add_events: bool = False, - add_vending_machines: bool = False, - add_team_positions: bool = False, - override_images: dict = None, - add_grid: bool = False, - ) -> Image.Image: - if override_images is None: - override_images = {} - - map_size = int((await self.get_info()).size) - - await self._handle_ratelimit( - 5 - + ( - 1 - if [add_icons, add_events, add_vending_machines].count(True) >= 1 - else 0 + 1 if add_team_positions else 0 - ) - ) - - app_request = self._generate_protobuf() - app_request.get_map = AppEmpty() - - await self.remote.send_message(app_request) - - app_message = await self.remote.get_response(app_request.seq, app_request) - - game_map = app_message.response.map - monuments = list(game_map.monuments) - - try: - image = Image.open(BytesIO(game_map.jpg_image)) - except Exception: - raise ImageError("Invalid bytes for the image") - - if not self.use_test_server: - image = image.crop((500, 500, game_map.height - 500, game_map.width - 500)) - - game_map = image.resize((map_size, map_size), Image.LANCZOS).convert("RGBA") - - if add_grid: - grid = generate_grid(map_size) - - game_map.paste(grid, (5, 5), grid) - - if add_icons or add_events or add_vending_machines: - map_markers = ( - await self.get_markers() if add_events or add_vending_machines else [] - ) - - if add_icons: - for monument in monuments: - if str(monument.token) == "DungeonBase": - continue - icon = convert_monument(monument.token, override_images) - if monument.token in override_images: - icon = icon.resize((150, 150)) - if str(monument.token) == "train_tunnel_display_name": - icon = icon.resize((100, 125)) - game_map.paste( - icon, - (format_coord(int(monument.x), int(monument.y), map_size)), - icon, - ) - - if add_vending_machines: - with resources.path( - "rustplus.api.icons", "vending_machine.png" - ) as path: - vending_machine = Image.open(path).convert("RGBA") - vending_machine = vending_machine.resize((100, 100)) - - for marker in map_markers: - if add_events: - if ( - marker.type == 2 - or marker.type == 4 - or marker.type == 5 - or marker.type == 6 - or marker.type == 8 - ): - icon = convert_marker(str(marker.type), marker.rotation) - if marker.type == 6: - x = marker.x - y = marker.y - if y > map_size: - y = map_size - if y < 0: - y = 100 - if x > map_size: - x = map_size - 75 - if x < 0: - x = 50 - game_map.paste(icon, (int(x), map_size - int(y)), icon) - else: - game_map.paste( - icon, - (format_coord(int(marker.x), int(marker.y), map_size)), - icon, - ) - if add_vending_machines and marker.type == 3: - game_map.paste( - vending_machine, - (int(marker.x) - 50, map_size - int(marker.y) - 50), - vending_machine, - ) - if add_team_positions: - team = await self.get_team_info() - for member in team.members: - if member.is_alive: - avatar = ( - Image.open( - requests.get( - f"https://companion-rust.facepunch.com/api/avatar/{member.steam_id}", - stream=True, - ).raw - ) - .resize((100, 100), Image.LANCZOS) - .convert("RGBA") - ) - - player_avatar = avatar_processing(avatar, 5, member.is_online) - - game_map.paste( - player_avatar, - (format_coord(int(member.x), int(member.y), map_size)), - player_avatar, - ) - - return game_map.resize((2000, 2000), Image.LANCZOS) - - async def get_entity_info(self, eid: int = None) -> RustEntityInfo: - await self._handle_ratelimit() - - if eid is None: - raise ValueError("EID cannot be None") - - app_request = self._generate_protobuf() - app_request.entity_id = eid - app_request.get_entity_info = AppEmpty() - - await self.remote.send_message(app_request) - - app_message = await self.remote.get_response(app_request.seq, app_request) - - return RustEntityInfo(app_message.response.entity_info) - - async def _update_smart_device(self, eid: int, value: bool) -> None: - await self._handle_ratelimit() - - entity_value = AppSetEntityValue() - entity_value.value = value - - app_request = self._generate_protobuf() - - app_request.entity_id = eid - app_request.set_entity_value = entity_value - - await self.remote.add_ignored_response(app_request.seq) - - await self.remote.send_message(app_request) - - async def turn_on_smart_switch(self, eid: int = None) -> None: - if eid is None: - raise ValueError("EID cannot be None") - - await self._update_smart_device(eid, True) - - async def turn_off_smart_switch(self, eid: int = None) -> None: - if eid is None: - raise ValueError("EID cannot be None") - - await self._update_smart_device(eid, False) - - async def promote_to_team_leader(self, steam_id: int = None) -> None: - if steam_id is None: - raise ValueError("SteamID cannot be None") - - await self._handle_ratelimit() - - leader_packet = AppPromoteToLeader() - leader_packet.steam_id = steam_id - - app_request = self._generate_protobuf() - app_request.promote_to_leader = leader_packet - - await self.remote.add_ignored_response(app_request.seq) - - await self.remote.send_message(app_request) - - async def get_current_events(self) -> List[RustMarker]: - return [ - marker - for marker in (await self.get_markers()) - if marker.type == 2 - or marker.type == 4 - or marker.type == 5 - or marker.type == 6 - or marker.type == 8 - ] - - async def get_contents( - self, eid: int = None, combine_stacks: bool = False - ) -> RustContents: - if eid is None: - raise ValueError("EID cannot be None") - - returned_data = await self.get_entity_info(eid) - - target_time = datetime.utcfromtimestamp(int(returned_data.protection_expiry)) - difference = target_time - datetime.utcnow() - - items = [] - - for item in returned_data.items: - items.append( - RustItem( - translate_id_to_stack(item.item_id), - item.item_id, - item.quantity, - item.item_is_blueprint, - ) - ) - - if combine_stacks: - merged_map = defaultdict(tuple) - - for item in items: - data = merged_map[str(item.item_id)] - if data: - count = int(data[0]) + int(item.quantity) - merged_map[str(item.item_id)] = (count, bool(item.is_blueprint)) - else: - merged_map[str(item.item_id)] = ( - int(item.quantity), - bool(item.is_blueprint), - ) - - items = [] - for key in merged_map.keys(): - items.append( - RustItem( - translate_id_to_stack(key), - key, - int(merged_map[key][0]), - bool(merged_map[key][1]), - ) - ) - - return RustContents(difference, bool(returned_data.has_protection), items) - - @deprecated("Use RustSocket.get_contents") - async def get_tc_storage_contents( - self, eid: int = None, combine_stacks: bool = False - ) -> RustContents: - return await self.get_contents(eid=eid, combine_stacks=combine_stacks) - - async def get_camera_manager(self, cam_id: str) -> CameraManager: - return await self.remote.create_camera_manager(cam_id) diff --git a/rustplus/commands/__init__.py b/rustplus/commands/__init__.py index 3cc2bff..b45ef7c 100644 --- a/rustplus/commands/__init__.py +++ b/rustplus/commands/__init__.py @@ -1,3 +1,3 @@ +from .chat_command_data import ChatCommandData +from .chat_command import ChatCommand, ChatCommandTime from .command_options import CommandOptions -from .command import Command, CommandTime -from .command_handler import CommandHandler diff --git a/rustplus/commands/chat_command.py b/rustplus/commands/chat_command.py new file mode 100644 index 0000000..c7cbce7 --- /dev/null +++ b/rustplus/commands/chat_command.py @@ -0,0 +1,33 @@ +import dataclasses +from collections import defaultdict +from typing import Dict, List + +from .chat_command_data import ChatCommandData +from ..identification import ServerDetails + + +@dataclasses.dataclass +class ChatCommandTime: + formatted_time: str + raw_time: int + + +class ChatCommand: + + REGISTERED_COMMANDS: Dict[ServerDetails, Dict[str, ChatCommandData]] = defaultdict( + dict + ) + + def __init__( + self, + sender_name: str, + sender_steam_id: int, + time: ChatCommandTime, + command: str, + args: List[str], + ) -> None: + self.sender_name = sender_name + self.sender_steam_id = sender_steam_id + self.time = time + self.command = command + self.args = args diff --git a/rustplus/commands/command_data.py b/rustplus/commands/chat_command_data.py similarity index 57% rename from rustplus/commands/command_data.py rename to rustplus/commands/chat_command_data.py index e9158c4..eaddd04 100644 --- a/rustplus/commands/command_data.py +++ b/rustplus/commands/chat_command_data.py @@ -1,21 +1,22 @@ -from typing import Callable, List +from typing import Callable -class CommandData: - def __init__(self, coro, aliases, callable_func) -> None: - self.coro = coro +class ChatCommandData: + + def __init__(self, coroutine: Callable, aliases=None, callable_func=None) -> None: + self.coroutine = coroutine self._aliases = aliases self._callable_func = callable_func @property - def aliases(self) -> List[str]: + def aliases(self): if self._aliases is None: return [] return self._aliases @property - def callable_func(self) -> Callable: + def callable_func(self): if self._callable_func is None: return lambda x: False diff --git a/rustplus/commands/command.py b/rustplus/commands/command.py deleted file mode 100644 index 2b7d3a6..0000000 --- a/rustplus/commands/command.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import List - - -class CommandTime: - def __init__(self, formatted_time, raw_time) -> None: - self.formatted_time = formatted_time - self.raw_time = raw_time - - -class Command: - def __init__( - self, - sender_name: str, - sender_steam_id: int, - time: CommandTime, - command: str, - args: List[str], - ) -> None: - self.sender_name = sender_name - self.sender_steam_id = sender_steam_id - self.time = time - self.command = command - self.args = args diff --git a/rustplus/commands/command_handler.py b/rustplus/commands/command_handler.py deleted file mode 100644 index 47d5107..0000000 --- a/rustplus/commands/command_handler.py +++ /dev/null @@ -1,68 +0,0 @@ -import asyncio -import shlex -from datetime import datetime - -from . import Command, CommandTime -from ..api.structures import RustChatMessage -from ..commands.command_options import CommandOptions -from ..commands.command_data import CommandData -from ..api.remote.events import RegisteredListener - - -class CommandHandler: - def __init__(self, command_options: CommandOptions, api) -> None: - self.command_options = command_options - self.commands = {} - self.api = api - - def register_command(self, data: CommandData) -> None: - if not asyncio.iscoroutinefunction(data.coro): - raise TypeError("The event registered must be a coroutine") - - self.commands[data.coro.__name__] = data - - async def run_command(self, message: RustChatMessage, prefix) -> None: - if prefix == self.command_options.prefix: - command = shlex.split(message.message)[0][len(prefix) :] - else: - command = prefix - - if command in self.commands: - data = self.commands[command] - - await data.coro( - Command( - message.name, - message.steam_id, - CommandTime(datetime.utcfromtimestamp(message.time), message.time), - command, - shlex.split(message.message)[1:], - ) - ) - else: - for command_name, data in self.commands.items(): - # Loop through all the commands and see if the command is in the data aliases list - # or if it matches the callable function - - if command in data.aliases or data.callable_func(command): - await data.coro( - Command( - message.name, - message.steam_id, - CommandTime( - datetime.utcfromtimestamp(message.time), message.time - ), - command, - shlex.split(message.message)[1:], - ), - ) - break - - def has_command(self, listener: RegisteredListener) -> bool: - return listener.listener_id in self.commands - - def remove_command(self, listener: RegisteredListener) -> None: - try: - del self.commands[listener.listener_id] - except KeyError: - pass diff --git a/rustplus/commands/command_options.py b/rustplus/commands/command_options.py index c12d463..232fd23 100644 --- a/rustplus/commands/command_options.py +++ b/rustplus/commands/command_options.py @@ -1,17 +1,9 @@ -from typing import List - from ..exceptions import PrefixNotDefinedError class CommandOptions: - def __init__( - self, prefix: str = None, overruling_commands: List[str] = None - ) -> None: + def __init__(self, prefix: str = None) -> None: if prefix is None: raise PrefixNotDefinedError("No prefix") - if overruling_commands is None: - overruling_commands = [] - self.prefix = prefix - self.overruling_commands = overruling_commands diff --git a/rustplus/conversation/__init__.py b/rustplus/conversation/__init__.py deleted file mode 100644 index fb32bed..0000000 --- a/rustplus/conversation/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .conversation import Conversation -from .conversation_prompt import ConversationPrompt -from .conversation_factory import ConversationFactory diff --git a/rustplus/conversation/conversation.py b/rustplus/conversation/conversation.py deleted file mode 100644 index 87042cb..0000000 --- a/rustplus/conversation/conversation.py +++ /dev/null @@ -1,62 +0,0 @@ -import asyncio -from typing import List, Any -from .conversation_prompt import ConversationPrompt -from ..api.remote.events import EventLoopManager - - -class Conversation: - def __init__( - self, - api, - target: int = None, - prompts: List[ConversationPrompt] = None, - register=None, - ) -> None: - if target is None: - raise ValueError("target must be specified") - self._target = target - - if prompts is None: - self._prompts = [] - else: - self._prompts = prompts - - self._answers = [] - self._seq = 0 - self._api = api - self._register = register - - def add_prompt(self, prompt: ConversationPrompt) -> None: - super(type(prompt), prompt).__init__(self) - self._prompts.append(prompt) - - def add_all_prompts(self, prompts: List[ConversationPrompt]) -> None: - for prompt in prompts: - super(ConversationPrompt, prompt).__init__(self) - self.add_prompt(prompt) - - def has_next(self) -> bool: - return self._seq + 1 < len(self._prompts) - - def get_current_prompt(self) -> ConversationPrompt: - return self._prompts[self._seq] - - def increment_prompt(self) -> None: - self._seq += 1 - - async def send_prompt(self, message: str) -> None: - if self._api.remote.ws is not None: - self._api.remote.ws.outgoing_conversation_messages.append(message) - await self._api.send_team_message(message) - - async def start(self) -> None: - self._register(self._target, self) - await self.send_prompt(await self._prompts[0].prompt()) - - def run_coro(self, coro, args) -> Any: - return asyncio.run_coroutine_threadsafe( - coro(*args), EventLoopManager.get_loop(self._api.server_id) - ).result() - - def get_answers(self) -> List[str]: - return self._answers diff --git a/rustplus/conversation/conversation_factory.py b/rustplus/conversation/conversation_factory.py deleted file mode 100644 index 33b3b85..0000000 --- a/rustplus/conversation/conversation_factory.py +++ /dev/null @@ -1,50 +0,0 @@ -import time -from threading import Thread - -from .conversation import Conversation - - -class ConversationFactory: - def __init__(self, api) -> None: - self.api = api - self.conversations = {} - self.expires = {} - self.gc_thread = Thread(target=self.garbage_collect, daemon=True) - self.gc_thread.start() - - def create_conversation(self, steamid: int) -> Conversation: - if steamid in self.conversations: - raise ValueError("Conversation already exists") - - return Conversation( - api=self.api, target=steamid, register=self._register_conversation - ) - - def _register_conversation(self, steamid, convo: Conversation) -> None: - self.conversations[steamid] = convo - self.expires[steamid] = time.time() + 60 * 5 - - def has_conversation(self, steamid: int) -> bool: - return steamid in self.conversations - - def get_conversation(self, steamid: int) -> Conversation: - return self.conversations[steamid] - - def abort_conversation(self, steamid: int) -> None: - try: - del self.conversations[steamid] - del self.expires[steamid] - except KeyError: - pass - - def garbage_collect(self) -> None: - while True: - to_remove = [] - for steamid, expire_time in self.expires.items(): - if expire_time < time.time(): - to_remove.append(steamid) - - for steamid in to_remove: - self.abort_conversation(steamid) - - time.sleep(5) diff --git a/rustplus/conversation/conversation_prompt.py b/rustplus/conversation/conversation_prompt.py deleted file mode 100644 index dc73678..0000000 --- a/rustplus/conversation/conversation_prompt.py +++ /dev/null @@ -1,12 +0,0 @@ -class ConversationPrompt: - def __init__(self, conversation) -> None: - self.conversation = conversation - - async def prompt(self) -> str: - return "" - - async def on_response(self, response) -> None: - pass - - async def on_finish(self) -> str: - return "" diff --git a/rustplus/events/__init__.py b/rustplus/events/__init__.py new file mode 100644 index 0000000..954af77 --- /dev/null +++ b/rustplus/events/__init__.py @@ -0,0 +1,4 @@ +from .chat_event import ChatEventPayload +from .entity_event import EntityEventPayload +from .team_event import TeamEventPayload +from .protobuf_event import ProtobufEventPayload diff --git a/rustplus/events/chat_event.py b/rustplus/events/chat_event.py new file mode 100644 index 0000000..d061f93 --- /dev/null +++ b/rustplus/events/chat_event.py @@ -0,0 +1,13 @@ +from rustplus.identification.handler_list import HandlerList +from ..structs import RustChatMessage + + +class ChatEventPayload: + HANDLER_LIST = HandlerList() + + def __init__(self, message: RustChatMessage) -> None: + self._message = message + + @property + def message(self) -> RustChatMessage: + return self._message diff --git a/rustplus/events/entity_event.py b/rustplus/events/entity_event.py new file mode 100644 index 0000000..b7e434a --- /dev/null +++ b/rustplus/events/entity_event.py @@ -0,0 +1,60 @@ +from typing import List + +from rustplus.identification.handler_list import EntityHandlerList +from ..remote.rustplus_proto import AppEntityPayloadItem, AppEntityChanged + + +class Item: + def __init__(self, app_message: AppEntityPayloadItem) -> None: + self._item_id: int = app_message.item_id + self._quantity: int = app_message.quantity + self._item_is_blueprint: bool = app_message.item_is_blueprint + + @property + def item_id(self) -> int: + return self._item_id + + @property + def quantity(self) -> int: + return self._quantity + + @property + def item_is_blueprint(self) -> bool: + return self._item_is_blueprint + + +class EntityEventPayload: + HANDLER_LIST = EntityHandlerList() + + def __init__(self, entity_changed: AppEntityChanged) -> None: + self._entity_id: int = entity_changed.entity_id + self._value: bool = entity_changed.payload.value + self._capacity: int = entity_changed.payload.capacity + self._has_protection: bool = entity_changed.payload.has_protection + self._protection_expiry: int = entity_changed.payload.protection_expiry + + self._items: List[Item] = [Item(item) for item in entity_changed.payload.items] + + @property + def entity_id(self) -> int: + return self._entity_id + + @property + def value(self) -> bool: + return self._value + + @property + def capacity(self) -> int: + return self._capacity + + @property + def has_protection(self) -> bool: + return self._has_protection + + @property + def protection_expiry(self) -> int: + return self._protection_expiry + + @property + def items(self) -> List[Item]: + return self._items diff --git a/rustplus/events/protobuf_event.py b/rustplus/events/protobuf_event.py new file mode 100644 index 0000000..06c7a68 --- /dev/null +++ b/rustplus/events/protobuf_event.py @@ -0,0 +1,12 @@ +from rustplus.identification.handler_list import HandlerList + + +class ProtobufEventPayload: + HANDLER_LIST = HandlerList() + + def __init__(self, message: bytes) -> None: + self._message = message + + @property + def message(self) -> bytes: + return self._message diff --git a/rustplus/events/team_event.py b/rustplus/events/team_event.py new file mode 100644 index 0000000..ef2cfc5 --- /dev/null +++ b/rustplus/events/team_event.py @@ -0,0 +1,18 @@ +from rustplus.identification.handler_list import HandlerList +from ..structs import RustTeamInfo + + +class TeamEventPayload: + HANDLER_LIST = HandlerList() + + def __init__(self, player_id: int, team_info: RustTeamInfo) -> None: + self._player_id = player_id + self._team_info = team_info + + @property + def player_id(self) -> int: + return self._player_id + + @property + def team_info(self) -> RustTeamInfo: + return self._team_info diff --git a/rustplus/api/icons/__init__.py b/rustplus/icons/__init__.py similarity index 100% rename from rustplus/api/icons/__init__.py rename to rustplus/icons/__init__.py diff --git a/rustplus/api/icons/airfield.png b/rustplus/icons/airfield.png similarity index 100% rename from rustplus/api/icons/airfield.png rename to rustplus/icons/airfield.png diff --git a/rustplus/api/icons/arctic_base.png b/rustplus/icons/arctic_base.png similarity index 100% rename from rustplus/api/icons/arctic_base.png rename to rustplus/icons/arctic_base.png diff --git a/rustplus/api/icons/bandit.png b/rustplus/icons/bandit.png similarity index 100% rename from rustplus/api/icons/bandit.png rename to rustplus/icons/bandit.png diff --git a/rustplus/api/icons/barn.png b/rustplus/icons/barn.png similarity index 100% rename from rustplus/api/icons/barn.png rename to rustplus/icons/barn.png diff --git a/rustplus/api/icons/cargo.png b/rustplus/icons/cargo.png similarity index 100% rename from rustplus/api/icons/cargo.png rename to rustplus/icons/cargo.png diff --git a/rustplus/api/icons/chinook.png b/rustplus/icons/chinook.png similarity index 100% rename from rustplus/api/icons/chinook.png rename to rustplus/icons/chinook.png diff --git a/rustplus/api/icons/chinook_blades.png b/rustplus/icons/chinook_blades.png similarity index 100% rename from rustplus/api/icons/chinook_blades.png rename to rustplus/icons/chinook_blades.png diff --git a/rustplus/api/icons/crate.png b/rustplus/icons/crate.png similarity index 100% rename from rustplus/api/icons/crate.png rename to rustplus/icons/crate.png diff --git a/rustplus/api/icons/desert_base.png b/rustplus/icons/desert_base.png similarity index 100% rename from rustplus/api/icons/desert_base.png rename to rustplus/icons/desert_base.png diff --git a/rustplus/api/icons/dome.png b/rustplus/icons/dome.png similarity index 100% rename from rustplus/api/icons/dome.png rename to rustplus/icons/dome.png diff --git a/rustplus/api/icons/excavator.png b/rustplus/icons/excavator.png similarity index 100% rename from rustplus/api/icons/excavator.png rename to rustplus/icons/excavator.png diff --git a/rustplus/api/icons/explosion.png b/rustplus/icons/explosion.png similarity index 100% rename from rustplus/api/icons/explosion.png rename to rustplus/icons/explosion.png diff --git a/rustplus/api/icons/ferryterminal.png b/rustplus/icons/ferryterminal.png similarity index 100% rename from rustplus/api/icons/ferryterminal.png rename to rustplus/icons/ferryterminal.png diff --git a/rustplus/api/icons/fishing.png b/rustplus/icons/fishing.png similarity index 100% rename from rustplus/api/icons/fishing.png rename to rustplus/icons/fishing.png diff --git a/rustplus/api/icons/harbour.png b/rustplus/icons/harbour.png similarity index 100% rename from rustplus/api/icons/harbour.png rename to rustplus/icons/harbour.png diff --git a/rustplus/api/icons/icon.png b/rustplus/icons/icon.png similarity index 100% rename from rustplus/api/icons/icon.png rename to rustplus/icons/icon.png diff --git a/rustplus/api/icons/junkyard.png b/rustplus/icons/junkyard.png similarity index 100% rename from rustplus/api/icons/junkyard.png rename to rustplus/icons/junkyard.png diff --git a/rustplus/api/icons/large_oil_rig.png b/rustplus/icons/large_oil_rig.png similarity index 100% rename from rustplus/api/icons/large_oil_rig.png rename to rustplus/icons/large_oil_rig.png diff --git a/rustplus/api/icons/launchsite.png b/rustplus/icons/launchsite.png similarity index 100% rename from rustplus/api/icons/launchsite.png rename to rustplus/icons/launchsite.png diff --git a/rustplus/api/icons/lighthouse.png b/rustplus/icons/lighthouse.png similarity index 100% rename from rustplus/api/icons/lighthouse.png rename to rustplus/icons/lighthouse.png diff --git a/rustplus/api/icons/military_tunnels.png b/rustplus/icons/military_tunnels.png similarity index 100% rename from rustplus/api/icons/military_tunnels.png rename to rustplus/icons/military_tunnels.png diff --git a/rustplus/api/icons/mining_outpost.png b/rustplus/icons/mining_outpost.png similarity index 100% rename from rustplus/api/icons/mining_outpost.png rename to rustplus/icons/mining_outpost.png diff --git a/rustplus/api/icons/mining_quarry_hqm.png b/rustplus/icons/mining_quarry_hqm.png similarity index 100% rename from rustplus/api/icons/mining_quarry_hqm.png rename to rustplus/icons/mining_quarry_hqm.png diff --git a/rustplus/api/icons/mining_quarry_stone.png b/rustplus/icons/mining_quarry_stone.png similarity index 100% rename from rustplus/api/icons/mining_quarry_stone.png rename to rustplus/icons/mining_quarry_stone.png diff --git a/rustplus/api/icons/mining_quarry_sulfur.png b/rustplus/icons/mining_quarry_sulfur.png similarity index 100% rename from rustplus/api/icons/mining_quarry_sulfur.png rename to rustplus/icons/mining_quarry_sulfur.png diff --git a/rustplus/api/icons/missile_silo.png b/rustplus/icons/missile_silo.png similarity index 100% rename from rustplus/api/icons/missile_silo.png rename to rustplus/icons/missile_silo.png diff --git a/rustplus/api/icons/outpost.png b/rustplus/icons/outpost.png similarity index 100% rename from rustplus/api/icons/outpost.png rename to rustplus/icons/outpost.png diff --git a/rustplus/api/icons/oxums.png b/rustplus/icons/oxums.png similarity index 100% rename from rustplus/api/icons/oxums.png rename to rustplus/icons/oxums.png diff --git a/rustplus/api/icons/patrol.png b/rustplus/icons/patrol.png similarity index 100% rename from rustplus/api/icons/patrol.png rename to rustplus/icons/patrol.png diff --git a/rustplus/api/icons/power_plant.png b/rustplus/icons/power_plant.png similarity index 100% rename from rustplus/api/icons/power_plant.png rename to rustplus/icons/power_plant.png diff --git a/rustplus/api/icons/satellite.png b/rustplus/icons/satellite.png similarity index 100% rename from rustplus/api/icons/satellite.png rename to rustplus/icons/satellite.png diff --git a/rustplus/api/icons/sewer.png b/rustplus/icons/sewer.png similarity index 100% rename from rustplus/api/icons/sewer.png rename to rustplus/icons/sewer.png diff --git a/rustplus/api/icons/small_oil_rig.png b/rustplus/icons/small_oil_rig.png similarity index 100% rename from rustplus/api/icons/small_oil_rig.png rename to rustplus/icons/small_oil_rig.png diff --git a/rustplus/api/icons/stables.png b/rustplus/icons/stables.png similarity index 100% rename from rustplus/api/icons/stables.png rename to rustplus/icons/stables.png diff --git a/rustplus/api/icons/supermarket.png b/rustplus/icons/supermarket.png similarity index 100% rename from rustplus/api/icons/supermarket.png rename to rustplus/icons/supermarket.png diff --git a/rustplus/api/icons/swamp.png b/rustplus/icons/swamp.png similarity index 100% rename from rustplus/api/icons/swamp.png rename to rustplus/icons/swamp.png diff --git a/rustplus/api/icons/train.png b/rustplus/icons/train.png similarity index 100% rename from rustplus/api/icons/train.png rename to rustplus/icons/train.png diff --git a/rustplus/api/icons/train_yard.png b/rustplus/icons/train_yard.png similarity index 100% rename from rustplus/api/icons/train_yard.png rename to rustplus/icons/train_yard.png diff --git a/rustplus/api/icons/underwater_lab.png b/rustplus/icons/underwater_lab.png similarity index 100% rename from rustplus/api/icons/underwater_lab.png rename to rustplus/icons/underwater_lab.png diff --git a/rustplus/api/icons/vending_machine.png b/rustplus/icons/vending_machine.png similarity index 100% rename from rustplus/api/icons/vending_machine.png rename to rustplus/icons/vending_machine.png diff --git a/rustplus/api/icons/water_treatment.png b/rustplus/icons/water_treatment.png similarity index 100% rename from rustplus/api/icons/water_treatment.png rename to rustplus/icons/water_treatment.png diff --git a/rustplus/identification/__init__.py b/rustplus/identification/__init__.py new file mode 100644 index 0000000..4d3c7aa --- /dev/null +++ b/rustplus/identification/__init__.py @@ -0,0 +1,2 @@ +from .registered_listener import RegisteredListener +from .server_details import ServerDetails diff --git a/rustplus/identification/handler_list.py b/rustplus/identification/handler_list.py new file mode 100644 index 0000000..06bee4e --- /dev/null +++ b/rustplus/identification/handler_list.py @@ -0,0 +1,76 @@ +from collections import defaultdict +from typing import Set, Dict +from rustplus.identification import ( + RegisteredListener, + ServerDetails, +) + + +class HandlerList: + def __init__(self) -> None: + self._handlers: Dict[ServerDetails, Set[RegisteredListener]] = defaultdict(set) + + def unregister( + self, listener: RegisteredListener, server_details: ServerDetails + ) -> None: + self._handlers[server_details].remove(listener) + + def register( + self, listener: RegisteredListener, server_details: ServerDetails + ) -> None: + self._handlers[server_details].add(listener) + + def has(self, listener: RegisteredListener, server_details: ServerDetails) -> bool: + return listener in self._handlers[server_details] + + def unregister_all(self) -> None: + self._handlers.clear() + + def get_handlers(self, server_details: ServerDetails) -> Set[RegisteredListener]: + return self._handlers.get(server_details, set()) + + +class EntityHandlerList(HandlerList): + def __init__(self) -> None: + super().__init__() + self._handlers: Dict[ServerDetails, Dict[str, Set[RegisteredListener]]] = ( + defaultdict(dict) + ) + + def unregister( + self, listener: RegisteredListener, server_details: ServerDetails + ) -> None: + if listener.listener_id in self._handlers.get(server_details): + self._handlers.get(server_details).get(listener.listener_id).remove( + listener + ) + + def register( + self, listener: RegisteredListener, server_details: ServerDetails + ) -> None: + if server_details not in self._handlers: + self._handlers[server_details] = defaultdict(set) + + if listener.listener_id not in self._handlers.get(server_details): + self._handlers.get(server_details)[listener.listener_id] = set() + + self._handlers.get(server_details).get(listener.listener_id).add(listener) + + def has(self, listener: RegisteredListener, server_details: ServerDetails) -> bool: + if ( + server_details in self._handlers + and listener.listener_id in self._handlers.get(server_details) + ): + return listener in self._handlers.get(server_details).get( + listener.listener_id + ) + + return False + + def unregister_all(self) -> None: + self._handlers.clear() + + def get_handlers( + self, server_details: ServerDetails + ) -> Dict[str, Set[RegisteredListener]]: + return self._handlers.get(server_details, dict()) diff --git a/rustplus/api/remote/events/registered_listener.py b/rustplus/identification/registered_listener.py similarity index 54% rename from rustplus/api/remote/events/registered_listener.py rename to rustplus/identification/registered_listener.py index 8ec6843..bc12fa3 100644 --- a/rustplus/api/remote/events/registered_listener.py +++ b/rustplus/identification/registered_listener.py @@ -1,23 +1,18 @@ -from typing import Union, Coroutine +from typing import Coroutine class RegisteredListener: def __init__( self, - listener_id: Union[str, int], + listener_id: str, coroutine: Coroutine, - entity_type: int = None, ) -> None: - self.listener_id = str(listener_id) + self.listener_id = listener_id self._coroutine = coroutine - self._entity_type = entity_type def get_coro(self): return self._coroutine - def get_entity_type(self): - return self._entity_type - def __eq__(self, other) -> bool: if not isinstance(other, RegisteredListener): return False @@ -25,8 +20,7 @@ def __eq__(self, other) -> bool: return ( self.listener_id == other.listener_id and self._coroutine == other.get_coro() - and self._entity_type == other.get_entity_type() ) def __hash__(self): - return hash((self.listener_id, self._coroutine, self._entity_type)) + return hash((self.listener_id, self._coroutine)) diff --git a/rustplus/identification/server_details.py b/rustplus/identification/server_details.py new file mode 100644 index 0000000..a73272a --- /dev/null +++ b/rustplus/identification/server_details.py @@ -0,0 +1,39 @@ +from typing import Union + + +class ServerDetails: + def __init__( + self, + ip: str, + port: Union[str, int, None], + player_id: int, + player_token: int, + secure: bool = False, + ) -> None: + self.ip = str(ip) + self.port = str(port) + self.player_id = int(player_id) + self.player_token = int(player_token) + self.secure = secure + + def get_server_string(self) -> str: + if self.port is None: + return f"{self.ip}" + return f"{self.ip}:{self.port}" + + def __str__(self) -> str: + return f"{self.ip}:{self.port} {self.player_id} {self.player_token}" + + def __hash__(self): + return hash(self.__str__()) + + def __eq__(self, o: object) -> bool: + if not isinstance(o, ServerDetails): + return False + + return ( + self.ip == o.ip + and self.port == o.port + and self.player_id == o.player_id + and self.player_token == o.player_token + ) diff --git a/rustplus/remote/__init__.py b/rustplus/remote/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rustplus/api/remote/camera/__init__.py b/rustplus/remote/camera/__init__.py similarity index 100% rename from rustplus/api/remote/camera/__init__.py rename to rustplus/remote/camera/__init__.py diff --git a/rustplus/api/remote/camera/camera_constants.py b/rustplus/remote/camera/camera_constants.py similarity index 100% rename from rustplus/api/remote/camera/camera_constants.py rename to rustplus/remote/camera/camera_constants.py diff --git a/rustplus/api/remote/camera/camera_manager.py b/rustplus/remote/camera/camera_manager.py similarity index 85% rename from rustplus/api/remote/camera/camera_manager.py rename to rustplus/remote/camera/camera_manager.py index 799ca77..5b854f6 100644 --- a/rustplus/api/remote/camera/camera_manager.py +++ b/rustplus/remote/camera/camera_manager.py @@ -1,6 +1,5 @@ import time from typing import Iterable, Union, List, Coroutine, Set, Callable - from PIL import Image from .camera_parser import Parser @@ -10,12 +9,17 @@ AppEmpty, AppRequest, AppCameraInfo, + AppCameraRays, + AppCameraSubscribe, ) -from ...structures import Vector +from ...structs import Vector from .structures import CameraInfo, Entity, LimitedQueue class CameraManager: + + ACTIVE_INSTANCE: Union["CameraManager", None] = None + def __init__( self, rust_socket, cam_id: str, cam_info_message: AppCameraInfo ) -> None: @@ -29,8 +33,9 @@ def __init__( ) self.time_since_last_subscribe: float = time.time() self.frame_callbacks: Set[Callable[[Image.Image], Coroutine]] = set() + CameraManager.ACTIVE_INSTANCE = self - async def add_packet(self, packet) -> None: + async def add_packet(self, packet: AppCameraRays) -> None: self._last_packets.add(packet) if len(self.frame_callbacks) == 0: @@ -122,8 +127,7 @@ async def send_combined_movement( for movement in movements: value = value | movement - await self.rust_socket._handle_ratelimit(0.01) - app_request: AppRequest = self.rust_socket._generate_protobuf() + app_request: AppRequest = await self.rust_socket._generate_request(0.01) cam_input = AppCameraInput() cam_input.buttons = value @@ -133,25 +137,29 @@ async def send_combined_movement( cam_input.mouse_delta = vector app_request.camera_input = cam_input - await self.rust_socket.remote.send_message(app_request) - await self.rust_socket.remote.add_ignored_response(app_request.seq) + await self.rust_socket.ws.send_message(app_request, True) async def exit_camera(self) -> None: - await self.rust_socket._handle_ratelimit() - app_request: AppRequest = self.rust_socket._generate_protobuf() + app_request: AppRequest = await self.rust_socket._generate_request() app_request.camera_unsubscribe = AppEmpty() - await self.rust_socket.remote.send_message(app_request) - await self.rust_socket.remote.add_ignored_response(app_request.seq) + await self.rust_socket.ws.send_message(app_request, True) self._open = False self._last_packets.clear() async def resubscribe(self) -> None: - await self.rust_socket.remote.subscribe_to_camera(self._cam_id, True) + + packet = await self.rust_socket._generate_request() + subscribe = AppCameraSubscribe() + subscribe.camera_id = self._cam_id + packet.camera_subscribe = subscribe + + self.rust_socket.ws.send_message(packet, True) + self.time_since_last_subscribe = time.time() self._open = True - self.rust_socket.remote.camera_manager = self + self.ACTIVE_INSTANCE = self async def get_entities_in_frame(self) -> List[Entity]: if self._last_packets is None: diff --git a/rustplus/api/remote/camera/camera_parser.py b/rustplus/remote/camera/camera_parser.py similarity index 98% rename from rustplus/api/remote/camera/camera_parser.py rename to rustplus/remote/camera/camera_parser.py index 6fde4cf..d729e46 100644 --- a/rustplus/api/remote/camera/camera_parser.py +++ b/rustplus/remote/camera/camera_parser.py @@ -318,9 +318,9 @@ def handle_entity( entity_size = np.array([entity.size.x, entity.size.y, entity.size.z, 0]) vertices = ( - MathUtils.get_player_vertices(entity.size) + MathUtils.get_player_vertices(Vector3(vector3=entity.size)) if entity.type == 2 - else MathUtils.get_tree_vertices(entity.size) + else MathUtils.get_tree_vertices(Vector3(vector3=entity.size)) ) # Add the position for the name tag to the vertices vertices = np.append(vertices, [np.array([0, 1.3, 0, 1])], axis=0) @@ -367,7 +367,9 @@ def handle_entity( colour = ( (PLAYER_COLOUR if not entity.name.isdigit() else SCIENTIST_COLOUR) if entity.type == 2 - else MathUtils.get_slightly_random_colour(TREE_COLOUR, entity.position) + else MathUtils.get_slightly_random_colour( + TREE_COLOUR, Vector3(vector3=entity.position) + ) ) MathUtils.set_polygon_with_depth( diff --git a/rustplus/api/remote/camera/structures.py b/rustplus/remote/camera/structures.py similarity index 90% rename from rustplus/api/remote/camera/structures.py rename to rustplus/remote/camera/structures.py index a04034c..b0f8e80 100644 --- a/rustplus/api/remote/camera/structures.py +++ b/rustplus/remote/camera/structures.py @@ -16,8 +16,8 @@ def is_move_option_permissible(self, value: int) -> bool: def __str__(self) -> str: return ( - f"CameraInfo(width={self.width}, height={self.height}, near_plane={self.near_plane}, " - f"far_plane={self.far_plane}, control_flags={self.control_flags})" + f"CameraInfo[width={self.width}, height={self.height}, near_plane={self.near_plane}, " + f"far_plane={self.far_plane}, control_flags={self.control_flags}]" ) @@ -32,8 +32,8 @@ def __init__(self, entity_data: AppCameraRaysEntity) -> None: def __str__(self) -> str: return ( - f"Entity(entity_id={self.entity_id}, type={self.type}, position={self.position}, " - f"rotation={self.rotation}, size={self.size}, name={self.name})" + f"Entity[entity_id={self.entity_id}, type={self.type}, position={self.position}, " + f"rotation={self.rotation}, size={self.size}, name={self.name}]" ) def __repr__(self): @@ -56,7 +56,7 @@ def __hash__(self): return hash((self.x, self.y, self.z)) def __str__(self) -> str: - return f"Vector3(x={self.x}, y={self.y}, z={self.z})" + return f"Vector3[x={self.x}, y={self.y}, z={self.z}]" class RayPacket: @@ -69,8 +69,8 @@ def __init__(self, ray_packet: AppCameraRays) -> None: def __str__(self) -> str: return ( - f"RayPacket(vertical_fov={self.vertical_fov}, sample_offset={self.sample_offset}, " - f"ray_data={self.ray_data}, distance={self.distance}, entities={self.entities})" + f"RayPacket[vertical_fov={self.vertical_fov}, sample_offset={self.sample_offset}, " + f"ray_data={self.ray_data}, distance={self.distance}, entities={self.entities}]" ) diff --git a/rustplus/remote/fcm/__init__.py b/rustplus/remote/fcm/__init__.py new file mode 100644 index 0000000..e1ccc3d --- /dev/null +++ b/rustplus/remote/fcm/__init__.py @@ -0,0 +1 @@ +from .fcm_listener import FCMListener diff --git a/rustplus/api/remote/fcm_listener.py b/rustplus/remote/fcm/fcm_listener.py similarity index 100% rename from rustplus/api/remote/fcm_listener.py rename to rustplus/remote/fcm/fcm_listener.py diff --git a/rustplus/remote/proxy/__init__.py b/rustplus/remote/proxy/__init__.py new file mode 100644 index 0000000..5679479 --- /dev/null +++ b/rustplus/remote/proxy/__init__.py @@ -0,0 +1 @@ +from .proxy_value_grabber import ProxyValueGrabber diff --git a/rustplus/remote/proxy/proxy_value_grabber.py b/rustplus/remote/proxy/proxy_value_grabber.py new file mode 100644 index 0000000..ef0cef7 --- /dev/null +++ b/rustplus/remote/proxy/proxy_value_grabber.py @@ -0,0 +1,32 @@ +import requests +import logging +import time + + +class ProxyValueGrabber: + + VALUE = -1 + LAST_FETCHED = -1 + + @staticmethod + def get_value() -> int: + + if ( + ProxyValueGrabber.VALUE != -1 + and ProxyValueGrabber.LAST_FETCHED >= time.time() - 600 + ): + return ProxyValueGrabber.VALUE + + data = requests.get("https://companion-rust.facepunch.com/api/version") + + if data.status_code == 200: + publish_time = data.json().get("minPublishedTime", None) + if publish_time is not None: + ProxyValueGrabber.VALUE = publish_time + 1 + ProxyValueGrabber.LAST_FETCHED = time.time() + return ProxyValueGrabber.VALUE + + logging.getLogger("rustplus.py").warning( + "Failed to get magic value from RustPlus Server" + ) + return 9999999999999 diff --git a/rustplus/api/remote/ratelimiter.py b/rustplus/remote/ratelimiter/__init__.py similarity index 65% rename from rustplus/api/remote/ratelimiter.py rename to rustplus/remote/ratelimiter/__init__.py index 3a89f7f..16bd653 100644 --- a/rustplus/api/remote/ratelimiter.py +++ b/rustplus/remote/ratelimiter/__init__.py @@ -4,7 +4,7 @@ from typing import Dict from ...exceptions.exceptions import RateLimitError -from ...utils import ServerID +from ...identification import ServerDetails class TokenBucket: @@ -19,21 +19,18 @@ def __init__( self.refresh_per_second = self.refresh_amount / self.refresh_rate def can_consume(self, amount) -> bool: - if (self.current - amount) >= 0: - return True - - return False + return (self.current - amount) >= 0 def consume(self, amount: int = 1) -> None: self.current -= amount def refresh(self) -> None: time_now = time.time() - time_delta = time_now - self.last_update self.last_update = time_now - - self.current = min([self.current + time_delta * self.refresh_amount, self.max]) + self.current = min( + self.current + time_delta * self.refresh_per_second, self.max + ) class RateLimiter: @@ -48,60 +45,60 @@ def default(cls) -> "RateLimiter": return cls() def __init__(self) -> None: - self.socket_buckets: Dict[ServerID, TokenBucket] = {} + self.socket_buckets: Dict[ServerDetails, TokenBucket] = {} self.server_buckets: Dict[str, TokenBucket] = {} self.lock = asyncio.Lock() def add_socket( self, - server_id: ServerID, + server_details: ServerDetails, current: float, maximum: float, refresh_rate: float, refresh_amount: float, ) -> None: - self.socket_buckets[server_id] = TokenBucket( + self.socket_buckets[server_details] = TokenBucket( current, maximum, refresh_rate, refresh_amount ) - if server_id.get_server_string() not in self.server_buckets: - self.server_buckets[server_id.get_server_string()] = TokenBucket( + if server_details.get_server_string() not in self.server_buckets: + self.server_buckets[server_details.get_server_string()] = TokenBucket( self.SERVER_LIMIT, self.SERVER_LIMIT, 1, self.SERVER_REFRESH_AMOUNT ) - async def can_consume(self, server_id: ServerID, amount: int = 1) -> bool: + async def can_consume(self, server_details: ServerDetails, amount: int = 1) -> bool: """ Returns whether the user can consume the amount of tokens provided """ async with self.lock: - can_consume = True - for bucket in [ - self.socket_buckets.get(server_id), - self.server_buckets.get(server_id.get_server_string()), + self.socket_buckets.get(server_details), + self.server_buckets.get(server_details.get_server_string()), ]: bucket.refresh() if not bucket.can_consume(amount): - can_consume = False - - return can_consume + return False + return True - async def consume(self, server_id: ServerID, amount: int = 1) -> None: + async def consume(self, server_details: ServerDetails, amount: int = 1) -> None: """ Consumes an amount of tokens from the bucket. You should first check to see whether it is possible with can_consume """ async with self.lock: for bucket in [ - self.socket_buckets.get(server_id), - self.server_buckets.get(server_id.get_server_string()), + self.socket_buckets.get(server_details), + self.server_buckets.get(server_details.get_server_string()), ]: bucket.refresh() if not bucket.can_consume(amount): - self.lock.release() raise RateLimitError("Not Enough Tokens") + for bucket in [ + self.socket_buckets.get(server_details), + self.server_buckets.get(server_details.get_server_string()), + ]: bucket.consume(amount) async def get_estimated_delay_time( - self, server_id: ServerID, target_cost: int + self, server_details: ServerDetails, target_cost: int ) -> float: """ Returns how long until the amount of tokens needed will be available @@ -109,8 +106,8 @@ async def get_estimated_delay_time( async with self.lock: delay = 0 for bucket in [ - self.socket_buckets.get(server_id), - self.server_buckets.get(server_id.get_server_string()), + self.socket_buckets.get(server_details), + self.server_buckets.get(server_details.get_server_string()), ]: val = ( math.ceil( @@ -126,9 +123,9 @@ async def get_estimated_delay_time( delay = val return delay - async def remove(self, server_id: ServerID) -> None: + async def remove(self, server_details: ServerDetails) -> None: """ Removes the limiter """ async with self.lock: - del self.socket_buckets[server_id] + del self.socket_buckets[server_details] diff --git a/rustplus/api/remote/rustplus_proto/__init__.py b/rustplus/remote/rustplus_proto/__init__.py similarity index 100% rename from rustplus/api/remote/rustplus_proto/__init__.py rename to rustplus/remote/rustplus_proto/__init__.py diff --git a/rustplus/api/remote/rustplus_proto/rustplus.py b/rustplus/remote/rustplus_proto/rustplus.py similarity index 99% rename from rustplus/api/remote/rustplus_proto/rustplus.py rename to rustplus/remote/rustplus_proto/rustplus.py index 8d67945..0c06efb 100644 --- a/rustplus/api/remote/rustplus_proto/rustplus.py +++ b/rustplus/remote/rustplus_proto/rustplus.py @@ -427,7 +427,7 @@ class AppClanChat(betterproto.Message): @dataclass class AppNexusAuth(betterproto.Message): - server_id: str = betterproto.string_field(1) + server_details: str = betterproto.string_field(1) player_token: int = betterproto.int32_field(2) diff --git a/rustplus/remote/websocket/__init__.py b/rustplus/remote/websocket/__init__.py new file mode 100644 index 0000000..bc2f0a5 --- /dev/null +++ b/rustplus/remote/websocket/__init__.py @@ -0,0 +1 @@ +from .ws import RustWebsocket diff --git a/rustplus/remote/websocket/ws.py b/rustplus/remote/websocket/ws.py new file mode 100644 index 0000000..1dd6a59 --- /dev/null +++ b/rustplus/remote/websocket/ws.py @@ -0,0 +1,348 @@ +import shlex +import base64 +import betterproto +from websockets.exceptions import InvalidURI, InvalidHandshake, ConnectionClosedError +from websockets.legacy.client import WebSocketClientProtocol +from websockets.client import connect +from asyncio import TimeoutError, Task, AbstractEventLoop +from typing import Union, Coroutine, Optional, Set, Dict +import logging +import asyncio + +from ..camera import CameraManager +from ..proxy import ProxyValueGrabber +from ..rustplus_proto import AppMessage, AppRequest +from ...commands import CommandOptions, ChatCommand, ChatCommandTime +from ...events import ( + ProtobufEventPayload, + EntityEventPayload, + TeamEventPayload, + ChatEventPayload, +) +from ...exceptions import ClientNotConnectedError, RequestError +from ...identification import ServerDetails, RegisteredListener +from ...structs import RustChatMessage, RustTeamInfo +from ...utils import YieldingEvent, convert_time + + +class RustWebsocket: + RESPONSE_TIMEOUT = 5 + + def __init__( + self, + server_details: ServerDetails, + command_options: Union[CommandOptions, None], + use_fp_proxy: bool, + use_test_server: bool, + debug: bool, + ) -> None: + self.server_details: ServerDetails = server_details + self.command_options: Union[CommandOptions, None] = command_options + self.connection: Union[WebSocketClientProtocol, None] = None + self.logger: logging.Logger = logging.getLogger("rustplus.py") + self.task: Union[Task, None] = None + self.debug: bool = debug + self.use_test_server: bool = use_test_server + self.use_fp_proxy: bool = use_fp_proxy + + self.responses: Dict[int, YieldingEvent] = {} + self.open = False + + async def connect(self) -> bool: + + address = ( + ( + f"{'wss' if self.server_details.secure else 'ws'}://" + + self.server_details.get_server_string() + ) + if not self.use_fp_proxy + else f"wss://companion-rust.facepunch.com/game/{self.server_details.ip}/{self.server_details.port}" + ) + f"?v={ProxyValueGrabber.get_value()}" + + try: + self.connection = await connect( + address, + close_timeout=0, + ping_interval=None, + max_size=1_000_000_000, + ) + except (InvalidURI, OSError, InvalidHandshake, TimeoutError) as err: + self.logger.warning("WebSocket connection error: %s", err) + return False + + if self.debug: + self.logger.info("Websocket connection established to %s", address) + + self.task = asyncio.create_task( + self.run(), name="[RustPlus.py] Websocket Polling Task" + ) + + self.open = True + + return True + + async def disconnect(self) -> None: + if self.task and self.connection: + self.task.cancel() + try: + await self.task + except asyncio.CancelledError: + pass # Ignore the cancellation error + + self.task = None + + self.open = False + await self.connection.close() + self.connection = None + + async def run(self) -> None: + while self.open: + try: + data = await self.connection.recv() + + await self.run_coroutine_non_blocking( + self.run_proto_event(data, self.server_details) + ) + + if self.use_test_server: + data = base64.b64decode(data) + + app_message = AppMessage() + app_message.parse(data) + + except ConnectionClosedError as e: + if self.debug: + self.logger.exception("Connection Interrupted: %s", e) + else: + self.logger.warning("Connection Interrupted: %s", e) + break + + except Exception as e: + self.logger.exception( + "An Error occurred whilst parsing the message from the server: %s", + e, + ) + continue + + try: + await self.run_coroutine_non_blocking(self.handle_message(app_message)) + except Exception as e: + self.logger.exception( + "An Error occurred whilst handling the message from the server %s", + e, + ) + + async def send_and_get(self, request: AppRequest) -> AppMessage: + await self.send_message(request) + return await self.get_response(request.seq) + + async def send_message( + self, request: AppRequest, ignore_response: bool = False + ) -> None: + if self.connection is None: + raise ClientNotConnectedError("No Current Websocket Connection") + + if self.debug: + self.logger.info(f"Sending Message with seq {request.seq}: {request}") + + if not ignore_response: + self.responses[request.seq] = YieldingEvent() + + try: + if self.use_test_server: + await self.connection.send( + base64.b64encode(bytes(request)).decode("utf-8") + ) + else: + await self.connection.send(bytes(request)) + except Exception as err: + self.logger.warning("WebSocket connection error: %s", err) + + async def get_response(self, seq: int) -> Union[AppMessage, None]: + + response = await self.responses[seq].wait(timeout=self.RESPONSE_TIMEOUT) + del self.responses[seq] + + return response + + async def handle_message(self, app_message: AppMessage) -> None: + if self.debug: + self.logger.info( + f"Received Message with seq {app_message.response.seq}: {app_message}" + ) + + if self.error_present(app_message.response.error.error): + raise RequestError(app_message.response.error.error) + + prefix = self.get_prefix( + str(app_message.broadcast.team_message.message.message) + ) + + if prefix is not None: + # Command + + if self.debug: + self.logger.info(f"Attempting to run Command: {app_message}") + + message = RustChatMessage(app_message.broadcast.team_message.message) + + parts = shlex.split(message.message) + command = parts[0][len(prefix) :] + + data = ChatCommand.REGISTERED_COMMANDS[self.server_details].get( + command, None + ) + + dao = ChatCommand( + message.name, + message.steam_id, + ChatCommandTime( + convert_time(message.time), + message.time, + ), + command, + parts[1:], + ) + + if data is not None: + await data.coroutine(dao) + else: + for command_name, data in ChatCommand.REGISTERED_COMMANDS[ + self.server_details + ].items(): + if command in data.aliases or data.callable_func(command): + await data.coroutine(dao) + break + + if self.is_entity_broadcast(app_message): + # Entity Event + if self.debug: + self.logger.info(f"Running Entity Event: {app_message}") + + handlers = EntityEventPayload.HANDLER_LIST.get_handlers( + self.server_details + ).get(str(app_message.broadcast.entity_changed.entity_id), []) + for handler in handlers: + await handler.get_coro()( + EntityEventPayload( + entity_changed=app_message.broadcast.entity_changed, + ) + ) + + elif self.is_camera_broadcast(app_message): + # Pipe packet into Camera Manager + if self.debug: + self.logger.info(f"Updating Camera Packet: {app_message}") + + if CameraManager.ACTIVE_INSTANCE is not None: + await CameraManager.ACTIVE_INSTANCE.add_packet( + app_message.broadcast.camera_rays + ) + + elif self.is_team_broadcast(app_message): + # Team Event + if self.debug: + self.logger.info(f"Running Team Event: {app_message}") + + # This means that the team of the current player has changed + handlers = TeamEventPayload.HANDLER_LIST.get_handlers(self.server_details) + team_event = TeamEventPayload( + app_message.broadcast.team_changed.player_id, + RustTeamInfo(app_message.broadcast.team_changed.team_info), + ) + for handler in handlers: + await handler.get_coro()(team_event) + + elif self.is_message(app_message): + # Chat message event + if self.debug: + self.logger.info(f"Running Chat Event: {app_message}") + + handlers = ChatEventPayload.HANDLER_LIST.get_handlers(self.server_details) + chat_event = ChatEventPayload( + RustChatMessage(app_message.broadcast.team_message.message) + ) + for handler in handlers: + await handler.get_coro()(chat_event) + + else: + # This means that it wasn't sent by the server and is a message from the server in response to an action + event: YieldingEvent = self.responses.get(app_message.response.seq, None) + if event is not None: + if self.debug: + self.logger.info(f"Running Response Event: {app_message}") + + event.set_with_value(app_message) + + def get_prefix(self, message: str) -> Optional[str]: + + if self.command_options is None: + return None + + if message.startswith(self.command_options.prefix): + return self.command_options.prefix + else: + return None + + @staticmethod + async def run_proto_event( + data: Union[str, bytes], server_details: ServerDetails + ) -> None: + handlers: Set[RegisteredListener] = ( + ProtobufEventPayload.HANDLER_LIST.get_handlers(server_details) + ) + for handler in handlers: + await handler.get_coro()(data) + + @staticmethod + def is_message(app_message: AppMessage) -> bool: + return betterproto.serialized_on_wire( + app_message.broadcast.team_message.message + ) + + @staticmethod + def is_camera_broadcast(app_message: AppMessage) -> bool: + return betterproto.serialized_on_wire(app_message.broadcast.camera_rays) + + @staticmethod + def is_entity_broadcast(app_message: AppMessage) -> bool: + return betterproto.serialized_on_wire(app_message.broadcast.entity_changed) + + @staticmethod + def is_team_broadcast(app_message: AppMessage) -> bool: + return betterproto.serialized_on_wire(app_message.broadcast.team_changed) + + @staticmethod + def get_proto_cost(app_request: AppRequest) -> int: + """ + Gets the cost of an AppRequest + """ + costs = [ + (app_request.get_time, 1), + (app_request.send_team_message, 2), + (app_request.get_info, 1), + (app_request.get_team_chat, 1), + (app_request.get_team_info, 1), + (app_request.get_map_markers, 1), + (app_request.get_map, 5), + (app_request.set_entity_value, 1), + (app_request.get_entity_info, 1), + (app_request.promote_to_leader, 1), + ] + for request, cost in costs: + if betterproto.serialized_on_wire(request): + return cost + + raise ValueError() + + @staticmethod + def error_present(message) -> bool: + """ + Checks message for error + """ + return message != "" + + @staticmethod + async def run_coroutine_non_blocking(coroutine: Coroutine) -> Task: + loop: AbstractEventLoop = asyncio.get_event_loop_policy().get_event_loop() + return loop.create_task(coroutine) diff --git a/rustplus/rust_api.py b/rustplus/rust_api.py new file mode 100644 index 0000000..398387f --- /dev/null +++ b/rustplus/rust_api.py @@ -0,0 +1,524 @@ +import asyncio +from collections import defaultdict +from datetime import datetime +from importlib import resources +from io import BytesIO +from typing import List, Union +import logging +from PIL import Image + +from .commands import CommandOptions +from .identification import ServerDetails +from .remote.camera import CameraManager +from .remote.rustplus_proto import ( + AppRequest, + AppEmpty, + AppSendMessage, + AppSetEntityValue, + AppPromoteToLeader, + AppCameraSubscribe, + AppMapMonument, + AppFlag, +) +from .remote.websocket import RustWebsocket +from .structs import ( + RustTime, + RustInfo, + RustChatMessage, + RustTeamInfo, + RustMarker, + RustMap, + RustEntityInfo, + RustContents, + RustItem, +) +from .utils import ( + convert_time, + translate_id_to_stack, + generate_grid, + fetch_avatar_icon, + format_coord, + convert_marker, + convert_monument, +) +from .remote.ratelimiter import RateLimiter + + +class RustSocket: + + def __init__( + self, + server_details: ServerDetails, + ratelimiter: Union[RateLimiter, None] = None, + command_options: Union[CommandOptions, None] = None, + use_fp_proxy: bool = False, + use_test_server: bool = False, + debug: bool = False, + ) -> None: + self.server_details = server_details + self.command_options = command_options + self.logger = logging.getLogger("rustplus.py") + + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.DEBUG) + console_handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + ) + self.logger.addHandler(console_handler) + self.logger.setLevel(logging.DEBUG) + + self.ws = RustWebsocket( + self.server_details, + self.command_options, + use_fp_proxy, + use_test_server, + debug, + ) + self.seq = 1 + + if ratelimiter: + self.ratelimiter = ratelimiter + else: + self.ratelimiter = RateLimiter() + + self.ratelimiter.add_socket( + self.server_details, + RateLimiter.SERVER_LIMIT, + RateLimiter.SERVER_LIMIT, + 1, + RateLimiter.SERVER_REFRESH_AMOUNT, + ) + + async def _handle_ratelimit(self, tokens) -> None: + while True: + if await self.ratelimiter.can_consume(self.server_details, tokens): + await self.ratelimiter.consume(self.server_details, tokens) + break + + await asyncio.sleep( + await self.ratelimiter.get_estimated_delay_time( + self.server_details, tokens + ) + ) + + async def _generate_request(self, tokens=1) -> AppRequest: + await self._handle_ratelimit(tokens) + + app_request = AppRequest() + app_request.seq = self.seq + self.seq += 1 + app_request.player_id = self.server_details.player_id + app_request.player_token = self.server_details.player_token + + return app_request + + async def connect(self) -> bool: + if await self.ws.connect(): + await self.get_time() # Wake up the connection + return True + return False + + async def disconnect(self) -> None: + await self.ws.disconnect() + + @staticmethod + async def hang() -> None: + """ + This Will permanently put your script into a state of 'hanging' Cannot be Undone. Only do this in scripts + using commands + + :returns Nothing, This will never return + """ + + while True: + await asyncio.sleep(1) + + async def get_time(self) -> Union[RustTime, None]: + """ + Gets the current in-game time from the server. + + :returns RustTime: The Time + """ + + packet = await self._generate_request() + packet.get_time = AppEmpty() + packet = await self.ws.send_and_get(packet) + + if packet is None: + return None + + return RustTime( + packet.response.time.day_length_minutes, + convert_time(packet.response.time.sunrise), + convert_time(packet.response.time.sunset), + convert_time(packet.response.time.time), + packet.response.time.time, + packet.response.time.time_scale, + ) + + async def send_team_message(self, message: str) -> None: + """ + Sends a message to the in-game team chat + + :param message: The string message to send + """ + + packet = await self._generate_request(tokens=2) + send_message = AppSendMessage() + send_message.message = message + packet.send_team_message = send_message + + await self.ws.send_message(packet, True) + + async def get_info(self) -> Union[RustInfo, None]: + """ + Gets information on the Rust Server + :return: RustInfo - The info of the server + """ + packet = await self._generate_request() + packet.get_info = AppEmpty() + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return RustInfo(response.response.info) + + async def get_team_chat(self) -> Union[List[RustChatMessage], None]: + """ + Gets the team chat from the server + + :return List[RustChatMessage]: The chat messages in the team chat + """ + packet = await self._generate_request() + packet.get_team_chat = AppEmpty() + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return [ + RustChatMessage(message) for message in response.response.team_chat.messages + ] + + async def get_team_info(self) -> Union[RustTeamInfo, None]: + """ + Gets Information on the members of your team + + :return RustTeamInfo: The info of your team + """ + packet = await self._generate_request() + packet.get_team_info = AppEmpty() + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return RustTeamInfo(response.response.team_info) + + async def get_markers(self) -> Union[List[RustMarker], None]: + """ + Gets all the map markers from the server + + :return List[RustMarker]: All the markers on the map + """ + packet = await self._generate_request() + packet.get_map_markers = AppEmpty() + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return [RustMarker(marker) for marker in response.response.map_markers.markers] + + async def get_map( + self, + add_icons: bool = False, + add_events: bool = False, + add_vending_machines: bool = False, + add_team_positions: bool = False, + override_images: dict = None, + add_grid: bool = False, + ) -> Union[Image.Image, None]: + """ + Gets an image of the map from the server with the specified additions + + :param add_icons: To add the monument icons + :param add_events: To add the Event icons + :param add_vending_machines: To add the vending icons + :param add_team_positions: To add the team positions + :param override_images: To override the images pre-supplied with RustPlus.py + :param add_grid: To add the grid to the map + :return Image: PIL Image + """ + + if override_images is None: + override_images = {} + + server_info = await self.get_info() + if server_info is None: + return None + + map_size = server_info.size + + packet = await self._generate_request(5) + packet.get_map = AppEmpty() + response = await self.ws.send_and_get(packet) + if response is None: + return None + + map_packet = response.response.map + monuments: List[AppMapMonument] = map_packet.monuments + + try: + output = Image.open(BytesIO(map_packet.jpg_image)) + except Exception as e: + self.logger.error(f"Error opening image: {e}") + return None + + output = output.crop( + (500, 500, map_packet.height - 500, map_packet.width - 500) + ) + output = output.resize((map_size, map_size), Image.LANCZOS).convert("RGBA") + + if add_grid: + output.paste(grid := generate_grid(map_size), (5, 5), grid) + + if add_icons or add_events or add_vending_machines: + map_markers = ( + await self.get_markers() if add_events or add_vending_machines else [] + ) + + if add_icons: + for monument in monuments: + if str(monument.token) == "DungeonBase": + continue + icon = convert_monument(monument.token, override_images) + if monument.token in override_images: + icon = icon.resize((150, 150)) + if str(monument.token) == "train_tunnel_display_name": + icon = icon.resize((100, 125)) + output.paste( + icon, + (format_coord(int(monument.x), int(monument.y), map_size)), + icon, + ) + + if add_vending_machines: + with resources.path("rustplus.icons", "vending_machine.png") as path: + vending_machine = Image.open(path).convert("RGBA") + vending_machine = vending_machine.resize((100, 100)) + + for marker in map_markers: + if add_events: + if marker.type in [2, 4, 5, 6, 8]: + icon = convert_marker(str(marker.type), marker.rotation) + if marker.type == 6: + x, y = marker.x, marker.y + y = min(max(y, 0), map_size) + x = min(max(x, 0), map_size - 75 if x > map_size else x) + output.paste(icon, (int(x), map_size - int(y)), icon) + else: + output.paste( + icon, + (format_coord(int(marker.x), int(marker.y), map_size)), + icon, + ) + + if add_vending_machines and marker.type == 3: + output.paste( + vending_machine, + (int(marker.x) - 50, map_size - int(marker.y) - 50), + vending_machine, + ) + + if add_team_positions: + team = await self.get_team_info() + if team is not None: + for member in team.members: + if not member.is_alive: + continue + + output.paste( + avatar := await fetch_avatar_icon( + member.steam_id, member.is_online + ), + format_coord(int(member.x), int(member.y), server_info.size), + avatar, + ) + + return output + + async def get_map_info(self) -> Union[RustMap, None]: + """ + Gets the raw map data from the server + + :return RustMap: The raw map of the server + """ + packet = await self._generate_request(tokens=5) + packet.get_map = AppEmpty() + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return RustMap(response.response.map) + + async def get_entity_info(self, eid: int = None) -> Union[RustEntityInfo, None]: + """ + Gets entity info from the server + + :param eid: The Entities ID + :return RustEntityInfo: The entity Info + """ + packet = await self._generate_request() + packet.get_entity_info = AppEmpty() + packet.entity_id = eid + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return RustEntityInfo(response.response.entity_info) + + async def set_entity_value(self, eid: int, value: bool = False) -> None: + """ + Turns off a given smart switch by entity ID + + :param eid: The Entities ID + :param value: The value to set + :return None: + """ + packet = await self._generate_request() + set_value = AppSetEntityValue() + set_value.value = value + packet.set_entity_value = AppEmpty() + packet.entity_id = eid + + await self.ws.send_message(packet, True) + + async def set_subscription_to_entity(self, eid: int, value: bool = True) -> None: + """ + Subscribes to an entity for events + + :param eid: The Entities ID + :param value: The value to set the subscription to + :return None: + """ + packet = await self._generate_request() + flag = AppFlag() + flag.value = value + packet.set_subscription = flag + packet.entity_id = eid + + await self.ws.send_message(packet, True) + + async def check_subscription_to_entity(self, eid: int) -> Union[bool, None]: + """ + Checks if you are subscribed to an entity + + :param eid: The Entities ID + :return bool: If you are subscribed + """ + packet = await self._generate_request() + packet.check_subscription = AppEmpty() + packet.entity_id = eid + + response = await self.ws.send_and_get(packet) + if response is None: + return None + + return response.response.flag.value + + async def promote_to_team_leader(self, steamid: int = None) -> None: + """ + Promotes a given user to the team leader by their 64-bit Steam ID + + :param steamid: The SteamID of the player to promote + :return None: + """ + packet = await self._generate_request() + promote_packet = AppPromoteToLeader() + promote_packet.steam_id = steamid + packet.promote_to_leader = promote_packet + + await self.ws.send_message(packet, True) + + async def get_contents( + self, eid: int = None, combine_stacks: bool = False + ) -> Union[RustContents, None]: + """ + Gets the contents of a storage monitor-attached container + + :param eid: The EntityID Of the storage Monitor + :param combine_stacks: Whether to combine alike stacks together + :return RustContents: The contents on the monitor + """ + returned_data = await self.get_entity_info(eid) + + if returned_data is None: + return None + + target_time = datetime.utcfromtimestamp(int(returned_data.protection_expiry)) + difference = target_time - datetime.utcnow() + + items = [] + + for item in returned_data.items: + items.append( + RustItem( + translate_id_to_stack(item.item_id), + item.item_id, + item.quantity, + item.item_is_blueprint, + ) + ) + + if combine_stacks: + merged_map = defaultdict(tuple) + + for item in items: + data = merged_map[str(item.item_id)] + if data: + count = int(data[0]) + int(item.quantity) + merged_map[str(item.item_id)] = (count, bool(item.is_blueprint)) + else: + merged_map[str(item.item_id)] = ( + int(item.quantity), + bool(item.is_blueprint), + ) + + items = [] + for key in merged_map.keys(): + items.append( + RustItem( + translate_id_to_stack(key), + key, + int(merged_map[key][0]), + bool(merged_map[key][1]), + ) + ) + + return RustContents(difference, bool(returned_data.has_protection), items) + + async def get_camera_manager(self, cam_id: str) -> Union[CameraManager, None]: + """ + Gets a camera manager for a given camera ID + + NOTE: This will override the current camera manager if one exists for the given ID so you cannot have multiple + + :param cam_id: The ID of the camera + :return CameraManager: The camera manager + :raises RequestError: If the camera is not found, or you cannot access it. See reason for more info + """ + packet = await self._generate_request() + subscribe = AppCameraSubscribe() + subscribe.camera_id = cam_id + packet.camera_subscribe = subscribe + + response = await self.ws.send_and_get(packet) + + if response is None: + return None + + return CameraManager(self, cam_id, response.response.camera_subscribe_info) diff --git a/rustplus/api/structures/__init__.py b/rustplus/structs/__init__.py similarity index 100% rename from rustplus/api/structures/__init__.py rename to rustplus/structs/__init__.py diff --git a/rustplus/api/structures/rust_chat_message.py b/rustplus/structs/rust_chat_message.py similarity index 97% rename from rustplus/api/structures/rust_chat_message.py rename to rustplus/structs/rust_chat_message.py index bc3cef5..c63d228 100644 --- a/rustplus/api/structures/rust_chat_message.py +++ b/rustplus/structs/rust_chat_message.py @@ -30,7 +30,7 @@ def colour(self) -> str: def time(self) -> int: return self._time - def __repr__(self): + def __str__(self): return "RustChatMessage[steam_id={}, sender_name={}, message={}, colour={}, time={}]".format( self._steam_id, self._name, self._message, self._colour, self._time ) diff --git a/rustplus/api/structures/rust_contents.py b/rustplus/structs/rust_contents.py similarity index 100% rename from rustplus/api/structures/rust_contents.py rename to rustplus/structs/rust_contents.py diff --git a/rustplus/api/structures/rust_entity_info.py b/rustplus/structs/rust_entity_info.py similarity index 96% rename from rustplus/api/structures/rust_entity_info.py rename to rustplus/structs/rust_entity_info.py index 563f55e..7f3faba 100644 --- a/rustplus/api/structures/rust_entity_info.py +++ b/rustplus/structs/rust_entity_info.py @@ -1,6 +1,6 @@ from typing import List from .serialization import Serializable -from ..remote import AppEntityInfo, AppEntityPayloadItem +from ..remote.rustplus_proto import AppEntityInfo, AppEntityPayloadItem class RustEntityInfoItem(Serializable): diff --git a/rustplus/api/structures/rust_info.py b/rustplus/structs/rust_info.py similarity index 100% rename from rustplus/api/structures/rust_info.py rename to rustplus/structs/rust_info.py diff --git a/rustplus/api/structures/rust_item.py b/rustplus/structs/rust_item.py similarity index 100% rename from rustplus/api/structures/rust_item.py rename to rustplus/structs/rust_item.py diff --git a/rustplus/api/structures/rust_map.py b/rustplus/structs/rust_map.py similarity index 100% rename from rustplus/api/structures/rust_map.py rename to rustplus/structs/rust_map.py diff --git a/rustplus/api/structures/rust_marker.py b/rustplus/structs/rust_marker.py similarity index 100% rename from rustplus/api/structures/rust_marker.py rename to rustplus/structs/rust_marker.py diff --git a/rustplus/api/structures/rust_team_info.py b/rustplus/structs/rust_team_info.py similarity index 100% rename from rustplus/api/structures/rust_team_info.py rename to rustplus/structs/rust_team_info.py diff --git a/rustplus/api/structures/rust_time.py b/rustplus/structs/rust_time.py similarity index 100% rename from rustplus/api/structures/rust_time.py rename to rustplus/structs/rust_time.py diff --git a/rustplus/api/structures/serialization.py b/rustplus/structs/serialization.py similarity index 100% rename from rustplus/api/structures/serialization.py rename to rustplus/structs/serialization.py diff --git a/rustplus/api/structures/util.py b/rustplus/structs/util.py similarity index 100% rename from rustplus/api/structures/util.py rename to rustplus/structs/util.py diff --git a/rustplus/utils/__init__.py b/rustplus/utils/__init__.py index f0acb2b..82fc751 100644 --- a/rustplus/utils/__init__.py +++ b/rustplus/utils/__init__.py @@ -1,6 +1,13 @@ -from .rust_utils import * from .deprecated import deprecated -from .grab_items import translate_id_to_stack -from .server_id import ServerID +from .utils import ( + convert_time, + generate_grid, + fetch_avatar_icon, + format_coord, + convert_marker, + convert_monument, + convert_event_type_to_name, +) +from .grab_items import translate_stack_to_id, translate_id_to_stack from .yielding_event import YieldingEvent from .emojis import Emoji diff --git a/rustplus/utils/emojis.py b/rustplus/utils/emojis.py index cdb55c8..333ce7a 100644 --- a/rustplus/utils/emojis.py +++ b/rustplus/utils/emojis.py @@ -2,21 +2,21 @@ class Emoji(Enum): - ANGRY = ":angry+0:" - COFFEE_CAN = ":coffeecan+0:" - EYEBROW = ":eyebrow+0:" - FUNNY = ":funny+0:" - HAPPY = ":happy+0:" - LAUGH = ":laugh+0:" - LOVE = ":love+0:" - MASK = ":mask+0:" - NERVOUS = ":nervous+0:" - NEUTRAL = ":neutral+0:" - HEART_ROCK = ":heartrock+0:" - SMILE_CRY = ":smilecry+0:" - COOL = ":cool+0:" - LIGHT = ":light+0:" - WORRIED = ":worried+0:" + ANGRY = ":angry:" + COFFEE_CAN = ":coffeecan:" + EYEBROW = ":eyebrow:" + FUNNY = ":funny:" + HAPPY = ":happy:" + LAUGH = ":laugh:" + LOVE = ":love:" + MASK = ":mask:" + NERVOUS = ":nervous:" + NEUTRAL = ":neutral:" + HEART_ROCK = ":heartrock:" + SMILE_CRY = ":smilecry:" + COOL = ":cool:" + LIGHT = ":light:" + WORRIED = ":worried:" HEART = ":heart:" SKULL = ":skull:" EYES = ":eyes:" @@ -25,7 +25,7 @@ class Emoji(Enum): TRUMPET = ":trumpet:" WAVE = ":wave:" YELLOW_PIN = ":yellowpin:" - SCIENTIST = ":scientist+0:" + SCIENTIST = ":scientist:" def __str__(self): return self.value diff --git a/rustplus/utils/grab_items.py b/rustplus/utils/grab_items.py index 45e9976..ea95ee0 100644 --- a/rustplus/utils/grab_items.py +++ b/rustplus/utils/grab_items.py @@ -13,6 +13,7 @@ "-2094954543": "Wood Armor Helmet", "-2086926071": "Potato", "-2084071424": "Potato Seed", + "-2073432256": "Skinning Knife", "-2072273936": "Bandage", "-2069578888": "M249", "-2067472972": "Sheet Metal Door", @@ -31,6 +32,7 @@ "-1997543660": "Horse Saddle", "-1994909036": "Sheet Metal", "-1992717673": "Large Furnace", + "-1989600732": "Hot Air Balloon Armor", "-1985799200": "Rug", "-1982036270": "High Quality Metal Ore", "-1978999529": "Salvaged Cleaver", @@ -42,6 +44,7 @@ "-1941646328": "Can of Tuna", "-1938052175": "Charcoal", "-1916473915": "Chinese Lantern", + "-1913996738": "Fish Trophy", "-1904821376": "Orange Roughy", "-1903165497": "Bone Helmet", "-1899491405": "Glue", @@ -50,6 +53,7 @@ "-1878764039": "Small Trout", "-1878475007": "Satchel Charge", "-1863559151": "Water Barrel", + "-1863063690": "Rocking Chair", "-1861522751": "Research Table", "-1850571427": "Silencer", "-1848736516": "Cooked Chicken", @@ -80,6 +84,7 @@ "-1709878924": "Raw Human Meat", "-1707425764": "Fishing Tackle", "-1698937385": "Herring", + "-1696379844": "Hazmat Youtooz", "-1695367501": "Shorts", "-1693832478": "Large Flatbed Vehicle Module", "-1691396643": "HV Pistol Ammo", @@ -110,6 +115,7 @@ "-1549739227": "Boots", "-1539025626": "Miners Hat", "-1538109120": "Violet Volcano Firework", + "-1536855921": "Shovel", "-1535621066": "Stone Fireplace", "-1530414568": "Cassette Recorder", "-1520560807": "Raw Bear Meat", @@ -130,16 +136,19 @@ "-1469578201": "Longsword", "-1448252298": "Electrical Branch", "-1442559428": "Hobo Barrel", + "-1442496789": "Pinata", "-1440987069": "Raw Chicken Breast", "-1432674913": "Anti-Radiation Pills", "-1429456799": "Prison Cell Wall", "-1423304443": "Medium Neon Sign", + "-1421257350": "Storage Barrel Horizontal", "-1405508498": "Muzzle Boost", "-1379835144": "Festive Window Garland", "-1379036069": "Canbourine", "-1370759135": "Portrait Picture Frame", "-1368584029": "Sickle", "-1367281941": "Waterpipe Shotgun", + "-1344017968": "Wanted Poster", "-1336109173": "Wood Double Door", "-1331212963": "Star Tree Topper", "-1330640246": "Junkyard Drum Kit", @@ -152,6 +161,7 @@ "-1293296287": "Small Oil Refinery", "-1286302544": "OR Switch", "-1284169891": "Water Pump", + "-1274093662": "Bath Tub Planter", "-1273339005": "Bed", "-1266045928": "Bunny Onesie", "-1262185308": "Binoculars", @@ -174,6 +184,7 @@ "-1162759543": "Cooked Horse Meat", "-1160621614": "Red Industrial Wall Light", "-1157596551": "Sulfur Ore", + "-1151332840": "Wooden Frontier Bar Doors", "-1138208076": "Small Wooden Sign", "-1137865085": "Machete", "-1130709577": "Pump Jack", @@ -217,6 +228,7 @@ "-965336208": "Chocolate Bar", "-961457160": "New Year Gong", "-956706906": "Prison Cell Gate", + "-948291630": "Seismic Sensor", "-946369541": "Low Grade Fuel", "-939424778": "Flasher Light", "-936921910": "Flashbang", @@ -224,15 +236,20 @@ "-930193596": "Fertilizer", "-929092070": "Basic Healing Tea", "-912398867": "Cassette - Medium", + "-907422733": "Large Backpack", "-904863145": "Semi-Automatic Rifle", + "-901370585": "Twitch Rivals Trophy 2023", "-888153050": "Halloween Candy", "-886280491": "Hemp Clone", + "-885833256": "Vampire Stake", + "-869598982": "Small Hunting Trophy", "-858312878": "Cloth", "-855748505": "Simple Handmade Sight", "-854270928": "Dragon Door Knocker", "-852563019": "M92 Pistol", "-851988960": "Salmon", "-850982208": "Key Lock", + "-849373693": "Frontier Horseshoe Single Item Rack", "-845557339": "Landscape Picture Frame", "-842267147": "Snowman Helmet", "-819720157": "Metal Window Bars", @@ -244,6 +261,7 @@ "-781014061": "Sprinkler", "-778875547": "Corn Clone", "-778367295": "L96 Rifle", + "-770304148": "Chinese Lantern White", "-769647921": "Skull Trophy", "-765183617": "Double Barrel Shotgun", "-763071910": "Lumberjack Hoodie", @@ -255,6 +273,7 @@ "-742865266": "Rocket", "-733625651": "Paddling Pool", "-727717969": "12 Gauge Slug", + "-722629980": "Heavy Scientist Youtooz", "-722241321": "Small Present", "-702051347": "Bandana Mask", "-700591459": "Can of Beans", @@ -312,6 +331,7 @@ "-363689972": "Snowball", "-343857907": "Sound Light", "-335089230": "High External Wooden Gate", + "-333406828": "Sled", "-324675402": "Reindeer Antlers", "-321733511": "Crude Oil", "-321431890": "Beach Chair", @@ -323,9 +343,12 @@ "-265876753": "Gun Powder", "-265292885": "Fluid Combiner", "-262590403": "Salvaged Axe", + "-258574361": "Dracula Cape", "-253079493": "Scientist Suit", + "-246672609": "Horizontal Weapon Rack", "-242084766": "Cooked Pork", "-237809779": "Hemp Seed", + "-218009552": "Homing Missile Launcher", "-216999575": "Counter", "-216116642": "Skull Door Knocker", "-211235948": "Xylobone", @@ -340,6 +363,7 @@ "-173268128": "Rustig\u00e9 Egg - White", "-173268126": "Rustig\u00e9 Egg - Ivory", "-173268125": "Rustig\u00e9 Egg - Green", + "-152332823": "Chicken Costume", "-151838493": "Wood", "-151387974": "Deluxe Christmas Lights", "-148794216": "Garage Door", @@ -347,7 +371,6 @@ "-144513264": "Pipe Tool", "-144417939": "Wire Tool", "-143132326": "Huge Wooden Sign", - "-135252633": "Sled", "-134959124": "Light Frankenstein Head", "-132516482": "Weapon Lasersight", "-132247350": "Small Water Catcher", @@ -359,10 +382,12 @@ "-99886070": "Violet Roman Candle", "-97956382": "Tool Cupboard", "-97459906": "Jumpsuit", + "-96256997": "Wide Weapon Rack", "-92759291": "Wooden Floor Spikes", "-89874794": "Low Quality Spark Plugs", "-78533081": "Burnt Deer Meat", "-75944661": "Eoka Pistol", + "-52398594": "Frontier Horns Single Item Rack", "-48090175": "Snow Jacket", "-44876289": "Igniter", "-41896755": "Workbench Level 2", @@ -412,9 +437,11 @@ "200773292": "Hammer", "204391461": "Coal :(", "204970153": "Wrapped Gift", + "209218760": "Head Bag", "215754713": "Bone Arrow", "223891266": "T-Shirt", "237239288": "Pants", + "240752557": "Tall Weapon Rack", "254522515": "Large Medkit", "261913429": "White Volcano Firework", "263834859": "Basic Scrap Tea", @@ -442,9 +469,11 @@ "442289265": "Holosight", "442886268": "Rocket Launcher", "443432036": "Fluid Switch & Pump", + "446206234": "Torch Holder", "476066818": "Cassette - Long", "479143914": "Gears", "479292118": "Large Loot Bag", + "486661382": "Clan Table", "492357192": "RAND Switch", "524678627": "Advanced Scrap Tea", "528668503": "Flame Turret", @@ -459,12 +488,15 @@ "588596902": "Handmade Shell", "593465182": "Table", "596469572": "RF Transmitter", + "602628465": "Parachute", "602741290": "Burlap Shirt", "603811464": "Advanced Max Health Tea", "605467368": "Incendiary 5.56 Rifle Ammo", + "607400343": "Legacy Wood Shelter", "609049394": "Battery - Small", "610102428": "Industrial Conveyor", "613961768": "Bota Bag", + "615112838": "Rail Road Planter", "621915341": "Raw Pork", "634478325": "CCTV Camera", "642482233": "Sticks", @@ -477,9 +509,11 @@ "671063303": "Riot Helmet", "671706427": "Reinforced Glass Window", "674734128": "Festive Doorway Garland", + "678698219": "M4 Shotgun", "680234026": "Yellow Perch", "696029452": "Paper Map", "699075597": "Wooden Cross", + "703057617": "Military Flame Thrower", "709206314": "Tiger Mask", "722955039": "Water Gun", "742745918": "Industrial Splitter", @@ -518,9 +552,11 @@ "895374329": "Passenger Vehicle Module", "915408809": "40mm Smoke Grenade", "926800282": "Medium Quality Valves", + "935606207": "Minigun", "935692442": "Longsleeve T-Shirt", "936496778": "Floor grill", "952603248": "Weapon flashlight", + "960673498": "Large Hunting Trophy", "963906841": "Rock", "968019378": "Clatter Helmet", "968421290": "Connected Speaker", @@ -546,6 +582,7 @@ "1110385766": "Metal Chest Plate", "1112162468": "Blue Berry", "1121925526": "Candy Cane", + "1132603396": "Weapon Rack Stand", "1142993169": "Ceiling Light", "1149964039": "Storage Monitor", "1153652756": "Large Wooden Sign", @@ -554,6 +591,7 @@ "1158340334": "Low Quality Crankshaft", "1159991980": "Code Lock", "1160881421": "Hitch & Trough", + "1168856825": "Metal Detector", "1171735914": "AND Switch", "1177596584": "Elevator", "1181207482": "Heavy Plate Helmet", @@ -578,7 +616,9 @@ "1272430949": "Wheelbarrow Piano", "1272768630": "Spoiled Human Meat", "1293102274": "XOR Switch", + "1296788329": "Homing Missile", "1305578813": "Small Neon Sign", + "1307626005": "Storage Barrel Vertical", "1315082560": "Ox Mask", "1318558775": "MP5A4", "1319617282": "Small Loot Bag", @@ -589,6 +629,7 @@ "1346158228": "Pumpkin Bucket", "1353298668": "Armored Door", "1358643074": "Snow Machine", + "1361520181": "Minecart Planter", "1366282552": "Leather Gloves", "1367190888": "Corn", "1371909803": "Tesla Coil", @@ -644,6 +685,7 @@ "1588492232": "Drone", "1601468620": "Blue Jumpsuit", "1602646136": "Stone Spear", + "1603174987": "Confetti Cannon", "1608640313": "Tank Top", "1614528785": "Heavy Frankenstein Torso", "1623701499": "Industrial Wall Light", @@ -693,6 +735,7 @@ "1849887541": "Small Generator", "1850456855": "Road Sign Kilt", "1856217390": "Egg Basket", + "1865253052": "Dracula Mask", "1873897110": "Cooked Bear Meat", "1874610722": "Armored Cockpit Vehicle Module", "1877339384": "Burlap Headwrap", @@ -718,6 +761,7 @@ "1965232394": "Crossbow", "1973165031": "Birthday Cake", "1973684065": "Burnt Chicken", + "1973949960": "Frontier Bolts Single Item Rack", "1975934948": "Survey Charge", "1983621560": "Floor triangle grill", "1989785143": "High Quality Horse Shoes", @@ -732,6 +776,7 @@ "2041899972": "Triangle Ladder Hatch", "2048317869": "Wolf Skull", "2063916636": "Advanced Ore Tea", + "2068884361": "Small Backpack", "2070189026": "Large Banner on pole", "2087678962": "Search Light", "2090395347": "Large Solar Panel", diff --git a/rustplus/utils/server_id.py b/rustplus/utils/server_id.py deleted file mode 100644 index 699f5a3..0000000 --- a/rustplus/utils/server_id.py +++ /dev/null @@ -1,26 +0,0 @@ -class ServerID: - def __init__(self, ip, port, player_id, player_token) -> None: - self.ip = ip - self.port = port - self.player_id = player_id - self.player_token = player_token - - def __str__(self) -> str: - return f"{self.ip}:{self.port} {self.player_id} {self.player_token}" - - def get_server_string(self) -> str: - return f"{self.ip}:{self.port}" - - def __hash__(self): - return hash(self.__str__()) - - def __eq__(self, o: object) -> bool: - if not isinstance(o, ServerID): - return False - - return ( - self.ip == o.ip - and self.port == o.port - and self.player_id == o.player_id - and self.player_token == o.player_token - ) diff --git a/rustplus/utils/rust_utils.py b/rustplus/utils/utils.py similarity index 57% rename from rustplus/utils/rust_utils.py rename to rustplus/utils/utils.py index f2520a3..7680e3d 100644 --- a/rustplus/utils/rust_utils.py +++ b/rustplus/utils/utils.py @@ -1,60 +1,133 @@ -from importlib import resources -from typing import Tuple -from PIL import Image, ImageDraw, ImageFont import logging import string +from importlib import resources +from typing import Tuple, Dict -from ..api.remote.rustplus_proto import AppMessage -from ..api.structures import RustTime +import requests +from PIL import ImageFont, Image, ImageDraw -ICONS_PATH = "rustplus.api.icons" +ICONS_PATH = "rustplus.icons" FONT_PATH = "rustplus.utils.fonts" GRID_DIAMETER = 146.28571428571428 - PLAYER_MARKER_ONLINE_COLOR = (201, 242, 155, 255) PLAYER_MARKER_OFFLINE_COLOR = (128, 128, 128, 255) -def format_time(protobuf: AppMessage) -> RustTime: - def convert_time(time) -> str: - hours, minutes = divmod(time * 60, 60) +def convert_time(time) -> str: + hours, minutes = divmod(time * 60, 60) - return ( - f"{int(hours)}:0{int(minutes)}" - if minutes <= 9 - else f"{int(hours)}:{int(minutes)}" - ) + return ( + f"{int(hours)}:0{int(minutes)}" + if minutes <= 9 + else f"{int(hours)}:{int(minutes)}" + ) + + +def convert_event_type_to_name(event: int) -> str: + if event == 1: + return "Player" + elif event == 2: + return "Explosion" + elif event == 3: + return "Vending Machine" + elif event == 4: + return "CH47 Chinook" + elif event == 5: + return "Cargo Ship" + elif event == 6: + return "Locked Crate" + elif event == 7: + return "Generic Radius" + elif event == 8: + return "Patrol Helicopter" + + +def generate_grid( + map_size: int, + text_size: int = 20, + text_padding: int = 5, + color: str = "black", +) -> Image.Image: + img = Image.new("RGBA", (map_size, map_size), (0, 0, 0, 0)) + d = ImageDraw.Draw(img) + + with resources.path(FONT_PATH, "PermanentMarker.ttf") as path: + font = ImageFont.truetype(str(path), text_size) - sunrise = convert_time(protobuf.response.time.sunrise) - sunset = convert_time(protobuf.response.time.sunset) - parsed_time = convert_time(protobuf.response.time.time) - - return RustTime( - protobuf.response.time.day_length_minutes, - sunrise, - sunset, - parsed_time, - protobuf.response.time.time, - protobuf.response.time.time_scale, + letters = list(string.ascii_uppercase) + letters.extend( + a + b for a in string.ascii_uppercase for b in string.ascii_uppercase ) + num_cells = int(map_size / GRID_DIAMETER) -def format_coord(x, y, map_size) -> Tuple[int, int]: - y = map_size - y - 75 - x -= 75 + for i in range(num_cells): + for j in range(num_cells): + start = (i * GRID_DIAMETER, j * GRID_DIAMETER) + end = ((i + 1) * GRID_DIAMETER, (j + 1) * GRID_DIAMETER) + d.rectangle((start, end), outline=color) - if x < 0: - x = 0 - if x > map_size: - x = map_size - 150 - if y < 0: - y = 0 - if y > map_size: - y = map_size - 150 + text = letters[i] + str(j) + text_pos = (start[0] + text_padding, start[1] + text_padding) + d.text(text_pos, text, fill=color, font=font) + + return img + + +def format_coord(x: int, y: int, map_size: int) -> Tuple[int, int]: + # Adjust y and x coordinates with offsets + y = max(0, min(map_size - y - 75, map_size - 150)) + x = max(0, min(x - 75, map_size - 150)) return x, y +async def fetch_avatar_icon(steam_id: int, online: bool) -> Image.Image: + avatar = ( + Image.open( + requests.get( + f"https://companion-rust.facepunch.com/api/avatar/{steam_id}", + stream=True, + ).raw + ) + .resize((100, 100), Image.LANCZOS) + .convert("RGBA") + ) + + return await avatar_processing(avatar, 5, online) + + +async def avatar_processing( + image: Image.Image, border_size: int, player_online: bool = False +) -> Image.Image: + size_with_border = ( + image.size[0] + 2 * border_size, + image.size[1] + 2 * border_size, + ) + + border_image = Image.new("RGBA", size_with_border, (0, 0, 0, 0)) + + mask = Image.new("L", size_with_border, 0) + draw = ImageDraw.Draw(mask) + + draw.ellipse([0, 0, size_with_border[0], size_with_border[1]], fill=255) + + border_layer = Image.new( + "RGBA", + size_with_border, + PLAYER_MARKER_ONLINE_COLOR if player_online else PLAYER_MARKER_OFFLINE_COLOR, + ) + border_image.paste(border_layer, mask=mask) + + image_mask = Image.new("L", image.size, 0) + draw = ImageDraw.Draw(image_mask) + draw.ellipse([0, 0, image.size[0], image.size[1]], fill=255) + + border_image.paste(image, (border_size, border_size), image_mask) + + return border_image + + def convert_marker(name, angle) -> Image.Image: name_to_file = { "2": "explosion.png", @@ -66,27 +139,31 @@ def convert_marker(name, angle) -> Image.Image: with resources.path(ICONS_PATH, name_to_file[name]) as path: icon = Image.open(path).convert("RGBA") + if name == "6": icon = icon.resize((85, 85)) elif name == "2": icon = icon.resize((96, 96)) - elif name == "4": - with resources.path(ICONS_PATH, "chinook_blades.png") as path: - blades = Image.open(path).convert("RGBA") - blades = blades.resize((100, 100)) - icon.paste(blades, (64 - 50, 96 - 50), blades) - icon.paste(blades, (64 - 50, 32 - 50), blades) - elif name == "8": - icon = icon.resize((200, 200)) - with resources.path(ICONS_PATH, "chinook_blades.png") as path: + elif name in ["4", "8"]: + blades_file = "chinook_blades.png" + blades_size = (100, 100) if name == "4" else (200, 200) + + with resources.path(ICONS_PATH, blades_file) as path: blades = Image.open(path).convert("RGBA") - blades = blades.resize((200, 200)) - icon.paste(blades, (0, 0), blades) + blades = blades.resize(blades_size) + + if name == "4": + icon.paste(blades, (64 - 50, 96 - 50), blades) + icon.paste(blades, (64 - 50, 32 - 50), blades) + else: + icon = icon.resize((200, 200)) + icon.paste(blades, (0, 0), blades) + icon = icon.rotate(angle) return icon -def convert_monument(name: str, override_images: dict) -> Image.Image: +def convert_monument(name: str, override_images: Dict[str, Image.Image]) -> Image.Image: name_to_file = { "supermarket": "supermarket.png", "mining_outpost_display_name": "mining_outpost.png", @@ -144,154 +221,5 @@ def convert_monument(name: str, override_images: dict) -> Image.Image: ) with resources.path(ICONS_PATH, "icon.png") as path: icon = Image.open(path).convert("RGBA") - return icon - - -def entity_type_to_string(id) -> str: - if id == 1: - return "Switch" - elif id == 2: - return "Alarm" - elif id == 3: - return "Storage Monitor" - else: - raise ValueError("Not Valid type") - - -def _get_grid_x(x): - counter = 1 - start_grid = 0 - while start_grid < x + GRID_DIAMETER: - if start_grid <= x <= (start_grid + GRID_DIAMETER): - # We're at the correct grid! - return _number_to_letters(counter) - counter += 1 - start_grid += GRID_DIAMETER - - -def _get_grid_y(y, map_size): - counter = 1 - number_of_grids = map_size // GRID_DIAMETER - start_grid = 0 - while start_grid < y + GRID_DIAMETER: - if start_grid <= y <= (start_grid + GRID_DIAMETER): - return number_of_grids - counter - counter += 1 - start_grid += GRID_DIAMETER - - -def _number_to_letters(num): - power, mod = divmod(num, 26) - out = chr(64 + mod) if mod else (power, "Z") - return _number_to_letters(power) + out if power else out - - -def _get_corrected_map_size(map_size): - remainder = map_size % GRID_DIAMETER - offset = GRID_DIAMETER - remainder - return map_size - remainder if remainder < 120 else map_size + offset - - -def _is_outside_grid_system(x, y, map_size, offset=0): - return ( - x < -offset or x > (map_size + offset) or y < -offset or y > (map_size + offset) - ) - - -class HackyBackwardsCompatCoordClass: - def __init__(self, x, y): - self.x = x - self.y = y - - def __getitem__(self, item): - if item == 0: - return self.x - elif item == 1: - return self.y - else: - raise IndexError("Index out of range") - - def __iter__(self): - yield self.x - yield self.y - - def __repr__(self): - return f"{self.x}{self.y}" - - def __str__(self): - return self.__repr__() - - -def convert_xy_to_grid( - coords: tuple, map_size: float, catch_out_of_bounds: bool = True -) -> HackyBackwardsCompatCoordClass: - corrected_map_size = _get_corrected_map_size(map_size) - - grid_pos_letters = _get_grid_x(coords[0]) - grid_pos_number = str(int(_get_grid_y(coords[1], corrected_map_size))) - - return HackyBackwardsCompatCoordClass(grid_pos_letters, grid_pos_number) - - -def generate_grid( - map_size: int, - text_size: int = 20, - text_padding: int = 5, - color: str = "black", -) -> Image.Image: - img = Image.new("RGBA", (map_size, map_size), (0, 0, 0, 0)) - d = ImageDraw.Draw(img) - - with resources.path(FONT_PATH, "PermanentMarker.ttf") as path: - font = ImageFont.truetype(str(path), text_size) - - letters = list(string.ascii_uppercase) - letters.extend( - a + b for a in string.ascii_uppercase for b in string.ascii_uppercase - ) - - num_cells = int(map_size / GRID_DIAMETER) - - for i in range(num_cells): - for j in range(num_cells): - start = (i * GRID_DIAMETER, j * GRID_DIAMETER) - end = ((i + 1) * GRID_DIAMETER, (j + 1) * GRID_DIAMETER) - d.rectangle((start, end), outline=color) - - text = letters[i] + str(j) - text_pos = (start[0] + text_padding, start[1] + text_padding) - d.text(text_pos, text, fill=color, font=font) - - return img - - -def avatar_processing( - image: Image.Image, border_size: int, player_online: bool = False -) -> Image.Image: - size_with_border = ( - image.size[0] + 2 * border_size, - image.size[1] + 2 * border_size, - ) - - border_image = Image.new("RGBA", size_with_border, (0, 0, 0, 0)) - - mask = Image.new("L", size_with_border, 0) - draw = ImageDraw.Draw(mask) - - draw.ellipse([0, 0, size_with_border[0], size_with_border[1]], fill=255) - - border_layer = Image.new( - "RGBA", - size_with_border, - PLAYER_MARKER_ONLINE_COLOR if player_online else PLAYER_MARKER_OFFLINE_COLOR, - ) - border_image.paste(border_layer, mask=mask) - - image_mask = Image.new("L", image.size, 0) - draw = ImageDraw.Draw(image_mask) - draw.ellipse([0, 0, image.size[0], image.size[1]], fill=255) - - border_image.paste(image, (border_size, border_size), image_mask) - - return border_image + return icon diff --git a/rustplus/utils/yielding_event.py b/rustplus/utils/yielding_event.py index a10d21e..2fa8500 100644 --- a/rustplus/utils/yielding_event.py +++ b/rustplus/utils/yielding_event.py @@ -1,6 +1,6 @@ import asyncio import contextlib -from typing import Any, Union +from typing import Any, Union, Optional class YieldingEvent(asyncio.Event): @@ -16,12 +16,14 @@ def clear(self) -> None: self.value = None super().clear() - async def wait(self) -> Any: - await super().wait() - return self.value + async def wait(self, timeout: Optional[float] = None) -> Any: + if timeout is not None: + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(super().wait(), timeout) + else: + await super().wait() - async def event_wait_for(self, timeout) -> Any: - # suppress TimeoutError because we'll return False in case of timeout - with contextlib.suppress(asyncio.TimeoutError): - await asyncio.wait_for(self.wait(), timeout) return self.value if self.is_set() else None + + async def event_wait_for(self, timeout: float) -> Any: + return await self.wait(timeout)