pw_env_setup: Add PW_CIPD_SERVICE_ACCOUNT variable

This is expected to be useful when bootstrapping a Pigweed project on
a GCE VM outside of a LUCI context.

Bug: b/244166911
Change-Id: I3b42c0734a9635bab8bfcf41b3f8889a2d3eb10b
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/109272
Commit-Queue: Ted Pudlik <tpudlik@google.com>
Reviewed-by: Joe Brennan <jmbrenna@google.com>
Reviewed-by: Rob Mohr <mohrr@google.com>
diff --git a/pw_env_setup/docs.rst b/pw_env_setup/docs.rst
index f85a647..7ae6782 100644
--- a/pw_env_setup/docs.rst
+++ b/pw_env_setup/docs.rst
@@ -419,6 +419,11 @@
   Python executable to be used, for example "python2" or "python3". Defaults to
   "python".
 
+``PW_CIPD_SERVICE_ACCOUNT_JSON``
+  Value to pass as ``-service-account-json`` to CIPD invocations. This should
+  point either to a service account JSON key file, or be the magical value
+  ``:gce`` to tell the tool to fetch tokens from GCE metadata server.
+
 ``PW_ENVIRONMENT_ROOT``
   Location to which packages are installed. Defaults to ``.environment`` folder
   within the checkout root.
diff --git a/pw_env_setup/py/pw_env_setup/cipd_setup/update.py b/pw_env_setup/py/pw_env_setup/cipd_setup/update.py
index d5ef2c0..efae350 100755
--- a/pw_env_setup/py/pw_env_setup/cipd_setup/update.py
+++ b/pw_env_setup/py/pw_env_setup/cipd_setup/update.py
@@ -30,8 +30,11 @@
 import sys
 
 
-def check_auth(cipd, package_files, spin):
+def check_auth(cipd, package_files, cipd_service_account, spin):
     """Check have access to CIPD pigweed directory."""
+    cmd = [cipd]
+    if cipd_service_account:
+        cmd.extend(['-service-account-json', cipd_service_account])
 
     paths = []
     for package_file in package_files:
@@ -48,7 +51,7 @@
 
     username = None
     try:
-        output = subprocess.check_output([cipd, 'auth-info'],
+        output = subprocess.check_output(cmd + ['auth-info'],
                                          stderr=subprocess.STDOUT).decode()
         logged_in = True
 
@@ -66,7 +69,7 @@
             # Not catching CalledProcessError because 'cipd ls' seems to never
             # return an error code unless it can't reach the CIPD server.
             output = subprocess.check_output(
-                [cipd, 'ls', path], stderr=subprocess.STDOUT).decode()
+                cmd + ['ls', path], stderr=subprocess.STDOUT).decode()
             if 'No matching packages' not in output:
                 continue
 
@@ -75,7 +78,7 @@
             # 'cipd instances' does use an error code if there's no such package
             # or that package is inaccessible.
             try:
-                subprocess.check_output([cipd, 'instances', path],
+                subprocess.check_output(cmd + ['instances', path],
                                         stderr=subprocess.STDOUT)
             except subprocess.CalledProcessError:
                 inaccessible_paths.append(path)
@@ -95,7 +98,8 @@
             stderr()
             stderr('Attempting CIPD login')
             try:
-                subprocess.check_call([cipd, 'auth-login'])
+                # Note that with -service-account-json, auth-login is a no-op.
+                subprocess.check_call(cmd + ['auth-login'])
             except subprocess.CalledProcessError:
                 stderr('CIPD login failed')
                 return False
@@ -284,6 +288,14 @@
         '-max-threads', '0',  # 0 means use CPU count.
     ]  # yapf: disable
 
+    cipd_service_account = None
+    if env_vars:
+        cipd_service_account = env_vars.get('PW_CIPD_SERVICE_ACCOUNT_JSON')
+    if not cipd_service_account:
+        cipd_service_account = os.environ.get('PW_CIPD_SERVICE_ACCOUNT_JSON')
+    if cipd_service_account:
+        cmd.extend(['-service-account-json', cipd_service_account])
+
     hasher = hashlib.sha256()
     encoded = '\0'.join(cmd)
     if hasattr(encoded, 'encode'):
@@ -308,7 +320,7 @@
                 if digest == digest_file:
                     return True
 
-    if not check_auth(cipd, package_files, spin):
+    if not check_auth(cipd, package_files, cipd_service_account, spin):
         return False
 
     # TODO(pwbug/135) Use function from common utility module.