AdbCommands structure change, Pycryptodome support, python3 fixes

.gitignore

@@ -5,3 +5,5 @@ adb.egg-info/
 .tox/
 /adb.zip
 /fastboot.zip
+.idea/
+*.DS_Store*

CONTRIBUTORS

@@ -6,3 +6,5 @@ Marc-Antoine Ruel <maruel@chromium.org>
 Max Borghino <fmborghino@gmail.com>
 Mohammad Abu-Garbeyyeh <github@mohammadag.com>
 Josip Delic <delijati@gmail.com>
+Greg E. <greateggsgreg@gmail.com>
+

adb/adb_commands.py

@@ -25,6 +25,7 @@
 import io
 import os
 import socket
+import posixpath
 
 from adb import adb_protocol
 from adb import common
@@ -54,52 +55,127 @@ class AdbCommands(object):
   protocol_handler = adb_protocol.AdbMessage
   filesync_handler = filesync_protocol.FilesyncProtocol
 
-  @classmethod
-  def ConnectDevice(
-      cls, port_path=None, serial=None, default_timeout_ms=None, **kwargs):
-    """Convenience function to get an adb device from usb path or serial.
+  def __init__(self):
+
+    self.__reset()
+
+  def __reset(self):
+    self.build_props = None
+    self._handle = None
+    self._device_state = None
+
+    # Connection table tracks each open AdbConnection objects per service type for program functions
+    # that choose to persist an AdbConnection object for their functionality, using
+    # self._get_service_connection
+    self._service_connections = {}
+
+  def _get_service_connection(self, service, service_command=None, create=True, timeout_ms=None):
+    """
+    Based on the service, get the AdbConnection for that service or create one if it doesnt exist
+
+    :param service:
+    :param service_command: Additional service parameters to append
+    :param create: If False, dont create a connection if it does not exist
+    :return:
+    """
+
+    connection = self._service_connections.get(service, None)
+
+    if connection:
+      return connection
+
+    if not connection and not create:
+      return None
+
+    if service_command:
+      destination_str = b'%s:%s' % (service, service_command)
+    else:
+      destination_str = service
+
+    connection = self.protocol_handler.Open(
+      self._handle, destination=destination_str, timeout_ms=timeout_ms)
+
+    self._service_connections.update({service: connection})
+
+    return connection
+
+  def ConnectDevice(self, port_path=None, serial=None, default_timeout_ms=None, **kwargs):
+    """Convenience function to setup a transport handle for the adb device from 
+         usb path or serial then connect to it.
 
     Args:
       port_path: The filename of usb port to use.
       serial: The serial number of the device to use.
       default_timeout_ms: The default timeout in milliseconds to use.
+      kwargs: handle: Device handle to use (instance of common.TcpHandle or common.UsbHandle)
+              banner: Connection banner to pass to the remote device
+              rsa_keys: List of AuthSigner subclass instances to be used for
+                  authentication. The device can either accept one of these via the Sign
+                  method, or we will send the result of GetPublicKey from the first one
+                  if the device doesn't accept any of them.
+              auth_timeout_ms: Timeout to wait for when sending a new public key. This
+                  is only relevant when we send a new public key. The device shows a
+                  dialog and this timeout is how long to wait for that dialog. If used
+                  in automation, this should be low to catch such a case as a failure
+                  quickly; while in interactive settings it should be high to allow
+                  users to accept the dialog. We default to automation here, so it's low
+                  by default.
 
     If serial specifies a TCP address:port, then a TCP connection is
     used instead of a USB connection.
     """
-    if serial and b':' in serial:
-        handle = common.TcpHandle(serial, timeout_ms=default_timeout_ms)
+
+    # If there isnt a handle override (used by tests), build one here
+    if 'handle' in kwargs:
+      self._handle = kwargs.pop('handle')
+    elif serial and b':' in serial:
+      self._handle = common.TcpHandle(serial, timeout_ms=default_timeout_ms)
     else:
-        handle = common.UsbHandle.FindAndOpen(
+      self._handle = common.UsbHandle.FindAndOpen(
             DeviceIsAvailable, port_path=port_path, serial=serial,
             timeout_ms=default_timeout_ms)
-    return cls.Connect(handle, **kwargs)
 
-  def __init__(self, handle, device_state):
-    self.handle = handle
-    self._device_state = device_state
+    self._Connect(**kwargs)
+
+    return self
 
   def Close(self):
-    self.handle.Close()
+    for conn in list(self._service_connections.values()):
+      if conn:
+        try:
+          conn.Close()
+        except:
+          pass
 
-  @classmethod
-  def Connect(cls, usb, banner=None, **kwargs):
+    if self._handle:
+      self._handle.Close()
+
+    self.__reset()
+
+  def _Connect(self, banner=None, **kwargs):
     """Connect to the device.
 
     Args:
