You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
557 lines
21 KiB
Python
557 lines
21 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
# Name: usb.py
|
|
# Purpose: Module to list USB devices and get details under Linux and Windows
|
|
# Authors: Sundar
|
|
# Licence: This file is a part of multibootusb package. You can redistribute it or modify
|
|
# under the terms of GNU General Public License, v.2 or above
|
|
|
|
import collections
|
|
import ctypes
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
if platform.system()=='Linux':
|
|
import dbus
|
|
|
|
from . import config
|
|
from . import gen
|
|
from . import osdriver
|
|
|
|
if platform.system() == 'Linux':
|
|
from . import udisks
|
|
UDISKS = udisks.get_udisks(ver=None)
|
|
if platform.system() == 'Windows':
|
|
import psutil
|
|
import win32com.client
|
|
# import wmi
|
|
import pythoncom
|
|
|
|
|
|
def is_block(usb_disk):
|
|
"""
|
|
Function to detect if the USB is block device
|
|
:param usb_disk: USB disk path
|
|
:return: True is devie is block device else False
|
|
"""
|
|
import stat
|
|
if platform.system() == 'Linux':
|
|
if len(usb_disk) != 9:
|
|
return False
|
|
elif platform.system() == 'Windows':
|
|
if len(usb_disk) != 2:
|
|
return False
|
|
else:
|
|
return True
|
|
try:
|
|
mode = os.stat(usb_disk).st_mode
|
|
gen.log(mode)
|
|
gen.log(stat.S_ISBLK(mode))
|
|
except:
|
|
return False
|
|
|
|
return stat.S_ISBLK(mode)
|
|
|
|
|
|
def disk_usage(mount_path):
|
|
"""
|
|
Return disk usage statistics about the given path as a (total, used, free)
|
|
namedtuple. Values are expressed in bytes.
|
|
"""
|
|
# Author: Giampaolo Rodola' <g.rodola [AT] gmail [DOT] com>
|
|
# License: MIT
|
|
_ntuple_diskusage = collections.namedtuple('usage', 'total used free')
|
|
|
|
if platform.system() == "Linux":
|
|
st = os.statvfs(mount_path)
|
|
free = st.f_bavail * st.f_frsize
|
|
total = st.f_blocks * st.f_frsize
|
|
used = (st.f_blocks - st.f_bfree) * st.f_frsize
|
|
|
|
return _ntuple_diskusage(total, used, free)
|
|
|
|
elif platform.system() == "Windows":
|
|
|
|
_, total, free = ctypes.c_ulonglong(), ctypes.c_ulonglong(), \
|
|
ctypes.c_ulonglong()
|
|
if sys.version_info >= (3,) or isinstance(mount_path, unicode):
|
|
fun = ctypes.windll.kernel32.GetDiskFreeSpaceExW
|
|
else:
|
|
fun = ctypes.windll.kernel32.GetDiskFreeSpaceExA
|
|
ret = fun(mount_path, ctypes.byref(_), ctypes.byref(total), ctypes.byref(free))
|
|
if ret == 0:
|
|
raise ctypes.WinError()
|
|
used = total.value - free.value
|
|
|
|
return _ntuple_diskusage(total.value, used, free.value)
|
|
else:
|
|
raise NotImplementedError("Platform not supported.")
|
|
|
|
|
|
def list_devices(fixed=False):
|
|
"""
|
|
List inserted USB devices.
|
|
:return: USB devices as list.
|
|
"""
|
|
devices = []
|
|
if platform.system() == "Linux":
|
|
try:
|
|
# pyudev is good enough to detect USB devices on modern Linux machines.
|
|
gen.log("Using pyudev for detecting USB drives...")
|
|
try:
|
|
import pyudev
|
|
except Exception as e:
|
|
gen.log('PyUdev is not installed on host system, using built-in.')
|
|
from . import pyudev
|
|
context = pyudev.Context()
|
|
|
|
for device in context.list_devices(subsystem='block'):
|
|
if fixed is True:
|
|
if device.get('DEVTYPE') in ['disk', 'partition'] and device.get('ID_PART_TABLE_TYPE'):
|
|
devices.append(str(device.get('DEVNAME')))
|
|
gen.log("\t" + device.get('DEVNAME'))
|
|
else:
|
|
if device.get('ID_BUS') in ['usb'] and device.get('ID_PART_TABLE_TYPE'):
|
|
devices.append(str(device.get('DEVNAME')))
|
|
gen.log("\t" + device.get('DEVNAME'))
|
|
|
|
except Exception as e:
|
|
gen.log(e)
|
|
bus = dbus.SystemBus()
|
|
try:
|
|
# You should come here only if your system does'nt have udev installed.
|
|
# We will use udiskd2 for now.
|
|
gen.log("Falling back to Udisks2..")
|
|
ud_manager_obj = bus.get_object(
|
|
'org.freedesktop.UDisks2', '/org/freedesktop/UDisks2')
|
|
ud_manager = dbus.Interface(
|
|
ud_manager_obj, 'org.freedesktop.DBus.ObjectManager')
|
|
for k, v in ud_manager.GetManagedObjects().items():
|
|
drive_info = v.get('org.freedesktop.UDisks2.Block', {})
|
|
if fixed is True:
|
|
if drive_info.get('IdUsage') == "filesystem" and not drive_info.get('ReadOnly'):
|
|
device = drive_info.get('Device')
|
|
device = bytearray(device).replace(b'\x00', b'').decode('utf-8')
|
|
devices.append(device)
|
|
else:
|
|
if drive_info.get('IdUsage') == "filesystem" and not drive_info.get(
|
|
'HintSystem') and not drive_info.get('ReadOnly'):
|
|
device = drive_info.get('Device')
|
|
device = bytearray(device).replace(
|
|
b'\x00', b'').decode('utf-8')
|
|
devices.append(device)
|
|
|
|
except Exception as e:
|
|
gen.log(e, error=True)
|
|
try:
|
|
# You must be using really old distro. Otherwise, the code
|
|
# should not reach here.
|
|
gen.log("Falling back to Udisks1...")
|
|
ud_manager_obj = bus.get_object(
|
|
"org.freedesktop.UDisks", "/org/freedesktop/UDisks")
|
|
ud_manager = dbus.Interface(
|
|
ud_manager_obj, 'org.freedesktop.UDisks')
|
|
for dev in ud_manager.EnumerateDevices():
|
|
device_obj = bus.get_object(
|
|
"org.freedesktop.UDisks", dev)
|
|
device_props = dbus.Interface(
|
|
device_obj, dbus.PROPERTIES_IFACE)
|
|
if device_props.Get('org.freedesktop.UDisks.Device',
|
|
"DriveConnectionInterface") == "usb" and device_props.Get(
|
|
'org.freedesktop.UDisks.Device', "DeviceIsPartition"):
|
|
if device_props.Get('org.freedesktop.UDisks.Device', "DeviceIsMounted"):
|
|
device_file = device_props.Get(
|
|
'org.freedesktop.UDisks.Device', "DeviceFile")
|
|
devices.append(device_file)
|
|
except Exception as e:
|
|
gen.log(e, error=True)
|
|
gen.log("No USB device found...")
|
|
|
|
devices.sort()
|
|
|
|
elif platform.system() == "Windows":
|
|
if fixed is True:
|
|
for drive in psutil.disk_partitions():
|
|
if 'cdrom' in drive.opts or drive.fstype == '':
|
|
# Skip cdrom drives or the disk with no filesystem
|
|
continue
|
|
devices.append(drive[0][:-1])
|
|
else:
|
|
try:
|
|
# Try new method using psutil. It should also detect USB 3.0 (but not tested by me)
|
|
for drive in psutil.disk_partitions():
|
|
if 'cdrom' in drive.opts or drive.fstype == '':
|
|
# Skip cdrom drives or the disk with no filesystem
|
|
continue
|
|
if 'removable' in drive.opts:
|
|
devices.append(drive[0][:-1])
|
|
except:
|
|
# Revert back to old method if psutil fails (which is unlikely)
|
|
oFS = win32com.client.Dispatch("Scripting.FileSystemObject")
|
|
oDrives = oFS.Drives
|
|
for drive in oDrives:
|
|
if drive.DriveType == 1 and drive.IsReady:
|
|
devices.append(drive)
|
|
|
|
if devices:
|
|
return devices
|
|
else:
|
|
gen.log("No USB device found...")
|
|
return None
|
|
|
|
|
|
def parent_partition(partition):
|
|
exploded = [c for c in partition]
|
|
while exploded[-1].isdigit():
|
|
exploded.pop()
|
|
return ''.join(exploded)
|
|
|
|
|
|
def details_udev(usb_disk_part):
|
|
"""
|
|
Get details of USB partition using udev
|
|
"""
|
|
assert usb_disk_part is not None
|
|
assert platform.system() == "Linux"
|
|
|
|
try:
|
|
import pyudev
|
|
except:
|
|
from . import pyudev
|
|
# Try with PyUdev to get the details of USB disks.
|
|
# This is the easiest and reliable method to find USB details.
|
|
# Also, it is a standalone package and no dependencies are required.
|
|
context = pyudev.Context()
|
|
try:
|
|
device = pyudev.Device.from_device_file(context, usb_disk_part)
|
|
except:
|
|
gen.log("ERROR: Unknown disk/partition (%s)" % str(usb_disk_part))
|
|
return None
|
|
|
|
try:
|
|
ppart = parent_partition(usb_disk_part)
|
|
fdisk_cmd_out = subprocess.check_output(
|
|
'LANG=C fdisk -l ' + ppart, shell=True)
|
|
partition_prefix = bytes(usb_disk_part, 'utf-8') + b' '
|
|
out = []
|
|
for l in fdisk_cmd_out.split(b'\n'):
|
|
if not l.startswith(b'/dev/'):
|
|
out.append(l)
|
|
continue
|
|
if l.startswith(partition_prefix):
|
|
out.append(l)
|
|
# filter out non-relevant partition info
|
|
fdisk_cmd_out = b'\n'.join(out)
|
|
|
|
except subprocess.CalledProcessError:
|
|
gen.log("ERROR: fdisk failed on disk/partition (%s)" %
|
|
str(usb_disk_part))
|
|
return None
|
|
|
|
detected_type = None
|
|
for keyword, ptype in [(b'Extended', 'extended partition'),
|
|
(b'swap', 'swap partition'),
|
|
(b'Linux LVM', 'lvm partition'),]:
|
|
if keyword in fdisk_cmd_out:
|
|
detected_type = ptype
|
|
break
|
|
if detected_type:
|
|
mount_point = ''
|
|
uuid = ''
|
|
file_system = ''
|
|
vendor = ''
|
|
model = ''
|
|
label = ''
|
|
devtype = detected_type
|
|
elif device.get('DEVTYPE') == "partition":
|
|
uuid = device.get('ID_FS_UUID') or ""
|
|
file_system = device.get('ID_FS_TYPE') or ""
|
|
label = device.get('ID_FS_LABEL') or ""
|
|
remounted = []
|
|
mount_point = UDISKS.mount(usb_disk_part, remounted) or ""
|
|
if remounted and remounted[0]:
|
|
config.add_remounted(usb_disk_part)
|
|
mount_point = mount_point.replace('\\x20', ' ')
|
|
vendor = device.get('ID_VENDOR') or ""
|
|
model = device.get('ID_MODEL') or ""
|
|
devtype = "partition"
|
|
|
|
elif device.get('DEVTYPE') == "disk":
|
|
mount_point = ""
|
|
uuid = ""
|
|
file_system = ""
|
|
label = device.get('ID_FS_LABEL') or ""
|
|
vendor = device.get('ID_VENDOR') or ""
|
|
model = device.get('ID_MODEL') or ""
|
|
devtype = "disk"
|
|
|
|
if mount_point not in ["", "None"]:
|
|
size_total = shutil.disk_usage(mount_point)[0]
|
|
size_used = shutil.disk_usage(mount_point)[1]
|
|
size_free = shutil.disk_usage(mount_point)[2]
|
|
|
|
else:
|
|
fdisk_cmd = 'LANG=C fdisk -l ' + usb_disk_part + \
|
|
' | grep "^Disk /" | sed -re "s/.*\s([0-9]+)\sbytes.*/\\1/"'
|
|
size_total = subprocess.check_output(fdisk_cmd, shell=True).strip()
|
|
size_used = ""
|
|
size_free = 0
|
|
mount_point = ""
|
|
|
|
return {'uuid': uuid, 'file_system': file_system, 'label': label, 'mount_point': mount_point,
|
|
'size_total': size_total, 'size_used': size_used, 'size_free': size_free,
|
|
'vendor': vendor, 'model': model, 'devtype': devtype}
|
|
|
|
|
|
def details_udisks2(usb_disk_part):
|
|
"""
|
|
Get details of USB disk detail.
|
|
usb_disk_part: It is the partition of an USB removable disk.
|
|
"""
|
|
import dbus
|
|
bus = dbus.SystemBus()
|
|
|
|
mount_point = ''
|
|
uuid = ''
|
|
file_system = ''
|
|
vendor = ''
|
|
model = ''
|
|
label = ''
|
|
devtype = "disk"
|
|
|
|
bd = bus.get_object('org.freedesktop.UDisks2', '/org/freedesktop/UDisks2/block_devices%s'%usb_disk_part[4:])
|
|
device = bd.Get('org.freedesktop.UDisks2.Block', 'Device', dbus_interface='org.freedesktop.DBus.Properties')
|
|
device = bytearray(device).replace(b'\x00', b'').decode('utf-8')
|
|
if device[-1].isdigit() is True:
|
|
uuid = bd.Get('org.freedesktop.UDisks2.Block', 'IdUUID', dbus_interface='org.freedesktop.DBus.Properties')
|
|
file_system = bd.Get('org.freedesktop.UDisks2.Block', 'IdType', dbus_interface='org.freedesktop.DBus.Properties')
|
|
mount_point = bd.Get('org.freedesktop.UDisks2.Filesystem', 'MountPoints', dbus_interface='org.freedesktop.DBus.Properties')
|
|
devtype = 'partition'
|
|
else:
|
|
devtype = 'disk'
|
|
file_system = 'No_File_System'
|
|
if mount_point:
|
|
# mount_point = str(bytearray(mount_point[0]).decode('utf-8').replace(b'\x00', b''))
|
|
mount_point = bytearray(mount_point[0]).replace(b'\x00', b'').decode('utf-8')
|
|
else:
|
|
try:
|
|
mount_point = UDISKS.mount(usb_disk_part)
|
|
config.add_remounted(usb_disk_part)
|
|
except:
|
|
mount_point = "No_Mount"
|
|
try:
|
|
label = bd.Get('org.freedesktop.UDisks2.Block', 'IdLabel', dbus_interface='org.freedesktop.DBus.Properties')
|
|
except:
|
|
label = "No_Label"
|
|
usb_drive_id = (bd.Get('org.freedesktop.UDisks2.Block', 'Drive', dbus_interface='org.freedesktop.DBus.Properties'))
|
|
bd1 = bus.get_object('org.freedesktop.UDisks2', usb_drive_id)
|
|
try:
|
|
vendor = bd1.Get('org.freedesktop.UDisks2.Drive', 'Vendor', dbus_interface='org.freedesktop.DBus.Properties')
|
|
except:
|
|
vendor = str('No_Vendor')
|
|
try:
|
|
model = bd1.Get('org.freedesktop.UDisks2.Drive', 'Model', dbus_interface='org.freedesktop.DBus.Properties')
|
|
except:
|
|
model = str('No_Model')
|
|
if not mount_point == "No_Mount":
|
|
size_total, size_used, size_free = \
|
|
shutil.disk_usage(mount_point)[:3]
|
|
else:
|
|
size_total = str('No_Mount')
|
|
size_used = str('No_Mount')
|
|
size_free = 0
|
|
|
|
return {'uuid': uuid, 'file_system': file_system, 'label': label, 'mount_point': mount_point,
|
|
'size_total': size_total, 'size_used': size_used, 'size_free': size_free,
|
|
'vendor': vendor, 'model': model, 'devtype': devtype}
|
|
|
|
|
|
def bytes2human(n):
|
|
"""
|
|
Convert the size to human readable format
|
|
Authored by 'Giampaolo Rodolà' and original link is:-
|
|
http://code.activestate.com/recipes/577972-disk-usage/
|
|
"""
|
|
try:
|
|
n = int(n)
|
|
except:
|
|
return 'Unknown'
|
|
symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
|
|
prefix = {}
|
|
for i, s in enumerate(symbols):
|
|
prefix[s] = 1 << (i + 1) * 10
|
|
for s in reversed(symbols):
|
|
if n >= prefix[s]:
|
|
value = float(n) / prefix[s]
|
|
return '%.1f%s' % (value, s)
|
|
return "%sB" % n
|
|
|
|
|
|
def gpt_device(dev_name):
|
|
"""
|
|
Find if the device inserted is GPT or not. We will just change the variable parameter in config file for later use
|
|
:param dev_name:
|
|
:return: True if GPT else False
|
|
"""
|
|
if platform.system() == 'Windows':
|
|
partition, disk = gen.wmi_get_drive_info(dev_name)
|
|
is_gpt = partition.Type.startswith('GPT:')
|
|
gen.log('Device %s is a %s disk...' %
|
|
(dev_name, is_gpt and 'GPT' or 'MBR'))
|
|
config.usb_gpt = is_gpt
|
|
return is_gpt
|
|
if platform.system() == "Linux":
|
|
if gen.has_digit(dev_name):
|
|
_cmd_out = subprocess.check_output("parted " + dev_name[:-1] + " print", shell=True)
|
|
else:
|
|
_cmd_out = subprocess.check_output("parted " + dev_name + " print", shell=True)
|
|
if b'msdos' in _cmd_out:
|
|
config.usb_gpt = False
|
|
gen.log('Device ' + dev_name + ' is a MBR disk...')
|
|
return False
|
|
elif b'gpt' in _cmd_out:
|
|
config.usb_gpt = True
|
|
gen.log('Device ' + dev_name + ' is a GPT disk...')
|
|
return True
|
|
|
|
|
|
def unmount(usb_disk):
|
|
UDISKS.unmount(usb_disk)
|
|
|
|
|
|
class RemountError(Exception):
|
|
def __init__(self, caught_exception, *args, **kw):
|
|
super(RemountError, self).__init__(*args, **kw)
|
|
self.caught_exception = caught_exception
|
|
|
|
def __str__(self):
|
|
return "%s due to '%s'" % (
|
|
self.__class__.__name__, self.caught_exception)
|
|
|
|
|
|
class UnmountError(RemountError):
|
|
def __init__(self, *args, **kw):
|
|
super(UnmountError, self).__init__(*args, **kw)
|
|
|
|
|
|
class MountError(RemountError):
|
|
def __init__(self, *args, **kw):
|
|
super(MountError, self).__init__(*args, **kw)
|
|
|
|
|
|
class UnmountedContext:
|
|
def __init__(self, usb_disk, exit_callback):
|
|
self.usb_disk = usb_disk
|
|
self.exit_callback = exit_callback
|
|
self.is_relevant = platform.system() != 'Windows' and \
|
|
self.usb_disk[-1:].isdigit()
|
|
|
|
def assert_no_access(self):
|
|
p = subprocess.Popen(['lsof', self.usb_disk],
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
stdin=subprocess.PIPE)
|
|
output = p.communicate()
|
|
if len(output[0].strip()) != 0:
|
|
gen.log("Open handle exists.")
|
|
gen.log(output[0])
|
|
raise UnmountError(Exception('open handle exists.'))
|
|
|
|
def __enter__(self):
|
|
if not self.is_relevant:
|
|
return
|
|
self.assert_no_access()
|
|
try:
|
|
gen.log("Unmounting %s" % self.usb_disk)
|
|
os.sync() # This is needed because UDISK.unmount() can timeout.
|
|
UDISKS.unmount(self.usb_disk)
|
|
except dbus.exceptions.DBusException as e:
|
|
gen.log("Unmount of %s has failed." % self.usb_disk)
|
|
# This may get the partition mounted. Don't call!
|
|
# self.exit_callback(details(self.usb_disk))
|
|
raise UnmountError(e)
|
|
gen.log("Unmounted %s" % self.usb_disk)
|
|
return self
|
|
|
|
def __exit__(self, type_, value, traceback_):
|
|
if not self.is_relevant:
|
|
return
|
|
os.sync() # This should not be strictly necessary
|
|
time.sleep(1) # Yikes, mount always fails without this sleep().
|
|
try:
|
|
mount_point = UDISKS.mount(self.usb_disk)
|
|
config.add_remounted(self.usb_disk)
|
|
self.exit_callback(details(self.usb_disk))
|
|
except dbus.exceptions.DBusException as e:
|
|
raise MountError(e)
|
|
gen.log("Mounted %s" % (self.usb_disk))
|
|
|
|
|
|
def check_vfat_filesystem(usb_disk, result=None):
|
|
p = subprocess.Popen(['fsck.vfat', '-n', usb_disk],
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
|
stdin=subprocess.PIPE)
|
|
output = p.communicate()
|
|
gen.log("fsck.vfat -n returned %d" % p.returncode)
|
|
gen.log(b"fsck.vfat -n said:" + b'\n---\n'.join(f for f in output if f))
|
|
if result is not None:
|
|
result.append((p.returncode, output, 'fsck.vfat -n'))
|
|
return len(output[0].split(b'\n'))==3 and output[1]==b'' \
|
|
and p.returncode==0
|
|
|
|
|
|
def repair_vfat_filesystem(usb_disk, result=None):
|
|
for args, input_ in [
|
|
(['-a', usb_disk], None, ),
|
|
(['-r', usb_disk], b'1\ny\n', ),
|
|
]:
|
|
cmd = ['fsck.vfat'] + args
|
|
p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE, stdin=subprocess.PIPE)
|
|
output = p.communicate(input=input_)
|
|
gen.log("%s returned %d" % (' '.join(cmd), p.returncode))
|
|
gen.log(b"It said:" + b'\n---\n'.join(f for f in output if f))
|
|
if result is not None:
|
|
result.append((p.returncode, output, ' '.join(cmd)))
|
|
return None
|
|
|
|
def details(usb_disk_part):
|
|
"""
|
|
Populate and get details of an USB disk.
|
|
:param usb_disk_part: USB disk. Example.. "/dev/sdb1" on Linux and "D:\" on Windows.
|
|
:return: label == > returns name/label of an inserted USB device.
|
|
mount_point == > returns mount path of an inserted USB device.
|
|
uuid == > returns uuid of an inserted USB device.
|
|
file_system == > returns type of filesystem of an inserted USB device.
|
|
device == > returns device path of an inserted USB device.
|
|
size_total == > returns total size in MB/GB of an inserted USB device.
|
|
size_free == > returns free size in MB/GB of an inserted USB device.
|
|
size_used == > returns used size in MB/GB of an inserted USB device.
|
|
vendor == > returns the name of the manufacturer.
|
|
model == > returns the model name of the USB.
|
|
"""
|
|
|
|
assert usb_disk_part is not None
|
|
|
|
details = {}
|
|
|
|
if platform.system() == 'Linux':
|
|
try:
|
|
details = details_udev(usb_disk_part)
|
|
except:
|
|
details = details_udisks2(usb_disk_part)
|
|
elif platform.system() == 'Windows':
|
|
details = osdriver.wmi_get_drive_info_ex(usb_disk_part)
|
|
|
|
return details
|
|
|
|
|
|
if __name__ == '__main__':
|
|
usb_devices = list_devices()
|
|
if usb_devices is not None:
|
|
for dev in usb_devices:
|
|
gen.log(details(dev))
|