WebSockets bugfix, updated Readme

Author: DarkCat09

README.md

@@ -62,9 +62,9 @@ if testserv != None:
 |v0.3|Implemented files API, added typization.|
 |v0.4|Implemented configuration API, some bugfixes.|
 |v0.5|The API was updated corresponding to new Aternos security methods. Huge thanks to [lusm554](https://github.com/lusm554).|
-|v0.6|Code refactoring, unit-tests, websocket API and session saving to prevent detecting automation access.|
-|v0.7|Full implementation of config API and Google Drive backups is planned.|
-|v0.8|Shared access API and permission management is planned.|
+|v0.6|Code refactoring, websockets API and session saving to prevent detecting automation access.|
+|v0.7|Full implementation of config and software API, unit tests and documentation is planned.|
+|v0.8|Shared access API and Google Drive backups is planned.|
 |v0.9.x|A long debugging before stable release, SemVer version code.|
 
 ## License

examples/websocket_example.py

@@ -1,3 +1,4 @@
+import asyncio
 from getpass import getpass
 from python_aternos import Client, atwss
 
@@ -10,6 +11,10 @@
 
 @socket.wssreceiver(atwss.Streams.console)
 async def console(msg):
-    print('Received: ' + msg)
+    print('Received:', msg)
 
-s.start()
+async def main():
+    s.start()
+    await socket.connect()
+
+asyncio.run(main())

python_aternos/atserver.py

@@ -17,7 +17,7 @@ class Edition(enum.IntEnum):
 class Status(enum.IntEnum):
 	off = 0
 	on = 1
-	loading = 2
+	starting = 2
 	shutdown = 3
 	unknown = 6
 	error = 7

python_aternos/atwss.py

@@ -1,6 +1,7 @@
 import enum
 import json
 import asyncio
+import logging
 import websockets
 from typing import Union, Any, Dict, Callable, Coroutine
 from typing import TYPE_CHECKING
@@ -9,12 +10,17 @@
 if TYPE_CHECKING:
 	from .atserver import AternosServer
 
-class Streams(enum.IntEnum):
-	status = 0
-	queue = 1
-	console = 2
-	ram = 3
-	tps = 4
+class Streams(enum.Enum):
+
+	status = (0,None)
+	queue = (1,None)
+	console = (2,'console')
+	ram = (3,'heap')
+	tps = (4,'tick')
+
+	def __init__(self, num:int, stream:str):
+		self.num = num
+		self.stream = stream
 
 class AternosWss:
 
@@ -23,7 +29,7 @@ def __init__(self, atserv:'AternosServer', autoconfirm:bool=False) -> None:
 		self.atserv = atserv
 		self.cookies = atserv.atconn.session.cookies
 		self.session = self.cookies['ATERNOS_SESSION']
-		self.servid = self.cookies['ATERNOS_SERVER']
+		self.servid = atserv.servid
 		self.recv = {}
 		self.autoconfirm = autoconfirm
 		self.confirmed = False
@@ -32,7 +38,7 @@ async def confirm(self) -> None:
 
 		self.atserv.confirm()
 
-	def wssreceiver(self, stream:int) -> Callable[[Callable[[Any],Coroutine[Any,Any,None]]],Any]:
+	def wssreceiver(self, stream:Streams) -> Callable[[Callable[[Any],Coroutine[Any,Any,None]]],Any]:
 		def decorator(func:Callable[[Any],Coroutine[Any,Any,None]]) -> None:
 			self.recv[stream] = func
 		return decorator
@@ -53,6 +59,30 @@ async def connect(self) -> None:
 			origin='https://aternos.org',
 			extra_headers=headers
 		)
+
+		@self.wssreceiver(Streams.status)
+		async def confirmfunc(msg):
+			# Autoconfirm
+			if not self.autoconfirm:
+				return
+			if msg['class'] == 'queueing' \
+			and msg['queue']['pending'] == 'pending'\
+			and not self.confirmed:
+				self.confirm()
+		
+		@self.wssreceiver(Streams.status)
+		async def streamsfunc(msg):
+			if msg['status'] == 2:
+				# Automatically start streams
+				for strm in self.recv:
+					if not isinstance(strm,Streams):
+						continue
+					if strm.stream:
+						logging.debug(f'Enabling {strm.stream} stream')
+						await self.send({
+							'stream': strm.stream,
+							'type': 'start'
+						})
 
 		await self.wssworker()
 
@@ -66,7 +96,7 @@ async def send(self, obj:Union[Dict[str, Any],str]) -> None:
 		if isinstance(obj, dict):
 			obj = json.dumps(obj)
 
-		self.socket.send(obj)
+		await self.socket.send(obj)
 
 	async def wssworker(self) -> None:
 
@@ -86,6 +116,7 @@ async def receiver(self) -> None:
 		while True:
 			data = await self.socket.recv()
 			obj = json.loads(data)
+			msgtype = -1
 
 			if obj['type'] == 'line':
 				msgtype = Streams.console
@@ -104,16 +135,6 @@ async def receiver(self) -> None:
 				msgtype = Streams.status
 				msg = json.loads(obj['message'])
 
-				if not self.autoconfirm:
-					continue
-				if msg['class'] == 'queueing' \
-				and msg['queue']['pending'] == 'pending'\
-				and not self.confirmed:
-					t = asyncio.create_task(
-						self.confirm()
-					)
-					await t
-
 			if msgtype in self.recv:
 				t = asyncio.create_task(
 					self.recv[msgtype](msg)