alune/adb.py

@@ -2,33 +2,41 @@
 Module for all ADB (Android Debug Bridge) related methods.
 """
 
+import asyncio
+import atexit
 import os.path
 import random
 
 from adb_shell.adb_device import AdbDeviceTcp
 from adb_shell.adb_device_async import AdbDeviceTcpAsync
 from adb_shell.auth.keygen import keygen
 from adb_shell.auth.sign_pythonrsa import PythonRSASigner
+from adb_shell.exceptions import TcpTimeoutException
+import av
+from av.error import InvalidDataError  # pylint: disable=no-name-in-module
 import cv2
 from loguru import logger
 import numpy
 from numpy import ndarray
 import psutil
 
 from alune import helpers
+from alune.config import AluneConfig
 from alune.images import ClickButton
 from alune.images import ImageButton
 from alune.screen import BoundingBox
 from alune.screen import ImageSearchResult
 
 
-class ADB:
+# The amount of attributes is fine in my opinion.
+# We could split off screen recording into its own class, but I don't see the need to.
+class ADB:  # pylint: disable=too-many-instance-attributes
     """
     Class to hold the connection to an ADB connection via TCP.
     USB connection is possible, but not supported at the moment.
     """
 
-    def __init__(self):
+    def __init__(self, config: AluneConfig):
         """
         Initiates base values for the ADB instance.
         """
@@ -37,16 +45,23 @@ def __init__(self):
         self._random = random.Random()
         self._rsa_signer = None
         self._device = None
+        self._config = config
+        self._default_port = config.get_adb_port()
 
-    async def load(self, port: int):
+        if not config.should_use_screen_record():
+            return
+
+        self._video_codec = av.codec.CodecContext.create("h264", "r")
+        self._is_screen_recording = False
+        self._should_stop_screen_recording = False
+        self._latest_frame = None
+
+    async def load(self):
         """
         Load the RSA signer and attempt to connect to a device via ADB TCP.
-
-        Args:
-            port: The port to attempt to connect to.
         """
         await self._load_rsa_signer()
-        await self._connect_to_device(port)
+        await self._connect_to_device(self._default_port)
 
     async def _load_rsa_signer(self):
         """
@@ -97,6 +112,25 @@ async def scan_localhost_devices(self) -> int | None:
         logger.warning("No local device was found. Make sure ADB is enabled in your emulator's settings.")
         return None
 
+    def mark_screen_record_for_close(self):
+        """
+        Tells the screen recording to close itself when possible.
+        """
+        if self._is_screen_recording:
+            self._should_stop_screen_recording = True
+            self._is_screen_recording = False
+
+    def create_screen_record_task(self):
+        """
+        Create the screen recording task. Will not start recording if there's already a recording.
+        """
+        if self._is_screen_recording:
+            return
+
+        self._should_stop_screen_recording = False
+        asyncio.create_task(self.__screen_record())
+        atexit.register(self.mark_screen_record_for_close)
+
     async def _connect_to_device(self, port: int, retry_with_scan: bool = True):
         """
         Connect to the device via TCP.
@@ -136,7 +170,7 @@ async def get_screen_size(self) -> str:
         Returns:
              A string containing 'WIDTHxHEIGHT'.
         """
-        shell_output = await self._device.shell("wm size | awk 'END{print $3}'")
+        shell_output = await self._wrap_shell_call("wm size | awk 'END{print $3}'")
         return shell_output.replace("\n", "")
 
     async def get_screen_density(self) -> str:
@@ -146,20 +180,20 @@ async def get_screen_density(self) -> str:
         Returns:
              A string containing the pixel density.
         """
-        shell_output = await self._device.shell("wm density | awk 'END{print $3}'")
+        shell_output = await self._wrap_shell_call("wm density | awk 'END{print $3}'")
         return shell_output.replace("\n", "")
 
     async def set_screen_size(self):
         """
         Set the screen size to 1280x720.
         """
-        await self._device.shell("wm size 1280x720")
+        await self._wrap_shell_call("wm size 1280x720")
 
     async def set_screen_density(self):
         """
         Set the screen pixel density to 240.
         """
-        await self._device.shell("wm density 240")
+        await self._wrap_shell_call("wm density 240")
 
     async def get_memory(self) -> int:
         """
