cr50_rma_open: fix linter errors

Fix errors from pylint. Use logging to print messages instead of print.

BUG=b:149405690
BRANCH=none
TEST=RMA open a device and enable testlab mode.

Change-Id: I63a8b108b0f436291659f2f787ebcf72e8f4a523
Signed-off-by: Mary Ruthven <mruthven@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2052269
Reviewed-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Wai-Hong Tam <waihong@google.com>
(cherry picked from commit 63e6072b454d41e85fd3763abff412d73090ce2d)
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2080634

1 files changed, 178 insertions, 193 deletions

diff --git a/extra/cr50_rma_open/cr50_rma_open.py b/extra/cr50_rma_open/cr50_rma_open.py
index 11fc90a6b9..8e1ea833ad 100755
--- a/extra/cr50_rma_open/cr50_rma_open.py
+++ b/extra/cr50_rma_open/cr50_rma_open.py
@@ -5,32 +5,8 @@
 # found in the LICENSE file.
 
 # Used to access the cr50 console and handle RMA Open
+"""Open cr50 using RMA authentication.
 
-import argparse
-import glob
-import re
-import serial
-import subprocess
-import sys
-import time
-
-SCRIPT_VERSION = 5
-CCD_IS_UNRESTRICTED = 1 << 0
-WP_IS_DISABLED = 1 << 1
-TESTLAB_IS_ENABLED = 1 << 2
-RMA_OPENED = CCD_IS_UNRESTRICTED | WP_IS_DISABLED
-URL = 'https://www.google.com/chromeos/partner/console/cr50reset?' \
-    'challenge=%s&hwid=%s'
-RMA_SUPPORT_PROD = '0.3.3'
-RMA_SUPPORT_PREPVT = '0.4.5'
-DEV_MODE_OPEN_PROD = '0.3.9'
-DEV_MODE_OPEN_PREPVT = '0.4.7'
-TESTLAB_PROD = '0.3.10'
-CR50_USB = '18d1:5014'
-CR50_LSUSB_CMD = ['lsusb', '-vd', CR50_USB]
-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.
@@ -66,6 +42,32 @@ prod image, you can skip this step.
     sudo python cr50_rma_open.py -t
 """
 
+import argparse
+import glob
+import logging
+import re
+import subprocess
+import sys
+import time
+
+import serial
+
+SCRIPT_VERSION = 5
+CCD_IS_UNRESTRICTED = 1 << 0
+WP_IS_DISABLED = 1 << 1
+TESTLAB_IS_ENABLED = 1 << 2
+RMA_OPENED = CCD_IS_UNRESTRICTED | WP_IS_DISABLED
+URL = ('https://www.google.com/chromeos/partner/console/cr50reset?'
+       'challenge=%s&hwid=%s')
+RMA_SUPPORT_PROD = '0.3.3'
+RMA_SUPPORT_PREPVT = '0.4.5'
+DEV_MODE_OPEN_PROD = '0.3.9'
+DEV_MODE_OPEN_PREPVT = '0.4.7'
+TESTLAB_PROD = '0.3.10'
+CR50_USB = '18d1:5014'
+CR50_LSUSB_CMD = ['lsusb', '-vd', CR50_USB]
+ERASED_BID = 'ffffffff'
+
 DEBUG_MISSING_USB = """
 Unable to find Cr50 Device 18d1:5014
 
@@ -126,48 +128,13 @@ DEBUG_DUT_CONTROL_OSERROR = """
 Run from chroot if you are trying to use a /dev/pts ccd servo console
 """
 
-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('-p', '--print_caps', action='store_true',
-        help='Print the ccd output when checking the capabilities')
-parser.add_argument('-t', '--enable_testlab', action='store_true',
-        help='enable testlab mode')
-parser.add_argument('-w', '--wp_disable', action='store_true',
-        help='Disable write protect')
-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')
-parser.add_argument('-P', '--servo_port', type=str, default='',
-        help='the servo port')
-parser.add_argument('-I', '--ip', type=str, default='',
-        help='The DUT IP. Necessary to do ccd open')
-
-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"""
 
     ENABLE_TESTLAB_CMD = 'ccd testlab enabled\n'
 
-    def __init__(self, device=None, usb_serial=None, print_caps=False,
-            servo_port=None, ip=None):
+    def __init__(self, device=None, usb_serial=None, servo_port=None, ip=None):
         self.servo_port = servo_port if servo_port else '9999'
