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