#!/usr/bin/python

import os, signal, sys
from subprocess import check_output, call, DEVNULL, STDOUT

signal.signal(signal.SIGINT, lambda x,y: sys.exit(0))

def main():

    # listing iso files
    f_dir = '$(pwd)'
    iso_list = check_output('find %s -name *.iso' % f_dir, shell=True) \
            .decode('utf-8').split('\n')[:-1]

    iso_list.sort(key=os.path.getmtime, reverse=True)

    print('List of your ISO files:')

    if len(iso_list) is 0:
        print('\tIt seems you don\'t have any ISO files!')
        exit()

    for i, iso in enumerate(iso_list):
        print('\t(%d) %s' % (i, iso))

    selected_iso = input('Select your ISO (default: 0) ')
    if selected_iso == '':
        selected_iso = iso_list[0]
    elif selected_iso.isdigit() and int(selected_iso) < len(iso_list):
        selected_iso = iso_list[int(selected_iso)]
    else:
        print('This is not a valid number!\n')
        exit()

    # listing drives
    drive_list = check_output('ls -l /dev/disk/by-id/', shell=True) \
            .decode('utf-8').split('\n')[:-1]

    labels = []
    print('\nList of your USB drives:')
    i = 0
    has_usb = False
    for d in drive_list:
        if not d[-1].isdigit() and 'usb' in d:
            name = ' '.join(d.split('usb-')[1].split('_')[:-1])
            label = d.split('/')[-1]
            labels.append(label)
            size = int(check_output('cat /sys/class/block/%s/size'
                    % label, shell=True)) / 2 / 1024 / 1024
            print('\t(%d) %s - %.2f GB' % (i, name, size))
            i += 1
            has_usb = True

    if not has_usb:
        print('\tIt seems you don\'t have any USB drive connected!')
        exit()

    # drive selection
    selected_drive = input('Select your USB drive! (default: 0) ')

    if selected_drive.isdigit() and int(selected_drive) < len(labels):
        selected_drive = '/dev/' + labels[int(selected_drive)]
    elif selected_drive == '':
        selected_drive = '/dev/' + labels[0]
    else:
        print('This is not a valid number!\n')
        exit()


    sure = True if input('This will erease %s Are you sure? (y/n) '
            % selected_drive) == 'y' else False

    if sure:
        is_mounted = True if call('mount | grep %s' % selected_drive, shell=True,
                stdout=DEVNULL, stderr=STDOUT) == 0 else False
        if is_mounted:
            mount_part = check_output('mount | grep %s' % selected_drive,
                    shell=True).decode('utf-8').split(' on')[0]
            call('umount %s' % mount_part, shell=True)
            print("drive unmounted!")

        call('sudo dd if=%s of=%s status=progress bs=4M'
                % (selected_iso, selected_drive), shell=True)
        call('sync', shell=True)
        print('Done! Now you can remove your drive!')
    else:
        exit()

if __name__ == "__main__":
    main()
