forked from mirrors/ranger_udisk_menu
e0db765f81
The JSON output from lsblk seems to have some issues that interfere with the `json.loads()` function. Specifically, the value 0B, which is expected to be a string. By doing a direct replace on the lsblk output for `0B` on `menu.py`, the JSON is parsed and the menu appears.
238 lines
7.8 KiB
Python
238 lines
7.8 KiB
Python
#!/usr/bin/python3
|
|
# coding: utf-8
|
|
# License: The MIT License
|
|
# Author: Alexander Lutsai <sl_ru@live.com>
|
|
# Year: 2021
|
|
# Description: This script draws menu to choose, mount and unmount drives
|
|
|
|
import curses
|
|
import curses.ascii
|
|
import subprocess
|
|
import json
|
|
import sys
|
|
|
|
|
|
class ChoosePartition:
|
|
blkinfo = None
|
|
screen = None
|
|
selected_partn = 1
|
|
selected_mountpoint = None
|
|
partn = 1
|
|
help_message = ["Press 'm' to mount, 'u' to unmount, 'g' to refresh",
|
|
" and 'e' to unmount, 'p' to poweroff drive, 'enter' to cd"]
|
|
message = ""
|
|
|
|
def __init__(self):
|
|
self.screen = curses.initscr()
|
|
curses.start_color()
|
|
curses.curs_set(0)
|
|
curses.noecho()
|
|
curses.cbreak()
|
|
self.screen.keypad(True)
|
|
curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_WHITE)
|
|
self.selected_partn = 1
|
|
self._read_partitions()
|
|
|
|
def _read_partitions(self):
|
|
r = subprocess.check_output(['lsblk', '--all', '--json', '-O'])
|
|
r = r.decode().replace('0B,', '\"0B\",')
|
|
self.blkinfo = json.loads(r.encode())
|
|
partn = 0
|
|
for bd in self.blkinfo['blockdevices']:
|
|
if 'children' not in bd:
|
|
continue
|
|
for part in bd['children']:
|
|
partn += 1
|
|
|
|
self.partn = partn
|
|
if self.selected_partn > self.partn:
|
|
self.selected_partn = self.partn
|
|
if self.selected_partn <= 0:
|
|
self.selected_partn = 1
|
|
|
|
def _get_part_by_partn(self):
|
|
partn = 0
|
|
for bd in self.blkinfo['blockdevices']:
|
|
if 'children' not in bd:
|
|
continue
|
|
for part in bd['children']:
|
|
partn += 1
|
|
if self.selected_partn == partn:
|
|
return part
|
|
return None
|
|
|
|
def _get_drive_by_partn(self):
|
|
partn = 0
|
|
for bd in self.blkinfo['blockdevices']:
|
|
if 'children' not in bd:
|
|
continue
|
|
for part in bd['children']:
|
|
partn += 1
|
|
if self.selected_partn == partn:
|
|
return bd
|
|
return None
|
|
|
|
def _select_print_part(self, part, is_selected, i):
|
|
if not ('mountpoint' in part and
|
|
'name' in part and
|
|
'size' in part):
|
|
raise Exception('Wrong lsblk json format.' +
|
|
'No mountpoint, name or size in the partition')
|
|
label = ""
|
|
label_fields = ['label', 'partlabel', 'parttypename', 'fstype']
|
|
for f in label_fields:
|
|
if f not in part:
|
|
continue
|
|
if part[f] is not None:
|
|
label = part[f]
|
|
break
|
|
|
|
mp = None
|
|
if part['mountpoint'] is not None:
|
|
mp = part['mountpoint']
|
|
if is_selected:
|
|
self.selected_mountpoint = mp
|
|
if mp is None:
|
|
mp = "Not mounted"
|
|
|
|
s = "{name:<12} {size:<8} {label:<16} {mp}".format(
|
|
name=part['name'] if part['name'] is not None else "None",
|
|
label=label if label is not None else "None",
|
|
size=part['size'] if part['size'] is not None else "None",
|
|
mp=mp
|
|
)
|
|
self.screen.addstr(2 + i, 4, s, curses.color_pair(is_selected))
|
|
|
|
def _select_print_block_device(self, bd, i):
|
|
if not ('model' in bd and
|
|
'size' in bd or
|
|
'name' in bd):
|
|
raise Exception('Wrong lsblk json format. ' +
|
|
'No model, size or name in blockdevice')
|
|
|
|
model = bd['model'] if bd['model'] is not None else ""
|
|
size = bd['size'] if bd['size'] is not None else ""
|
|
self.screen.addstr(2 + i, 2, bd['name'] + " " + model + " " + size)
|
|
|
|
def _select_print(self, x):
|
|
self.screen.clear()
|
|
self.screen.border(0)
|
|
self.screen.addstr(1, 2, self.help_message[0])
|
|
self.screen.addstr(2, 2, self.help_message[1])
|
|
|
|
partn = 0
|
|
i = 0
|
|
if 'blockdevices' not in self.blkinfo:
|
|
raise Exception('Wrong lsblk json format. No field "blockdevices"')
|
|
for bd in self.blkinfo['blockdevices']:
|
|
i += 1
|
|
bd_selected = False
|
|
bd_i = i
|
|
self._select_print_block_device(bd, bd_i)
|
|
if 'children' not in bd:
|
|
continue
|
|
for part in bd['children']:
|
|
i += 1
|
|
partn += 1
|
|
is_selected = 0 if self.selected_partn != partn else 1
|
|
if is_selected:
|
|
bd_selected = True
|
|
self._select_print_part(part, is_selected, i)
|
|
if bd_selected:
|
|
self.screen.addstr(2 + bd_i, 1, ">")
|
|
self.screen.addstr(2 + i + 2, 4, self.message)
|
|
|
|
def _eject_all(self):
|
|
blk = None
|
|
partn = 0
|
|
for bd in self.blkinfo['blockdevices']:
|
|
if 'children' not in bd:
|
|
continue
|
|
for part in bd['children']:
|
|
partn += 1
|
|
if self.selected_partn == partn:
|
|
blk = bd
|
|
if blk is None:
|
|
return
|
|
for part in blk['children']:
|
|
self.unmount(part)
|
|
|
|
def select(self):
|
|
sel = None
|
|
x = 0
|
|
# quit when pressed `q` or `Esc` or `Ctrl+g`
|
|
while x not in (ord('q'), curses.ascii.ESC,
|
|
curses.ascii.BEL, curses.ascii.NL):
|
|
self._select_print(x)
|
|
x = self.screen.getch()
|
|
if x in (ord('j'), curses.ascii.SO, curses.KEY_DOWN):
|
|
# down
|
|
self.selected_partn += 1
|
|
if self.selected_partn > self.partn:
|
|
self.selected_partn = self.partn
|
|
elif x in (ord('k'), curses.ascii.DLE, curses.KEY_UP):
|
|
# up
|
|
self.selected_partn -= 1
|
|
if self.selected_partn <= 0:
|
|
self.selected_partn = 1
|
|
elif x == ord('e'):
|
|
sel = self._eject_all()
|
|
elif x == ord('m'):
|
|
sel = self._get_part_by_partn()
|
|
if sel is not None:
|
|
self.mount(sel)
|
|
elif x == ord('u'):
|
|
sel = self._get_part_by_partn()
|
|
if sel is not None:
|
|
self.unmount(sel)
|
|
elif x == ord('p'):
|
|
sel_drive = self._get_drive_by_partn()
|
|
if sel_drive is not None:
|
|
self.poweroff(sel_drive)
|
|
elif x == ord('g') or x == ord('r'):
|
|
self._read_partitions()
|
|
curses.endwin()
|
|
if self.selected_mountpoint is not None and x == curses.ascii.NL:
|
|
return self.selected_mountpoint
|
|
return ""
|
|
|
|
def _udisk_mount_unmount(self, cmd, dev):
|
|
r = ""
|
|
try:
|
|
r = subprocess.run(
|
|
['udisksctl', cmd, '-b', dev], capture_output=True)
|
|
r = (r.stdout.decode(encoding="utf-8") +
|
|
r.stderr.decode(encoding="utf-8"))
|
|
self.message = r
|
|
except Exception as e:
|
|
self.message = cmd + " error: " + r + str(e)
|
|
self._read_partitions()
|
|
|
|
def get_drive_path(self, drive):
|
|
if 'path' not in drive:
|
|
drive['path'] = '/dev/' + drive['kname']
|
|
return drive['path']
|
|
|
|
def unmount(self, dev):
|
|
p = self.get_drive_path(dev)
|
|
self._udisk_mount_unmount("unmount", p)
|
|
|
|
def poweroff(self, dev):
|
|
p = self.get_drive_path(dev)
|
|
self._udisk_mount_unmount("power-off", p)
|
|
|
|
def mount(self, dev):
|
|
p = self.get_drive_path(dev)
|
|
self._udisk_mount_unmount("mount", p)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cp = ChoosePartition()
|
|
sel = cp.select()
|
|
# print(sel)
|
|
# print(len(sys.argv))
|
|
if len(sys.argv) >= 2:
|
|
# print(sys.argv[1])
|
|
with open(sys.argv[1], 'w') as f:
|
|
f.write(sel)
|