-      usb: UsbHandle or TcpHandle instance to use.
       banner: See protocol_handler.Connect.
-      **kwargs: See protocol_handler.Connect for kwargs. Includes rsa_keys,
-          and auth_timeout_ms.
+      **kwargs: See protocol_handler.Connect and adb_commands.ConnectDevice for kwargs.
+           Includes handle, rsa_keys, and auth_timeout_ms.
     Returns:
       An instance of this class if the device connected successfully.
     """
+
     if not banner:
       banner = socket.gethostname().encode()
-    device_state = cls.protocol_handler.Connect(usb, banner=banner, **kwargs)
+
+    conn_str = self.protocol_handler.Connect(self._handle, banner=banner, **kwargs)
+
     # Remove banner and colons after device state (state::banner)
-    device_state = device_state.split(b':')[0]
-    return cls(usb, device_state)
+    parts = conn_str.split(b'::')
+    self._device_state = parts[0]
+
+    # Break out the build prop info
+    self.build_props = str(parts[1].split(b';'))
+
+    return True
 
   @classmethod
   def Devices(cls):
@@ -129,13 +205,13 @@ def Install(self, apk_path, destination_dir='', timeout_ms=None, replace_existin
     if not destination_dir:
       destination_dir = '/data/local/tmp/'
     basename = os.path.basename(apk_path)
-    destination_path = destination_dir + basename
+    destination_path = posixpath.join(destination_dir, basename)
     self.Push(apk_path, destination_path, timeout_ms=timeout_ms, progress_callback=transfer_progress_callback)
 
     cmd = ['pm install']
     if replace_existing:
       cmd.append('-r')
-    cmd.append('"%s"' % destination_path)
+    cmd.append('"{}"'.format(destination_path))
     return self.Shell(' '.join(cmd), timeout_ms=timeout_ms)
 
   def Uninstall(self, package_name, keep_data=False, timeout_ms=None):
@@ -167,50 +243,56 @@ def Push(self, source_file, device_filename, mtime='0', timeout_ms=None, progres
       progress_callback: callback method that accepts filename, bytes_written and total_bytes,
                          total_bytes will be -1 for file-like objects
     """
-    should_close = False
     if isinstance(source_file, str):
       if os.path.isdir(source_file):
         self.Shell("mkdir " + device_filename)
         for f in os.listdir(source_file):
           self.Push(os.path.join(source_file, f), device_filename + '/' + f, progress_callback=progress_callback)
         return
       source_file = open(source_file, "rb")
-      should_close = True
 
