1"""Disk utility module, no mixins here!
2
3    examples:
4    1) get disk size
5    from mozharness.base.diskutils import DiskInfo, DiskutilsError
6    ...
7    try:
8        DiskSize().get_size(path='/', unit='Mb')
9    except DiskutilsError as e:
10        # manage the exception e.g: log.error(e)
11        pass
12    log.info("%s" % di)
13
14
15    2) convert disk size:
16    from mozharness.base.diskutils import DiskutilsError, convert_to
17    ...
18    file_size = <function that gets file size in bytes>
19    # convert file_size to GB
20    try:
21        file_size = convert_to(file_size, from_unit='bytes', to_unit='GB')
22    except DiskutilsError as e:
23        # manage the exception e.g: log.error(e)
24        pass
25
26"""
27import ctypes
28import os
29import sys
30import logging
31from mozharness.base.log import INFO, numeric_log_level
32
33# use mozharness log
34log = logging.getLogger(__name__)
35
36
37class DiskutilsError(Exception):
38    """Exception thrown by Diskutils module"""
39    pass
40
41
42def convert_to(size, from_unit, to_unit):
43    """Helper method to convert filesystem sizes to kB/ MB/ GB/ TB/
44       valid values for source_format and destination format are:
45           * bytes
46           * kB
47           * MB
48           * GB
49           * TB
50        returns: size converted from source_format to destination_format.
51    """
52    sizes = {'bytes': 1,
53             'kB': 1024,
54             'MB': 1024 * 1024,
55             'GB': 1024 * 1024 * 1024,
56             'TB': 1024 * 1024 * 1024 * 1024}
57    try:
58        df = sizes[to_unit]
59        sf = sizes[from_unit]
60        return size * sf / df
61    except KeyError:
62        raise DiskutilsError('conversion error: Invalid source or destination format')
63    except TypeError:
64        raise DiskutilsError('conversion error: size (%s) is not a number' % size)
65
66
67class DiskInfo(object):
68    """Stores basic information about the disk"""
69    def __init__(self):
70        self.unit = 'bytes'
71        self.free = 0
72        self.used = 0
73        self.total = 0
74
75    def __str__(self):
76        string = ['Disk space info (in %s)' % self.unit]
77        string += ['total: %s' % self.total]
78        string += ['used: %s' % self.used]
79        string += ['free: %s' % self.free]
80        return " ".join(string)
81
82    def _to(self, unit):
83        from_unit = self.unit
84        to_unit = unit
85        self.free = convert_to(self.free, from_unit=from_unit, to_unit=to_unit)
86        self.used = convert_to(self.used, from_unit=from_unit, to_unit=to_unit)
87        self.total = convert_to(self.total, from_unit=from_unit, to_unit=to_unit)
88        self.unit = unit
89
90
91class DiskSize(object):
92    """DiskSize object
93    """
94    @staticmethod
95    def _posix_size(path):
96        """returns the disk size in bytes
97           disk size is relative to path
98        """
99        # we are on a POSIX system
100        st = os.statvfs(path)
101        disk_info = DiskInfo()
102        disk_info.free = st.f_bavail * st.f_frsize
103        disk_info.used = (st.f_blocks - st.f_bfree) * st.f_frsize
104        disk_info.total = st.f_blocks * st.f_frsize
105        return disk_info
106
107    @staticmethod
108    def _windows_size(path):
109        """returns size in bytes, works only on windows platforms"""
110        # we're on a non POSIX system (windows)
111        # DLL call
112        disk_info = DiskInfo()
113        dummy = ctypes.c_ulonglong()  # needed by the dll call but not used
114        total = ctypes.c_ulonglong()  # stores the total space value
115        free = ctypes.c_ulonglong()   # stores the free space value
116        # depending on path format (unicode or not) and python version (2 or 3)
117        # we need to call GetDiskFreeSpaceExW or GetDiskFreeSpaceExA
118        called_function = ctypes.windll.kernel32.GetDiskFreeSpaceExA
119        if isinstance(path, unicode) or sys.version_info >= (3,):
120            called_function = ctypes.windll.kernel32.GetDiskFreeSpaceExW
121        # we're ready for the dll call. On error it returns 0
122        if called_function(path,
123                           ctypes.byref(dummy),
124                           ctypes.byref(total),
125                           ctypes.byref(free)) != 0:
126            # success, we can use the values returned by the dll call
127            disk_info.free = free.value
128            disk_info.total = total.value
129            disk_info.used = total.value - free.value
130        return disk_info
131
132    @staticmethod
133    def get_size(path, unit, log_level=INFO):
134        """Disk info stats:
135                total => size of the disk
136                used  => space used
137                free  => free space
138          In case of error raises a DiskutilError Exception
139        """
140        try:
141            # let's try to get the disk size using os module
142            disk_info = DiskSize()._posix_size(path)
143        except AttributeError:
144            try:
145                # os module failed. let's try to get the size using
146                # ctypes.windll...
147                disk_info = DiskSize()._windows_size(path)
148            except AttributeError:
149                # No luck! This is not a posix nor window platform
150                # raise an exception
151                raise DiskutilsError('Unsupported platform')
152
153        disk_info._to(unit)
154        lvl = numeric_log_level(log_level)
155        log.log(lvl, msg="%s" % disk_info)
156        return disk_info
157