Migrate USB driver check to python

rootfs/standard/etc/systemd/system/usb_driver_check.service

@@ -8,7 +8,7 @@ After=mynode.service
 
 [Service]
 Type=oneshot
-ExecStart=/usr/bin/mynode_usb_driver_check.sh
+ExecStart=/usr/local/bin/python3 /usr/bin/mynode_usb_driver_check.py
 User=root
 Group=root
 StandardOutput=syslog

rootfs/standard/usr/bin/mount_drive.tcl

@@ -79,7 +79,7 @@ proc createMyNodeFsOnBlockDevice {blockDevice} {
 
     if [catch {
         # Run USB check to make sure we are using a good driver
-        runCommand /usr/bin/mynode_usb_driver_check.sh > /dev/null
+        runCommand /usr/local/bin/python3 /usr/bin/mynode_usb_driver_check.py > /dev/null
 
         puts "Waiting on format confirmation..."
         runCommand echo "drive_format_confirm" > /tmp/.mynode_status

rootfs/standard/usr/bin/mynode_startup.sh

@@ -137,6 +137,8 @@ do
     chmod 644 /home/bitcoin/.mynode/mynode_serial
 done
 
+# Check for USB driver updates before mount or clone tool opening
+/usr/local/bin/python3 /usr/bin/mynode_usb_driver_check.py
 
 # Clone tool was opened
 if [ -f /home/bitcoin/open_clone_tool ]; then

rootfs/standard/usr/bin/mynode_usb_driver_check.py

@@ -0,0 +1,208 @@
+#!/usr/local/bin/python3
+
+import os
+import re
+import logging
+import argparse
+from systemd import journal
+from utilities import *
+from device_info import *
+
+args = None
+def get_args():
+    global args
+    return args
+def set_args(a):
+    global args
+    args = a
+
+def is_uas_forced():
+    if os.path.isfile("/home/bitcoin/.mynode/.uas_usb_enabled") or \
+       os.path.isfile("/mnt/hdd/mynode/settings/.uas_usb_enabled"):
+        return True
+    return False
+
+def device_checks_uas():
+    d = get_device_type()
+    if (d == "raspi4" or d == "rockpi4"):
+        return True
+    return False
+
+def has_quirks_setting():
+    d = get_device_type()
+    if d == "raspi4":
+        line = run_linux_cmd("head -n 1 /boot/cmdline.txt")
+        if "usb-storage.quirks=" in line:
+            return True
+    elif d == "rockpi4":
+        lines = run_linux_cmd("cat /boot/armbianEnv.txt")
+        if "usbstoragequirks=" in lines:
+            return True
+    else:
+        raise Exception("Unexpected Device Type")
+    return False
+
+def get_current_usb_quirks():
+    d = get_device_type()
+    quirks = []
+    try:
+        if d == "raspi4":
+            if os.path.isfile("/boot/cmdline.txt"):
+                content = run_linux_cmd("head -n 1 /boot/cmdline.txt")
+                m = re.search("usb-storage.quirks=(.+)", content)
+                quirks_string = m.group(1)
+                quirks = quirks_string.split(",")
+            else:
+                raise Exception("Missing file: /boot/cmdline.txt")
+        elif d == "rockpi4":
+            if os.path.isfile("/boot/armbianEnv.txt"):
+                content = run_linux_cmd("cat /boot/armbianEnv.txt")
+                m = re.search("usbstoragequirks=(.+)", content)
+                quirks_string = m.group(1)
+                quirks = quirks_string.split(",")
+            else:
+                raise Exception("Missing file: /boot/armbianEnv.txt")
+        else:
+            raise Exception("Unexpected Device Type")
+    except Exception as e:
+        log_message("Exception in get_current_usb_quirks: {}".format(str(e)))
+    return quirks
+
+def generate_quirks_string(quirks_list):
+    quirks_string = ""
+    d = get_device_type()
+    if d == "raspi4":
+        quirks_string += "usb-storage.quirks="
+    elif d == "rockpi4":
+        quirks_string += "usbstoragequirks="
+    else:
+        raise Exception("Unexpected Device Type")
+
+    for q in quirks_list:
+        quirks_string += q + ","
+    quirks_string = quirks_string.rstrip(",")
+    return quirks_string
+
+def update_usb_quirks(quirks_list):
+    d = get_device_type()
+    quirks_string = generate_quirks_string(quirks_list)
+    log_message("Updating Quirks: {}".format(quirks_string))
+    if has_quirks_setting():
+        # Update Quirks
+        if d == "raspi4":
+            run_linux_cmd("sed -i \"s/usb-storage.quirks=.*/"+quirks_string+"/g\" /boot/cmdline.txt")
+        elif d == "rockpi4":
+            run_linux_cmd("sed -i \"s/usbstoragequirks=.*/"+quirks_string+"/g\" /boot/armbianEnv.txt")
+            run_linux_cmd("mkimage -C none -A arm -T script -d /boot/boot.cmd /boot/boot.scr")
+        else:
+            raise Exception("Unexpected Device Type")
+    else:
+        # Add Quirks
+        if d == "raspi4":
+            contents = run_linux_cmd("head -n 1 /boot/cmdline.txt").strip()
+            contents += " " + quirks_string
+            set_file_contents("/boot/cmdline.txt", contents)
+        elif d == "rockpi4":
+            # Rock pi 4 normally never has this missing, if so, need updates
+            run_linux_cmd("sed -i \"s/usbstoragequirks=.*/"+quirks_string+"/g\" /boot/armbianEnv.txt")
+            run_linux_cmd("mkimage -C none -A arm -T script -d /boot/boot.cmd /boot/boot.scr")
+        else:
+            raise Exception("Unexpected Device Type")
+
+def remove_usb_quirks():
+    if has_quirks_setting():
+        current = get_current_usb_quirks()
+        if current != ["none"]:
+            update_usb_quirks(["none"])
+            reboot()
+    else:
+        # No quirks setting, nothing to do
+        return
+
+def get_quirks_from_uas_devices():
+    quirks = []
+    lsusbt_output = run_linux_cmd("lsusb -t")
+    if "Driver=uas" in lsusbt_output:
+        log_message("UAS Driver in use! Looking for likely candidates...")
+        lsusb_output = run_linux_cmd("lsusb")
+        lsusb_lines = lsusb_output.splitlines()
+        for line in lsusb_lines:
+            try:
+                if "SATA 6Gb/s bridge" in line:
+                    m = re.search("usb-storage.quirks=(.+)", line)
+                    m = re.search("Bus ([0-9+]) Device ([0-9]+): ID (.+)", line)
+                    bus = m.group(1)
+                    dev = m.group(2)
+                    id = m.group(3)
+                    quirks.append(id+":u")
+            except Exception as e:
+                log_message("Unable to scan USB device: {}".format(line))
+    return quirks
+
+def get_required_usb_quirks():
+    required_quirks = []
+
+    # Add known devices
+    required_quirks.append("174c:55aa:u") # ASMedia Technology Inc.
+    required_quirks.append("152d:1561:u") # JMicron Technology Corp.
+
+    # Add any devices currently found as UAS
+    required_quirks += get_quirks_from_uas_devices()
+
+    # Remove duplicates
+    required_quirks = list(set(required_quirks))
+
+    return required_quirks
+
+def reboot():
+    if not get_args().no_reboot:
+        log_message("Rebooting...")
+        run_linux_cmd("sync")
+        time.sleep(3)
+        run_linux_cmd("mynode-reboot")
+    else:
+        log_message("Skipping reboot.")
+
+def main():
+    parser = argparse.ArgumentParser(description='Check and Update USB Drivers')
+    #parser.add_argument('--no-modify', action='store_true', help="Do not modify any files")
+    parser.add_argument('--no-reboot', action='store_true', help="Do not reboot the device after updating files")
+    args = parser.parse_args()
+    set_args(args)
+
+    if not device_checks_uas():
+        log_message("Device doesn't need to check UAS. Exiting.")
+        return
+
+    if is_uas_forced():
+        log_message("UAS is allowed by settings. Removing USB quirks.")
+        remove_usb_quirks()
+        return
+
+    current_quirks = get_current_usb_quirks()
+    print("Current Quirks:  {}".format(current_quirks))
+    required_quirks = get_required_usb_quirks()
+    print("Required Quirks: {}".format(required_quirks))
+    updated_quirks = current_quirks.copy()
+    for r in required_quirks:
+        if r not in current_quirks:
+            updated_quirks.append(r)
+
+    if updated_quirks != current_quirks:
+        update_usb_quirks(updated_quirks)
+        reboot()
+    else:
+        log_message("No update necessary. Exiting.")
+
+
+# This is the main entry point for the program
+if __name__ == "__main__":
+    try:
+        log = logging.getLogger('usb_driver_check')
+        log.addHandler(journal.JournaldLogHandler())
+        log.setLevel(logging.INFO)
+        set_logger(log)
+
+        main()
+    except Exception as e:
+        log_message("Exception: {}".format(str(e)))

rootfs/standard/usr/bin/mynode_usb_driver_check.sh

@@ -2,81 +2,83 @@
 
 set -x
 
-source /usr/share/mynode/mynode_device_info.sh
+# NO LONGER USED
 
-# Allow UAS
+# source /usr/share/mynode/mynode_device_info.sh
-if [ -f /mnt/hdd/mynode/settings/.uas_usb_enabled ] || [ -f /home/bitcoin/.mynode/.uas_usb_enabled ]; then
-    cat /boot/cmdline.txt | grep "usb-storage.quirks="
-    if [ $? -eq 0 ]; then
-        cat /boot/cmdline.txt | grep "usb-storage.quirks=none"
-        if [ $? -eq 0 ]; then
-            exit 0
-        else
-            sed -i "s/usb-storage.quirks=.*/usb-storage.quirks=none/g" /boot/cmdline.txt
-            sync
-            /usr/bin/mynode-reboot
-        fi
-    fi
-    exit 0
-else
-    # Disable UAS
-    lsusb -t | grep "Driver=uas"
-    if [ $? -eq 0 ]; then
-        echo "UAS FOUND"
-        USBINFO=$(lsusb | grep "SATA 6Gb/s bridge")
-        DEVID=$(egrep -o '[0-9a-f]+:[0-9a-f]+' <<< $USBINFO)
-        QUIRK="${DEVID}:u"
-        echo "UAS IN USE ON $DEVID"
 
-        # Raspberry Pi
+# # Allow UAS
-        if [ $IS_RASPI -eq 1 ]; then
+# if [ -f /mnt/hdd/mynode/settings/.uas_usb_enabled ] || [ -f /home/bitcoin/.mynode/.uas_usb_enabled ]; then
-            echo "IS RASPBERRY PI"
+#     cat /boot/cmdline.txt | grep "usb-storage.quirks="
-            if [ -f /boot/cmdline.txt ]; then
+#     if [ $? -eq 0 ]; then
-                CMDLINE=$(head -n 1 /boot/cmdline.txt)
+#         cat /boot/cmdline.txt | grep "usb-storage.quirks=none"
-                cat /boot/cmdline.txt | grep "usb-storage.quirks"
+#         if [ $? -eq 0 ]; then
-                if [ $? -eq 0 ]; then
+#             exit 0
-                    cat /boot/cmdline.txt | grep "usb-storage.quirks=${QUIRK}"
+#         else
-                    if [ $? -eq 0 ]; then
+#             sed -i "s/usb-storage.quirks=.*/usb-storage.quirks=none/g" /boot/cmdline.txt
-                        # Quirk already added, exit 0
+#             sync
-                        echo "QUIRK ALREADY EXISTS, EXITING"
+#             /usr/bin/mynode-reboot
-                        exit 0
+#         fi
-                    else
+#     fi
-                        # Different quirk exists, update and reboot
+#     exit 0
-                        echo "DIFFERENT QUIRK FOUND, UPDATING AND REBOOTING"
+# else
-                        sed -i "s/usb-storage.quirks=.*/usb-storage.quirks=${QUIRK}/g" /boot/cmdline.txt
+#     # Disable UAS
-                    fi
+#     lsusb -t | grep "Driver=uas"
-                else
+#     if [ $? -eq 0 ]; then
-                    # No quirk found, add it and reboot
+#         echo "UAS FOUND"
-                    echo "NO QUIRK FOUND, ADDING AND REBOOTING"
+#         USBINFO=$(lsusb | grep "SATA 6Gb/s bridge")
-                    echo "${CMDLINE} usb-storage.quirks=${QUIRK}" > /boot/cmdline.txt
+#         DEVID=$(egrep -o '[0-9a-f]+:[0-9a-f]+' <<< $USBINFO)
-                fi
+#         QUIRK="${DEVID}:u"
+#         echo "UAS IN USE ON $DEVID"
 
-                sync
+#         # Raspberry Pi
-                sleep 5s
+#         if [ $IS_RASPI -eq 1 ]; then
-                /usr/bin/mynode-reboot
+#             echo "IS RASPBERRY PI"
-            fi
+#             if [ -f /boot/cmdline.txt ]; then
-        fi
+#                 CMDLINE=$(head -n 1 /boot/cmdline.txt)
+#                 cat /boot/cmdline.txt | grep "usb-storage.quirks"
+#                 if [ $? -eq 0 ]; then
+#                     cat /boot/cmdline.txt | grep "usb-storage.quirks=${QUIRK}"
+#                     if [ $? -eq 0 ]; then
+#                         # Quirk already added, exit 0
+#                         echo "QUIRK ALREADY EXISTS, EXITING"
+#                         exit 0
+#                     else
+#                         # Different quirk exists, update and reboot
+#                         echo "DIFFERENT QUIRK FOUND, UPDATING AND REBOOTING"
+#                         sed -i "s/usb-storage.quirks=.*/usb-storage.quirks=${QUIRK}/g" /boot/cmdline.txt
+#                     fi
+#                 else
+#                     # No quirk found, add it and reboot
+#                     echo "NO QUIRK FOUND, ADDING AND REBOOTING"
+#                     echo "${CMDLINE} usb-storage.quirks=${QUIRK}" > /boot/cmdline.txt
+#                 fi
 
-        # RockPi 4
+#                 sync
-        if [ $IS_ROCKPI4 -eq 1 ]; then
+#                 sleep 5s
-            echo "IS ROCKPI 4"
+#                 /usr/bin/mynode-reboot
-            if [ -f /boot/armbianEnv.txt ]; then
+#             fi
-                cat /boot/armbianEnv.txt | grep "${QUIRK}"
+#         fi
-                if [ $? -eq 0 ]; then
+
-                    # Quirk already added, exit 0
+#         # RockPi 4
-                    echo "QUIRK ALREADY EXISTS, EXITING"
+#         if [ $IS_ROCKPI4 -eq 1 ]; then
-                    exit 0
+#             echo "IS ROCKPI 4"
-                else
+#             if [ -f /boot/armbianEnv.txt ]; then
-                    echo "ADDING QUIRK AND REBOOTING"
+#                 cat /boot/armbianEnv.txt | grep "${QUIRK}"
-                    sed -i "s/usbstoragequirks=.*/&,${QUIRK}/g" cat /boot/armbianEnv.txt
+#                 if [ $? -eq 0 ]; then
-                    mkimage -C none -A arm -T script -d /boot/boot.cmd /boot/boot.scr
+#                     # Quirk already added, exit 0
-                    sync
+#                     echo "QUIRK ALREADY EXISTS, EXITING"
-                    sleep 5s
+#                     exit 0
-                    /usr/bin/mynode-reboot
+#                 else
-                fi
+#                     echo "ADDING QUIRK AND REBOOTING"
-            fi
+#                     sed -i "s/usbstoragequirks=.*/&,${QUIRK}/g" /boot/armbianEnv.txt
-        fi
+#                     mkimage -C none -A arm -T script -d /boot/boot.cmd /boot/boot.scr
-    else
+#                     sync
-        echo "No UAS devices found"
+#                     sleep 5s
-    fi
+#                     /usr/bin/mynode-reboot
-fi
+#                 fi
+#             fi
+#         fi
+#     else
+#         echo "No UAS devices found"
+#     fi
+# fi