import asyncio
import os
import os.path
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from contextlib import suppress
from typing import Any, Dict, List, Optional, Tuple, Union
import aiohttp
import portpicker
from worlds._sc2common.bot import logger
from . import paths, wsl
from .controller import Controller
from .paths import Paths
from .versions import VERSIONS
class kill_switch:
_to_kill: List[Any] = []
@classmethod
def add(cls, value):
logger.debug("kill_switch: Add switch")
cls._to_kill.append(value)
@classmethod
def kill_all(cls):
logger.info(f"kill_switch: Process cleanup for {len(cls._to_kill)} processes")
for p in cls._to_kill:
# pylint: disable=W0212
p._clean(verbose=False)
class SC2Process:
"""
A class for handling SCII applications.
:param host: hostname for the url the SCII application will listen to
:param port: the websocket port the SCII application will listen to
:param fullscreen: whether to launch the SCII application in fullscreen or not, defaults to False
:param resolution: (window width, window height) in pixels, defaults to (1024, 768)
:param placement: (x, y) the distances of the SCII app's top left corner from the top left corner of the screen
e.g. (20, 30) is 20 to the right of the screen's left border, and 30 below the top border
:param render:
:param sc2_version:
:param base_build:
:param data_hash:
"""
def __init__(
self,
host: Optional[str] = None,
port: Optional[int] = None,
fullscreen: bool = False,
resolution: Optional[Union[List[int], Tuple[int, int]]] = None,
placement: Optional[Union[List[int], Tuple[int, int]]] = None,
render: bool = False,
sc2_version: str = None,
base_build: str = None,
data_hash: str = None,
) -> None:
assert isinstance(host, str) or host is None
assert isinstance(port, int) or port is None
self._render = render
self._arguments: Dict[str, str] = {"-displayMode": str(int(fullscreen))}
if not fullscreen:
if resolution and len(resolution) == 2:
self._arguments["-windowwidth"] = str(resolution[0])
self._arguments["-windowheight"] = str(resolution[1])
if placement and len(placement) == 2:
self._arguments["-windowx"] = str(placement[0])
self._arguments["-windowy"] = str(placement[1])
self._host = host or os.environ.get("SC2CLIENTHOST", "127.0.0.1")
self._serverhost = os.environ.get("SC2SERVERHOST", self._host)
if port is None:
self._port = portpicker.pick_unused_port()
else:
self._port = port
self._used_portpicker = bool(port is None)
self._tmp_dir = tempfile.mkdtemp(prefix="SC2_")
self._process: subprocess = None
self._session = None
self._ws = None
self._sc2_version = sc2_version
self._base_build = base_build
self._data_hash = data_hash
async def __aenter__(self) -> Controller:
kill_switch.add(self)
def signal_handler(*_args):
# unused arguments: signal handling library expects all signal
# callback handlers to accept two positional arguments
kill_switch.kill_all()
signal.signal(signal.SIGINT, signal_handler)
try:
self._process = self._launch()
self._ws = await self._connect()
except:
await self._close_connection()
self._clean()
raise
return Controller(self._ws, self)
async def __aexit__(self, *args):
logger.exception("async exit")
await self._close_connection()
kill_switch.kill_all()
signal.signal(signal.SIGINT, signal.SIG_DFL)
@property
def ws_url(self):
return f"ws://{self._host}:{self._port}/sc2api"
@property
def versions(self):
"""Opens the versions.json file which origins from
https://github.com/Blizzard/s2client-proto/blob/master/buildinfo/versions.json"""
return VERSIONS
def find_data_hash(self, target_sc2_version: str) -> Optional[str]:
""" Returns the data hash from the matching version string. """
version: dict
for version in self.versions:
if version["label"] == target_sc2_version:
return version["data-hash"]
return None
def _launch(self):
if self._base_build:
executable = str(paths.latest_executeble(Paths.BASE / "Versions", self._base_build))
else:
executable = str(Paths.EXECUTABLE)
if self._port is None:
self._port = portpicker.pick_unused_port()
self._used_portpicker = True
args = paths.get_runner_args(Paths.CWD) + [
executable,
"-listen",
self._serverhost,
"-port",
str(self._port),
"-dataDir",
str(Paths.BASE),
"-tempDir",
self._tmp_dir,
]
for arg, value in self._arguments.items():
args.append(arg)
args.append(value)
if self._sc2_version:
def special_match(strg: str):
""" Tests if the specified version is in the versions.py dict. """
for version in self.versions:
if version["label"] == strg:
return True
return False
valid_version_string = special_match(self._sc2_version)
if valid_version_string:
self._data_hash = self.find_data_hash(self._sc2_version)
assert (
self._data_hash is not None
), f"StarCraft 2 Client version ({self._sc2_version}) was not found inside sc2/versions.py file. Please check your spelling or check the versions.py file."
else:
logger.warning(
f'The submitted version string in sc2.rungame() function call (sc2_version="{self._sc2_version}") was not found in versions.py. Running latest version instead.'
)
if self._data_hash:
args.extend(["-dataVersion", self._data_hash])
if self._render:
args.extend(["-eglpath", "libEGL.so"])
# if logger.getEffectiveLevel() <= logging.DEBUG:
args.append("-verbose")
sc2_cwd = str(Paths.CWD) if Paths.CWD else None
if paths.PF in {"WSL1", "WSL2"}:
return wsl.run(args, sc2_cwd)
return subprocess.Popen(
args,
cwd=sc2_cwd,
# Suppress Wine error messages
stderr=subprocess.DEVNULL
# , env=run_config.env
)
async def _connect(self):
# How long it waits for SC2 to start (in seconds)
for i in range(180):
if self._process is None:
# The ._clean() was called, clearing the process
logger.debug("Process cleanup complete, exit")
sys.exit()
await asyncio.sleep(1)
try:
self._session = aiohttp.ClientSession()
ws = await self._session.ws_connect(self.ws_url, timeout=120)
# FIXME fix deprecation warning in for future aiohttp version
# ws = await self._session.ws_connect(
# self.ws_url, timeout=aiohttp.client_ws.ClientWSTimeout(ws_close=120)
# )
logger.debug("Websocket connection ready")
return ws
except aiohttp.client_exceptions.ClientConnectorError:
await self._session.close()
if i > 15:
logger.debug("Connection refused (startup not complete (yet))")
logger.debug("Websocket connection to SC2 process timed out")
raise TimeoutError("Websocket")
async def _close_connection(self):
logger.info(f"Closing connection at {self._port}...")
if self._ws is not None:
await self._ws.close()
if self._session is not None:
await self._session.close()
# pylint: disable=R0912
def _clean(self, verbose=True):
if verbose:
logger.info("Cleaning up...")
if self._process is not None:
if paths.PF in {"WSL1", "WSL2"}:
if wsl.kill(self._process):
logger.error("KILLED")
elif self._process.poll() is None:
for _ in range(3):
self._process.terminate()
time.sleep(0.5)
if not self._process or self._process.poll() is not None:
break
else:
self._process.kill()
self._process.wait()
logger.error("KILLED")
# Try to kill wineserver on linux
if paths.PF in {"Linux", "WineLinux"}:
# Command wineserver not detected
with suppress(FileNotFoundError):
with subprocess.Popen(["wineserver", "-k"]) as p:
p.wait()
if os.path.exists(self._tmp_dir):
shutil.rmtree(self._tmp_dir)
self._process = None
self._ws = None
if self._used_portpicker and self._port is not None:
portpicker.return_port(self._port)
self._port = None
if verbose:
logger.info("Cleanup complete")