Add more shopfloor flexibility.

BUG=None
TEST=Manual on device

Change-Id: Icef86f40a131dfe6b16cde03b3c631c70c6353bd
Reviewed-on: https://gerrit.chromium.org/gerrit/45067
Reviewed-by: Shuo-Peng Liao <deanliao@chromium.org>
Tested-by: Jon Salz <jsalz@chromium.org>
Commit-Queue: Jon Salz <jsalz@chromium.org>
diff --git a/py/gooftool/gooftool.py b/py/gooftool/gooftool.py
index 96d3bc2..65232dd 100755
--- a/py/gooftool/gooftool.py
+++ b/py/gooftool/gooftool.py
@@ -38,7 +38,7 @@
 from cros.factory.event_log import TimedUuid
 from cros.factory.test.factory import FACTORY_LOG_PATH
 from cros.factory.utils.process_utils import Spawn
-from cros.factory.system.vpd import FilterVPD
+from cros.factory.privacy import FilterDict
 
 
 # Use a global event log, so that only a single log is created when
@@ -345,7 +345,7 @@
       value = rw_vpd[key]
       if (known_valid_values is not None) and (value not in known_valid_values):
         sys.exit('Invalid RW VPD entry : key %r, value %r' % (key, value))
-    _event_log.Log('vpd', ro_vpd=FilterVPD(ro_vpd), rw_vpd=FilterVPD(rw_vpd))
+    _event_log.Log('vpd', ro_vpd=FilterDict(ro_vpd), rw_vpd=FilterDict(rw_vpd))
   map(hwid_tool.Validate.Status, options.status)
 
   if not options.hwid or not options.probe_results:
diff --git a/py/privacy.py b/py/privacy.py
new file mode 100644
index 0000000..5c870be
--- /dev/null
+++ b/py/privacy.py
@@ -0,0 +1,27 @@
+#!/usr/bin/python
+# pylint: disable=W0212
+#
+# Copyright (c) 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.
+
+
+# Keys that may not be logged (in VPDs or device data).
+BLACKLIST_KEYS = [
+  'ubind_attribute',
+  'gbind_attribute'
+]
+
+
+def FilterDict(data):
+  """Redacts values of any keys in BLACKLIST_KEYS.
+
+  Args:
+    data: A dictionary to redact.
+  """
+  def FilterItem(k, v):
+    if v is None:
+      return None
+    return '<redacted %d chars>' % len(v) if k in BLACKLIST_KEYS else v
+
+  return dict((k, FilterItem(k, v)) for k, v in data.iteritems())
diff --git a/py/privacy_unittest.py b/py/privacy_unittest.py
new file mode 100755
index 0000000..c5a848c
--- /dev/null
+++ b/py/privacy_unittest.py
@@ -0,0 +1,23 @@
+#!/usr/bin/python
+# pylint: disable=W0212
+#
+# Copyright (c) 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.
+
+import unittest2
+
+import factory_common  # pylint: disable=W0611
+from cros.factory import privacy
+
+class PrivacyTest(unittest2.TestCase):
+  def testFilterDict(self):
+    self.assertEquals(
+        dict(a='A', b='B',
+             ubind_attribute='<redacted 1 chars>',
+             gbind_attribute='<redacted 2 chars>'),
+        privacy.FilterDict(
+            dict(a='A', b='B', ubind_attribute='U', gbind_attribute='GG')))
+
+if __name__ == '__main__':
+  unittest2.main()
diff --git a/py/system/vpd.py b/py/system/vpd.py
index bef5654..ebe5a6c 100644
--- a/py/system/vpd.py
+++ b/py/system/vpd.py
@@ -1,4 +1,3 @@
-
 #!/usr/bin/python
 # pylint: disable=W0212
 #
@@ -12,6 +11,7 @@
 
 
 import factory_common  # pylint: disable=W0611
+from cros.factory import privacy
 from cros.factory.utils.process_utils import Spawn
 
 
@@ -26,21 +26,6 @@
 VPD_VALUE_PATTERN = re.compile(r'^[ !#-~]*$')
 
 
