gooftool: add a flag to enable set sn bits

This is the the fundation of approach (1) and (2) listed in the issue.
We also removed CheckCr50SetSnBitsDependency() because the script is now
always included in the test image.

BUG=b:157210082
TEST=make test

Change-Id: I7d566e9a354cc881639fc3c87552a6f4b5a78c2a
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/factory/+/2212021
Reviewed-by: Stimim Chen <stimim@chromium.org>
Tested-by: Stimim Chen <stimim@chromium.org>
Commit-Queue: Stimim Chen <stimim@chromium.org>
diff --git a/py/gooftool/commands.py b/py/gooftool/commands.py
index c889d9b..cf97522 100755
--- a/py/gooftool/commands.py
+++ b/py/gooftool/commands.py
@@ -267,6 +267,10 @@
     '--no_generate_mfg_date', action='store_false', dest='generate_mfg_date',
     help='Do not generate manufacturing date nor write mfg_date into VPD.')
 
+_enable_zero_touch_cmd_arg = CmdArg(
+    '--enable_zero_touch', action='store_true',
+    help='Set attested_device_id for zero-touch feature.')
+
 
 @Command(
     'verify_ec_key',
@@ -364,9 +368,11 @@
   return GetGooftool(options).VerifyCrosConfig()
 
 
-@Command('verify-sn-bits')
+@Command('verify-sn-bits',
+         _enable_zero_touch_cmd_arg)
 def VerifySnBits(options):
-  GetGooftool(options).VerifySnBits()
+  if options.enable_zero_touch:
+    GetGooftool(options).VerifySnBits()
 
 
 @Command('write_protect')
@@ -466,17 +472,11 @@
 
 @Command('cr50_write_flash_info',
          _rma_mode_cmd_arg,
-         CmdArg('--expect_zero_touch', action='store_true',
-                help='zero touch feature is expected, the command will fail '
-                     'immediately if required dependencies are not found.'))
+         _enable_zero_touch_cmd_arg)
 def Cr50WriteFlashInfo(options):
   """Set the serial number bits, board id and flags on the Cr50 chip."""
-  # The '--expect_zero_touch' argument is for testing purpose, therefore, the
-  # argument can only be specified by directly using `cr50_write_flash_info`
-  # subcommand.  And the `expect_zero_touch` attribute won't exist when this
-  # function is invoked by other subcommands, e.g. `finalize`.
-  expect_zero_touch = getattr(options, 'expect_zero_touch', False)
-  GetGooftool(options).Cr50WriteFlashInfo(expect_zero_touch, options.rma_mode)
+  GetGooftool(options).Cr50WriteFlashInfo(
+      options.enable_zero_touch, options.rma_mode)
   event_log.Log('cr50_write_flash_info')
 
 
@@ -566,7 +566,8 @@
          _enforced_release_channels_cmd_arg,
          _waive_list_cmd_arg,
          _skip_list_cmd_arg,
-         _no_ectool_cmd_arg)
+         _no_ectool_cmd_arg,
+         _enable_zero_touch_cmd_arg)
 def Verify(options):
   """Verifies if whole factory process is ready for finalization.
 
@@ -812,7 +813,8 @@
          _rlz_embargo_end_date_offset_cmd_arg,
          _waive_list_cmd_arg,
          _skip_list_cmd_arg,
-         _no_generate_mfg_date_cmd_arg)
+         _no_generate_mfg_date_cmd_arg,
+         _enable_zero_touch_cmd_arg)
 def Finalize(options):
   """Verify system readiness and trigger transition into release state.
 
diff --git a/py/gooftool/core.py b/py/gooftool/core.py
index 00a153b..cfd8a2a 100644
--- a/py/gooftool/core.py
+++ b/py/gooftool/core.py
@@ -629,8 +629,6 @@
       raise Error('write protectioin switch of EC is disabled.')
 
   def VerifySnBits(self):
-    if not self.CheckCr50SetSnBitsDependency():
-      return
     # Add '-n' to dry run.
     result = self._util.shell(['/usr/share/cros/cr50-set-sn-bits.sh', '-n'])
     stdout = result.stdout.strip()
@@ -640,9 +638,6 @@
     logging.info('stderr: %s', stderr)
 
     if result.status != 0:
-      # TODO(b/157210082): In the future, it should be okay to have
-      # cr50-set-sn-bits.sh, but not enabling zero-touch.
-
       # Fail reason, either:
       # - attested_device_id is not set
       # - SN bits has been set differently
@@ -864,10 +859,6 @@
            codecs.encode(secret_bytes, 'hex').decode('utf-8')},
           partition=vpd.VPD_READONLY_PARTITION_NAME)
 
-  def CheckCr50SetSnBitsDependency(self):
-    script_path = '/usr/share/cros/cr50-set-sn-bits.sh'
-    return os.path.exists(script_path)
-
   def Cr50SetSnBits(self):
     """Set the serial number bits on the Cr50 chip.
 
@@ -883,14 +874,6 @@
     vpd_key = 'attested_device_id'
     has_vpd_key = self._vpd.GetValue(vpd_key) is not None
 
-    # If the script does not exist, that board is not able to do Zero-Touch.
-    if not self.CheckCr50SetSnBitsDependency():
-      logging.warning('The Cr50 script to set serial number bits is not found, '
-                      'those bits will not be set on this device.')
-      if has_vpd_key:
-        raise Error('Zero-Touch is not enabled, but %r is set.' % vpd_key)
-      return
-
     # The script exists, Zero-Touch is enabled.
     if not has_vpd_key:
       # TODO(stimim): What if Zero-Touch is enabled on a program (e.g. hatch),
@@ -963,14 +946,8 @@
       logging.exception('Failed to set Cr50 Board ID.')
       raise
 
-  def Cr50WriteFlashInfo(self, expect_zero_touch=False, rma_mode=False):
+  def Cr50WriteFlashInfo(self, enable_zero_touch=False, rma_mode=False):
     """Write device info into cr50 flash."""
-    if expect_zero_touch and not self.CheckCr50SetSnBitsDependency():
-      logging.error('zero_touch feature is expected, but we cannot find '
-                    'required dependencies.  Please check if USE flag '
-                    '`zero_touch` is set when building the test image.')
-      return
-
     cros_config = cros_config_module.CrosConfig(self._util.shell)
     is_whitelabel, whitelabel_tag = cros_config.GetWhiteLabelTag()
 
@@ -992,7 +969,7 @@
         raise Error('whitelabel_tag reported by cros_config and VPD does not '
                     'match.  Have you reboot the device after updating VPD '
                     'fields?')
-    if not rma_mode:
+    if not rma_mode and enable_zero_touch:
       self.Cr50SetSnBits()
     self.Cr50SetBoardId(is_whitelabel)