diskdump.py

#!/usr/bin/env python3

import subprocess
import json
import os
from datetime import datetime
import sys
import re


def sanitize_name(name):
    """Convert disk model to a safe filename."""
    if not name:
        return "diskimage"
    name = name.strip().lower()
    name = re.sub(r'\s+', '_', name)
    name = re.sub(r'[^a-z0-9._-]', '', name)
    return name


def get_disks():
    result = subprocess.run(
        ["lsblk", "-J", "-o", "NAME,SIZE,TYPE,MODEL"],
        capture_output=True,
        text=True,
        check=True
    )

    data = json.loads(result.stdout)

    disks = []
    for device in data["blockdevices"]:
        if device["type"] == "disk":
            disks.append(device)

    return disks


def choose_disk(disks):
    print("\nDetected Drives:\n")

    for i, d in enumerate(disks):
        model = (d.get("model") or "Unknown").strip()
        print(f"{i+1}) /dev/{d['name']} | {d['size']} | {model}")

    while True:
        try:
            choice = int(input("\nChoose a disk number: "))
            if 1 <= choice <= len(disks):
                return disks[choice-1]
        except ValueError:
            pass

        print("Invalid selection.")


def run_dd(source, output_file):
    print("\nStarting backup...\n")

    cmd = [
        "sudo",
        "dd",
        f"if={source}",
        f"of={output_file}",
        "bs=4M",
        "status=progress",
        "conv=fsync"
    ]

    process = subprocess.run(cmd)

    if process.returncode != 0:
        print("\nBackup failed.")
        sys.exit(1)


def main():

    disks = get_disks()
    disk = choose_disk(disks)

    model = sanitize_name(disk.get("model"))
    timestamp = datetime.now().strftime("%d-%m-%Y_%Hh%M")

    directory = input("\nEnter destination directory: ").strip()
    if not os.path.isdir(directory):
        print("Directory does not exist.")
        return

    fmt = input("\nChoose image format (.img or .iso): ").strip().lower()
    if fmt not in [".img", ".iso"]:
        print("Invalid format.")
        return

    filename = f"{model}_{timestamp}{fmt}"
    output_file = os.path.join(directory, filename)

    source = f"/dev/{disk['name']}"

    print(f"\nSource: {source}")
    print(f"Destination: {output_file}")

    confirm = input("\nProceed with backup? (yes/no): ").lower()
    if confirm not in ["yes", "y"]:
        print("Cancelled.")
        return

    run_dd(source, output_file)

    print("\nBackup completed successfully!")
    print(f"Image saved to: {output_file}")


if __name__ == "__main__":
    main()