client/site_tests/factory_VPD/factory_VPD.py

# 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 os
import serial_task
import region_task
import registration_codes

from autotest_lib.client.bin import test
from autotest_lib.client.common_lib import error, utils
from autotest_lib.client.cros import factory
from autotest_lib.client.cros.factory import shopfloor
from autotest_lib.client.cros.factory import task
from autotest_lib.client.cros.factory import ui

_MESSAGE_FETCH_FROM_SHOP_FLOOR = "Fetching VPD from shop floor server..."
_MESSAGE_WRITING = "Writing VPD:"


class WriteVpdTask(task.FactoryTask):

    def __init__(self, vpd):
        self.vpd = vpd

    def mask_bind_attributes(self):
        """Mask sensitive bind attributes for display"""
        mask_keys = ['ubind_attribute', 'gbind_attribute']
        disp_dict = {'ro':{}, 'rw':{}}

        for vpd_type in ('ro', 'rw'):
            for key in self.vpd[vpd_type]:
                if key in mask_keys:
                    disp_dict[vpd_type][key] = '********'
                else:
                    disp_dict[vpd_type][key] = self.vpd[vpd_type][key]

        return disp_dict

    def write_vpd(self):
        """Writes a VPD structure into system.

        @param vpd: A dictionary with 'ro' and 'rw' keys, each associated with a
          key-value VPD data set.
        """
        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
        disp_vpd = self.mask_bind_attributes()
        VPD_LIST = (('RO_VPD', 'ro'), ('RW_VPD', 'rw'))
        for (section, vpd_type) in VPD_LIST:
            if not vpd.get(vpd_type, None):
                continue
            parameter = format_vpd_parameter(vpd[vpd_type])
            disp_param = format_vpd_parameter(disp_vpd[vpd_type])
            ret = os.system('vpd -i %s %s' % (section, parameter))
            factory.log('vpd -i %s %s' % (section, disp_param))
            factory.log('vpd returned: %d' % (ret))
        self.stop()

    def start(self):
        # Flatten key-values in VPD dictionary.
        vpd = self.vpd
        vpd_list = []
        for vpd_type in ('ro', 'rw'):
            for key in sorted(vpd[vpd_type]):
                vpd_list += ['%s: %s = %s' %
                             (vpd_type, key, vpd[vpd_type][key])]
                if key in ['ubind_attribute', 'gbind_attribute']:
                    registration_codes.CheckRegistrationCode(vpd[vpd_type][key])

        self.add_widget(ui.make_label("%s\n%s" % (_MESSAGE_WRITING,
                                                  '\n'.join(vpd_list))))
        task.schedule(self.write_vpd)


class ShopFloorVpdTask(task.FactoryTask):

    def __init__(self, vpd):
        self.vpd = vpd

    def start(self):
        self.add_widget(ui.make_label(_MESSAGE_FETCH_FROM_SHOP_FLOOR))
        task.schedule(self.fetch_vpd)

    def fetch_vpd(self):
        self.vpd.update(shopfloor.get_vpd())
        self.stop()


class factory_VPD(test.test):
    version = 5

    SERIAL_TASK_NAME = 'serial'
    REGION_TASK_NAME = 'region'

    def run_once(self, override_vpd=None,
                 task_list=[SERIAL_TASK_NAME, REGION_TASK_NAME]):
        factory.log('%s run_once' % self.__class__)
        self.tasks = []
        self.vpd = override_vpd or {'ro': {}, 'rw': {}}

        if not override_vpd:
            if shopfloor.is_enabled():
                self.tasks += [ShopFloorVpdTask(self.vpd)]
            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)]
        task.run_factory_tasks(self.job, self.tasks)

        factory.log('%s run_once finished' % repr(self.__class__))