diff --git a/alune/adb.py b/alune/adb.py index d2adfaa..e583808 100644 --- a/alune/adb.py +++ b/alune/adb.py @@ -2,6 +2,8 @@ Module for all ADB (Android Debug Bridge) related methods. """ +import asyncio +import atexit import os.path import random @@ -9,6 +11,9 @@ 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 @@ -16,19 +21,22 @@ 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,7 +335,7 @@ 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" ) @@ -297,4 +343,83 @@ 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 diff --git a/alune/config.py b/alune/config.py index a60e2e7..cd789f4 100644 --- a/alune/config.py +++ b/alune/config.py @@ -212,3 +212,12 @@ def get_queue_timeout(self) -> int: 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"] diff --git a/alune/resources/config.yaml b/alune/resources/config.yaml index 212bd96..e6a40db 100644 --- a/alune/resources/config.yaml +++ b/alune/resources/config.yaml @@ -31,6 +31,14 @@ surrender_early: false # Default value : 0 (disabled = surrender as fast as possible) surrender_random_delay: 0 +# Screen recording settings. +screen_record: + # Whether we should use screen recording. + # If false, the bot will fall back to screen capture. + # Screen recording saves ~600ms per screen check, but may occasionally cause a timeout. + # The timeout will be caught and handled. + enabled: true + # Configuration for chance events in the bot, in percent. chances: # The chance % to buy experience per check - configuring 50 means it's a 50% chance. @@ -44,6 +52,6 @@ adb_port: 5555 # Changing these below values manually can potentially break the bot, so don't! # Version of the YAML. -version: 7 +version: 9 # Version of the TFT set. set: 13 diff --git a/main.py b/main.py index 1361067..1253cb9 100644 --- a/main.py +++ b/main.py @@ -291,7 +291,9 @@ async def loop_disconnect_wrapper(adb_instance: ADB, config: AluneConfig): await loop(adb_instance, config) except TcpTimeoutException: logger.warning("ADB device was disconnected, attempting one reconnect...") - await adb_instance.load(config.get_adb_port()) + adb_instance.mark_screen_record_for_close() + await adb_instance.load() + adb_instance.create_screen_record_task() if not adb_instance.is_connected(): raise_and_exit("Could not reconnect. Please check your emulator for any errors. Exiting.") logger.info("Reconnected to device, continuing main loop.") @@ -539,15 +541,24 @@ async def main(): await check_version() - adb_instance = ADB() + adb_instance = ADB(config) - await adb_instance.load(config.get_adb_port()) + await adb_instance.load() if not adb_instance.is_connected(): logger.error("There is no ADB device ready. Exiting.") return logger.debug("ADB is connected, checking phone and app details") await check_phone_preconditions(adb_instance) + + if config.should_use_screen_record(): + logger.info("The bot will use live screen recording for image searches.") + adb_instance.create_screen_record_task() + while await adb_instance.get_screen() is None: + logger.debug("Waiting for frame data to become available...") + await asyncio.sleep(0.5) + logger.debug("Frames are now available.") + logger.info("Connected to ADB and device is set up correctly, starting main loop.") if config.should_surrender(): diff --git a/pyproject.toml b/pyproject.toml index 5b68539..f289e49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "psutil==6.0.0", "keyboard==0.13.5", "imutils==0.5.4", + "av==14.0.0", ] [project.optional-dependencies]