servo_updater: add better default args

Let servo_updater look for configs and firmwares
in the default installed locations.

BUG=b:69016431
BRANCH=None
TEST=sudo servo_updater -b servo_v4

Change-Id: I069b96044ce17992e41a8fc66931477823eef986
Signed-off-by: Nick Sanders <nsanders@google.com>
Reviewed-on: https://chromium-review.googlesource.com/792603
Commit-Ready: Nick Sanders <nsanders@chromium.org>
Tested-by: Nick Sanders <nsanders@chromium.org>
Reviewed-by: Wai-Hong Tam <waihong@google.com>

1 files changed, 55 insertions, 5 deletions

diff --git a/extra/usb_updater/servo_updater.py b/extra/usb_updater/servo_updater.py
index ad453a5c20..c142f51c38 100755
--- a/extra/usb_updater/servo_updater.py
+++ b/extra/usb_updater/servo_updater.py
@@ -15,8 +15,14 @@ import json
 import fw_update
 import ecusb.tiny_servo_common as c
 
+FIRMWARE_PATH = "/usr/share/servo_updater/firmware/"
+CONFIGS_PATH = "/usr/share/servo_updater/configs/"
+
+class ServoUpdaterException(Exception):
+  """Raised on exceptions generated by servo_updater."""
 
 def flash(brdfile, serialno, binfile):
+  """Call fw_update to upload to updater USB endpoint."""
   p = fw_update.Supdate()
   p.load_board(brdfile)
   p.connect_usb(serialname=serialno)
@@ -34,6 +40,8 @@ def flash(brdfile, serialno, binfile):
 
 
 def select(vidpid, iface, serialno, region, debuglog=False):
+  """Ensure the servo is in the expected ro/rw partition."""
+
   if region not in ["rw", "ro"]:
     raise Exception("Region must be ro or rw")
 
@@ -50,21 +58,64 @@ def select(vidpid, iface, serialno, region, debuglog=False):
   pty.close()
 
 
+def findfiles(cname, fname):
+  """Select config and firmware binary files.
+
+  This checks default file names and paths.
+  In: /usr/share/servo_updater/[firmware|configs]
+  check for board.json, board.bin
+
+  Args:
+    cname: config name, or board name. eg. "servo_v4".
+    fname: firmware binary name. Can be None to try default.
+  Returns:
+    cname, fname: validated filenames selected from the path.
+  """
+  if not os.path.isfile(cname):
+    # If not an existing file, try checking on the default path.
+    newname = CONFIGS_PATH + cname
+    if os.path.isfile(newname):
+      cname = newname
+    elif os.path.isfile(newname + ".json"):
+      # Try appending ".json" to convert board name to config file.
+      cname = newname + ".json"
+    else:
+      raise ServoUpdaterException("Can't find file: %s." % cname)
+
+  if not fname:
+    # If None, try defaults.
+    dev = None
+    if "servo_v4" in cname:
+      dev = "servo_v4"
+    if "servo_micro" in cname:
+      dev = "servo_micro"
+    if os.path.isfile(FIRMWARE_PATH + dev + ".bin"):
+      fname = FIRMWARE_PATH + dev + ".bin"
+  elif not os.path.isfile(fname):
+    # If a name is specified but not found, try the default path.
+    if os.path.isfile(FIRMWARE_PATH + fname):
+      fname = FIRMWARE_PATH + fname
+    else:
+      raise ServoUpdaterException("Can't find file: %s." % fname)
+
+  return cname, fname
+
+
 def main():
   parser = argparse.ArgumentParser(description="Image a servo micro device")
   parser.add_argument('-s', '--serialno', type=str,
       help="serial number to program", default=None)
   parser.add_argument('-b', '--board', type=str,
-      help="Board configuration json file", default="servo_v4.json")
+      help="Board configuration json file", default="servo_v4")
   parser.add_argument('-f', '--file', type=str,
-      help="Complete ec.bin file", default="servo_v4.bin")
+      help="Complete ec.bin file", default=None)
   parser.add_argument('-v', '--verbose', action="store_true",
       help="Chatty output")
 
   args = parser.parse_args()
 
-  brdfile = args.board
-  binfile = args.file
+  brdfile, binfile = findfiles(args.board, args.file)
+
   serialno = args.serialno
   debuglog = (args.verbose is True)
 
@@ -84,4 +135,3 @@ def main():
 
 if __name__ == "__main__":
   main()
-