check for nvme ses capability

ssd-eraser

@@ -13,8 +13,9 @@
 __version__ = '0.1'
 
 import os
+import json
 from time import sleep
-from subprocess import run, CalledProcessError
+from subprocess import run, CalledProcessError, PIPE
 try:
     import pyudev
 except ImportError:
@@ -148,7 +149,14 @@ def unfreeze():
 def erase(disk):
     print(YELLOW+"Erasing "+disk+WHITE)
     if DINFO[disk][1] == 'NVMe':
-        cmd = ['nvme', 'format', disk, '--ses=2']
+        checkcap = ['nvme', 'id-ctrl', disk, '-o', 'json']
+        getcap = run(checkcap, stdout=PIPE, check=True)
+        nvmecap = json.loads(getcap.stdout)
+        if nvmecap['fna'] == 5:
+            ses = '--ses=2'
+        else:
+            ses = '--ses=1'
+        cmd = ['nvme', 'format', disk, ses]
         try:
             run(cmd, check=True)
             sleep(20)