Add APK upload/install endpoints and androguard (#643)

Author: mahbd

pyproject.toml

@@ -103,6 +103,7 @@ dependencies = [
     "pymysql>=1.1.2",
     "mysql-connector-python>=9.5.0",
     "oracledb>=3.4.1",
+    "androguard>=4.1.2",
 ]
 
 [dependency-groups]

server/mobile.py

@@ -1,15 +1,18 @@
+import asyncio
+import base64
 import hashlib
 import os
+import shutil
 import subprocess
-import base64
 from typing import Literal
-import asyncio
 
 import requests
-from fastapi import APIRouter
+from androguard.core.apk import APK
+from fastapi import APIRouter, UploadFile, File
 from pydantic import BaseModel
 
-from Framework.Utilities import ConfigModule, CommonUtil
+from Framework.Utilities import CommonUtil, ConfigModule
+from settings import ZEUZ_NODE_DOWNLOADS_DIR
 
 ADB_PATH = "adb"  # Ensure ADB is in PATH
 UI_XML_PATH = "ui.xml"
@@ -29,12 +32,14 @@ class InspectorResponse(BaseModel):
 
 class DeviceInfo(BaseModel):
     """Model for device information."""
+
     serial: str
     status: str
     name: str | None = None
     # model: str | None = None
     # product: str | None = None
 
+
 @router.get("/devices", response_model=list[DeviceInfo])
 def get_devices():
     """Get list of connected Android devices."""
@@ -72,29 +77,28 @@ def inspect(device_serial: str | None = None):
         capture_screenshot(device_serial=device_serial)
 
         # Read XML file
-        with open(UI_XML_PATH, 'r') as xml_file:
+        with open(UI_XML_PATH, "r") as xml_file:
             xml_content = xml_file.read()
-            
+
         # Read and encode screenshot
-        with open(SCREENSHOT_PATH, 'rb') as img_file:
+        with open(SCREENSHOT_PATH, "rb") as img_file:
             screenshot_bytes = img_file.read()
-            screenshot_base64 = base64.b64encode(screenshot_bytes).decode('utf-8')
-            
+            screenshot_base64 = base64.b64encode(screenshot_bytes).decode("utf-8")
+
         return InspectorResponse(
-            status="ok",
-            ui_xml=xml_content,
-            screenshot=screenshot_base64
+            status="ok", ui_xml=xml_content, screenshot=screenshot_base64
         )
     except Exception as e:
-        return InspectorResponse(
-            status="error",
-            error=str(e)
-        )
+        return InspectorResponse(status="error", error=str(e))
+
 
 @router.get("/dump/driver")
 def dump_driver():
     """Dump the current driver."""
-    from Framework.Built_In_Automation.Mobile.CrossPlatform.Appium.BuiltInFunctions import appium_driver
+    from Framework.Built_In_Automation.Mobile.CrossPlatform.Appium.BuiltInFunctions import (
+        appium_driver,
+    )
+
     if appium_driver is None:
         return
     return appium_driver.page_source
@@ -103,7 +107,14 @@ def dump_driver():
 def run_adb_command(command):
     """Run an ADB command and return the output."""
     try:
-        result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        result = subprocess.run(
+            command,
+            shell=True,
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
         return result.stdout.strip()
     except subprocess.CalledProcessError as e:
         return f"Error: {e.stderr.strip()}"
@@ -116,11 +127,14 @@ def capture_ui_dump(device_serial: str | None = None):
         f"{ADB_PATH} {device_flag} shell uiautomator dump /sdcard/ui.xml".strip()
     )
     if out.startswith("Error:"):
-        from Framework.Built_In_Automation.Mobile.CrossPlatform.Appium.BuiltInFunctions import appium_driver
+        from Framework.Built_In_Automation.Mobile.CrossPlatform.Appium.BuiltInFunctions import (
+            appium_driver,
+        )
+
         if appium_driver is None:
             return
         page_src = appium_driver.page_source
-        with open(UI_XML_PATH, 'w') as xml_file:
+        with open(UI_XML_PATH, "w") as xml_file:
             xml_file.write(page_src)
     else:
         out = run_adb_command(
@@ -137,7 +151,10 @@ def capture_screenshot(device_serial: str | None = None):
         f"{ADB_PATH} {device_flag} shell screencap -p /sdcard/screen.png".strip()
     )
     if out.startswith("Error:"):
-        from Framework.Built_In_Automation.Mobile.CrossPlatform.Appium.BuiltInFunctions import appium_driver
+        from Framework.Built_In_Automation.Mobile.CrossPlatform.Appium.BuiltInFunctions import (
+            appium_driver,
+        )
+
         if appium_driver is None:
             return
         full_screenshot_path = os.path.join(os.getcwd(), SCREENSHOT_PATH)
@@ -156,10 +173,16 @@ async def upload_android_ui_dump():
         try:
             capture_ui_dump()
             try:
-                with open(UI_XML_PATH, 'r') as xml_file:
+                with open(UI_XML_PATH, "r") as xml_file:
                     xml_content = xml_file.read()
-                    xml_content = xml_content.replace("<?xml version='1.0' encoding='UTF-8' standalone='yes' ?>", "", 1)
-                    new_xml_hash = hashlib.sha256(xml_content.encode('utf-8')).hexdigest()
+                    xml_content = xml_content.replace(
+                        "<?xml version='1.0' encoding='UTF-8' standalone='yes' ?>",
+                        "",
+                        1,
+                    )
+                    new_xml_hash = hashlib.sha256(
+                        xml_content.encode("utf-8")
+                    ).hexdigest()
                     # Don't upload if the content hasn't changed
                     if prev_xml_hash == new_xml_hash:
                         await asyncio.sleep(5)
@@ -169,17 +192,59 @@ async def upload_android_ui_dump():
             except FileNotFoundError:
                 await asyncio.sleep(5)
                 continue
-            url = ConfigModule.get_config_value("Authentication", "server_address").strip() + "/node_ai_contents/"
+            url = (
+                ConfigModule.get_config_value(
+                    "Authentication", "server_address"
+                ).strip()
+                + "/node_ai_contents/"
+            )
             apiKey = ConfigModule.get_config_value("Authentication", "api-key").strip()
             res = requests.post(
                 url,
                 headers={"X-Api-Key": apiKey},
                 json={
                     "dom_mob": {"dom": xml_content},
-                    "node_id": CommonUtil.MachineInfo().getLocalUser().lower()
-                })
+                    "node_id": CommonUtil.MachineInfo().getLocalUser().lower(),
+                },
+            )
             if res.ok:
                 CommonUtil.ExecLog("", "UI dump uploaded successfully", iLogLevel=1)
         except Exception as e:
             CommonUtil.ExecLog("", f"Error uploading UI dump: {str(e)}", iLogLevel=3)
         await asyncio.sleep(5)
