1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
|
"""BattlEye RCon client."""
from logging import getLogger
from socket import SOCK_DGRAM
from typing import Callable
from rcon.battleye.proto import HEADER_SIZE
from rcon.battleye.proto import RESPONSE_TYPES
from rcon.battleye.proto import CommandRequest
from rcon.battleye.proto import CommandResponse
from rcon.battleye.proto import Header
from rcon.battleye.proto import LoginRequest
from rcon.battleye.proto import LoginResponse
from rcon.battleye.proto import Request
from rcon.battleye.proto import Response
from rcon.battleye.proto import ServerMessage
from rcon.battleye.proto import ServerMessageAck
from rcon.client import BaseClient
from rcon.exceptions import WrongPassword
__all__ = ["Client"]
MessageHandler = Callable[[ServerMessage], None]
def log_message(server_message: ServerMessage) -> None:
"""Default handler, logging the server message."""
getLogger("Server message").info(server_message.message)
class Client(BaseClient, socket_type=SOCK_DGRAM):
"""BattlEye RCon client."""
def __init__(
self,
*args,
max_length: int = 4096,
message_handler: MessageHandler = log_message,
**kwargs,
):
super().__init__(*args, **kwargs)
self.max_length = max_length
self.message_handler = message_handler
def handle_server_message(self, message: ServerMessage) -> None:
"""Handle the respective server message."""
with self._socket.makefile("wb") as file:
file.write(bytes(ServerMessageAck(message.seq)))
self.message_handler(message)
def receive(self) -> Response:
"""Receive a packet."""
return RESPONSE_TYPES[
(
header := Header.from_bytes(
(data := self._socket.recv(self.max_length))[:HEADER_SIZE]
)
).type
].from_bytes(header, data[HEADER_SIZE:])
def receive_transaction(self):
command_responses = []
login_response = None
seq = 0
while True:
response = self.receive()
if isinstance(response, LoginResponse) and login_response is None:
login_response = response
continue
if isinstance(response, CommandResponse):
command_responses.append(response)
seq = response.seq
continue
if isinstance(response, ServerMessage):
self.handle_server_message(response)
if login_response is not None:
return login_response
if len(command_responses) >= seq:
break
return "".join(
command_response.message
for command_response in sorted(command_responses, key=lambda cr: cr.seq)
)
def communicate(self, request: Request) -> Response | str:
"""Send a request and receive a response."""
with self._socket.makefile("wb") as file:
file.write(bytes(request))
return self.receive_transaction()
def login(self, passwd: str) -> bool:
"""Log-in the user."""
if not self.communicate(LoginRequest(passwd)).success:
raise WrongPassword()
return True
def run(self, command: str, *args: str) -> str:
"""Execute a command and return the text message."""
return self.communicate(CommandRequest.from_command(command, *args))
|
