servo: create common ap-driver

Beating the login hurdle is one aspect that several
servo automations might want to share. This CL introduces
a common driver for AP console communication that houses
tools to beat the login hurdle on the AP console.
Going forward common utilities to do work on the AP console
should be housed in ap.py, and automation drivers that
interact with the AP console should inherit from or use
an instance of the ap driver.

CL also shores up the logic and increases the speed to beat
the login hurdle.

BUG=chromium:760267
TEST=manual
sudo servod -b $BOARD
dut-control login
login:no
dut-control login:yes
dut-control login
login:yes

Change-Id: I253abf30e4f0c90bc915b31b8fe57d9b3a9ed4c0
Reviewed-on: https://chromium-review.googlesource.com/870028
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>
Reviewed-by: Todd Broch <tbroch@chromium.org>
diff --git a/servo/data/uart_common.xml b/servo/data/uart_common.xml
index 126e266..eb8d07a 100644
--- a/servo/data/uart_common.xml
+++ b/servo/data/uart_common.xml
@@ -44,19 +44,19 @@
     <name>login</name>
     <doc>Query if AP UART session is logged in. Ask to login/logout of a session
     by calling login:yes and login:no</doc>
-    <params map="yesno" subtype="login" interface="8" drv="login">
+    <params map="yesno" subtype="login" interface="8" drv="ap">
     </params>
   </control>
   <control>
     <name>login_username</name>
     <doc>Username used to log into a session.</doc>
-    <params input_type="str" subtype="username" interface="8" drv="login">
+    <params input_type="str" subtype="username" interface="8" drv="ap">
     </params>
   </control>
   <control>
     <name>login_password</name>
     <doc>Password used to log into a session.</doc>
-    <params input_type="str" subtype="password" interface="8" drv="login">
+    <params input_type="str" subtype="password" interface="8" drv="ap">
     </params>
   </control>
   <control>
diff --git a/servo/drv/__init__.py b/servo/drv/__init__.py
index 053bc36..beb3391 100644
--- a/servo/drv/__init__.py
+++ b/servo/drv/__init__.py
@@ -39,7 +39,7 @@
 import larvae_adc
 import lcm2004
 import link_power
-import login
+import ap
 import loglevel
 import ltc1663
 import lumpy_power