-        self.print_caps = print_caps
         self.ip = ip
         if device:
             self.set_cr50_device(device)
@@ -175,24 +142,22 @@ class RMAOpen(object):
             self.find_cr50_servo_uart()
         else:
             self.find_cr50_device(usb_serial)
-        info('DEVICE: ' + self.device)
+        logging.info('DEVICE: %s', self.device)
         self.check_version()
         self.print_platform_info()
-        info('Cr50 setup ok')
+        logging.info('Cr50 setup ok')
         self.update_ccd_state()
         self.using_ccd = self.device_is_running_with_servo_ccd()
 
-
     def _dut_control(self, control):
         """Run dut-control and return the response"""
         try:
-            return subprocess.check_output(['dut-control', '-p',
-                    self.servo_port, control], encoding='utf-8').strip()
+            cmd = ['dut-control', '-p', self.servo_port, control]
+            return subprocess.check_output(cmd, encoding='utf-8').strip()
         except OSError:
-            debug(DEBUG_DUT_CONTROL_OSERROR)
+            logging.warning(DEBUG_DUT_CONTROL_OSERROR)
             raise
 
-
     def find_cr50_servo_uart(self):
         """Save the device used for the console.
 
@@ -201,12 +166,10 @@ class RMAOpen(object):
         self._dut_control('cr50_uart_timestamp:off')
         self.device = self._dut_control('cr50_uart_pty').split(':')[-1]
 
-
     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
 
@@ -220,15 +183,15 @@ class RMAOpen(object):
         try:
             ser = serial.Serial(self.device, timeout=1)
         except OSError:
-            debug('Permission denied ' + self.device)
-            debug('Try running cr50_rma_open with sudo')
+            logging.warning('Permission denied %s', self.device)
+            logging.warning('Try running cr50_rma_open with sudo')
             raise
         write_cmd = cmd + '\n\n'
         ser.write(write_cmd.encode('utf-8'))
         if nbytes:
-            output = ser.read(nbytes).strip()
+            output = ser.read(nbytes)
         else:
-            output = ser.readall().strip()
+            output = ser.readall()
         ser.close()
 
         output = output.decode('utf-8').strip() if output else ''
@@ -238,7 +201,6 @@ class RMAOpen(object):
             return ''.join(output.rpartition(split_cmd)[1::]).split('>')[0]
         return output
 
-
     def device_is_running_with_servo_ccd(self):
         """Return True if the device is a servod ccd console"""
         # servod uses /dev/pts consoles. Non-servod uses /dev/ttyUSBX
@@ -252,10 +214,9 @@ class RMAOpen(object):
         # is run in the chroot, so run it last.
         if 'ccd_cr50' not in self._dut_control('servo_type'):
             return False
-        info('running through servod ccd')
+        logging.info('running through servod ccd')
         return True
 
-
     def get_rma_challenge(self):
         """Get the rma_auth challenge
 
@@ -279,14 +240,13 @@ class RMAOpen(object):
             The RMA challenge with all whitespace removed.
         """
         output = self.send_cmd_get_output('rma_auth').strip()
-        print('rma_auth output:\n', output)
+        logging.info('rma_auth output:\n%s', output)
         # Extract the challenge from the console output
         if 'generated challenge:' in output:
             return output.split('generated challenge:')[-1].strip()
-        challenge = ''.join(re.findall(' \S{5}' * 4, output))
+        challenge = ''.join(re.findall(r' \S{5}' * 4, output))
         # Remove all whitespace
-        return re.sub('\s', '', challenge)
-
+        return re.sub(r'\s', '', challenge)
 
     def generate_challenge_url(self, hwid):
         """Get the rma_auth challenge
@@ -297,13 +257,12 @@ class RMAOpen(object):
 
         challenge = self.get_rma_challenge()
         self.print_platform_info()
-        info('CHALLENGE: ' + challenge)
-        info('HWID:' + hwid)
+        logging.info('CHALLENGE: %s', challenge)
+        logging.info('HWID: %s', hwid)
         url = URL % (challenge, hwid)
