Fix AndroidDevice#take_screenshot failure caused by bracket characters

mobly/controllers/android_device.py

@@ -1076,7 +1076,7 @@ def take_screenshot(
     filename_no_extension, _ = os.path.splitext(filename)
     device_path = os.path.join('/storage/emulated/0/', filename)
     self.adb.shell(
-        ['screencap', '-p', '-a' if all_displays else '', device_path],
+        ['screencap', '-p', '-a' if all_displays else '', f'"{device_path}"'],
         timeout=TAKE_SCREENSHOT_TIMEOUT_SECOND,
     )
     utils.create_dir(destination)
@@ -1099,13 +1099,13 @@ def take_screenshot(
             os.path.join(destination, os.path.basename(device_path))
         )
         self.log.debug('Screenshot taken, saved on the host: %s', pic_paths[-1])
-        self.adb.shell(['rm', device_path])
+        self.adb.shell(['rm', f'"{device_path}"'])
       return pic_paths
     # handle single screenshot when all_displays=False
     self.adb.pull([device_path, destination])
     pic_path = os.path.join(destination, filename)
     self.log.debug('Screenshot taken, saved on the host: %s', pic_path)
-    self.adb.shell(['rm', device_path])
+    self.adb.shell(['rm', f'"{device_path}"'])
     return pic_path
 
   def run_iperf_client(self, server_host, extra_args=''):

tests/lib/mock_android_device.py

@@ -34,11 +34,18 @@
     'ro.hardware': 'marlin',
 }
 
+_ESCAPE_CHARACTER = '\\'
+_SPECIAL_CHARS_IN_FILE_PATH = ('(', ')')
+
 
 class Error(Exception):
   pass
 
 
+class AdbShellCmdInvalidPathError(Error):
+  pass
+
+
 def get_mock_ads(num):
   """Generates a list of mock AndroidDevice objects.
 
@@ -76,6 +83,35 @@ def list_adb_devices():
   return [ad.serial for ad in get_mock_ads(5)]
 
 
+def _assert_valid_path_in_adb_shell_cmd(path: str):
+  """Asserts that file paths passed to adb shell commands are valid.
+
+  File paths that contain special characters should be quoted, or the special
+  characters should be escaped.
+  """
+  if not path:
+    return
+  if not any(ch in path for ch in _SPECIAL_CHARS_IN_FILE_PATH):
+    return
+  if any(
+      [
+          path[0] == '"' and path[-1] == '"',
+          path[0] == "'" and path[-1] == "'",
+      ]
+  ):
+    return
+  for idx, ch in enumerate(path):
+    if ch not in _SPECIAL_CHARS_IN_FILE_PATH:
+      continue
+    if idx > 0 and path[idx - 1] == _ESCAPE_CHARACTER:
+      continue
+    raise AdbShellCmdInvalidPathError(
+        'Got invalid file path in adb shell command. The path is not quoted '
+        f'and special character, character "{ch}" at index {idx}, is not'
+        f'escaped. Full path: {path}'
+    )
+
+
 class MockAdbProxy:
   """Mock class that swaps out calls to adb with mock calls."""
 
@@ -137,6 +173,14 @@ def shell(self, params, timeout=None):
       )
     elif 'which' in params:
       return b''
+    elif isinstance(params, list) and params and params[0] == 'rm':
+      for arg in params[1:]:
+        if arg and not arg.startswith('-'):
+          _assert_valid_path_in_adb_shell_cmd(arg)
+    elif isinstance(params, list) and params and params[0] == 'screencap':
+      for arg in params[1:]:
+        if arg and not arg.startswith('-'):
+          _assert_valid_path_in_adb_shell_cmd(arg)
 
   def getprop(self, params):
     if params in self.mock_properties:

tests/mobly/controllers/android_device_test.py

@@ -1159,6 +1159,33 @@ def test_AndroidDevice_take_screenshot(
         ),
     )
 
+  @mock.patch(
+      'mobly.controllers.android_device_lib.adb.AdbProxy',
+      return_value=mock_android_device.MockAdbProxy('1'),
+  )
+  @mock.patch(
+      'mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+      return_value=mock_android_device.MockFastbootProxy('1'),
+  )
+  @mock.patch('mobly.utils.create_dir')
+  @mock.patch('mobly.logger.get_log_file_timestamp')
+  def test_AndroidDevice_take_screenshot_when_debug_tag_has_special_characters(
+      self,
+      get_log_file_timestamp_mock,
+      create_dir_mock,
+      FastbootProxy,
+      MockAdbProxy,
+  ):
+    get_log_file_timestamp_mock.return_value = '07-22-2019_17-53-34-450'
+    mock_serial = '1'
+    ad = android_device.AndroidDevice(serial=mock_serial)
+    ad.debug_tag = '1(Tester)'
+
+    try:
+      ad.take_screenshot(self.tmp_dir)
+    except Exception:
+      self.fail('take_screenshot method should not raise an exception.')
+
   @mock.patch(
       'mobly.controllers.android_device_lib.adb.AdbProxy',
       return_value=mock_android_device.MockAdbProxy('1'),
@@ -1301,6 +1328,46 @@ def custom_shell_for_screenshot(params, timeout=None):
         ],
     )
 
+  @mock.patch(
+      'mobly.controllers.android_device_lib.adb.AdbProxy',
+      return_value=mock_android_device.MockAdbProxy('1'),
+  )
+  @mock.patch(
+      'mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+      return_value=mock_android_device.MockFastbootProxy('1'),
+  )
+  @mock.patch('mobly.utils.create_dir')
+  @mock.patch('mobly.logger.get_log_file_timestamp')
+  def test_AndroidDevice_take_screenshot_all_displays_when_debug_tag_has_special_char(
+      self,
+      get_log_file_timestamp_mock,
+      create_dir_mock,
+      FastbootProxy,
+      MockAdbProxy,
+  ):
+    test_adb_proxy = MockAdbProxy.return_value
+    original_mock_adb_instance_shell = test_adb_proxy.shell
+
+    def custom_shell_for_screenshot(params, timeout=None):
+      if f'ls /storage/emulated/0/*.png' in params:
+        return (
+            b'/storage/emulated/0/screenshot,1(Tester),1,fakemodel,07-22-2019_17-53-34-450_0.png\n'
+            b'/storage/emulated/0/screenshot,1(Tester),1,fakemodel,07-22-2019_17-53-34-450_1.png\n'
+        )
+      return original_mock_adb_instance_shell(params, timeout)
+
+    test_adb_proxy.shell = custom_shell_for_screenshot
+
+    get_log_file_timestamp_mock.return_value = '07-22-2019_17-53-34-450'
+    mock_serial = '1'
+    ad = android_device.AndroidDevice(serial=mock_serial)
+    ad.debug_tag = '1(Tester)'
+
+    try:
+      ad.take_screenshot(self.tmp_dir, all_displays=True)
+    except Exception:
+      self.fail('take_screenshot method should not raise an exception.')
+
   @mock.patch(
       'mobly.controllers.android_device_lib.adb.AdbProxy',
       return_value=mock_android_device.MockAdbProxy('1'),