uiautodev 0.3.3__py3-none-any.whl

uiautodev/utils/common.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import datetime
+import json as sysjson
+import platform
+import re
+import socket
+import sys
+import typing
+import uuid
+from http.client import HTTPConnection, HTTPResponse
+from typing import Optional, TypeVar, Union
+
+from pydantic import BaseModel
+from pygments import formatters, highlight, lexers
+
+from uiautodev.exceptions import RequestError
+from uiautodev.model import Node
+
+
+def is_output_terminal() -> bool:
+    """
+    Check if the standard output is attached to a terminal.
+    """
+    return sys.stdout.isatty()
+
+
+def enable_windows_ansi_support():
+    if platform.system().lower() == "windows" and is_output_terminal():
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+        kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
+
+
+def default_json_encoder(obj):
+    if isinstance(obj, bytes):
+        return f'<{obj.hex()}>'
+    if isinstance(obj, datetime.datetime):
+        return str(obj)
+    if isinstance(obj, uuid.UUID):
+        return str(obj)
+    if isinstance(obj, BaseModel):
+        return obj.model_dump()
+    raise TypeError()
+
+
+def print_json(buf, colored=None, default=default_json_encoder):
+    """ copy from pymobiledevice3 """
+    formatted_json = sysjson.dumps(buf, sort_keys=True, indent=4, default=default)
+    if colored is None:
+        if is_output_terminal():
+            colored = True
+            enable_windows_ansi_support()
+        else:
+            colored = False
+
+    if colored:
+        colorful_json = highlight(formatted_json, lexers.JsonLexer(),
+                                  formatters.TerminalTrueColorFormatter(style='stata-dark'))
+        print(colorful_json)
+    else:
+        print(formatted_json)
+    
+
+_T = TypeVar("_T")
+
+def convert_to_type(value: str, _type: _T) -> _T:
+    """ usage example:
+    convert_to_type("123", int)
+    """
+    if _type in (int, float, str):
+        return _type(value)
+    if _type == bool:
+        return value.lower() in ("true", "1")
+    if _type == Union[int, float]:
+        return float(value) if "." in value else int(value)
+    if _type == re.Pattern:
+        return re.compile(value)
+    raise NotImplementedError(f"convert {value} to {_type}")
+    
+
+def convert_params_to_model(params: list[str], model: BaseModel) -> BaseModel:
+    """ used in cli.py """
+    assert len(params) > 0
+    if len(params) == 1:
+        try:
+            return model.model_validate_json(params)
+        except Exception as e:
+            print("module_parse_error", e)
+
+    value = {}
+    type_hints = typing.get_type_hints(model)
+    for p in params:
+        if "=" not in p:
+            _type = type_hints.get(p)
+            if _type == bool:
+                value[p] = True
+                continue
+            elif _type is None:
+                print(f"unknown key: {p}")
+                continue
+            raise ValueError(f"missing value for {p}")
+        k, v = p.split("=", 1)
+        _type = type_hints.get(k)
+        if _type is None:
+            print(f"unknown key: {k}")
+            continue
+        value[k] = convert_to_type(v, _type)
+    return model.model_validate(value)
+
+
+class SocketHTTPConnection(HTTPConnection):
+    def __init__(self, conn: socket.socket, timeout: float):
+        super().__init__("localhost", timeout=timeout)
+        self.__conn = conn
+        
+    def connect(self):
+        self.sock = self.__conn
+
+    def __enter__(self) -> HTTPConnection:
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+
+class MySocketHTTPConnection(SocketHTTPConnection):
+    def connect(self):
+        super().connect()
+        self.sock.settimeout(self.timeout)
+
+
+def fetch_through_socket(sock: socket.socket, path: str, method: str = "GET", json: Optional[dict] = None, timeout: float = 60) -> bytearray:
+    """ usage example:
+    with socket.create_connection((host, port)) as s:
+        request_through_socket(s, "GET", "/")
+    """
+    conn = MySocketHTTPConnection(sock, timeout)
+    try:
+        if json is None:
+            conn.request(method, path)
+        else:
+            conn.request(method, path, body=sysjson.dumps(json), headers={"Content-Type": "application/json"})
+        response = conn.getresponse()
+        if response.getcode() != 200:
+            raise RequestError(f"request {method} {path}, status: {response.getcode()}")
+        content = bytearray()
+        while chunk := response.read(40960):
+            content.extend(chunk)
+        return content
+    finally:
+        conn.close()
+
+
+def node_travel(node: Node, dfs: bool = True):
+    """ usage example:
+    for n in node_travel(node):
+        print(n)
+    """
+    if not dfs:
+        yield node
+    for child in node.children:
+        yield from node_travel(child, dfs)
+    if dfs:
+        yield node