-        info('GOTO:\n' + url)
-        print('If the server fails to debug the challenge make sure the RLZ '
-              'is whitelisted')
-
+        logging.info('GOTO:\n %s', url)
+        logging.info('If the server fails to debug the challenge make sure the '
+                     'RLZ is whitelisted')
 
     def try_authcode(self, authcode):
         """Try opening cr50 with the authcode
@@ -314,8 +273,8 @@ class RMAOpen(object):
         # 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')
+        logging.info('CR50 RESPONSE: %s', output)
+        logging.info('waiting for cr50 reboot')
         # Cr50 may be rebooting. Wait a bit
         time.sleep(5)
         if self.using_ccd:
@@ -327,41 +286,36 @@ class RMAOpen(object):
         authcode_match = 'process_response: success!' in output
         if not self.check(CCD_IS_UNRESTRICTED):
             if not authcode_match:
-                debug(DEBUG_AUTHCODE_MISMATCH)
+                logging.warning(DEBUG_AUTHCODE_MISMATCH)
                 message = 'Authcode mismatch. Check args and url'
             else:
                 message = 'Could not set all capability privileges to Always'
             raise ValueError(message)
 
-
     def wp_is_force_disabled(self):
         """Returns True if write protect is forced disabled"""
         output = self.send_cmd_get_output('wp')
         wp_state = output.split('Flash WP:', 1)[-1].split('\n', 1)[0].strip()
-        info('wp: ' + wp_state)
+        logging.info('wp: %s', wp_state)
         return wp_state == 'forced disabled'
 
-
     def testlab_is_enabled(self):
         """Returns True if testlab mode is enabled"""
         output = self.send_cmd_get_output('ccd testlab')
         testlab_state = output.split('mode')[-1].strip().lower()
-        info('testlab: ' + testlab_state)
+        logging.info('testlab: %s', testlab_state)
         return testlab_state == 'enabled'
 
-
     def ccd_is_restricted(self):
         """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 self.print_caps:
-            print('CURRENT CCD SETTINGS:\n', output)
+        logging.debug('CURRENT CCD SETTINGS:\n%s', output)
         restricted = 'IfOpened' in output or 'IfUnlocked' in output
-        info('ccd: ' + ('Restricted' if restricted else 'Unrestricted'))
+        logging.info('ccd: %srestricted', '' if restricted else 'Un')
         return restricted
 
-
     def update_ccd_state(self):
         """Get the wp and ccd state from cr50. Save it in _ccd_state"""
         self._ccd_state = 0
@@ -372,12 +326,10 @@ class RMAOpen(object):
         if self.testlab_is_enabled():
             self._ccd_state |= TESTLAB_IS_ENABLED
 
-
     def check(self, setting):
         """Returns true if the all of the 1s in setting are 1 in _ccd_state"""
         return self._ccd_state & setting == setting
 
-
     def _has_testlab_support(self):
         """Return True if you can enable testlab mode"""
         # all prepvt images can enable testlab
@@ -385,7 +337,6 @@ class RMAOpen(object):
             return True
         return not self._running_version_is_older(DEV_MODE_OPEN_PROD)
 
-
     def _capabilities_allow_open_from_console(self):
         """Return True if ccd open is Always allowed from usb"""
         output = self.send_cmd_get_output('ccd')
@@ -401,13 +352,11 @@ class RMAOpen(object):
             return True
         return not self._running_version_is_older(DEV_MODE_OPEN_PREPVT)
 
-
     def _run_on_dut(self, command):
         """Run the command on the DUT."""
         return subprocess.check_output(['ssh', self.ip, command],
                                        encoding='utf-8')
 
-
     def _open_in_dev_mode(self):
         """Open Cr50 when it's in dev mode"""
         output = self.send_cmd_get_output('ccd')
@@ -415,12 +364,12 @@ class RMAOpen(object):
         if 'State: Open' not in output:
             # Verify the device is in devmode before trying to run open.
             if 'dev_mode' not in output:
-                debug('Enter dev mode to open ccd or update to %s' %
-                      TESTLAB_PROD)
+                logging.warning('Enter dev mode to open ccd or update to %s',
+                                TESTLAB_PROD)
                 raise ValueError('DUT not in dev mode')
             if not self.ip:
-                debug("If your DUT doesn't have ssh support, run 'gsctool -a "
-                      "-o' from the AP")
+                logging.warning("If your DUT doesn't have ssh support, run "
+                                "'gsctool -a -o' from the AP")
                 raise ValueError('Cannot run ccd open without dut ip')
             self._run_on_dut('gsctool -a -o')
             # Wait >1 second for cr50 to update ccd state
@@ -428,22 +377,21 @@ class RMAOpen(object):
             output = self.send_cmd_get_output('ccd')
             if 'State: Open' not in output:
                 raise ValueError('Could not open cr50')
-        info('ccd is open')
-
+        logging.info('ccd is open')
 
     def enable_testlab(self):
         """Disable write protect"""
         if not self._has_testlab_support():
