diff --git a/rustplus/annotations/__init__.py b/rustplus/annotations/__init__.py new file mode 100644 index 0000000..fab6e56 --- /dev/null +++ b/rustplus/annotations/__init__.py @@ -0,0 +1 @@ +from .command import Command diff --git a/rustplus/annotations/command.py b/rustplus/annotations/command.py new file mode 100644 index 0000000..ceedfc3 --- /dev/null +++ b/rustplus/annotations/command.py @@ -0,0 +1,19 @@ +from typing import Callable + +from ..identification import RegisteredListener, ServerID +from ..commands import ChatCommand, ChatCommandData + + +def Command(server_id: ServerID, aliases: list = None, alias_func: Callable = None): + + 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_id][func.__name__] = command_data + + return RegisteredListener(func.__name__, func) + + return wrapper diff --git a/rustplus/commands/__init__.py b/rustplus/commands/__init__.py new file mode 100644 index 0000000..b45ef7c --- /dev/null +++ b/rustplus/commands/__init__.py @@ -0,0 +1,3 @@ +from .chat_command_data import ChatCommandData +from .chat_command import ChatCommand, ChatCommandTime +from .command_options import CommandOptions diff --git a/rustplus/commands/chat_command.py b/rustplus/commands/chat_command.py new file mode 100644 index 0000000..4a6d67b --- /dev/null +++ b/rustplus/commands/chat_command.py @@ -0,0 +1,31 @@ +import dataclasses +from collections import defaultdict +from typing import Dict, List + +from .chat_command_data import ChatCommandData +from ..identification import ServerID + + +@dataclasses.dataclass +class ChatCommandTime: + formatted_time: str + raw_time: int + + +class ChatCommand: + + REGISTERED_COMMANDS: Dict[ServerID, 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/chat_command_data.py b/rustplus/commands/chat_command_data.py new file mode 100644 index 0000000..eaddd04 --- /dev/null +++ b/rustplus/commands/chat_command_data.py @@ -0,0 +1,23 @@ +from typing import Callable + + +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): + if self._aliases is None: + return [] + + return self._aliases + + @property + def callable_func(self): + if self._callable_func is None: + return lambda x: False + + return self._callable_func diff --git a/rustplus/commands/command_options.py b/rustplus/commands/command_options.py new file mode 100644 index 0000000..2c37456 --- /dev/null +++ b/rustplus/commands/command_options.py @@ -0,0 +1,11 @@ +from ..exceptions import PrefixNotDefinedError + + +class CommandOptions: + def __init__( + self, prefix: str = None + ) -> None: + if prefix is None: + raise PrefixNotDefinedError("No prefix") + + self.prefix = prefix 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/remote/websocket/ws.py b/rustplus/remote/websocket/ws.py index 2c35cb3..2fcb9f5 100644 --- a/rustplus/remote/websocket/ws.py +++ b/rustplus/remote/websocket/ws.py @@ -1,3 +1,4 @@ +import shlex import betterproto from websockets.exceptions import InvalidURI, InvalidHandshake from websockets.legacy.client import WebSocketClientProtocol @@ -8,15 +9,17 @@ import asyncio from ..rustplus_proto import AppMessage, AppRequest -from ...exceptions import ClientNotConnectedError +from ...commands import CommandOptions, ChatCommand, ChatCommandTime +from ...exceptions import ClientNotConnectedError, RequestError from ...identification import ServerID from ...structs import RustChatMessage -from ...utils.yielding_event import YieldingEvent +from ...utils import YieldingEvent, convert_time class RustWebsocket: - def __init__(self, server_id: ServerID) -> None: + def __init__(self, server_id: ServerID, command_options: Union[None, CommandOptions]) -> None: self.server_id: ServerID = server_id + self.command_options: Union[None, CommandOptions] = command_options self.connection: Union[WebSocketClientProtocol, None] = None self.logger: logging.Logger = logging.getLogger("rustplus.py") self.task: Union[Task, None] = None @@ -107,7 +110,7 @@ async def handle_message(self, app_message: AppMessage) -> None: ) if self.error_present(app_message.response.error.error): - raise Exception(app_message.response.error.error) + raise RequestError(app_message.response.error.error) prefix = self.get_prefix( str(app_message.broadcast.team_message.message.message) @@ -122,7 +125,35 @@ async def handle_message(self, app_message: AppMessage) -> None: ) message = RustChatMessage(app_message.broadcast.team_message.message) - # TODO await self.remote.command_handler.run_command(message, prefix) + + parts = shlex.split(message.message) + command = parts[0][len(prefix) :] + + data = ChatCommand.REGISTERED_COMMANDS[self.server_id].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_id + ].items(): + if ( + command in data.aliases + or data.callable_func(command) + ): + await data.coroutine(dao) + break if self.is_entity_broadcast(app_message): # This means that an entity has changed state @@ -178,20 +209,14 @@ async def handle_message(self, app_message: AppMessage) -> None: def get_prefix(self, message: str) -> Optional[str]: - return None + if self.command_options is None: + return None - if self.remote.use_commands: - if message.startswith(self.remote.command_options.prefix): - return self.remote.command_options.prefix + if message.startswith(self.command_options.prefix): + return self.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( diff --git a/rustplus/rust_api.py b/rustplus/rust_api.py index a9ca60b..ab81af7 100644 --- a/rustplus/rust_api.py +++ b/rustplus/rust_api.py @@ -5,6 +5,7 @@ import logging from PIL import Image +from .commands import CommandOptions from .exceptions import RateLimitError from .identification import ServerID from .remote.camera import CameraManager @@ -27,9 +28,10 @@ class RustSocket: def __init__( - self, server_id: ServerID, ratelimiter: Union[None, RateLimiter] = None + self, server_id: ServerID, ratelimiter: Union[None, RateLimiter] = None, command_options: Union[None, CommandOptions] = None ) -> None: self.server_id = server_id + self.command_options = command_options self.logger = logging.getLogger("rustplus.py") console_handler = logging.StreamHandler() @@ -40,7 +42,7 @@ def __init__( self.logger.addHandler(console_handler) self.logger.setLevel(logging.DEBUG) - self.ws = RustWebsocket(self.server_id) + self.ws = RustWebsocket(self.server_id, self.command_options) self.seq = 1 if ratelimiter: @@ -79,6 +81,7 @@ async def __generate_request(self, tokens=1) -> AppRequest: async def connect(self) -> None: await self.ws.connect() + await self.get_time() # Wake up the connection @staticmethod async def hang() -> None: diff --git a/rustplus/utils/__init__.py b/rustplus/utils/__init__.py index 65d8243..4a9099f 100644 --- a/rustplus/utils/__init__.py +++ b/rustplus/utils/__init__.py @@ -1,3 +1,4 @@ from .deprecated import deprecated from .utils import convert_time from .grab_items import translate_stack_to_id, translate_id_to_stack +from .yielding_event import YieldingEvent