initial working version

ssd-eraser

+#!/usr/bin/python3
+
+import os
+from time import sleep
+from subprocess import run, CalledProcessError
+try:
+    import pyudev
+except ImportError:
+    print("Module pyudev missing. Install with:")
+    print("'pip install pyudev'")
+    exit(1)
+
+# Got root?
+if os.getuid() != 0:
+    print("\nOnly root can run this script.\n")
+    exit(1)
+
+# Main Vars
+C = pyudev.Context()
+DISKS = []
+DINFO = {}
+for device in C.list_devices(subsystem='block', DEVTYPE='disk'):
+    # Ignore USB and CDROM
+    if device.get('ID_BUS') == 'ata' and device.get('ID_TYPE') == 'disk':
+        if device.get('ID_ATA_ROTATION_RATE_RPM') == '0':
+            DISKS.append(device.device_node)
+        DINFO[device.device_node] = device.get('ID_MODEL'), 'SSD'
+    if 'nvme' in device.get('DEVNAME'):
+        DISKS.append(device.device_node)
+        DINFO[device.device_node] = device.get('ID_SERIAL'), 'NVMe'
+
+# Colours
+RED = '\033[91m'
+YELLOW = '\033[93m'
+GREEN = '\033[92m'
+WHITE = '\033[0m'
+
+# Functions
+def banner():
+    os.system('clear')
+    b = """
+       _/_/_/_/  _/_/_/      _/_/      _/_/_/  _/_/_/_/  _/_/_/
+      _/        _/    _/  _/    _/  _/        _/        _/    _/
+     _/_/_/    _/_/_/    _/_/_/_/    _/_/    _/_/_/    _/_/_/
+    _/        _/    _/  _/    _/        _/  _/        _/    _/
+   _/_/_/_/  _/    _/  _/    _/  _/_/_/    _/_/_/_/  _/    _/
+        """
+    print('\n')
+    print(RED+b+WHITE)
+
+def main_menu():
+    banner()
+    all_disks = 0
+    print("\nChoose disk for more info:\n")
+    for i in range(len(DISKS)):
+        print('{0}) {1} [{2}]'.format(i, DISKS[i], DINFO[DISKS[i]][0]))
+    print("\n")
+    if len(DISKS) > 1:
+        all_disks = 1
+        print("a) All disks")
+    print("q) Quit\n")
+    choice = input(">>  ")
+    disk_select(choice, all_disks)
+    return
+
+def disk_select(choice, all_disks):
+    try:
+        if choice == '':
+            main_menu()
+        elif choice =='a' and all_disks:
+            erase_menu(DISKS)
+        elif choice =='q':
+            exit()
+        else:
+            try:
+                erase_menu(DISKS[int(choice)])
+            except IndexError:
+                main_menu()
+    except ValueError:
+        main_menu()
+    return
+
+def erase_menu(disk):
+    banner()
+    if type(disk) is list:
+        for d in disk:
+            cmd = ['lsblk','-no', 'NAME,SIZE,TYPE,FSTYPE', d]
+            print("\n")
+            print("[{0}] = [{1}]".format(DINFO[d][0], DINFO[d][1]))
+            print("-"*(len(DINFO[d][0])+2))
+            run(cmd)
+    else:
+        cmd = ['lsblk','-no', 'NAME,SIZE,TYPE,FSTYPE', disk]
+        print("\n")
+        print("[{0}] = [{1}]".format(DINFO[disk][0], DINFO[disk][1]))
+        print("-"*(len(DINFO[disk][0])+2))
+        run(cmd)
+
+    print(YELLOW+"\nReady to secure erase."+WHITE)
+    print("Computer will enter sleep-mode and proceed after 5sec.\n")
+    print("Continue?\n")
+    print("y) Yes")
+    print("c) Cancel")
+    print("q) Quit\n")
+    choice = input(">>  ")
+    confirm_select(choice, disk)
+
+def confirm_select(choice, disk):
+    try:
+        if choice == 'c':
+            main_menu()
+        elif choice == 'q':
+            exit()
+        elif choice == 'y':
+            try:
+                unfreeze()
+                if type(disk) is list:
+                    erase_all(disk)
+                else:
+                    erase(disk)
+            except:
+                print("Something went wrong...")
+                exit(1)
+        else:
+            erase_menu(disk)
+    except ValueError:
+        erase_menu(disk)
+    return
+
+def unfreeze():
+    cmd = ['rtcwake', '-m', 'mem', '-s', '5']
+    run(cmd)
+    
+def erase(disk):
+    print(YELLOW+"Erasing "+disk+WHITE)
+    if DINFO[disk][1] == 'NVMe':
+        cmd = ['nvme', 'format', disk, '--ses=2']
+        try:
+            run(cmd, check=True)
+            sleep(20)
+        except CalledProcessError:
+            print(RED+"Error erasing "+disk+WHITE)
+    else:
+        cmds = [['hdparm', '--security-set-pass', 'NULL', disk],
+                ['hdparm', '--security-erase', 'NULL', disk]]
+        for cmd in cmds:
+            try:
+                run(cmd, check=True)
+                sleep(20)
+            except CalledProcessError:
+                print(RED+"Error erasing "+disk+WHITE)
+    print(GREEN+"Done"+WHITE)
+
+def erase_all(disk):
+    for d in disk:
+        erase(d)
+
+# Main
+if __name__ == "__main__":
+    main_menu()