From f43bc6d7311b8a21282c86d081ec06c4b7eb804e Mon Sep 17 00:00:00 2001
From: Martin Miglio <marmig0404@gmail.com>
Date: Sat, 22 Jul 2023 14:54:57 -0400
Subject: [PATCH] Add get_adb_connection to get ADB connection info

---
 demo/demo.ipynb     | 16 ++++++++++++++++
 pymemuc/_command.py | 36 ++++++++++++++++++++++++++++++++++++
 pymemuc/pymemuc.py  |  1 +
 3 files changed, 53 insertions(+)

diff --git a/demo/demo.ipynb b/demo/demo.ipynb
index 68171ea..fa3cb2d 100644
--- a/demo/demo.ipynb
+++ b/demo/demo.ipynb
@@ -89,6 +89,22 @@
     "vm_index = memuc.create_vm()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get VM's ADB connection info (IP and port)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "memuc.get_adb_connection(vm_index)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/pymemuc/_command.py b/pymemuc/_command.py
index 0403f91..b9f2e43 100644
--- a/pymemuc/_command.py
+++ b/pymemuc/_command.py
@@ -1,6 +1,7 @@
 """This module contains functions for commanding running virtual machines with memuc.exe.
 Functions for interacting with running VMs are defined here."""
 from typing import TYPE_CHECKING, Literal, Tuple, Union
+from urllib.parse import urlparse
 
 from ._decorators import _retryable
 from .exceptions import PyMemucError, PyMemucIndexError, PyMemucTimeoutExpired
@@ -628,3 +629,38 @@ def send_adb_command_vm(
         return output
 
     raise PyMemucIndexError("Please specify either a vm index or a vm name")
+
+
+def get_adb_connection(
+    self: "PyMemuc",
+    vm_index: Union[int, None] = None,
+    vm_name: Union[str, None] = None,
+    timeout: Union[int, None] = None,
+) -> Tuple[Union[str, None], Union[int, None]]:
+    """Get the adb connection information for a VM
+
+    :param vm_index: VM index. Defaults to None.
+    :type vm_index: int, optional
+    :param vm_name: VM name. Defaults to None.
+    :type vm_name: str, optional
+    :param timeout: Timeout for the command. Defaults to None.
+    :type timeout: int, optional
+    :raises PyMemucIndexError: an error if neither a vm index or a vm name is specified
+    :raises PyMemucTimeoutExpired: an error if the command times out
+    :raises PyMemucError: an error if the command fails
+    :return: the ip and port of the adb connection as a tuple
+    :rtype: tuple[str | None, int | None]
+    """
+    adb_output = self.send_adb_command_vm(
+        ["shell", "ifconfig"],
+        vm_index=vm_index,
+        vm_name=vm_name,
+        timeout=timeout,
+    )
+    try:
+        adb_output = adb_output.split("\n")[0]
+        _, connection_string = adb_output.split("connected to ")
+        connection_string = urlparse(f"//{connection_string}")
+        return connection_string.hostname, connection_string.port
+    except ValueError as err:
+        raise PyMemucError(f"Failed to get adb connection: {adb_output}") from err
diff --git a/pymemuc/pymemuc.py b/pymemuc/pymemuc.py
index a3a9d6c..cd886af 100644
--- a/pymemuc/pymemuc.py
+++ b/pymemuc/pymemuc.py
@@ -25,6 +25,7 @@ class PyMemuc:
         create_app_shortcut_vm,
         disconnect_internet_vm,
         execute_command_vm,
+        get_adb_connection,
         get_app_info_list_vm,
         get_public_ip_vm,
         input_text_vm,