Python script to utilise ATA security with SATA SSDs

disk_utility.py

#!/usr/bin/python3

import argparse
import getpass
import os
import subprocess
import time
import typing

# Documentation:
# https://ata.wiki.kernel.org/index.php/ATA_Secure_Erase#Step_2_-_Enable_security_by_setting_a_user_password:
# https://ata.wiki.kernel.org/index.php/ATA_Secure_Erase
# http://chrisarges.net/2018/02/16/using-sed-encryption-on-disks.html
# https://www.zeitgeist.se/2014/09/07/enabling-ata-security-on-a-self-encrypting-ssd/

default_device = "/dev/sda"
default_partition = "/dev/sda1"
default_mount_path = "/media/data"

class DiskUtility:
    def __init__(self, device: str, partition: str):
        self.__device: str = device
        self.__partition: str = partition

    def configure(self, password: str):
        print("Setting user password...")
        args = ["hdparm", "--security-mode", "m", "--user-master", "u", "--security-set-pass", password, self.__device]
        subprocess.run(args, check=True)

    def unlock(self, password: str):
        print("Unlocking device...")
        args = ["hdparm", "--user-master", "u", "--security-unlock", password, self.__device]
        subprocess.run(args, check=True)

    def disable(self, password: str):
        print("Disabling device encryption on {0}...".format(self.__device))
        args = ["hdparm", "--user-master", "u", "--security-disable", password, self.__device]
        subprocess.run(args, check=True)

    def mount(self, path: str):
        print("Calling parted to refresh partitions in /dev/*...")
        os.system("parted -s -l 1> /dev//zero")

        # this step might not be necessary anymore
        print("Waiting for partition to become available...")
        while not os.path.exists(self.__partition):
            print("...")
            time.sleep(1)

        print("Mounting {0} in {1}...".format(self.__partition, path))
        args = ["mount", "-t", "ext4,rw", self.__partition, path]
        subprocess.run(args, check=True)

    def info(self):
        args = ["hdparm", "-I", self.__device]
        subprocess.run(args, check=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to simplify handling disk encryption.")
    parser.add_argument("-d", "--device", help="Device path: e.g. /dev/sda")
    parser.add_argument("-p", "--partition", help="Partition: e.g. /dev/sda1")
    parser.add_argument("-c", "--configure", action="store_true", help="mode to configure disk encryption")
    parser.add_argument("-m", "--mount", action="store_true", help="mode to mount the disk")
    parser.add_argument("-i", "--info", action="store_true", help="info")
    parser.add_argument("-r", "--remove_encryption", action="store_true", help="mode to disable encryption")
    parser.add_argument("-mp", "--mount_path", help="mount path")
    args = parser.parse_args()

    password = getpass.getpass(prompt="Type in your encryption password: ")
    password_confirmation = getpass.getpass(prompt="Confirm the encryption password: ")
    if password != password_confirmation:
        raise RuntimeError("Given passwords were not equal.")

    device = args.device
    partition = args.partition
    mount_path = args.mount_path

    if not device or len(device) == 0:
        device = default_device
    if not mount_path or len(mount_path) == 0:
        mount_path = default_mount_path
    if not partition or len(partition) == 0:
        partition = default_partition

    print("Selected disk: {0}".format(device))
    disk_util = DiskUtility(device, partition)

    if args.configure:
        disk_util.configure(password)
    elif args.info:
        disk_util.info()
    elif args.mount:
        disk_util.unlock(password)
        disk_util.mount(mount_path)
    elif args.remove_encryption:
        disk_util.disable(password)