#!/usr/bin/python3
import errno
import os
import sys
from sys import stdout
from sys import argv
import subprocess

##### global data #######
drv_lst = []  #list of removable "sdx" block devices found in /sys/block
partition_list0 = []
partition_list1 = []
partition_list2 = []
partition_list3 = []
partition_list4 = []
partition_list5 = []
partition_list6 = []
mounts = []  # list of mounts, indexed by drv_lst
vendors = []
models = []
drive_sizes = []

def have_exe(name):
    try:
        cp = subprocess.run(name, universal_newlines=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except OSError as e:
        if e.errno == os.errno.ENOENT:
            return False
    return(True)

def quote_shell(args):
    return " ".join(quote(arg) for arg in args)


def quote_applescript(string):
    charmap = {
        "\n": "\\n",
        "\r": "\\r",
        "\t": "\\t",
        "\"": "\\\"",
        "\\": "\\\\",
    }
    return '"%s"' % "".join(charmap.get(char, char) for char in string)


def elevate(show_console=True, graphical=False):
    if os.getuid() == 0:
        return
    print("format_fat32 needs to run as root, will ask for password")
    args = [sys.executable] + sys.argv
    commands = []


    commands.append(["sudo"] + args)

    for args in commands:
        try:
            os.execlp(args[0], *args)
        except OSError as e:
            if e.errno != errno.ENOENT or args[0] == "sudo":
                raise


def is_root():
    return os.getuid() == 0

#returns list
#example= "/proc/mounts"
#if error returns len() = 0
def file_read_to_list(fname):
    lst=[]
    try:
        f = open(fname,"r")
    except IOError:
        print("failed to open",fname)
        return('')
    for l in f:
        lst.append(l.strip())
    f.close()
    return(lst)

#do unix command, returns 0 if error
#example cmd = "ls -d /sys/block/*"  returns full path of dir entries
#example cmd = "ls -w1 /sys/block"  returns single dir/file per list entry
#if error returns len() = 0
def do_cmd(cmd):
    lst = []
    try:
        f = os.popen(cmd)
    except IOError:
        print("dir_read failed")
        return('')
    for l in f.readlines():
        lst.append(l.strip())
    return(lst)

def find_drives():
    global drv_lst
    drv_lst = []
    lst = do_cmd("ls -w1 /sys/block")
    if len(lst) == 0:
        print("error reading /sys/block")
        return()
    print("block devices found=",lst)
    for blk in lst:
        if "sd" in blk:
#            print("found-",blk)
            pathx = "/sys/block/"+blk+"/removable"
            rflag = file_read_to_list(pathx)
#            print("read ",pathx,"got-",rflag)
            if rflag[0] == '1':
#                print("found drive")
                drv_lst.append(blk)
    return(drv_lst) 

#find partitions for drive at indx in drv_lst 
def find_partitions_one_drive(indx):
    plist=[]
    cmd = "ls -w1 /sys/block/"+drv_lst[indx]
    lst = do_cmd(cmd)
    if len(lst) == 0:
        print("error reading ",cmd)
        return()
    for partit in lst:
        if drv_lst[indx] in partit:
            print("found partition ",partit)
            plist.append(partit)
    return(plist)

def find_all_partitions():
    global partition_list0,partition_list1,partition_list2
    indx = 0
    for drivex in drv_lst:
        x = find_partitions_one_drive(indx)
        globals()["partition_list"+str(indx)] = x
#        print("find_all_partitions found",x)
        indx = indx + 1

def get_mounts():
    global mounts
    added = 0
    for xdrive in drv_lst:
        m = file_read_to_list("/proc/mounts")
#        print("len(m)=",len(m))
#        print("m=",m)
        for line in m:
            if line.find("/dev/"+xdrive) != -1:
#                print("adding line-",line)
                #x = line.split()
                mounts.append(line)
#                print("len(line)=",len(line),"len(mounts)=",len(mounts))
                added = 1
                break
        if added == 0:
            mounts.append("0")
        added = 0

def get_drv_model():
    for xdrive in drv_lst:
        m = file_read_to_list("/sys/block/"+xdrive+"/device/vendor")
        vendors.append(m)
        m = file_read_to_list("/sys/block/"+xdrive+"/device/model")
        models.append(m)

def get_sizes():
    for xdrive in drv_lst:
        m = file_read_to_list("/sys/block/"+xdrive+"/size")
        x = int(m[0])
        xx = x/2000000
        xx = round(xx,0)
        if xx < 1:
            z = str(int(xx))+"Megs"
        else:
            z = str(int(xx))+"Gigs" 
        drive_sizes.append(z)

def get_key():
    import tty, termios
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setraw(sys.stdin.fileno())
        ch = sys.stdin.read(1)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    return(ch)

def sfdisk_part(y):
    print("using sfdisk")
    cmd = "sfdisk --delete /dev/"+str(drv_lst[y])
    do_cmd(cmd)
    cmd = "echo ',,c;' | sfdisk --force /dev/"+str(drv_lst[y])
    do_cmd(cmd)
    cmd = "mkfs.exfat -n fat32 /dev/"+str(drv_lst[y])+"1"
    do_cmd(cmd)
    print("drive formatted")   

def parted_part(y):
    print("using parted")
    indx = 1
    for partit in globals()["partition_list"+str(y)]:
        print("/dev/"+str(drv_lst[y]))
        cmd = "parted -s /dev/"+str(drv_lst[y])+" rm "+str(indx)
        print(cmd)
        do_cmd(cmd)
        indx = indx + 1
    cmd = "parted -s /dev/"+str(drv_lst[y])+" mklabel msdos mkpart primary fat32 0 100%"
    print("cmd=",cmd)
    do_cmd(cmd)
    cmd = "mkfs.exfat -n fat32 /dev/"+str(drv_lst[y])+"1"
    do_cmd(cmd)
    print("drive formatted")   

############# start of execution #################
xmessage = have_exe("xmessage")
notify = have_exe("notify-send")

if stdout.isatty() == False:
#    print('The output device is a teletype. Or something like a teletype.')
#else:
    if xmessage == True:
        os.popen("xmessage -center format-fat32 needs a terminal")
    elif notify == True:
        os.popen("notify-send '**whoops** format-fat32 needs a terminal'")
    sys.exit()


try:
    from shlex import quote
except ImportError:
    from pipes import quote

elevate()

sfdisk = have_exe(["sfdisk","-V"])
parted= have_exe(["parted","-v"])
mkfs = have_exe(["mkfs.fat","--help"])
#print("xmessage=",xmessage,"notify=",notify,"sfdisk=",sfdisk,"parted=",parted,"mkfs=",mkfs)

x = find_drives()
#print("drv_lst=",drv_lst)
find_all_partitions()
#print("partition_list0=",partition_list0)
#print("partition_list1=",partition_list1)

get_mounts()
#print("len(mounts)=",len(mounts))
#print("mounts=",mounts)

get_drv_model()
#print("vendors=",vendors)
#print("models=",models)

get_sizes()
#print("drive_sizes=",drive_sizes)

do_exit = 0
if sfdisk == False and parted == False:
    print("**** sfdisk or parted programs needed to continue *****")
    do_exit = 1
if mkfs == False:
    print("**** the mkfs program is needed to continue *****")
    do_exit = 1
if do_exit == 1:
    sys.exit()

if len(drv_lst) == 0:
    print("No removable drives found\nPress any key to continue")
    get_key()
    sys.exit()

menu = "Enter number of option\n 0 = exit\n"
indx = 1
x=""
for xdrive in drv_lst:
    y = str(vendors[indx -1])
    z = str(models[indx -1])
    w = str(drive_sizes[indx-1])

    x = " "+str(indx) +" format "+xdrive+y+z+" size-"+w+"\n"
    menu = menu + x
    indx = indx + 1
print(menu)
x = get_key()

if x == "0":
    sys.exit()
y = int(x)
print("unmounting all partition for drive "+str(drv_lst[y-1])+" "+str(globals()["partition_list"+str(y-1)]))
for partit in globals()["partition_list"+str(y-1)]:
    print("/dev/"+str(partit))
    do_cmd("umount /dev/"+str(partit))

if sfdisk == True:
#    sfdisk_part(y-1)
#elif partit == True:
    parted_part(y-1)

lst = do_cmd("fdisk -l /dev/"+str(drv_lst[y-1]))
print(lst)