Cura/util/removableStorage.py

   1 import platform
   2 import string
   3 import glob
   4 import os
   5 import stat
   6 import time
   7 import subprocess
   8 try:
   9         from xml.etree import cElementTree as ElementTree
  10 except:
  11         from xml.etree import ElementTree
  12
  13 _removeableCache = None
  14 _removeableCacheTime = None
  15
  16 def _parseStupidPListXML(e):
  17         if e.tag == 'plist':
  18                 return _parseStupidPListXML(list(e)[0])
  19         if e.tag == 'array':
  20                 ret = []
  21                 for c in list(e):
  22                         ret.append(_parseStupidPListXML(c))
  23                 return ret
  24         if e.tag == 'dict':
  25                 ret = {}
  26                 key = None
  27                 for c in list(e):
  28                         if c.tag == 'key':
  29                                 key = c.text
  30                         elif key is not None:
  31                                 ret[key] = _parseStupidPListXML(c)
  32                                 key = None
  33                 return ret
  34         if e.tag == 'true':
  35                 return True
  36         if e.tag == 'false':
  37                 return False
  38         return e.text
  39
  40 def _findInTree(t, n):
  41         ret = []
  42         if type(t) is dict:
  43                 if '_name' in t and t['_name'] == n:
  44                         ret.append(t)
  45                 for k, v in t.items():
  46                         ret += _findInTree(v, n)
  47         if type(t) is list:
  48                 for v in t:
  49                         ret += _findInTree(v, n)
  50         return ret
  51
  52 def getPossibleSDcardDrives():
  53         global _removeableCache, _removeableCacheTime
  54         if _removeableCache is not None and time.time() - _removeableCacheTime < 5.0:
  55                 return _removeableCache
  56
  57         drives = []
  58         if platform.system() == "Windows":
  59                 from ctypes import windll
  60                 import ctypes
  61                 bitmask = windll.kernel32.GetLogicalDrives()
  62                 for letter in string.uppercase:
  63                         if bitmask & 1 and windll.kernel32.GetDriveTypeA(letter + ':/') == 2:
  64                                 volumeName = ''
  65                                 nameBuffer = ctypes.create_unicode_buffer(1024)
  66                                 if windll.kernel32.GetVolumeInformationW(ctypes.c_wchar_p(letter + ':/'), nameBuffer, ctypes.sizeof(nameBuffer), None, None, None, None, 0) == 0:
  67                                         volumeName = nameBuffer.value
  68                                 if volumeName == '':
  69                                         volumeName = 'NO NAME'
  70
  71                                 freeBytes = ctypes.c_longlong(0)
  72                                 if windll.kernel32.GetDiskFreeSpaceExA(letter + ':/', ctypes.byref(freeBytes), None, None) == 0:
  73                                         continue
  74                                 if freeBytes.value < 1:
  75                                         continue
  76                                 drives.append(('%s (%s:)' % (volumeName, letter), letter + ':/', volumeName))
  77                         bitmask >>= 1
  78         elif platform.system() == "Darwin":
  79                 p = subprocess.Popen(['system_profiler', 'SPUSBDataType', '-xml'], stdout=subprocess.PIPE)
  80                 xml = ElementTree.fromstring(p.communicate()[0])
  81                 p.wait()
  82
  83                 xml = _parseStupidPListXML(xml)
  84                 for dev in _findInTree(xml, 'Mass Storage Device'):
  85                         if 'removable_media' in dev and dev['removable_media'] == 'yes' and 'volumes' in dev and len(dev['volumes']) > 0:
  86                                 for vol in dev['volumes']:
  87                                         if 'mount_point' in vol:
  88                                                 volume = vol['mount_point']
  89                                                 drives.append((os.path.basename(volume), volume, os.path.basename(volume)))
  90         else:
  91                 for volume in glob.glob('/media/*'):
  92                         drives.append((os.path.basename(volume), volume, os.path.basename(volume)))
  93
  94         _removeableCache = drives
  95         _removeableCacheTime = time.time()
  96         return drives
  97
  98 def ejectDrive(driveName):
  99         if platform.system() == "Windows":
 100                 cmd = '"%s" %s>NUL' % (os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'EjectMedia.exe')), driveName)
 101         elif platform.system() == "Darwin":
 102                 cmd = "diskutil eject '%s' > /dev/null 2>&1" % (driveName)
 103         else:
 104                 cmd = "umount '%s' > /dev/null 2>&1" % (driveName)
 105         if os.system(cmd):
 106                 return False
 107         else:
 108                 return True
 109
 110 if __name__ == '__main__':
 111         print getPossibleSDcardDrives()