@@ -168,13 +202,25 @@ async def get_memory(self) -> int:
         Returns:
             The memory of the device in kB.
         """
-        shell_output = await self._device.shell("grep MemTotal /proc/meminfo | awk '{print $2}'")
+        shell_output = await self._wrap_shell_call("grep MemTotal /proc/meminfo | awk '{print $2}'")
         return int(shell_output)
 
     async def get_screen(self) -> ndarray | None:
         """
         Gets a ndarray which contains the values of the gray-scaled pixels
-        currently on the screen.
+        currently on the screen. Uses buffered frames from screen recording, available instantly.
+
+        Returns:
+            The ndarray containing the gray-scaled pixels. Is None until the first screen record frame is processed.
+        """
+        if self._config.should_use_screen_record():
+            return self._latest_frame
+        return await self._get_screen_capture()
+
+    async def _get_screen_capture(self) -> ndarray | None:
+        """
+        Gets a ndarray which contains the values of the gray-scaled pixels
+        currently on the screen. Uses screencap, so will take some processing time.
 
         Returns:
             The ndarray containing the gray-scaled pixels.
@@ -237,7 +283,7 @@ async def click(self, x: int, y: int):
         """
         # input tap x y comes with the downtime of tapping too fast for the game sometimes,
         # so we swipe on the same coordinate to simulate a longer press with a random duration.
-        await self._device.shell(f"input swipe {x} {y} {x} {y} {self._random.randint(60, 120)}")
+        await self._wrap_shell_call(f"input swipe {x} {y} {x} {y} {self._random.randint(60, 120)}")
 
     async def is_tft_installed(self) -> bool:
         """
@@ -246,7 +292,7 @@ async def is_tft_installed(self) -> bool:
         Returns:
             Whether the TFT package is in the list of the installed packages.
         """
-        shell_output = await self._device.shell(f"pm list packages | grep {self.tft_package_name}")
+        shell_output = await self._wrap_shell_call(f"pm list packages | grep {self.tft_package_name}")
         if not shell_output:
             return False
 
@@ -273,14 +319,14 @@ async def is_tft_active(self) -> bool:
         Returns:
              Whether TFT is the currently active window.
         """
-        shell_output = await self._device.shell("dumpsys window | grep -E 'mCurrentFocus' | awk '{print $3}'")
+        shell_output = await self._wrap_shell_call("dumpsys window | grep -E 'mCurrentFocus' | awk '{print $3}'")
         return shell_output.split("/")[0].replace("\n", "") == self.tft_package_name
 
     async def start_tft_app(self):
         """
         Start TFT using the activity manager (am).
         """
-        await self._device.shell(f"am start -n {self.tft_package_name}/{self._tft_activity_name}")
+        await self._wrap_shell_call(f"am start -n {self.tft_package_name}/{self._tft_activity_name}")
 
     async def get_tft_version(self) -> str:
         """
@@ -289,12 +335,91 @@ async def get_tft_version(self) -> str:
         Returns:
             The versionName of the tft package.
         """
