Simple dmenu pmount/pumount helper for removable block devices

dmenu_mount.py

#!/usr/bin/env python3
# dmenu removable devie mount/umount helper
import subprocess as subp, glob, sys, os, os.path as p

def error(msg):
    sys.stderr.write('ERROR: %s\n' % msg)
    subp.call(['i3-nagbar', '-t', 'error', '-m', 'ERROR: ' + msg, '-f', 'Monospace'])
    sys.exit(1)

def get_removable_devices():
    removables = []
    for i in glob.glob('/sys/class/block/*/removable'):
            with open(i, 'rb') as f:
                if f.read(1) == b'1':
                    removables.append(p.join(p.sep + 'dev', p.basename(p.dirname(i))))
    return removables

def get_usb_partitions():
    partitions = []
    for i in glob.glob('/dev/disk/by-path/*-usb-*-part[0-9]'):
        partitions.append(p.abspath(p.join(p.dirname(i), os.readlink(i))))
    return partitions

def get_device_partitions(d):
    partitions = []
    for i in glob.glob(d + '*'):
        if i != d:
            partitions.append(i)
    return partitions

def get_mounted_removables(p):
    ms = {}
    with open('/proc/mounts', 'rb') as f:
        for d in f.read().decode('utf-8').split('\n'):
            if len(d) == 0:
                continue
            x=d.split(' ', 2)
            if x[0] in p:
                ms[x[0]] = x[1]
    return ms

def dmenu(options, dmenu):
    cmd = subp.Popen(dmenu, stdin=subp.PIPE, stdout=subp.PIPE)
    return cmd.communicate('\n'.join(options).encode('utf-8'))[0].decode('utf-8').strip('\n')

MOUNTED_AT = ' mounted at '
def run(cmd):
    if cmd is None or len(cmd) == 0:
        return
    cmd = 'p' + cmd # yeah, just prepend the p for pmount/pumount here
    if cmd.find(MOUNTED_AT):
        cmd = cmd.split(MOUNTED_AT)[0]
    subp.call(cmd.split(' '))

def is_in_path(x):
    for d in os.getenv('PATH').split(':'):
        if p.isfile(p.join(d, x)) and os.access(p.join(d, x), os.R_OK | os.X_OK):
            return True
    return False

def main():
    # override some options
    dmenu_cmd = ['dmenu'] + sys.argv[1:] + ['-l', '5', '-p', 'removable devices']

    if not (is_in_path('dmenu') and is_in_path('i3-nagbar') and is_in_path('pmount') and
            is_in_path('pumount')):
        error("Need dmenu, i3-nagbar ;), pmount and pumount in $PATH")

    try:
        p = []
        for d in get_removable_devices():
            ps = get_device_partitions(d)
            if len(ps) > 0:
                p += ps
        p += get_usb_partitions()
        p = list(set(p))
    except:
        error("Could not enumerate removable media/partitions")

    try:
        mounted = get_mounted_removables(set(p))
    except:
        error("Could not get mounted removable media");

    items = ['umount ' + d + MOUNTED_AT + mounted[d] if d in mounted else 'mount ' + d for d in p]

    # mounted first, then sort according to name
    items.sort(key=lambda x: ('a' if x[0] == 'u' else 'b') + x.split(' ')[1])

    try:
        run(dmenu(items, dmenu_cmd))
    except:
        error("Could not spawn dmenu or the actual mount command")

__name__ == '__main__' and main()