extra: add cr50 rma open utility

The testlab will be running rma open on a bunch of devices. This change
adds a script for them to use.

The script handles things like finding the cr50 console and converting
the rma_auth output into a full challenge url. The user can click on the
url and generate the authcode. They can input the generated authcode
back into the script. The script will check 'ccd' output to make sure
rma open succeeds.

BUG=none
BRANCH=none
TEST=on a reef and bob running 3.4, run the entire cr50_rma_open
process. Verify RMA open succeeds

Change-Id: I2c9aef565351b65848f29083cd1e2c1f6e77e2a4
Signed-off-by: Mary Ruthven <mruthven@google.com>
Reviewed-on: https://chromium-review.googlesource.com/994417
Commit-Ready: Mary Ruthven <mruthven@chromium.org>
Tested-by: Mary Ruthven <mruthven@chromium.org>
Reviewed-by: Wai-Hong Tam <waihong@google.com>

1 files changed, 376 insertions, 0 deletions

diff --git a/extra/cr50_rma_open/cr50_rma_open.py b/extra/cr50_rma_open/cr50_rma_open.py
new file mode 100644
index 0000000000..81ba82f5b5
--- /dev/null
+++ b/extra/cr50_rma_open/cr50_rma_open.py
@@ -0,0 +1,376 @@
+#!/usr/bin/python2
+# Copyright 2018 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.
+
+# Used to access the cr50 console and handle RMA Open
+
+import argparse
+import glob
+import re
+import serial
+import subprocess
+import sys
+import time
+
+URL = 'https://www.google.com/chromeos/partner/console/cr50reset?' \
+    'challenge=%s&hwid=%s'
+RMA_SUPPORT = '0.3.3'
+CR50_USB = '18d1:5014'
+ERASED_BID = 'ffffffff'
+
+HELP_INFO = """
+Run RMA Open to enable CCD on Cr50. The utility can be used to get a
+url that will generate an authcode to open cr50. It can also be used to
+try opening cr50 with the generated authcode.
+
+The last challenge is the only valid one, so don't generate a challenge
+10 times and then use the first URL. You can only use the last one.
+
+For RMA Open:
+Connect suzyq to the dut and your workstation.
+
+Check the basic setup with
+    sudo python cr50_rma_open.py -c
+
+If the setup is broken. Follow the debug print statements to try to fix
+the error. Rerun until the script says Cr50 setup ok.
+
+After the setup is verified, run the following command to generate the
+challenge url
+    sudo python cr50_rma_open.py -g -i $HWID
+
+Go to the URL from by that command to generate an authcode. Once you have
+the authcode, you can use it to open cr50.
+    sudo python cr50_rma_open.py -a $AUTHCODE
+"""
+
+DEBUG_MISSING_USB = """
+Unable to find Cr50 Device 18d1:5014
+
+DEBUG MISSING USB:
+    - Make sure suzyq is plugged into the correct DUT port
+    - Try flipping the cable
+    - unplug the cable for 5s then plug it back in
+"""
+
+DEBUG_DEVICE = """
+DEBUG DEVICE COMMUNICATION:
+Issues communicating with %s
+
+A 18d1:5014 device exists, so make sure you have selected the correct
+/dev/ttyUSB
+"""
+
+DEBUG_SERIALNAME = """
+DEBUG SERIALNAME:
+Found the USB device, but can't match the usb serialname. Check the
+serialname you passed into cr50_rma_open or try running without a
+serialname.
+"""
+
+DEBUG_CONNECTION = """
+DEBUG CONNECTION:
+Found the USB device but cant communicate with any of the consoles.
+
+Try Running cr50_rma_open again. If it still fails unplug the ccd cable
+for 5 seconds and plug it back in.
+"""
+
+DEBUG_TOO_MANY_USB_DEVICES = """
+DEBUG SELECT USB:
+More than one cr50 usb device was found. Disconnect all but one device
+or use the -s option with the correct usb serialname.
+"""
+
+DEBUG_ERASED_BOARD_ID = """
+DEBUG ERASED BOARD ID:
+If you are using a prePVT device run
+/usr/share/cros/cr50-set-board-id.sh proto
+
+If you are running a MP device, please talk to someone.
+"""
+
+DEBUG_AUTHCODE_MISMATCH = """
+DEBUG AUTHCODE MISMATCH:
+    - Check the URL matches the one generated by the last cr50_rma_open
+      run.
+    - Check you used the correct authcode.
+    - Make sure the cr50 version is greater than 3.3.
+    - try generating another URL by rerunning the generate command and
+      rerunning the process.
+"""
+
+parser = argparse.ArgumentParser(
+    description=HELP_INFO, formatter_class=argparse.RawTextHelpFormatter)
+parser.add_argument('-g', '--generate_challenge', action='store_true',
+        help='Generate Cr50 challenge. Must be used in combination with -i')
+parser.add_argument('-c', '--check_connection', action='store_true',
+        help='Check cr50 console connection works')
+parser.add_argument('-s', '--serialname', type=str, default='',
+        help='The cr50 usb serialname')
+parser.add_argument('-d', '--device', type=str, default='',
+        help='cr50 console device ex /dev/ttyUSB0')
+parser.add_argument('-i', '--hwid', type=str, default='',
+        help='The board hwid. Necessary to generate a challenge')
+parser.add_argument('-a', '--authcode', type=str, default='',
+        help='The authcode string generated from the challenge url')
+
+def debug(string):
+    """Print yellow string"""
+    print '\033[93m' + string + '\033[0m'
+
+def info(string):
+    """Print green string"""
+    print '\033[92m' + string + '\033[0m'
+
+class RMAOpen(object):
+    """Used to find the cr50 console and run RMA open"""
+
+    def __init__(self, device=None, usb_serial=None):
+        if device:
+            self.set_cr50_device(device)
+        else:
+            self.find_cr50_device(usb_serial)
+        info('DEVICE: ' + self.device)
+        self.check_version()
+        self.print_platform_info()
+        info('Cr50 setup ok')
+
+
+    def set_cr50_device(self, device):
+        """Save the device used for the console"""
+        self.device = device
+
+
+    def send_cmd_get_output(self, cmd, nbytes=0):
+        """Send a cr50 command and get the output
+
+        Args:
+            cmd: The cr50 command string
+            nbytes: The number of bytes to read from the console. If 0 read all
+                    of the console output.
+        Returns:
+            The command output
+        """
+        try:
+            ser = serial.Serial(self.device, timeout=1)
+        except OSError, e:
+            debug('Permission denied ' + self.device)
+            debug('Try running cr50_rma_open with sudo')
+            raise
+        ser.write(cmd + '\n\n')
+        if nbytes:
+            output = ser.read(nbytes).strip()
+        else:
+            output = ser.readall().strip()
+        ser.close()
+
+        # Return only the command output
+        split_cmd = cmd + '\r'
+        if output and cmd and split_cmd in output:
+            return ''.join(output.rpartition(split_cmd)[1::]).split('>')[0]
+        return output
+
+
+    def get_rma_challenge(self):
+        """Get the rma_auth challenge
+
+        Returns:
+            The RMA challenge with all whitespace removed.
+        """
+        output = self.send_cmd_get_output('rma_auth').strip()
+        print 'rma_auth output:\n', output
+        # Extract the challenge from the console output
+        challenge = ''.join(re.findall(' \S{5}' * 4, output))
+        # Remove all whitespace
+        return re.sub('\s', '', challenge)
+
+
+    def generate_challenge_url(self, hwid):
+        """Get the rma_auth challenge
+
+        Returns:
+            The RMA challenge with all whitespace removed.
+        """
+
+        challenge = self.get_rma_challenge()
+        self.print_platform_info()
+        info('CHALLENGE: ' + challenge)
+        info('HWID:' + hwid)
+        url = URL % (challenge, hwid)
+        info('GOTO:\n' + url)
+        print 'If the server fails to debug the challenge make sure the RLZ is '
+        print 'whitelisted'
+
+
+    def try_authcode(self, authcode):
+        """Try opening cr50 with the authcode"""
+        # rma_auth may cause the system to reboot. Don't wait to read all that
+        # output. Read the first 300 bytes and call it a day.
+        output = self.send_cmd_get_output('rma_auth ' + authcode, nbytes=300)
+        print 'CR50 RESPONSE:', output
+        print 'waiting for cr50 reboot'
+        # Cr50 may be rebooting. Wait a bit
+        time.sleep(3)
+        self.check_ccd_settings('process_response: success!' in output)
+
+
+    def ccd_is_closed(self, debug=False):
+        """Returns True if any of the capabilities are still restricted"""
+        output = self.send_cmd_get_output('ccd')
+        if 'Capabilities' not in output:
+            raise ValueError('Could not get ccd output')
+        if debug:
+            print 'CURRENT CCD SETTINGS:\n', output
+        is_closed = 'IfOpened' in output or 'IfUnlocked' in output
+        if not is_closed:
+            info('Cr50 ccd is Open. RMA Open complete')
+        return is_closed
+
+
+    def check_ccd_settings(self, authcode_match):
+        """Raise an error if Cr50 CCD is not open
+
+        Choose the error and debug messages based on if Cr50 successfully
+        matched the authcode.
+
+        Args:
+            authcode_match: True if Cr50 successfully processed the authcode.
+
+        Raises:
+            ValueError if 'ccd' does shows any capabilities are restricted.
+        """
+        if self.ccd_is_closed(debug=True):
+            if not authcode_match:
+                debug(DEBUG_AUTHCODE_MISMATCH)
+                raise ValueError('Authcode mismatch. Check args and url')
+            else:
+                raise ValueError('Could not set all capability privileges to '
+                        'Always')
+
+
+    def check_version(self):
+        """Make sure cr50 is running a version that supports RMA Open"""
+        output = self.send_cmd_get_output('version')
+        if not output.strip():
+            debug(DEBUG_DEVICE % self.device)
+            raise ValueError('Could not communicate with %s' % self.device)
+
+        version = re.search('RW.*\* ([\d\.]+)/', output).group(1)
+        print 'RMA support added in:', RMA_SUPPORT
+        print 'Running Cr50 Version:', version
+        fields = [int(field) for field in version.split('.')]
+        rma_fields = [int(field) for field in RMA_SUPPORT.split('.')]
+        for i, field in enumerate(fields):
+            if field < int(rma_fields[i]):
+                raise ValueError('%s does not have RMA support. Update to at '
+                        'least %s' % (version, RMA_SUPPORT))
+
+
+    def device_matches_devid(self, devid, device):
+        """Return True if the device matches devid.
+
+        Use the sysinfo output from device to determine if it matches devid
+
+        Returns:
+            True if sysinfo from device shows the given devid. False if there
+            is no output or sysinfo doesn't contain the devid.
+        """
+
+        self.set_cr50_device(device)
+        sysinfo = self.send_cmd_get_output('sysinfo')
+        # Make sure there is some output, and it shows it's from Cr50
+        if not sysinfo or 'cr50' not in sysinfo:
+            return False
+        print sysinfo
+        # The cr50 device id should be in the sysinfo output, if we found
+        # the right console. Make sure it is
+        return devid in sysinfo
+
+
+    def find_cr50_device(self, usb_serial):
+        """Find the cr50 console device
+
+        The Cr50 usb serialname matches the cr50 devid. Convert the serialname
+        to devid. Use that to check all of the consoles and find cr50's.
+
+        Args:
+            usb_serial: an optional string. The serialname of the cr50 usb
+                        device
+        Raises:
+            ValueError if the console can't be found with the given serialname
+        """
+        usb_serial = self.find_cr50_usb(usb_serial)
+        info('SERIALNAME: ' + usb_serial)
+        devid = '0x' + ' 0x'.join(usb_serial.lower().split('-'))
+        info('DEVID: ' + devid)
+
+        # Get all the usb devices
+        devices = glob.glob('/dev/ttyUSB*')
+        # Typically Cr50 has the lowest number. Sort the devices, so we're more
+        # likely to try the cr50 console first.
+        devices.sort()
+
+        # Find the one that is the cr50 console
+        for device in devices:
+            print 'testing', device
+            if self.device_matches_devid(devid, device):
+                print 'found device', device
+                return
+        debug(DEBUG_CONNECTION)
+        raise ValueError('Found USB device, but could not communicate with '
+                'cr50 console')
+
+
+    def print_platform_info(self):
+        """Print the cr50 BID RLZ code"""
+        bid_output = self.send_cmd_get_output('bid')
+        bid = re.search('Board ID: (\S+),', bid_output).group(1)
+        if bid == ERASED_BID:
+            debug(DEBUG_ERASED_BOARD_ID)
+            raise ValueError('Cannot run RMA Open when board id is erased')
+        bid = int(bid, 16)
+        chrs = [chr((bid >> (8 * i)) & 0xff) for i in range(4)]
+        info('RLZ: ' + ''.join(chrs[::-1]))
+
+
+    def find_cr50_usb(self, usb_serial):
+        """Make sure the Cr50 USB device exists"""
+        try:
+            output = subprocess.check_output(['lsusb', '-vd', CR50_USB])
+        except:
+            debug(DEBUG_MISSING_USB)
+            raise ValueError('Could not find Cr50 USB device')
+        serialnames = re.findall('iSerial +\d+ (\S+)\s', output)
+        if usb_serial:
+            if usb_serial not in serialnames:
+                debug(DEBUG_SERIALNAME)
+                raise ValueError('Could not find usb device "%s"' % usb_serial)
+            return usb_serial
+        if len(serialnames) > 1:
+            print 'Found Cr50 device serialnames ', ', '.join(serialnames)
+            debug(DEBUG_TOO_MANY_USB_DEVICES)
+            raise ValueError('Too many cr50 usb devices')
+        return serialnames[0]
+
+
+def main():
+    args = parser.parse_args()
+
+    cr50_rma_open = RMAOpen(args.device, args.serialname)
+    if args.check_connection or not cr50_rma_open.ccd_is_closed():
+        sys.exit(0)
+    elif args.authcode:
+        info('Using authcode: ' + args.authcode)
+        cr50_rma_open.try_authcode(args.authcode)
+        info('Successfully ran RMA Open')
+    elif args.generate_challenge:
+        if not args.hwid:
+            debug('--hwid necessary to generate challenge url')
+            sys.exit(0)
+        cr50_rma_open.generate_challenge_url(args.hwid)
+
+if __name__ == "__main__":
+    main()
+