-        return await self._device.shell(
+        return await self._wrap_shell_call(
             f"dumpsys package {self.tft_package_name} | grep versionName | sed s/[[:space:]]*versionName=//g"
         )
 
     async def go_back(self):
         """
         Send a back key press event to the device.
         """
-        await self._device.shell("input keyevent 4")
+        await self._wrap_shell_call("input keyevent 4")
+
+    async def _wrap_shell_call(self, shell_command: str, retries: int = 0):
+        """
+        Wrapper for shell commands to catch timeout exceptions.
+        Retries 3 times with incremental backoff.
+
+        Args:
+            shell_command: The shell command to call.
+            retries: Optional, the amount of attempted retries so far.
+
+        Returns:
+            The output of the shell command.
+        """
+        try:
+            return await self._device.exec_out(shell_command)
+        except TcpTimeoutException:
+            if retries == 3:
+                raise
+            logger.debug(f"Timed out while calling '{shell_command}', retrying {3 - retries} times.")
+            await asyncio.sleep(1 + (1 * retries))
+            return await self._wrap_shell_call(shell_command, retries=retries + 1)
+
+    async def __convert_frame_to_cv2(self, frame_bytes: bytes):
+        """
+        Convert frame bytes to a CV2 compatible gray image.
+
+        Args:
+            frame_bytes: Byte output of the screen record session.
+        """
+        packets = self._video_codec.parse(frame_bytes)
+        if not packets:
+            return
+
+        try:
+            frames = self._video_codec.decode(packets[0])
+        except InvalidDataError:
+            return
+        if not frames:
+            return
+
+        self._latest_frame = frames[0].to_ndarray(format="gray8").copy()  # Change to bgr24 if color is ever needed.
+
+    async def __write_frame_data(self):
+        """
+        Start a streaming shell that outputs screenrecord frame bytes and store it as a cv2 compatible image.
+        """
+        # output-format h264 > H264 is the only format that outputs to console which we can work with.
+        # time-limit 10 > Restarts screen recording every 10 seconds instead of every 180. Fixes compression artifacts.
+        # bit-rate 16M > 16_000_000 Mbps, could probably be lowered or made configurable, but works well.
+        # - at the end makes screenrecord output to console, if format is h264.
+        async for data in self._device.streaming_shell(
+            command="screenrecord --time-limit 8 --output-format h264 --bit-rate 16M --size 1280x720 -", decode=False
+        ):
+            if self._should_stop_screen_recording:
+                break
+
+            await self.__convert_frame_to_cv2(data)
+
+    async def __screen_record(self):
+        """
+        Start the screen record session. Restarts itself until an external value stops it.
+        """
+        if self._is_screen_recording:
+            return
+
+        await self._device.exec_out("pkill -2 screenrecord")
+        logger.debug("Screen record starting.")
+
+        self._is_screen_recording = True
+        while not self._should_stop_screen_recording:
+            try:
+                await self.__write_frame_data()
+            except TcpTimeoutException:
+                logger.warning("Timed out while re-/starting screen record, waiting 5 seconds.")
+                await asyncio.sleep(5)
+
+        logger.debug("Screen record stopped.")
+        await self._device.exec_out("pkill -2 screenrecord")
+        self._is_screen_recording = False

alune/config.py

@@ -108,32 +108,27 @@ def _sanitize_game_mode(self):
         Sanitize the user configured game mode by checking against valid values.
         """
         game_mode = self._config.get("game_mode", "normal")
-        if game_mode not in {"normal", "dawn of heroes"}:
+        if game_mode not in {"normal"}:
             logger.warning(f"The configured game mode '{game_mode}' does not exist. Playing 'normal' instead.")
             self._config["game_mode"] = "normal"
 
     def _sanitize_traits(self):
         """
         Sanitize the user configured traits by checking against currently implemented traits.
         """
-        if self._config["game_mode"] == "dawn of heroes":
-            trait_class = images.DawnOfHeroesTrait
-        else:
-            trait_class = images.Trait
-
-        current_traits = [trait.name for trait in list(trait_class)]
+        current_traits = [trait.name for trait in list(images.Trait)]
         configured_traits = self._config.get("traits", [])
 
         allowed_traits = []
         for trait in configured_traits:
             if trait.upper() not in current_traits:
                 logger.warning(f"The configured trait '{trait}' does not exist. Skipping it.")
                 continue
-            allowed_traits.append(trait_class[trait.upper()])
+            allowed_traits.append(images.Trait[trait.upper()])
 
         if len(allowed_traits) == 0:
-            logger.warning(f"No valid traits were configured. Falling back to {trait_class.get_default_traits()}.")
-            allowed_traits = trait_class.get_default_traits()
+            logger.warning(f"No valid traits were configured. Falling back to {images.Trait.get_default_traits()}.")
+            allowed_traits = images.Trait.get_default_traits()
 
         self._config["traits"] = allowed_traits
 
@@ -203,3 +198,21 @@ def get_chance_to_buy_xp(self) -> int:
             The chance in percent from 0 to 100.
         """
         return self._config["chances"]["buy_xp"]
+
+    def get_queue_timeout(self) -> int:
+        """
+        Get the queue timeout in seconds.
+
+        Returns:
+            The queue timeout in seconds.
+        """
+        return self._config["queue_timeout"]
+
+    def should_use_screen_record(self) -> bool:
+        """
+        Get the screen record setting the user wants.
+
+        Returns:
+            Whether we should use screen recording.
+        """
+        return self._config["screen_record"]["enabled"]