au_test_harness: fork sudo module here

Until this project is migrated to Python 3, fork modules here.

BUG=chromium:1170006, chromium:1171299
TEST=CQ passes

Change-Id: I06a8605d864ead4f89972af726f3e9ab3ed956ec
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crostestutils/+/2659358
Tested-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Sean Abraham <seanabraham@chromium.org>
diff --git a/au_test_harness/cros_au_test_harness.py b/au_test_harness/cros_au_test_harness.py
index 8a1ddc5..978381f 100755
--- a/au_test_harness/cros_au_test_harness.py
+++ b/au_test_harness/cros_au_test_harness.py
@@ -30,7 +30,7 @@
 from chromite.lib import cros_build_lib
 from chromite.lib import cros_logging as logging
 from chromite.lib import parallel
-from chromite.lib import sudo
+import sudo  # TODO(crbug.com/1170006): Move back to chromite.
 from chromite.lib import timeout_util
 
 from crostestutils.au_test_harness import au_test
diff --git a/au_test_harness/sudo.py b/au_test_harness/sudo.py
new file mode 100644
index 0000000..55adf06
--- /dev/null
+++ b/au_test_harness/sudo.py
@@ -0,0 +1,155 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2011-2012 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.
+
+"""Helper methods and classes related to managing sudo."""
+
+from __future__ import print_function
+
+import errno
+import os
+import signal
+import subprocess
+import sys
+
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_logging as logging
+
+
+class SudoKeepAlive(cros_build_lib.MasterPidContextManager):
+  """Keep sudo auth cookie fresh.
+
+  This refreshes the sudo auth cookie; this is implemented this
+  way to ensure that sudo has access to both invoking tty, and
+  will update the user's tty-less cookie.
+  see crosbug/18393.
+  """
+
+  def __init__(self, ttyless_sudo=True, repeat_interval=4):
+    """Run sudo with a noop, to reset the sudo timestamp.
+
+    Args:
+      ttyless_sudo: Whether to update the tty-less cookie.
+      repeat_interval: In minutes, the frequency to run the update.
+    """
+    cros_build_lib.MasterPidContextManager.__init__(self)
+    self._ttyless_sudo = ttyless_sudo
+    self._repeat_interval = repeat_interval
+    self._proc = None
+    self._existing_keepalive_value = None
+
+  @staticmethod
+  def _IdentifyTTY():
+    for source in (sys.stdin, sys.stdout, sys.stderr):
+      try:
+        return os.ttyname(source.fileno())
+      except EnvironmentError as e:
+        if e.errno not in (errno.EINVAL, errno.ENOTTY):
+          raise
+
+    return 'unknown'
+
+  def _DaemonNeeded(self):
+    """Discern which TTYs require sudo keep alive code.
+
+    Returns:
+      A string representing the set of ttys we need daemons for.
+      This will be the empty string if no daemon is needed.
+    """
+    existing = os.environ.get('CROS_SUDO_KEEP_ALIVE')
+    needed = set([self._IdentifyTTY()])
+    if self._ttyless_sudo:
+      needed.add('unknown')
+    if existing is not None:
+      needed -= set(existing.split(':'))
+    return ':'.join(needed)
+
+  def _enter(self):
+    if os.getuid() == 0:
+      cros_build_lib.Die('This script cannot be run as root.')
+
+    start_for_tty = self._DaemonNeeded()
+    if not start_for_tty:
+      # Daemon is already started.
+      return
+
+    # Note despite the impulse to use 'sudo -v' instead of 'sudo true', the
+    # builder's sudoers configuration is slightly whacked resulting in it
+    # asking for password everytime.  As such use 'sudo true' instead.
+    cmds = ['sudo -n true 2>/dev/null',
+            'sudo -n true < /dev/null > /dev/null 2>&1']
+
+    # First check to see if we're already authed.  If so, then we don't
+    # need to prompt the user for their password.
+    for idx, cmd in enumerate(cmds):
+      ret = cros_build_lib.run(
+          cmd, print_cmd=False, shell=True, check=False)
+
+      if ret.returncode != 0:
+        tty_msg = 'Please disable tty_tickets using these instructions: %s'
+        if os.path.exists('/etc/goobuntu'):
+          url = 'https://goto.google.com/chromeos-sudoers'
+        else:
+          url = 'https://goo.gl/fz9YW'
+
+        # If ttyless sudo is not strictly required for this script, don't
+        # prompt for a password a second time. Instead, just complain.
+        if idx > 0:
+          logging.error(tty_msg, url)
+          if not self._ttyless_sudo:
+            break
+
+        # We need to go interactive and allow sudo to ask for credentials.
+        interactive_cmd = cmd.replace(' -n', '')
+        cros_build_lib.run(interactive_cmd, shell=True, print_cmd=False)
+
+        # Verify that sudo access is set up properly.
+        try:
+          cros_build_lib.run(cmd, shell=True, print_cmd=False)
+        except cros_build_lib.RunCommandError:
+          if idx == 0:
+            raise
+          cros_build_lib.Die('tty_tickets must be disabled. ' + tty_msg, url)
+
+    # Anything other than a timeout results in us shutting down.
+    repeat_interval = self._repeat_interval * 60
+    cmd = ('while :; do read -t %i; [ $? -le 128 ] && exit; %s; done' %
+           (repeat_interval, '; '.join(cmds)))
+
+    def ignore_sigint():
+      # We don't want our sudo process shutdown till we shut it down;
+      # since it's part of the session group it however gets SIGINT.
+      # Thus suppress it (which bash then inherits).
+      signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    # We don't use threads here.
+    # pylint: disable=bad-option-value,subprocess-popen-preexec-fn
+    self._proc = subprocess.Popen(['bash', '-c', cmd], shell=False,
+                                  close_fds=True, preexec_fn=ignore_sigint,
+                                  stdin=subprocess.PIPE)
+
+    self._existing_keepalive_value = os.environ.get('CROS_SUDO_KEEP_ALIVE')
+    os.environ['CROS_SUDO_KEEP_ALIVE'] = start_for_tty
+
+  def _exit(self, exc_type, exc, exc_tb):
+    if self._proc is None:
+      return
+
+    try:
+      self._proc.terminate()
+      self._proc.wait()
+    except EnvironmentError as e:
+      if e.errno != errno.ESRCH:
+        raise
+
+    if self._existing_keepalive_value is not None:
+      os.environ['CROS_SUDO_KEEP_ALIVE'] = self._existing_keepalive_value
+    else:
+      os.environ.pop('CROS_SUDO_KEEP_ALIVE', None)
+
+
+def SetFileContents(path, value, cwd=None):
+  """Set a given filepath contents w/ the passed in value."""
+  cros_build_lib.sudo_run(['tee', path], stdout=True,
+                          print_cmd=False, input=value, cwd=cwd)