-            debug('Testlab mode is not supported in prod iamges')
+            logging.warning('Testlab mode is not supported in prod iamges')
             return
         # Some cr50 images need to be in dev mode before they can be opened.
         if self._requires_dev_mode_open():
             self._open_in_dev_mode()
         else:
             self.send_cmd_get_output('ccd open')
-        print('Enabling testlab mode reqires pressing the power button.')
-        print('Once the process starts keep tapping the power button for 10 '
-              'seconds.')
+        logging.info('Enabling testlab mode reqires pressing the power button.')
+        logging.info('Once the process starts keep tapping the power button '
+                     'for 10 seconds.')
         input("Press Enter when you're ready to start...")
         end_time = time.time() + 15
 
@@ -452,7 +400,7 @@ class RMAOpen(object):
         output = ''
         # start ccd testlab enable
         ser.write(self.ENABLE_TESTLAB_CMD.encode('utf-8'))
-        print('start pressing the power button\n\n')
+        logging.info('start pressing the power button\n\n')
         # Print all of the cr50 output as we get it, so the user will have more
         # information about pressing the power button. Tapping the power button
         # a couple of times should do it, but this will give us more confidence
@@ -463,8 +411,8 @@ class RMAOpen(object):
                 full_lines = output.rsplit('\n', 1)[0]
                 new_lines = full_lines
                 if printed_lines:
-                    new_lines = full_lines.split(printed_lines, 1)[-1]
-                print(new_lines, end=' ')
+                    new_lines = full_lines.split(printed_lines, 1)[-1].strip()
+                logging.info('\n%s', new_lines)
                 printed_lines = full_lines
 
                 # Make sure the process hasn't ended. If it has, print the last
@@ -472,11 +420,11 @@ class RMAOpen(object):
                 new_lines = output.split(printed_lines, 1)[-1]
                 if 'CCD test lab mode enabled' in output:
                     # print the last of the ou
-                    print(new_lines)
+                    logging.info(new_lines)
                     break
                 elif 'Physical presence check timeout' in output:
-                    print(new_lines)
-                    debug('Did not detect power button press in time')
+                    logging.info(new_lines)
+                    logging.warning('Did not detect power button press in time')
                     raise ValueError('Could not enable testlab mode try again')
         finally:
             ser.close()
@@ -487,42 +435,38 @@ class RMAOpen(object):
         if not self.check(TESTLAB_IS_ENABLED):
             raise ValueError('Could not enable testlab mode try again')
 
-
     def wp_disable(self):
         """Disable write protect"""
-        print('Disabling write protect')
+        logging.info('Disabling write protect')
         self.send_cmd_get_output('wp disable')
         # Update the state after attempting to disable write protect
         self.update_ccd_state()
         if not self.check(WP_IS_DISABLED):
             raise ValueError('Could not disable write protect')
 
-
     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)
+            logging.warning(DEBUG_DEVICE, self.device)
             raise ValueError('Could not communicate with %s' % self.device)
 
-        version = re.search('RW.*\* ([\d\.]+)/', output).group(1)
-        print('Running Cr50 Version:', version)
+        version = re.search(r'RW.*\* ([\d\.]+)/', output).group(1)
+        logging.info('Running Cr50 Version: %s', version)
         self.running_ver_fields = [int(field) for field in version.split('.')]
 
         # prePVT images have even major versions. Prod have odd
         self.is_prepvt = self.running_ver_fields[1] % 2 == 0
         rma_support = RMA_SUPPORT_PREPVT if self.is_prepvt else RMA_SUPPORT_PROD
 
-        print('%s RMA support added in: %s' %
-              ('prePVT' if self.is_prepvt else 'prod', rma_support))
-        if (not self.is_prepvt and
-            self._running_version_is_older(TESTLAB_PROD)):
-            debug('No testlab support in old prod images')
-            debug('Update to %s to enable testlab' % TESTLAB_PROD)
+        logging.info('%s RMA support added in: %s',
+                     'prePVT' if self.is_prepvt else 'prod', rma_support)
+        if not self.is_prepvt and self._running_version_is_older(TESTLAB_PROD):
+            logging.warning('No testlab support in old prod images')
+            logging.warning('Update to %s to enable testlab', TESTLAB_PROD)
         if self._running_version_is_older(rma_support):
             raise ValueError('%s does not have RMA support. Update to at '
-                    'least %s' % (version, rma_support))
-
+                             'least %s' % (version, rma_support))
 
     def _running_version_is_older(self, comp_ver):
         """Returns True if running version is older than comp_ver."""
