Add support for registration codes.

BUG=chrome-os-partner:11912
TEST=Manual on device

Change-Id: I2b5a2c7fea0f7defdce1ac2a7f74cb10e73ed400
Reviewed-on: https://gerrit.chromium.org/gerrit/28684
Tested-by: Jon Salz <jsalz@chromium.org>
Reviewed-by: Hung-Te Lin <hungte@chromium.org>
Commit-Ready: Jon Salz <jsalz@chromium.org>
diff --git a/client/site_tests/factory_VPD/factory_VPD.py b/client/site_tests/factory_VPD/factory_VPD.py
index e3fd600..7d49a10 100644
--- a/client/site_tests/factory_VPD/factory_VPD.py
+++ b/client/site_tests/factory_VPD/factory_VPD.py
@@ -10,6 +10,7 @@
 from autotest_lib.client.common_lib import error, utils
 from autotest_lib.client.cros import factory_setup_modules
 from cros.factory.test import factory
+from cros.factory.test import registration_codes
 from cros.factory.test import shopfloor
 from cros.factory.test import task
 from cros.factory.test import ui
@@ -18,10 +19,19 @@
 _MESSAGE_WRITING = "Writing VPD:"
 
 
+def format_vpd_parameter(vpd_dict):
+    """Formats a key-value dictionary into VPD syntax."""
+    # Writes in sorted ordering so the VPD structure will be more
+    # deterministic.
+    return ' '.join(('-s "%s"="%s"' % (key, vpd_dict[key])
+                     for key in sorted(vpd_dict)))
+
+
 class WriteVpdTask(task.FactoryTask):
 
-    def __init__(self, vpd):
+    def __init__(self, vpd, registration_code_map):
         self.vpd = vpd
+        self.registration_code_map = registration_code_map
 
     def write_vpd(self):
         """Writes a VPD structure into system.
@@ -33,13 +43,6 @@
             factory.log(command)
             utils.system(command)
 
-        def format_vpd_parameter(vpd_dict):
-            """Formats a key-value dictionary into VPD syntax."""
-            # Writes in sorted ordering so the VPD structure will be more
-            # deterministic.
-            return ' '.join(('-s "%s"="%s"' % (key, vpd_dict[key])
-                             for key in sorted(vpd_dict)))
-
         vpd = self.vpd
         VPD_LIST = (('RO_VPD', 'ro'), ('RW_VPD', 'rw'))
         for (section, vpd_type) in VPD_LIST:
@@ -47,6 +50,26 @@
                 continue
             parameter = format_vpd_parameter(vpd[vpd_type])
             shell('vpd -i %s %s' % (section, parameter))
+
+        if self.registration_code_map is not None:
+            # Check registration codes (fail test if invalid).
+            for k in ['user', 'group']:
+                registration_codes.CheckRegistrationCode(
+                    self.registration_code_map[k])
+
+            # Add registration codes, being careful not to log the command.
+            factory.log('Storing registration codes')
+            utils.system(
+                'vpd -i %s %s' % (
+                    'RW_VPD',
+                    format_vpd_parameter(
+                        # See <http://src.chromium.org/svn/trunk/src/chrome/
+                        # browser/chromeos/extensions/echo_private_api.cc>.
+                        {'ubind_attribute':
+                             self.registration_code_map['user'],
+                         'gbind_attribute':
+                             self.registration_code_map['group']})))
+
         self.stop()
 
     def start(self):
@@ -56,6 +79,8 @@
         for vpd_type in ('ro', 'rw'):
             vpd_list += ['%s: %s = %s' % (vpd_type, key, vpd[vpd_type][key])
                          for key in sorted(vpd[vpd_type])]
+        if self.registration_code_map:
+            vpd_list += ['rw: registration codes']
 
         self.add_widget(ui.make_label("%s\n%s" % (_MESSAGE_WRITING,
                                                   '\n'.join(vpd_list))))
@@ -64,8 +89,9 @@
 
 class ShopFloorVpdTask(task.FactoryTask):
 
-    def __init__(self, vpd):
+    def __init__(self, vpd, registration_code_map):
         self.vpd = vpd
+        self.registration_code_map = registration_code_map
 
     def start(self):
         self.add_widget(ui.make_label(_MESSAGE_FETCH_FROM_SHOP_FLOOR))
@@ -73,6 +99,9 @@
 
     def fetch_vpd(self):
         self.vpd.update(shopfloor.get_vpd())
+        if self.registration_code_map is not None:
+            self.registration_code_map.update(
+                shopfloor.get_registration_code_map())
         self.stop()
 
 
@@ -83,20 +112,23 @@
     REGION_TASK_NAME = 'region'
 
     def run_once(self, override_vpd=None,
+                 store_registration_codes=False,
                  task_list=[SERIAL_TASK_NAME, REGION_TASK_NAME]):
         factory.log('%s run_once' % self.__class__)
         self.tasks = []
         self.vpd = override_vpd or {'ro': {}, 'rw': {}}
+        self.registration_code_map = {} if store_registration_codes else None
 
         if not override_vpd:
             if shopfloor.is_enabled():
-                self.tasks += [ShopFloorVpdTask(self.vpd)]
+                self.tasks += [
+                    ShopFloorVpdTask(self.vpd, self.registration_code_map)]
             else:
                 if self.SERIAL_TASK_NAME in task_list:
                     self.tasks += [serial_task.SerialNumberTask(self.vpd)]
                 if self.REGION_TASK_NAME in task_list:
                     self.tasks += [region_task.SelectRegionTask(self.vpd)]
-        self.tasks += [WriteVpdTask(self.vpd)]
+        self.tasks += [WriteVpdTask(self.vpd, self.registration_code_map)]
         task.run_factory_tasks(self.job, self.tasks)
 
         factory.log('%s run_once finished' % repr(self.__class__))