Update Factorio, Scum, Palworld

discordgsm/protocols/factorio.py

@@ -1,4 +1,5 @@
 import re
+import time
 from typing import TYPE_CHECKING
 
 import aiohttp
@@ -12,18 +13,28 @@
 
 
 class Factorio(Protocol):
-    pre_query_required = True
+    pre_query_required = False
     name = "factorio"
     master_servers = None
 
-    async def pre_query(self):
+    def __init__(self, kv: dict):
+        super().__init__(kv)
+
         # Get FACTORIO_USERNAME and FACTORIO_TOKEN on https://www.factorio.com/profile
         # Notice: It requires to own Factorio on steam account
-        username, token = (
+        self.username, self.token = (
             str(env("FACTORIO_USERNAME")).strip(),
             str(env("FACTORIO_TOKEN")).strip(),
         )
-        url = f"https://multiplayer.factorio.com/get-games?username={username}&token={token}"
+
+        if self.username and self.token:
+            self.pre_query_required = True
+
+    async def pre_query(self):
+        if not self.pre_query_required:
+            return
+
+        url = f"https://multiplayer.factorio.com/get-games?username={self.username}&token={self.token}"
 
         async with aiohttp.ClientSession() as session:
             async with session.get(url) as response:
@@ -39,35 +50,47 @@ async def pre_query(self):
         Factorio.master_servers = master_servers
 
     async def query(self):
-        if Factorio.master_servers is None:
-            await self.pre_query()
-
         host, port = str(self.kv["host"]), int(str(self.kv["port"]))
         ip = await Socket.gethostbyname(host)
-        host_address = f"{ip}:{port}"
 
-        if host_address not in Factorio.master_servers:
-            raise Exception("Server not found")
+        if self.pre_query_required:
+            if Factorio.master_servers is None:
+                await self.pre_query()
+
+            host_address = f"{ip}:{port}"
+
+            if host_address not in Factorio.master_servers:
+                raise Exception("Server not found")
+
+            data = dict(Factorio.master_servers[host_address])
+            ping = 0
+        else:
+            url = f"https://master-server.opengsq.com/factorio/search?host={ip}&port={port}"
+            start = time.time()
 
-        server = dict(Factorio.master_servers[host_address])
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as response:
+                    response.raise_for_status()
+                    data: dict = await response.json()
+                    ping = int((time.time() - start) * 1000)
 
         # Remove the rich text formatting
         # https://wiki.factorio.com/Rich_text
-        name = re.sub(r"\[\w*=\w*\]|\[\/\w*\]", "", server["name"])
-        players = server["players"] if "players" in server else []
+        name = re.sub(r"\[\w*=\w*\]|\[\/\w*\]", "", data["name"])
+        players = list(data.get("players", []))
 
         result: GamedigResult = {
             "name": name,
             "map": "",
-            "password": server["has_password"],
+            "password": data["has_password"],
             "numplayers": len(players),
             "numbots": 0,
-            "maxplayers": server["max_players"],
+            "maxplayers": data["max_players"],
             "players": [{"name": player, "raw": player} for player in players],
             "bots": None,
-            "connect": server["host_address"],
-            "ping": 0,
-            "raw": server,
+            "connect": data["host_address"],
+            "ping": ping,
+            "raw": data,
         }
 
         return result

discordgsm/protocols/palworld.py

@@ -1,10 +1,8 @@
 import time
 from typing import TYPE_CHECKING
-import aiohttp
 
-import opengsq
+import aiohttp
 from opengsq.protocol_socket import Socket
-from opengsq.exceptions.server_not_found_exception import ServerNotFoundException
 
 from discordgsm.protocols.protocol import Protocol
 
@@ -24,6 +22,7 @@ async def query(self):
 
         async with aiohttp.ClientSession() as session:
             async with session.get(url) as response:
+                response.raise_for_status()
                 data: dict = await response.json()
                 ping = int((time.time() - start) * 1000)
 

discordgsm/protocols/scum.py

@@ -35,7 +35,7 @@ async def query(self):
             "maxplayers": data.get("maxplayers", 0),
             "players": None,
             "bots": None,
-            "connect": f"{host}:{data.get("port", port) - 2}",
+            "connect": f"{host}:{data.get('port', port) - 2}",
             "ping": ping,
             "raw": data,
         }