@@ -532,7 +476,6 @@ class RMAOpen(object):
                 return True
         return False
 
-
     def device_matches_devid(self, devid, device):
         """Return True if the device matches devid.
 
@@ -542,18 +485,16 @@ class RMAOpen(object):
             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)
+        logging.debug('Sysinfo output: %s', 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
 
@@ -567,9 +508,9 @@ class RMAOpen(object):
             ValueError if the console can't be found with the given serialname
         """
         usb_serial = self.find_cr50_usb(usb_serial)
-        info('SERIALNAME: ' + usb_serial)
+        logging.info('SERIALNAME: %s', usb_serial)
         devid = '0x' + ' 0x'.join(usb_serial.lower().split('-'))
-        info('DEVID: ' + devid)
+        logging.info('DEVID: %s', devid)
 
         # Get all the usb devices
         devices = glob.glob('/dev/ttyUSB*')
@@ -579,102 +520,146 @@ class RMAOpen(object):
 
         # Find the one that is the cr50 console
         for device in devices:
-            print('testing', device)
+            logging.info('testing %s', device)
             if self.device_matches_devid(devid, device):
-                print('found device', device)
+                logging.info('found device: %s', device)
                 return
-        debug(DEBUG_CONNECTION)
+        logging.warning(DEBUG_CONNECTION)
         raise ValueError('Found USB device, but could not communicate with '
-                'cr50 console')
-
+                         '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)
+        bid = re.search(r'Board ID: (\S+?)[:,]', bid_output).group(1)
         if bid == ERASED_BID:
-            debug(DEBUG_ERASED_BOARD_ID)
+            logging.warning(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]))
-
+        logging.info('RLZ: %s', ''.join(chrs[::-1]))
 
-    def find_cr50_usb(self, usb_serial):
+    @staticmethod
+    def find_cr50_usb(usb_serial):
         """Make sure the Cr50 USB device exists"""
         try:
             output = subprocess.check_output(CR50_LSUSB_CMD, encoding='utf-8')
         except:
-            debug(DEBUG_MISSING_USB)
+            logging.warning(DEBUG_MISSING_USB)
             raise ValueError('Could not find Cr50 USB device')
-        serialnames = re.findall('iSerial +\d+ (\S+)\s', output)
+        serialnames = re.findall(r'iSerial +\d+ (\S+)\s', output)
         if usb_serial:
             if usb_serial not in serialnames:
-                debug(DEBUG_SERIALNAME)
+                logging.warning(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)
+            logging.info('Found Cr50 device serialnames %s',
+                         ', '.join(serialnames))
+            logging.warning(DEBUG_TOO_MANY_USB_DEVICES)
             raise ValueError('Too many cr50 usb devices')
         return serialnames[0]
 
+    def print_dut_state(self):
+        """Print CCD RMA and testlab mode state."""
+        if not self.check(CCD_IS_UNRESTRICTED):
+            logging.info('CCD is still restricted.')
+            logging.info('Run cr50_rma_open.py -g -i $HWID to generate a url')
+            logging.info('Run cr50_rma_open.py -a $AUTHCODE to open cr50 with '
+                         'an authcode')
+        elif not self.check(WP_IS_DISABLED):
+            logging.info('WP is still enabled.')
+            logging.info('Run cr50_rma_open.py -w to disable write protect')
+        if self.check(RMA_OPENED):
+            logging.info('RMA Open complete')
+
+        if not self.check(TESTLAB_IS_ENABLED) and self.is_prepvt:
+            logging.info('testlab mode is disabled.')
+            logging.info('If you are prepping a device for the testlab, you '
+                         'should enable testlab mode.')
+            logging.info('Run cr50_rma_open.py -t to enable testlab mode')
+
+
+def parse_args(argv):
+    """Get cr50_rma_open args."""
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument('-g', '--generate_challenge', action='store_true',
+                        help='Generate Cr50 challenge. Must be used with -i')
+    parser.add_argument('-t', '--enable_testlab', action='store_true',
+                        help='enable testlab mode')
+    parser.add_argument('-w', '--wp_disable', action='store_true',
+                        help='Disable write protect')
+    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', '--debug', action='store_true',
+                        help='print debug messages')
+    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. Needed to generate a challenge')
+    parser.add_argument('-a', '--authcode', type=str, default='',
+                        help='The authcode string from the challenge url')
+    parser.add_argument('-P', '--servo_port', type=str, default='',
+                        help='the servo port')
+    parser.add_argument('-I', '--ip', type=str, default='',
+                        help='The DUT IP. Necessary to do ccd open')
+    return parser.parse_args(argv)
+
+
+def main(argv):
+    """Run cr50 rma open."""
+    opts = parse_args(argv)
+
+    loglevel = logging.INFO
+    log_format = '%(levelname)7s'
+    if opts.debug:
+        loglevel = logging.DEBUG
+        log_format += ' - %(lineno)3d:%(funcName)-15s'
+    log_format += ' - %(message)s'
+    logging.basicConfig(level=loglevel, format=log_format)
 
