From 3c5f9c1bbe76cac7040da09e786aa59f74f0488b Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Mon, 18 Nov 2019 21:47:23 -0700
Subject: [PATCH] Added workaround for MX500 SSD issues

---
 .bin/Scripts/functions/hw_diags.py | 39 ++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/.bin/Scripts/functions/hw_diags.py b/.bin/Scripts/functions/hw_diags.py
index 79123b3c..528040a6 100644
--- a/.bin/Scripts/functions/hw_diags.py
+++ b/.bin/Scripts/functions/hw_diags.py
@@ -22,6 +22,7 @@ TOP_PANE_TEXT = TOP_PANE_TEXT.format(**COLORS)
 
 # Regex
 REGEX_ERROR_STATUS = re.compile('|'.join(STATUSES['RED']))
+REGEX_MX500 = re.compile(r'CT(250|500|1000|2000)MX500SSD(1|4)', re.IGNORECASE)
 
 
 # Error Classes
@@ -87,6 +88,7 @@ class CpuObj():
 class DiskObj():
   """Object for tracking disk specific data."""
   def __init__(self, disk_path):
+    self.attributes = {}
     self.checkbox = None
     self.attr_type = 'UNKNOWN'
     self.disk_ok = True
@@ -119,6 +121,13 @@ class DiskObj():
     self.get_smart_details()
     self.description = '{size} ({tran}) {model} {serial}'.format(
       **self.lsblk)
+    self.update_attributes()
+
+  def update_attributes(self):
+    """HACK to adjust attribute thresholds for Crucial MX500 SSDs."""
+    self.attributes = ATTRIBUTES.copy()
+    if REGEX_MX500.search(self.lsblk['model']):
+      self.attributes['SMART'][197].update({'Warning': 1, 'Error': 2})
 
   def add_nvme_smart_note(self, note):
     """Add note that will be included in the NVMe / SMART report."""
@@ -189,8 +198,8 @@ class DiskObj():
     elif self.smart_attributes:
       poh = self.smart_attributes.get(9, {}).get('raw', -1)
 
-    error_thresh = ATTRIBUTES['SMART'][9]['Error']
-    max_thresh = ATTRIBUTES['SMART'][9]['Maximum']
+    error_thresh = self.attributes['SMART'][9]['Error']
+    max_thresh = self.attributes['SMART'][9]['Maximum']
 
     return error_thresh <= poh < max_thresh
 
@@ -211,23 +220,23 @@ class DiskObj():
     elif self.smart_attributes:
       items = self.smart_attributes.items()
     for k, v in items:
-      if k in ATTRIBUTES[attr_type]:
-        if not ATTRIBUTES[attr_type][k]['Error']:
+      if k in self.attributes[attr_type]:
+        if not self.attributes[attr_type][k]['Error']:
           # Informational attribute, skip
           continue
-        if ATTRIBUTES[attr_type][k]['Ignore']:
+        if self.attributes[attr_type][k]['Ignore']:
           # Attribute is non-failing, skip
           continue
-        if v['raw'] >= ATTRIBUTES[attr_type][k]['Error']:
-          if (ATTRIBUTES[attr_type][k]['Maximum']
-              and v['raw'] >= ATTRIBUTES[attr_type][k]['Maximum']):
+        if v['raw'] >= self.attributes[attr_type][k]['Error']:
+          if (self.attributes[attr_type][k]['Maximum']
+              and v['raw'] >= self.attributes[attr_type][k]['Maximum']):
             # Non-standard value, skip
             continue
           else:
             disk_ok = False
 
             # Disable override if necessary
-            if ATTRIBUTES[attr_type][k].get('Critical', False):
+            if self.attributes[attr_type][k].get('Critical', False):
               self.override_disabled = True
 
     # SMART overall assessment
@@ -298,7 +307,7 @@ class DiskObj():
       attr_type = 'SMART'
       items = self.smart_attributes.items()
     for k, v in items:
-      if k in ATTRIBUTES[attr_type]:
+      if k in self.attributes[attr_type]:
         _note = ''
         _color = COLORS['GREEN']
 
@@ -308,19 +317,23 @@ class DiskObj():
         else:
           _line = '  {i:>3} / {h}: {n:28}'.format(
             i=k,
-            h=ATTRIBUTES[attr_type][k]['Hex'],
+            h=self.attributes[attr_type][k]['Hex'],
             n=v['name'][:28])
 
         # Set color
         for _t, _c in ATTRIBUTE_COLORS:
-          if ATTRIBUTES[attr_type][k][_t]:
-            if v['raw'] >= ATTRIBUTES[attr_type][k][_t]:
+          if self.attributes[attr_type][k][_t]:
+            if v['raw'] >= self.attributes[attr_type][k][_t]:
               _color = COLORS[_c]
               if _t == 'Error':
                 _note = '(failed)'
               elif _t == 'Maximum':
                 _note = '(invalid?)'
 
+        # 197/C5 warning
+        if str(k) == '197' and REGEX_MX500.search(self.lsblk['model']):
+          _note = '(MX500 thresholds)'
+
         # 199/C7 warning
         if str(k) == '199' and v['raw'] > 0:
           _note = '(bad cable?)'