diff --git a/servo/drv/ap.py b/servo/drv/ap.py
new file mode 100644
index 0000000..996f338
--- /dev/null
+++ b/servo/drv/ap.py
@@ -0,0 +1,123 @@
+# Copyright 2017 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""
+    AP console communications driver. Automations that need to interact
+    with the AP console should inherit from this driver, or build a composition
+    containing this driver.
+    This driver exposes methods to login, logout, set password, and
+    username to login, and check if a session is logged in.
+"""
+import logging
+import pexpect
+import time
+
+import pty_driver
+
+
+class apError(Exception):
+  """Exception class for AP errors."""
+
+class ap(pty_driver.ptyDriver):
+  """ Wrapper class around ptyDriver to handle communication
+      with the AP console.
+  """
+
+  # This is a class level variable because each
+  # servo command runs on their own instance of
+  # dutMetadata. We share it across instances
+  # to allow for updating and retrieval.
+
+  # TODO(coconutruben): right now the defaults are hard coded
+  # into the driver. Evaluate if it makes more sense for the
+  # defaults to be set by the xml commands, or to allow servo
+  # invocation to overwrite them.
+  _login_info = {
+      'username'  : 'root',
+      'password'  : 'test0000'
+  }
+
+  def __init__(self, interface, params):
+    """Initializes the AP driver.
+
+    Args:
+      interface: A driver interface object. This is the AP uart interface.
+      params: A dictionary of parameters, but is ignored.
+    """
+    super(ap, self).__init__(interface, params)
+    self._logger.debug("")
+
+  def _Get_password(self):
+    """ Returns password currently used for login attempts. """
+    return self._login_info['password']
+
+  def _Set_password(self, value):
+    """ Set |value| as password for login attempts. """
+    self._login_info['password'] = value
+
+  def _Get_username(self):
+    """ Returns username currently used for login attempts. """
+    return self._login_info['username']
+
+  def _Set_username(self, value):
+    """ Set |value| as username for login attempts. """
+    self._login_info['username'] = value
+
+  def _Get_login(self):
+    """ Heuristic to determine if a session is logged in
+        on the CPU uart terminal.
+
+        Sends a newline to the terminal, and evaluates the output
+        to determine login status.
+
+        Returns:
+          True 1 a session is logged in, 0 otherwise.
+    """
+    try:
+      match = self._issue_cmd_get_results([''], [r"localhost\x1b\[01;34m\s"
+                                                 r"[^\s/]+\s[#$]|"
+                                                 r"localhost login:"])
+      return 0 if 'localhost login:' in match[0] else 1
+    except pty_driver.ptyError:
+      return 0
+
+  def _Set_login(self, value):
+    """ Login/out of a session on the CPU uart terminal.
+
+        Uses username and password stored in |_login_info| to
+        attempt login.
+    """
+    # TODO(coconutruben): the login/logout logic fails silently and user has
+    # to call login command again to verify. Consider if raising an error on
+    # failure here is appropiate.
+    if value == 1:
+      # 1 means login desired.
+      if not self._Get_login():
+        with self._open():
+          # Make sure uart capture does not interfere with matching the expected
+          # response
+          self._interface.pause_capture()
+          try:
+            self._send("")
+            self._child.expect('localhost login:', 3)
+            match = self._child.match
+            if not match:
+              raise apError("Username prompt did not show up on login attempt")
+            self._send(self._login_info['username'], flush=False)
+            self._child.expect('Password:', 2)
+            match = self._child.match
+            if not match:
+              raise apError("Password prompt did not show up on login attempt")
+            self._send(self._login_info['password'], flush=False)
+          except pexpect.TIMEOUT:
+            raise apError("Timeout waiting for response when attempting to "
+                          "log into AP console.")
+          finally:
+            # Reenable capturing the console output
+            self._interface.resume_capture()
+        time.sleep(0.1)
+    if value == 0:
+      # 0 means logout desired.
+      if self._Get_login():
+        self._issue_cmd(['exit'])
diff --git a/servo/drv/login.py b/servo/drv/login.py
deleted file mode 100644
index b5e2b09..0000000
--- a/servo/drv/login.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# Copyright 2017 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-"""
-    Login driver to handle the login hurdle on the CPU uart console.
-    This driver exposes methods to login, logout, set password, and
-    username to login, and check if a session is logged in.
-"""
-import logging
-import time
-
-import pty_driver
-
-
-class loginError(Exception):
-  """Exception class for login errors."""
-
-class login(pty_driver.ptyDriver):
-  """ Wrapper class around ptyDriver to handle communication
-      with the AP console.
-  """
-
-  # This is a class level variable because each
-  # servo command runs on their own instance of
-  # dutMetadata. We share it across instances
-  # to allow for updating and retrieval.
-
-  # TODO(coconutruben): right now the defaults are hard coded
-  # into the driver. Evaluate if it makes more sense for the
-  # defaults to be set by the xml commands.
-  _login_info = {
-      'username'  : 'root',
-      'password'  : 'test0000'
-  }
-
-  def __init__(self, interface, params):
-    """Initializes the login driver.
-
-    Args:
-      interface: A driver interface object. This is the AP uart interface.
-      params: A dictionary of parameters, but is ignored.
-    """
-    super(login, self).__init__(interface, params)
-    self._logger.debug("")
-
-  def _Get_password(self):
-    """ Returns password currently used for login attempts. """
-    return self._login_info['password']
-
-  def _Set_password(self, value):
-    """ Set |value| as password for login attempts. """
-    self._login_info['password'] = value
-
-  def _Get_username(self):
-    """ Returns username currently used for login attempts. """
-    return self._login_info['username']
-
-  def _Set_username(self, value):
-    """ Set |value| as username for login attempts. """
-    self._login_info['username'] = value
-
-  def _Get_login(self):
-    """ Heuristic to determine if a session is logged in
-        on the CPU uart terminal.
-
-        Sends a newline to the terminal, and evaluates the output
-        to determine login status.
-
-        Returns:
-          True 1 a session is logged in, 0 otherwise.
-    """
-    # TODO(coconutruben): currently, this assumes localhost as the
-    # host name, and a specific pattern for PS1. Make this more
-    # robust to OS changes by generating the regex based on the
-    # PS1 variable.
-    try:
-      self._issue_cmd_get_results([''], ['localhost.*\s+[#$]'])
-      return 1
-    except pty_driver.ptyError:
-      return 0
-
-  def _Set_login(self, value):
-    """ Login/out of a session on the CPU uart terminal.
-
-        Uses username and password stored in |_login_info| to
-        attempt login.
-    """
-    # TODO(coconutruben): the login/logout logic fails silently and user has
-    # to call login command again to verify. Consider if raising an error on
-    # failure here is appropiate.
-    # TODO(coconutruben): the login sequence relies on timing currently.
-    # Consider if it would be more robust to use fdexpect code directly
-    # to see when to send the username, and the password.
-    if value == 1:
-      # 1 means login desired.
-      if not self._Get_login():
-        with self._open():
-          self._send([self._login_info['username'],
-                      self._login_info['password']],
-                     0.1, False)
-      time.sleep(0.1)
-    if value == 0:
-      # 0 means logout desired.
-      if self._Get_login():
-        self._issue_cmd(['exit'])
diff --git a/servo/terminal_freezer.py b/servo/terminal_freezer.py
index c6bbce8..5770714 100644
--- a/servo/terminal_freezer.py
+++ b/servo/terminal_freezer.py
@@ -36,10 +36,11 @@
     CheckForPIDNamespace()
 
   def __enter__(self):
+    ret = ''
     try:
       ret = subprocess.check_output(['lsof', '-FR', self._tty],
                                     stderr=subprocess.STDOUT)
-    except subprocess.check_output:
+    except subprocess.CalledProcessError:
       # Ignore non-zero return codes.
       pass