-def main():
-    args = parser.parse_args()
     tried_authcode = False
-    info('Running cr50_rma_open version %s' % SCRIPT_VERSION)
+    logging.info('Running cr50_rma_open version %s', SCRIPT_VERSION)
 
-    cr50_rma_open = RMAOpen(args.device, args.serialname, args.print_caps,
-            args.servo_port, args.ip)
-    if args.check_connection:
+    cr50_rma_open = RMAOpen(opts.device, opts.serialname, opts.servo_port,
+                            opts.ip)
+    if opts.check_connection:
         sys.exit(0)
 
     if not cr50_rma_open.check(CCD_IS_UNRESTRICTED):
-        if args.generate_challenge:
-            if not args.hwid:
-                debug('--hwid necessary to generate challenge url')
+        if opts.generate_challenge:
+            if not opts.hwid:
+                logging.warning('--hwid necessary to generate challenge url')
                 sys.exit(0)
-            cr50_rma_open.generate_challenge_url(args.hwid)
+            cr50_rma_open.generate_challenge_url(opts.hwid)
             sys.exit(0)
-        elif args.authcode:
-            info('Using authcode: ' + args.authcode)
-            ccd_state = cr50_rma_open.try_authcode(args.authcode)
+        elif opts.authcode:
+            logging.info('Using authcode: %s', opts.authcode)
+            cr50_rma_open.try_authcode(opts.authcode)
             tried_authcode = True
 
     if not cr50_rma_open.check(WP_IS_DISABLED) and (tried_authcode or
-            args.wp_disable):
+                                                    opts.wp_disable):
         if not cr50_rma_open.check(CCD_IS_UNRESTRICTED):
             raise ValueError("Can't disable write protect unless ccd is "
-                    "open. Run through the rma open process first")
+                             "open. Run through the rma open process first")
         if tried_authcode:
-            debug("It's weird rma open didn't disable write protect. Trying to "
-                    "disable it manually")
+            logging.warning('RMA Open did not disable write protect. File a '
+                            'bug')
+            logging.warning('Trying to disable it manually')
         cr50_rma_open.wp_disable()
 
-    if not cr50_rma_open.check(TESTLAB_IS_ENABLED) and args.enable_testlab:
+    if not cr50_rma_open.check(TESTLAB_IS_ENABLED) and opts.enable_testlab:
         if not cr50_rma_open.check(CCD_IS_UNRESTRICTED):
-            raise ValueError("Can't enable testlab mode unless ccd is "
-                    "open. Run through the rma open process first")
+            raise ValueError("Can't enable testlab mode unless ccd is open."
+                             "Run through the rma open process first")
         cr50_rma_open.enable_testlab()
 
+    cr50_rma_open.print_dut_state()
 
-    if not cr50_rma_open.check(CCD_IS_UNRESTRICTED):
-        print('CCD is still restricted.')
-        print('Run cr50_rma_open.py -g -i $HWID to generate a url')
-        print('Run cr50_rma_open.py -a $AUTHCODE to open cr50 with an authcode')
-    elif not cr50_rma_open.check(WP_IS_DISABLED):
-        print('WP is still enabled.')
-        print('Run cr50_rma_open.py -w to disable write protect')
-    if cr50_rma_open.check(RMA_OPENED):
-        info('RMA Open complete')
-
-    if not cr50_rma_open.check(TESTLAB_IS_ENABLED) and cr50_rma_open.is_prepvt:
-        print('testlab mode is still disabled.')
-        print('If you are prepping a device for the testlab, you should enable '
-              'testlab mode')
-        print('Run cr50_rma_open.py -t to enable testlab mode')
-
-if __name__ == "__main__":
-    main()
 
+if __name__ == '__main__':
+    sys.exit(main(sys.argv[1:]))