+
+
+@router.post("/apk-upload")
+def handle_apk_upload(file: UploadFile = File(...)):
+    dir_path = f"{ZEUZ_NODE_DOWNLOADS_DIR}/apk"
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+    filename = file.filename or "uploaded.apk"
+    file_path = os.path.join(dir_path, filename)
+    with open(file_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+    return {"message": "APK uploaded successfully", "filename": filename}
+
+
+def get_package_name(file_path: str) -> str | None:
+    """Extract package name from APK using androguard."""
+    try:
+        apk = APK(file_path)
+        return apk.get_package()
+    except Exception:
+        return None
+
+
+@router.post("/apk-install")
+def handle_apk_install(filename: str, serial: str):
+    dir_path = f"{ZEUZ_NODE_DOWNLOADS_DIR}/apk"
+    file_path = os.path.join(dir_path, filename)
+    if not os.path.exists(file_path):
+        return {"message": "APK not found", "filename": filename}
+    package_name = get_package_name(file_path)
+    try:
+        subprocess.run([ADB_PATH, "-s", serial, "install", file_path], check=True)
+        return {"message": "APK installed successfully", "filename": filename, "package_name": package_name}
+    except Exception as e:
+        return {"message": f"Error installing APK: {str(e)}", "filename": filename, "package_name": package_name}
+