Migrate USB driver check to python

This commit is contained in:
Taylor Helsper 2022-06-17 22:42:53 -05:00
parent 79ebdfe62d
commit 290622b487
5 changed files with 288 additions and 76 deletions

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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)))

View File

@ -2,81 +2,83 @@
set -x
source /usr/share/mynode/mynode_device_info.sh
# NO LONGER USED
# Allow UAS
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"
# source /usr/share/mynode/mynode_device_info.sh
# Raspberry Pi
if [ $IS_RASPI -eq 1 ]; then
echo "IS RASPBERRY PI"
if [ -f /boot/cmdline.txt ]; then
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
# # Allow UAS
# 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"
sync
sleep 5s
/usr/bin/mynode-reboot
fi
fi
# # Raspberry Pi
# if [ $IS_RASPI -eq 1 ]; then
# echo "IS RASPBERRY PI"
# if [ -f /boot/cmdline.txt ]; then
# 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
if [ $IS_ROCKPI4 -eq 1 ]; then
echo "IS ROCKPI 4"
if [ -f /boot/armbianEnv.txt ]; then
cat /boot/armbianEnv.txt | grep "${QUIRK}"
if [ $? -eq 0 ]; then
# Quirk already added, exit 0
echo "QUIRK ALREADY EXISTS, EXITING"
exit 0
else
echo "ADDING QUIRK AND REBOOTING"
sed -i "s/usbstoragequirks=.*/&,${QUIRK}/g" cat /boot/armbianEnv.txt
mkimage -C none -A arm -T script -d /boot/boot.cmd /boot/boot.scr
sync
sleep 5s
/usr/bin/mynode-reboot
fi
fi
fi
else
echo "No UAS devices found"
fi
fi
# sync
# sleep 5s
# /usr/bin/mynode-reboot
# fi
# fi
# # RockPi 4
# if [ $IS_ROCKPI4 -eq 1 ]; then
# echo "IS ROCKPI 4"
# if [ -f /boot/armbianEnv.txt ]; then
# cat /boot/armbianEnv.txt | grep "${QUIRK}"
# if [ $? -eq 0 ]; then
# # Quirk already added, exit 0
# echo "QUIRK ALREADY EXISTS, EXITING"
# exit 0
# else
# echo "ADDING QUIRK AND REBOOTING"
# sed -i "s/usbstoragequirks=.*/&,${QUIRK}/g" /boot/armbianEnv.txt
# mkimage -C none -A arm -T script -d /boot/boot.cmd /boot/boot.scr
# sync
# sleep 5s
# /usr/bin/mynode-reboot
# fi
# fi
# fi
# else
# echo "No UAS devices found"
# fi
# fi