add finding default-device instead of using the hardcoded one

Author: sujinmkang

ssdutil/main.py

@@ -18,26 +18,44 @@
 DEFAULT_DEVICE="/dev/sda"
 SYSLOG_IDENTIFIER = "ssdutil"
 DISK_TYPE_SSD = "0"
+DISK_INVALID = "-1"
 
 # Global logger instance
 log = logger.Logger(SYSLOG_IDENTIFIER)
 
+def get_default_disk():
+    """Check default disk"""
+    default_device = DEFAULT_DEVICE
+    host_mnt = '/host'
+    cmd = "lsblk -l -n |grep {}".format(host_mnt)
+    proc = subprocess.Popen(cmd, shell=True, text=True, stdout=subprocess.PIPE)
+    out = proc.stdout.readline()
+    if host_mnt in out:
+        dev_nums = out.split()[1]
+        dev_maj_num = out.split(':')[0]
+
+        cmd = "lsblk -l -I {} |grep disk".format(dev_maj_num)
+        proc = subprocess.Popen(cmd, shell=True, text=True, stdout=subprocess.PIPE)
+        out = proc.stdout.readline()
+        if "disk" in out:
+            default_device = os.path.join("/dev/", out.split()[0])
+
+    return default_device
+
 
 def get_disk_type(diskdev):
     """Check disk type"""
     diskdev_name = diskdev.replace('/dev/','')
-    cmd = "lsblk -l -n |grep disk"
+    cmd = "lsblk -l -n |grep {}".format(diskdev_name)
     proc = subprocess.Popen(cmd, shell=True, text=True, stdout=subprocess.PIPE)
-    outs = proc.stdout.readlines()
-    for out in outs:
-        if out.split()[0] is diskdev_name:
-              cmd = "cat /sys/block/{}/queue/rotational".format(diskdev_name)
-              proc = subprocess.Popen(cmd, shell=True, text=True, stdout=subprocess.PIPE)
-              out = proc.stdout.readline()
-              return out.rstrip()
-
-    print("disk {} does not exist in the device".format(diskdev_name))
-    sys.exit(1)
+    out = proc.stdout.readline()
+    if diskdev_name not in out:
+        return DISK_INVAILD
+    cmd = "cat /sys/block/{}/queue/rotational".format(diskdev_name)
+    proc = subprocess.Popen(cmd, shell=True, text=True, stdout=subprocess.PIPE)
+    out = proc.stdout.readline()
+    disk_type = out.rstrip()
+    return disk_type
 
 
 def import_ssd_api(diskdev):
@@ -79,17 +97,22 @@ def ssdutil():
         sys.exit(1)
 
     parser = argparse.ArgumentParser()
-    parser.add_argument("-d", "--device", help="Device name to show health info", default=DEFAULT_DEVICE)
+    parser.add_argument("-d", "--device", help="Device name to show health info", default=None)
     parser.add_argument("-v", "--verbose", action="store_true", default=False, help="Show verbose output (some additional parameters)")
     parser.add_argument("-e", "--vendor", action="store_true", default=False, help="Show vendor output (extended output if provided by platform vendor)")
     args = parser.parse_args()
 
-    disk_type = get_disk_type(args.device)
+    if args.device:
+        disk_device = args.device
+    else:
+        disk_device = get_default_disk()
+
+    disk_type = get_disk_type(disk_device)
     if disk_type != DISK_TYPE_SSD:
         print("Disk is not SSD")
         sys.exit(1)
 
-    ssd = import_ssd_api(args.device)
+    ssd = import_ssd_api(disk_device)
 
     print("Device Model : {}".format(ssd.get_model()))
     if args.verbose: