#!/usr/bin/env python3

"""\
wipedisk - Wipe the header of a block device
Copyright (c) 2023 - Bart Sjerps <bart@dirty-cache.com>
URL: https://github.com/bsjerps/asmdisks
License: GPLv3+
"""

import os, sys, stat, argparse
from time import sleep
from subprocess import run, CalledProcessError

desc = """\
Wipes the header of a disk device with zeroes (by default, the first 1MiB).
The force (-f) flag is required for safety. A backup of the wiped section is made in /tmp.
A period of 5 seconds is waited before the actual wipe (to allow CTRL-C abort)
"""

def execute(cmd):
    try:
        print(f"Running {cmd}")
        r = run(cmd.split(), check=True)
    except CalledProcessError as e:
        sys.exit(10)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=desc, formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=40))
    parser.add_argument("-f", "--force", action='store_true', help='Required for safety')
    parser.add_argument("-n", "--dryrun", action='store_true', help='Skip actual wipe')
    parser.add_argument("-w", "--nowait", action='store_true', help='Skip the delay')
    parser.add_argument("-s", "--size", type=int, metavar='size', default=1, help='Overwrite <size> mib of the disk')
    parser.add_argument("dev", metavar='device')

    args = parser.parse_args()
    dev  = args.dev
    name = args.dev.replace('/','_')
    st   = os.stat(dev)

    if not stat.S_ISBLK(st.st_mode):
        print(f"{dev} is not a block device")
        sys.exit()

    if not args.force:
        print("Force flag (-f) required!")
        sys.exit()
    
    if not args.nowait:
        print("Waiting 5 seconds...")
        sleep(5)
    
    backupfile = f'/tmp/backup_{name}'

    print(f"Backing up to {backupfile}")
    execute(f"/usr/bin/dd if={dev} of={backupfile} bs=1M count={args.size}")

    if not args.dryrun:
        execute(f"/usr/bin/dd if=/dev/zero of={dev} bs=1M count={args.size}")