-    connection = self.protocol_handler.Open(
-        self.handle, destination=b'sync:', timeout_ms=timeout_ms)
-    self.filesync_handler.Push(connection, source_file, device_filename,
+    with source_file:
+      connection = self.protocol_handler.Open(
+        self._handle, destination=b'sync:', timeout_ms=timeout_ms)
+      self.filesync_handler.Push(connection, source_file, device_filename,
                                mtime=int(mtime), progress_callback=progress_callback)
-    if should_close:
-      source_file.close()
     connection.Close()
 
-  def Pull(self, device_filename, dest_file='', progress_callback=None, timeout_ms=None):
+  def Pull(self, device_filename, dest_file=None, timeout_ms=None, progress_callback=None):
     """Pull a file from the device.
 
     Args:
       device_filename: Filename on the device to pull.
       dest_file: If set, a filename or writable file-like object.
       timeout_ms: Expected timeout for any part of the pull.
+      progress_callback: callback method that accepts filename, bytes_written and total_bytes,
+                         total_bytes will be -1 for file-like objects
 
     Returns:
-      The file data if dest_file is not set.
+      The file data if dest_file is not set. Otherwise, True if the destination file exists
     """
     if not dest_file:
       dest_file = io.BytesIO()
     elif isinstance(dest_file, str):
-      dest_file = open(dest_file, 'wb')
-    connection = self.protocol_handler.Open(
-        self.handle, destination=b'sync:',
-        timeout_ms=timeout_ms)
-    self.filesync_handler.Pull(connection, device_filename, dest_file, progress_callback)
-    connection.Close()
+      dest_file = open(dest_file, 'w')
+    else:
+      raise ValueError("destfile is of unknown type")
+
+    conn = self.protocol_handler.Open(
+        self._handle, destination=b'sync:', timeout_ms=timeout_ms)
+
+    self.filesync_handler.Pull(conn, device_filename, dest_file, progress_callback)
+
+    conn.Close()
     if isinstance(dest_file, io.BytesIO):
       return dest_file.getvalue()
+    else:
+      dest_file.close()
+      return os.path.exists(dest_file)
 
   def Stat(self, device_filename):
     """Get a file's stat() information."""
-    connection = self.protocol_handler.Open(self.handle, destination=b'sync:')
+    connection = self.protocol_handler.Open(self._handle, destination=b'sync:')
     mode, size, mtime = self.filesync_handler.Stat(
         connection, device_filename)
     connection.Close()
@@ -222,7 +304,7 @@ def List(self, device_path):
     Args:
       device_path: Directory to list.
     """
-    connection = self.protocol_handler.Open(self.handle, destination=b'sync:')
+    connection = self.protocol_handler.Open(self._handle, destination=b'sync:')
     listing = self.filesync_handler.List(connection, device_path)
     connection.Close()
     return listing
@@ -233,32 +315,37 @@ def Reboot(self, destination=b''):
     Args:
       destination: Specify 'bootloader' for fastboot.
     """
-    self.protocol_handler.Open(self.handle, b'reboot:%s' % destination)
+    self.protocol_handler.Open(self._handle, b'reboot:%s' % destination)
 
   def RebootBootloader(self):
     """Reboot device into fastboot."""
     self.Reboot(b'bootloader')
 
   def Remount(self):
     """Remount / as read-write."""
-    return self.protocol_handler.Command(self.handle, service=b'remount')
+    return self.protocol_handler.Command(self._handle, service=b'remount')
 
   def Root(self):
     """Restart adbd as root on the device."""
-    return self.protocol_handler.Command(self.handle, service=b'root')
+    return self.protocol_handler.Command(self._handle, service=b'root')
 
   def EnableVerity(self):
     """Re-enable dm-verity checking on userdebug builds"""
-    return self.protocol_handler.Command(self.handle, service=b'enable-verity')
+    return self.protocol_handler.Command(self._handle, service=b'enable-verity')
 
   def DisableVerity(self):
     """Disable dm-verity checking on userdebug builds"""
-    return self.protocol_handler.Command(self.handle, service=b'disable-verity')
+    return self.protocol_handler.Command(self._handle, service=b'disable-verity')
 
   def Shell(self, command, timeout_ms=None):
-    """Run command on the device, returning the output."""
+    """Run command on the device, returning the output.
+
+    Args:
+      command: Shell command to run
+      timeout_ms: Maximum time to allow the command to run.
+    """
     return self.protocol_handler.Command(
-        self.handle, service=b'shell', command=command,
+        self._handle, service=b'shell', command=command,
         timeout_ms=timeout_ms)
 
   def StreamingShell(self, command, timeout_ms=None):
@@ -272,13 +359,34 @@ def StreamingShell(self, command, timeout_ms=None):
       The responses from the shell command.
     """
     return self.protocol_handler.StreamingCommand(
-        self.handle, service=b'shell', command=command,
+        self._handle, service=b'shell', command=command,
         timeout_ms=timeout_ms)
 
   def Logcat(self, options, timeout_ms=None):
     """Run 'shell logcat' and stream the output to stdout.
 
     Args:
       options: Arguments to pass to 'logcat'.
+      timeout_ms: Maximum time to allow the command to run.
     """
     return self.StreamingShell('logcat %s' % options, timeout_ms)
+
+  def InteractiveShell(self, cmd=None, strip_cmd=True, delim=None, strip_delim=True):
+    """Get stdout from the currently open interactive shell and optionally run a command
+        on the device, returning all output.
+
+    Args:
+      command: Optional. Command to run on the target.
+      strip_cmd: Optional (default True). Strip command name from stdout.
+      delim: Optional. Delimiter to look for in the output to know when to stop expecting more output
+      (usually the shell prompt)
+      strip_delim: Optional (default True): Strip the provided delimiter from the output
+
+    Returns:
+      The stdout from the shell command.
+    """
+    conn = self._get_service_connection(b'shell:')
+
+    return self.protocol_handler.InteractiveShellCommand(
+      conn, cmd=cmd, strip_cmd=strip_cmd,
+      delim=delim, strip_delim=strip_delim)