uiautodev/utils/exceptions.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""Created on Tue Mar 05 2024 10:18:09 by codeskyblue
+
+Copy from https://github.com/doronz88/pymobiledevice3
+"""
+
+
+from uiautodev.exceptions import IOSDriverException
+
+
+class NotPairedError(IOSDriverException):
+    pass
+
+
+
+class MuxException(IOSDriverException):
+    pass
+
+
+class MuxVersionError(MuxException):
+    pass
+
+
+class BadCommandError(MuxException):
+    pass
+
+
+class BadDevError(MuxException):
+    pass
+
+
+class ConnectionFailedError(MuxException):
+    pass
+
+
+class ConnectionFailedToUsbmuxdError(ConnectionFailedError):
+    pass
+
+
+class ArgumentError(IOSDriverException):
+    pass

uiautodev/utils/usbmux.py

@@ -0,0 +1,485 @@
+"""
+Copy from https://github.com/doronz88/pymobiledevice3
+
+Add http.client.HTTPConnection
+"""
+import abc
+import plistlib
+import socket
+import sys
+import time
+from dataclasses import dataclass
+from http.client import HTTPConnection
+from typing import List, Mapping, Optional
+
+from construct import Const, CString, Enum, FixedSized, GreedyBytes, Int16ul, Int32ul, Padding, Prefixed, StreamError, \
+    Struct, Switch, this
+
+from uiautodev.utils.exceptions import BadCommandError, BadDevError, ConnectionFailedError, \
+    ConnectionFailedToUsbmuxdError, MuxException, MuxVersionError, NotPairedError
+
+usbmuxd_version = Enum(Int32ul,
+                       BINARY=0,
+                       PLIST=1,
+                       )
+
+usbmuxd_result = Enum(Int32ul,
+                      OK=0,
+                      BADCOMMAND=1,
+                      BADDEV=2,
+                      CONNREFUSED=3,
+                      BADVERSION=6,
+                      )
+
+usbmuxd_msgtype = Enum(Int32ul,
+                       RESULT=1,
+                       CONNECT=2,
+                       LISTEN=3,
+                       ADD=4,
+                       REMOVE=5,
+                       PAIRED=6,
+                       PLIST=8,
+                       )
+
+usbmuxd_header = Struct(
+    'version' / usbmuxd_version,  # protocol version
+    'message' / usbmuxd_msgtype,  # message type
+    'tag' / Int32ul,  # responses to this query will echo back this tag
+)
+
+usbmuxd_request = Prefixed(Int32ul, Struct(
+    'header' / usbmuxd_header,
+    'data' / Switch(this.header.message, {
+        usbmuxd_msgtype.CONNECT: Struct(
+            'device_id' / Int32ul,
+            'port' / Int16ul,  # TCP port number
+            'reserved' / Const(0, Int16ul),
+        ),
+        usbmuxd_msgtype.PLIST: GreedyBytes,
+    }),
+), includelength=True)
+
+usbmuxd_device_record = Struct(
+    'device_id' / Int32ul,
+    'product_id' / Int16ul,
+    'serial_number' / FixedSized(256, CString('ascii')),
+    Padding(2),
+    'location' / Int32ul
+)
+
+usbmuxd_response = Prefixed(Int32ul, Struct(
+    'header' / usbmuxd_header,
+    'data' / Switch(this.header.message, {
+        usbmuxd_msgtype.RESULT: Struct(
+            'result' / usbmuxd_result,
+        ),
+        usbmuxd_msgtype.ADD: usbmuxd_device_record,
+        usbmuxd_msgtype.REMOVE: Struct(
+            'device_id' / Int32ul,
+        ),
+        usbmuxd_msgtype.PLIST: GreedyBytes,
+    }),
+), includelength=True)
+
+
+        
+
+@dataclass
+class MuxDevice:
+    devid: int
+    serial: str
+    connection_type: str
+
+    def connect(self, port: int, usbmux_address: Optional[str] = None) -> socket.socket:
+        mux = create_mux(usbmux_address=usbmux_address)
+        try:
+            return mux.connect(self, port)
+        except:  # noqa: E722
+            mux.close()
+            raise
+
+    @property
+    def is_usb(self) -> bool:
+        return self.connection_type == 'USB'
+
+    @property
+    def is_network(self) -> bool:
+        return self.connection_type == 'Network'
+
+    def matches_udid(self, udid: str) -> bool:
+        return self.serial.replace('-', '') == udid.replace('-', '')
+
+    def make_http_connection(self, port: int) -> HTTPConnection:
+        return USBMuxHTTPConnection(self, port)
+
+
+class SafeStreamSocket:
+    """ wrapper to native python socket object to be used with construct as a stream """
+
+    def __init__(self, address, family):
+        self._offset = 0
+        self.sock = socket.socket(family, socket.SOCK_STREAM)
+        self.sock.connect(address)
+
+    def send(self, msg: bytes) -> int:
+        self._offset += len(msg)
+        self.sock.sendall(msg)
+        return len(msg)
+
+    def recv(self, size: int) -> bytes:
+        msg = b''
+        while len(msg) < size:
+            chunk = self.sock.recv(size - len(msg))
+            self._offset += len(chunk)
+            if not chunk:
+                raise MuxException('socket connection broken')
+            msg += chunk
+        return msg
+
+    def close(self) -> None:
+        self.sock.close()
+
+    def settimeout(self, interval: float) -> None:
+        self.sock.settimeout(interval)
+
+    def setblocking(self, blocking: bool) -> None:
+        self.sock.setblocking(blocking)
+
+    def tell(self) -> int:
+        return self._offset
+
+    read = recv
+    write = send
+
+
+class MuxConnection:
+    # used on Windows
+    ITUNES_HOST = ('127.0.0.1', 27015)
+
+    # used for macOS and Linux
+    USBMUXD_PIPE = '/var/run/usbmuxd'
+
+    @staticmethod
+    def create_usbmux_socket(usbmux_address: Optional[str] = None) -> SafeStreamSocket:
+        try:
+            if usbmux_address is not None:
+                if ':' in usbmux_address:
+                    # assume tcp address
+                    hostname, port = usbmux_address.split(':')
+                    port = int(port)
+                    address = (hostname, port)
+                    family = socket.AF_INET
+                else:
+                    # assume unix domain address
+                    address = usbmux_address
+                    family = socket.AF_UNIX
+            else:
+                if sys.platform in ['win32', 'cygwin']:
+                    address = MuxConnection.ITUNES_HOST
+                    family = socket.AF_INET
+                else:
+                    address = MuxConnection.USBMUXD_PIPE
+                    family = socket.AF_UNIX
+            return SafeStreamSocket(address, family)
+        except ConnectionRefusedError:
+            raise ConnectionFailedToUsbmuxdError()
+
+    @staticmethod
+    def create(usbmux_address: Optional[str] = None):
+        # first attempt to connect with possibly the wrong version header (plist protocol)
+        sock = MuxConnection.create_usbmux_socket(usbmux_address=usbmux_address)
+
+        message = usbmuxd_request.build({
+            'header': {'version': usbmuxd_version.PLIST, 'message': usbmuxd_msgtype.PLIST, 'tag': 1},
+            'data': plistlib.dumps({'MessageType': 'ReadBUID'})
+        })
+        sock.send(message)
+        response = usbmuxd_response.parse_stream(sock)
+
+        # if we sent a bad request, we should re-create the socket in the correct version this time
+        sock.close()
+        sock = MuxConnection.create_usbmux_socket(usbmux_address=usbmux_address)
+
+        if response.header.version == usbmuxd_version.BINARY:
+            return BinaryMuxConnection(sock)
+        elif response.header.version == usbmuxd_version.PLIST:
+            return PlistMuxConnection(sock)
+
+        raise MuxVersionError(f'usbmuxd returned unsupported version: {response.version}')
+
+    def __init__(self, sock: SafeStreamSocket):
+        self._sock = sock
+
+        # after initiating the "Connect" packet, this same socket will be used to transfer data into the service
+        # residing inside the target device. when this happens, we can no longer send/receive control commands to
+        # usbmux on same socket
+        self._connected = False
+
+        # message sequence number. used when verifying the response matched the request
+        self._tag = 1
+
+        self.devices = []
+
+    @abc.abstractmethod
+    def _connect(self, device_id: int, port: int):
+        """ initiate a "Connect" request to target port """
+        pass
+
+    @abc.abstractmethod
+    def get_device_list(self, timeout: float = None):
+        """
+        request an update to current device list
+        """
+        pass
+
+    def connect(self, device: MuxDevice, port: int) -> socket.socket:
+        """ connect to a relay port on target machine and get a raw python socket object for the connection """
+        self._connect(device.devid, socket.htons(port))
+        self._connected = True
+        return self._sock.sock
+
+    def close(self):
+        """ close current socket """
+        self._sock.close()
+
+    def _assert_not_connected(self):
+        """ verify active state is in state for control messages """
+        if self._connected:
+            raise MuxException('Mux is connected, cannot issue control packets')
+
+    def _raise_mux_exception(self, result: int, message: str = None):
+        exceptions = {
+            int(usbmuxd_result.BADCOMMAND): BadCommandError,
+            int(usbmuxd_result.BADDEV): BadDevError,
+            int(usbmuxd_result.CONNREFUSED): ConnectionFailedError,
+            int(usbmuxd_result.BADVERSION): MuxVersionError,
+        }
+        exception = exceptions.get(result, MuxException)
+        raise exception(message)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+class BinaryMuxConnection(MuxConnection):
+    """ old binary protocol """
+
+    def __init__(self, sock: SafeStreamSocket):
+        super().__init__(sock)
+        self._version = usbmuxd_version.BINARY
+
+    def get_device_list(self, timeout: float = None):
+        """ use timeout to wait for the device list to be fully populated """
+        self._assert_not_connected()
+        end = time.time() + timeout
+        self.listen()
+        while time.time() < end:
+            self._sock.settimeout(end - time.time())
+            try:
+                self._receive_device_state_update()
+            except (BlockingIOError, StreamError):
+                continue
+            except IOError:
+                try:
+                    self._sock.setblocking(True)
+                    self.close()
+                except OSError:
+                    pass
+                raise MuxException('Exception in listener socket')
+
+    def listen(self):
+        """ start listening for events of attached and detached devices """
+        self._send_receive(usbmuxd_msgtype.LISTEN)
+
+    def _connect(self, device_id: int, port: int):
+        self._send({'header': {'version': self._version,
+                               'message': usbmuxd_msgtype.CONNECT,
+                               'tag': self._tag},
+                    'data': {'device_id': device_id, 'port': port},
+                    })
+        response = self._receive()
+        if response.header.message != usbmuxd_msgtype.RESULT:
+            raise MuxException(f'unexpected message type received: {response}')
+
+        if response.data.result != usbmuxd_result.OK:
+            raise self._raise_mux_exception(int(response.data.result),
+                                            f'failed to connect to device: {device_id} at port: {port}. reason: '
+                                            f'{response.data.result}')
+
+    def _send(self, data: Mapping):
+        self._assert_not_connected()
+        self._sock.send(usbmuxd_request.build(data))
+        self._tag += 1
+
+    def _receive(self, expected_tag: int = None):
+        self._assert_not_connected()
+        response = usbmuxd_response.parse_stream(self._sock)
+        if expected_tag and response.header.tag != expected_tag:
+            raise MuxException(f'Reply tag mismatch: expected {expected_tag}, got {response.header.tag}')
+        return response
+
+    def _send_receive(self, message_type: int):
+        self._send({'header': {'version': self._version, 'message': message_type, 'tag': self._tag},
+                    'data': b''})
+        response = self._receive(self._tag - 1)
+        if response.header.message != usbmuxd_msgtype.RESULT:
+            raise MuxException(f'unexpected message type received: {response}')
+
+        result = response.data.result
+        if result != usbmuxd_result.OK:
+            raise self._raise_mux_exception(int(result), f'{message_type} failed: error {result}')
+
+    def _add_device(self, device: MuxDevice):
+        self.devices.append(device)
+
+    def _remove_device(self, device_id: int):
+        self.devices = [device for device in self.devices if device.devid != device_id]
+
+    def _receive_device_state_update(self):
+        response = self._receive()
+        if response.header.message == usbmuxd_msgtype.ADD:
+            # old protocol only supported USB devices
+            self._add_device(MuxDevice(response.data.device_id, response.data.serial_number, 'USB'))
+        elif response.header.message == usbmuxd_msgtype.REMOVE:
+            self._remove_device(response.data.device_id)
+        else:
+            raise MuxException(f'Invalid packet type received: {response}')
+
+
+class PlistMuxConnection(BinaryMuxConnection):
+    def __init__(self, sock: SafeStreamSocket):
+        super().__init__(sock)
+        self._version = usbmuxd_version.PLIST
+
+    def listen(self) -> None:
+        self._send_receive({'MessageType': 'Listen'})
+
+    def get_pair_record(self, serial: str) -> Mapping:
+        # serials are saved inside usbmuxd without '-'
+        self._send({'MessageType': 'ReadPairRecord', 'PairRecordID': serial})
+        response = self._receive(self._tag - 1)
+        pair_record = response.get('PairRecordData')
+        if pair_record is None:
+            raise NotPairedError('device should be paired first')
+        return plistlib.loads(pair_record)
+
+    def get_device_list(self, timeout: float = None) -> None:
+        """ get device list synchronously without waiting the timeout """
+        self.devices = []
+        self._send({'MessageType': 'ListDevices'})
+        for response in self._receive(self._tag - 1)['DeviceList']:
+            if response['MessageType'] == 'Attached':
+                super()._add_device(MuxDevice(response['DeviceID'], response['Properties']['SerialNumber'],
+                                              response['Properties']['ConnectionType']))
+            elif response['MessageType'] == 'Detached':
+                super()._remove_device(response['DeviceID'])
+            else:
+                raise MuxException(f'Invalid packet type received: {response}')
+
+    def get_buid(self) -> str:
+        """ get SystemBUID """
+        self._send({'MessageType': 'ReadBUID'})
+        return self._receive(self._tag - 1)['BUID']
+
+    def save_pair_record(self, serial: str, device_id: int, record_data: bytes):
+        # serials are saved inside usbmuxd without '-'
+        self._send_receive({'MessageType': 'SavePairRecord',
+                            'PairRecordID': serial,
+                            'PairRecordData': record_data,
+                            'DeviceID': device_id})
+
+    def _connect(self, device_id: int, port: int):
+        self._send_receive({'MessageType': 'Connect', 'DeviceID': device_id, 'PortNumber': port})
+
+    def _send(self, data: Mapping):
+        request = {'ClientVersionString': 'qt4i-usbmuxd', 'ProgName': 'pymobiledevice3', 'kLibUSBMuxVersion': 3}
+        request.update(data)
+        super()._send({'header': {'version': self._version,
+                                  'message': usbmuxd_msgtype.PLIST,
+                                  'tag': self._tag},
+                       'data': plistlib.dumps(request),
+                       })
+
+    def _receive(self, expected_tag: int = None) -> Mapping:
+        response = super()._receive(expected_tag=expected_tag)
+        if response.header.message != usbmuxd_msgtype.PLIST:
+            raise MuxException(f'Received non-plist type {response}')
+        return plistlib.loads(response.data)
+
+    def _send_receive(self, data: Mapping):
+        self._send(data)
+        response = self._receive(self._tag - 1)
+        if response['MessageType'] != 'Result':
+            raise MuxException(f'got an invalid message: {response}')
+        if response['Number'] != 0:
+            raise self._raise_mux_exception(response['Number'], f'got an error message: {response}')
+
+
+def create_mux(usbmux_address: Optional[str] = None) -> MuxConnection:
+    return MuxConnection.create(usbmux_address=usbmux_address)
+
+
+def list_devices(usbmux_address: Optional[str] = None) -> List[MuxDevice]:
+    mux = create_mux(usbmux_address=usbmux_address)
+    mux.get_device_list(0.1)
+    devices = mux.devices
+    mux.close()
+    return devices
+
+
+def select_device(udid: str = None, connection_type: str = None, usbmux_address: Optional[str] = None) \
+        -> Optional[MuxDevice]:
+    """
+    select a UsbMux device according to given arguments.
+    if more than one device could be selected, always prefer the usb one.
+    """
+    tmp = None
+    for device in list_devices(usbmux_address=usbmux_address):
+        if connection_type is not None and device.connection_type != connection_type:
+            # if a specific connection_type was desired and not of this one then skip
+            continue
+
+        if udid is not None and not device.matches_udid(udid):
+            # if a specific udid was desired and not of this one then skip
+            continue
+
+        # save best result as a temporary
+        tmp = device
+
+        if device.is_usb:
+            # always prefer usb connection
+            return device
+
+    return tmp
+
+
+def select_devices_by_connection_type(connection_type: str, usbmux_address: Optional[str] = None) -> List[MuxDevice]:
+    """
+    select all UsbMux devices by connection type
+    """
+    tmp = []
+    for device in list_devices(usbmux_address=usbmux_address):
+        if device.connection_type == connection_type:
+            tmp.append(device)
+
+    return tmp
+
+
+
+class USBMuxHTTPConnection(HTTPConnection):
+    def __init__(self, device: MuxDevice, port=8100):
+        super().__init__("localhost", port)
+        self.__device = device
+        self.__port = port
+
+    def connect(self):
+        self.sock = self.__device.connect(self.__port)
+
+    def __enter__(self) -> HTTPConnection:
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()