Update disk.py from pymodules

This commit is contained in:
CorpNewt 2019-09-24 14:40:15 -05:00
parent 9d423bebf5
commit b6dfc2dd1f

View file

@ -1,6 +1,10 @@
import subprocess, plistlib, sys, os, time, json import subprocess, plistlib, sys, os, time, json
sys.path.append(os.path.abspath(os.path.dirname(os.path.realpath(__file__)))) sys.path.append(os.path.abspath(os.path.dirname(os.path.realpath(__file__))))
import run import run
if sys.version_info < (3,0):
# Force use of StringIO instead of cStringIO as the latter
# has issues with Unicode strings
from StringIO import StringIO
class Disk: class Disk:
@ -21,8 +25,8 @@ class Disk:
def _get_str(self, val): def _get_str(self, val):
# Helper method to return a string value based on input type # Helper method to return a string value based on input type
if (sys.version_info < (3,0) and isinstance(val, (str, unicode))) or (sys.version_info >= (3,0) and isinstance(val, str)): if (sys.version_info < (3,0) and isinstance(val, unicode)) or (sys.version_info >= (3,0) and isinstance(val, bytes)):
return val return val.encode("utf-8")
return str(val) return str(val)
def _get_plist(self, s): def _get_plist(self, s):
@ -31,8 +35,33 @@ class Disk:
if sys.version_info >= (3, 0): if sys.version_info >= (3, 0):
p = plistlib.loads(s.encode("utf-8")) p = plistlib.loads(s.encode("utf-8"))
else: else:
p = plistlib.readPlistFromString(s.encode("utf-8")) # p = plistlib.readPlistFromString(s)
except: # We avoid using readPlistFromString() as that uses
# cStringIO and fails when Unicode strings are detected
# Don't subclass - keep the parser local
from xml.parsers.expat import ParserCreate
# Create a new PlistParser object - then we need to set up
# the values and parse.
pa = plistlib.PlistParser()
# We also monkey patch this to encode unicode as utf-8
def end_string():
d = pa.getData()
if isinstance(d,unicode):
d = d.encode("utf-8")
pa.addObject(d)
pa.end_string = end_string
parser = ParserCreate()
parser.StartElementHandler = pa.handleBeginElement
parser.EndElementHandler = pa.handleEndElement
parser.CharacterDataHandler = pa.handleData
if isinstance(s, unicode):
# Encode unicode -> string; use utf-8 for safety
s = s.encode("utf-8")
# Parse the string
parser.Parse(s, 1)
p = pa.root
except Exception as e:
print(e)
pass pass
return p return p
@ -189,7 +218,7 @@ class Disk:
# Iterate!! # Iterate!!
if not disk or not len(self._get_str(disk)): if not disk or not len(self._get_str(disk)):
return None return None
disk = disk.lower() disk = self._get_str(disk).lower()
if disk.startswith("/dev/r"): if disk.startswith("/dev/r"):
disk = disk[len("/dev/r"):] disk = disk[len("/dev/r"):]
elif disk.startswith("/dev/"): elif disk.startswith("/dev/"):
@ -198,10 +227,10 @@ class Disk:
return disk return disk
for d in self.disks.get("AllDisksAndPartitions", []): for d in self.disks.get("AllDisksAndPartitions", []):
for a in d.get("APFSVolumes", []): for a in d.get("APFSVolumes", []):
if disk in [ a.get(x, "").lower() for x in ["DeviceIdentifier", "VolumeName", "VolumeUUID", "DiskUUID", "MountPoint"] ]: if disk in [ self._get_str(a.get(x, "")).lower() for x in ["DeviceIdentifier", "VolumeName", "VolumeUUID", "DiskUUID", "MountPoint"] ]:
return a.get("DeviceIdentifier", None) return a.get("DeviceIdentifier", None)
for a in d.get("Partitions", []): for a in d.get("Partitions", []):
if disk in [ a.get(x, "").lower() for x in ["DeviceIdentifier", "VolumeName", "VolumeUUID", "DiskUUID", "MountPoint"] ]: if disk in [ self._get_str(a.get(x, "")).lower() for x in ["DeviceIdentifier", "VolumeName", "VolumeUUID", "DiskUUID", "MountPoint"] ]:
return a.get("DeviceIdentifier", None) return a.get("DeviceIdentifier", None)
# At this point, we didn't find it # At this point, we didn't find it
return None return None