Add Commands

rustplus/annotations/__init__.py

@@ -0,0 +1 @@
+from .command import Command

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

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

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

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

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

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(

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:

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