-# Keys that may not be logged.
-VPD_BLACKLIST_KEYS = [
-  'ubind_attribute',
-  'gbind_attribute'
-]
-def FilterVPD(vpd_map):
-  """Redact values of any keys in VPD_BLACKLIST_KEYS."""
-  def FilterItem(k, v):
-    if v is None:
-      return None
-    return '<redacted %d chars>' % len(v) if k in VPD_BLACKLIST_KEYS else v
-
-  return dict((k, FilterItem(k, v)) for k, v in sorted(vpd_map.iteritems()))
-
-
 class Partition(object):
   """A VPD partition.
 
@@ -79,7 +64,7 @@
         with a redacted value.
     """
     if log:
-      logging.info('Updating %s: %s', self.name, FilterVPD(items))
+      logging.info('Updating %s: %s', self.name, privacy.FilterDict(items))
 
     command = ['vpd', '-i', self.name]
 
diff --git a/py/system/vpd_unittest.py b/py/system/vpd_unittest.py
index f51b2be..23192ea 100755
--- a/py/system/vpd_unittest.py
+++ b/py/system/vpd_unittest.py
@@ -22,14 +22,6 @@
   def tearDown(self):
     self.mox.UnsetStubs()
 
-  def testFilterVPD(self):
-    self.assertEquals(
-        dict(a='A', b='B',
-             ubind_attribute='<redacted 1 chars>',
-             gbind_attribute='<redacted 2 chars>'),
-        vpd.FilterVPD(
-            dict(a='A', b='B', ubind_attribute='U', gbind_attribute='GG')))
-
   def testGetAll(self):
     process = self.mox.CreateMockAnything()
     process.stdout_lines(strip=True).AndReturn(['"a"="b"',
diff --git a/py/test/pytests/start.py b/py/test/pytests/start.py
index 2d80a29..18ad40d 100755
--- a/py/test/pytests/start.py
+++ b/py/test/pytests/start.py
@@ -23,6 +23,7 @@
 import time
 import unittest
 
+from cros.factory.event_log import EventLog
 from cros.factory.test import factory
 from cros.factory.test import shopfloor
 from cros.factory.test import test_ui
@@ -31,7 +32,7 @@
 from cros.factory.test.args import Arg
 from cros.factory.test.factory_task import FactoryTask, FactoryTaskManager
 from cros.factory.test.event import Event
-from cros.factory.event_log import EventLog
+from cros.factory.test.utils import Enum
 from cros.factory.utils.process_utils import CheckOutput
 
 
@@ -304,9 +305,11 @@
     Arg('require_external_power', bool,
         'Prompts and waits for external power to be applied.',
         default=False, optional=True),
-    Arg('require_shop_floor', bool,
+    Arg('require_shop_floor', Enum([True, False, 'defer']),
         'Prompts and waits for serial number as input if no VPD keys are '
-        'provided as serial numbers, or reads serial numbers from VPD.',
+        'provided as serial numbers, or reads serial numbers from VPD. '
+        'This may be set to True, or "defer" to enable shopfloor but skip '
+        'reading the serial number.',
         default=None, optional=True),
     Arg('check_factory_install_complete', bool,
         'Check factory install process was complete.',
@@ -336,7 +339,8 @@
     if self.args.require_shop_floor is not None:
       shopfloor.set_enabled(self.args.require_shop_floor)
 
-    if self.args.require_shop_floor:
+    if (self.args.require_shop_floor and
+        self.args.require_shop_floor != 'defer'):
       if self.args.serial_number_vpd_keys:
         self._task_list.append(ReadVPDSerialTask(self))
       else:
diff --git a/py/test/pytests/vpd.py b/py/test/pytests/vpd.py
index 701db21..d774889 100644
--- a/py/test/pytests/vpd.py
+++ b/py/test/pytests/vpd.py
@@ -4,6 +4,27 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+"""Writes device VPD.
+
+This test can determine VPD values in several different ways based on the
+argument:
+
+- Manually.
+- Directly from shopfloor server.
+- From shopfloor device data.  If this option is selected with the
+  use_shopfloor_device_data arg, the following algorithm is applied:
+
+  - Locale fields (RO initial_locale, keyboard_layout, initial_timezone)
+    are set based on the 'locale' entry, which must be an item in the locale
+    database in locale.py.
+  - Registration codes are set based on the 'ubind_attribute' and
+    'gbind_attribute' entries.
+  - The RO 'serial_number' field is set based on the 'serial_number' entry.
+  - If the device data dictionary contains any keys of the format
+    'vpd.ro.xxx' or 'vpd.rw.xxx', the respective field in the RO/RW VPD
+    is set.
+"""
+
 import logging
 import re
 import unittest
@@ -333,6 +354,11 @@
         'Whether to store registration codes onto the machine.', default=False),
     Arg('task_list', list, 'A list of tasks to execute.',
         default=[VPDTasks.serial, VPDTasks.region]),
+    Arg('use_shopfloor_device_data', bool,
+        'If shopfloor is enabled, use accumulated data in shopfloor device '
+        'data dictionary instead of contacting shopfloor server again. '
+        'See file-level docs in vpd.py for more information.',
+        default=False),
     Arg('manual_input_fields', list, 'A list of tuples (vpd_region, key, '
         'en_display_name, zh_display_name, VALUE_CHECK) indicating the VPD '
         'fields that need to be manually entered.\n'
@@ -344,6 +370,39 @@
         'it.', default=[], optional=True)
   ]
 
+  def _ReadShopFloorDeviceData(self):
+    device_data = shopfloor.GetDeviceData()
+    required_keys = set(['serial_number', 'locale',
+                         'ubind_attribute', 'gbind_attribute'])
+    missing_keys = required_keys - set(device_data.keys())
+    if missing_keys:
+      self.fail('Missing keys in shopfloor device data: %r' %
+                sorted(missing_keys))
+
+    self.vpd['ro']['serial_number'] = device_data['serial_number']
+
+    locale_code = device_data['locale']
+    regions = [entry for entry in locale.DEFAULT_REGION_LIST
+               if entry[0] == locale_code]
+    if not regions:
+      logging.exception('Invalid locale %r', locale_code)
+    dummy_locale, layout, timezone, dummy_description = (
+        locale.BuildRegionInformation(regions[0]))
+
+    self.vpd['ro']['initial_locale'] = locale_code
+    self.vpd['ro']['keyboard_layout'] = layout
+    self.vpd['ro']['initial_timezone'] = timezone
+
+    for k, v in device_data.iteritems():
+      match = re.match(r'$vpd\.(ro|rw)\.(.+)^', k)
+      if match:
+        self.vpd[match.group(1)][match.group(2)] = v
+
+    self.registration_code_map = {
+        'user': device_data['ubind_attribute'],
+        'group': device_data['gbind_attribute'],
+        }
+
   def setUp(self):
     self.ui = test_ui.UI()
     self.template = OneSection(self.ui)
@@ -361,7 +420,10 @@
     if not (self.args.override_vpd and self.ui.InEngineeringMode()):
       if shopfloor.is_enabled():
         # Grab from ShopFloor, then input manual fields (if any).
-        self.tasks += [ShopFloorVPDTask(self)]
+        if self.args.use_shopfloor_device_data:
+          self._ReadShopFloorDeviceData()
+        else:
+          self.tasks += [ShopFloorVPDTask(self)]
         for v in self.args.manual_input_fields:
           self.tasks += [ManualInputTask(
               self, VPDInfo(v[0], v[1], v[2], v[3], v[4]))]
diff --git a/py/test/shopfloor.py b/py/test/shopfloor.py
index c5c170a..7040c0c 100644
--- a/py/test/shopfloor.py
+++ b/py/test/shopfloor.py
@@ -30,6 +30,7 @@
 from xmlrpclib import Binary
 
 import factory_common # pylint: disable=W0611
+from cros.factory import privacy
 from cros.factory.test import factory
 from cros.factory.utils import net_utils
 from cros.factory.utils.process_utils import Spawn
@@ -409,4 +410,4 @@
   data = factory.get_state_instance().update_shared_data_dict(
       KEY_DEVICE_DATA, new_device_data)
   logging.info('Updated device data; complete device data is now %s',
-               data)
+               privacy.FilterDict(data))