burn_to_usb.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import os
import sys


parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter
)

parser.add_argument(
    '--image-file',
    '-i',
    type=str,
    help='input image file'
)

parser.add_argument(
    '--target-device',
    '-d',
    type=str,
    help='target USB device',
    default='/dev/sdd'
)

parser.add_argument(
    '--whitelist',
    '-w',
    type=str,
    help='csv whitelist in device names',
    default='USB,DataTraveler,Patriot,usb'
)

parser.add_argument(
    '--test',
    '-t',
    help='test mode: do not do anything',
    action="store_true"
)


if len(sys.argv) == 1:
    parser.print_help()
    sys.exit(1)

args = parser.parse_args()

dev = args.target_device
image_file = args.image_file
whitelist = '\|'.join(args.whitelist.split(','))

is_usb = os.system('find /dev/disk/by-id/ -lname *%s | grep "%s"' % (os.path.basename(args.target_device), whitelist)) == 0

if is_usb:
    print('Device %s is a USB device, proceeding' % dev)
else:
    sys.exit('Device %s is not an USB device or is not found, adjust whitelist %s' % (dev, args.whitelist))

cmd = ''
ext = os.path.splitext(image_file)[1]
if image_file.endswith('.tar.bz2'):
    cmd += 'tar xj'
elif image_file.endswith('.tar.gz'):
    cmd += 'tar xz'
elif image_file.endswith('.tar'):
    cmd += 'tar x'
elif image_file.endswith('.xz'):
    cmd += 'xzcat'
elif ext in ['.raw', '.hddimg', '.iso']:
    pass
else:
    print('Unsupported extension %s' % ext)
    sys.exit(1)

if image_file.startswith('http'):
    cmd = "curl %s | %s" % (image_file, cmd)
else:
    if ext in ['.raw', '.hddimg', '.iso']:
        cmd += 'cat'
    cmd += " %s |" % image_file

cmd += ' sudo dd of=%s bs=4M conv=fsync status=progress' % dev
print(cmd)

if not args.test:
    os.system(cmd)
    print('Syncing')
    os.system('sync')
    print('All done')
else:
    print('Test mode, exiting without doing anything')