from __future__ import division

import os
import statvfs
import re


# List of filesystem types authorized when generating disk use statistics.
STABLE_FILESYSTEMS = frozenset(
["ext", "ext2", "ext3", "ext4", "reiserfs", "ntfs", "msdos", "dos", "vfat",
"xfs", "hpfs", "jfs", "ufs", "hfs", "hfsplus"])


EXTRACT_DEVICE = re.compile("([a-z]+)[0-9]*")


def get_mount_info(mounts_file, statvfs_,
filesystems_whitelist=STABLE_FILESYSTEMS):
"""
This is a generator that yields information about mounted filesystems.

@param mounts_file: A file with information about mounted filesystems,
such as C{/proc/mounts}.
@param statvfs_: A function to get file status information.
@param filesystems_whitelist: Optionally, a list of which filesystems to
stat.
@return: A C{dict} with C{device}, C{mount-point}, C{filesystem},
C{total-space} and C{free-space} keys. If the filesystem information
is not available, C{None} is returned. Both C{total-space} and
C{free-space} are in megabytes.
"""
for line in open(mounts_file):
try:
device, mount_point, filesystem = line.split()[:3]
mount_point = mount_point.decode("string-escape")
except ValueError:
continue
if (filesystems_whitelist is not None and
filesystem not in filesystems_whitelist):
continue
megabytes = 1024 * 1024
try:
stats = statvfs_(mount_point)
except OSError:
continue
block_size = stats[statvfs.F_BSIZE]
total_space = (stats[statvfs.F_BLOCKS] * block_size) // megabytes
free_space = (stats[statvfs.F_BFREE] * block_size) // megabytes
yield {"device": device, "mount-point": mount_point,
"filesystem": filesystem, "total-space": total_space,
"free-space": free_space}


def get_filesystem_for_path(path, mounts_file, statvfs_):
"""
Tries to determine to which of the mounted filesystem C{path} belongs to,
and then returns information about that filesystem or C{None} if it
couldn't be determined.

@param path: The path we want filesystem information about.
@param mounts_file: A file with information about mounted filesystems,
such as C{/proc/mounts}.
@param statvfs_: A function to get file status information.
@param filesystems_whitelist: Optionally, a list of which filesystems to
stat.
@return: A C{dict} with C{device}, C{mount-point}, C{filesystem},
C{total-space} and C{free-space} keys. If the filesystem information
is not available, C{None} is returned. Both C{total-space} and
C{free-space} are in megabytes.
"""
candidate = None
path = os.path.realpath(path)
path_segments = path.split("/")
for info in get_mount_info(mounts_file, statvfs_):
mount_segments = info["mount-point"].split("/")
if path.startswith(info["mount-point"]):
if ((not candidate)
or path_segments[:len(mount_segments)] == mount_segments):
candidate = info
return candidate


def is_device_removable(device):
"""
This function returns whether a given device is removable or not by looking
at the corresponding /sys/block/<device>/removable file

@param device: The filesystem path to the device, e.g. /dev/sda1
"""
# Shortcut the case where the device an SD card. The kernel/udev currently
# consider SD cards (mmcblk devices) to be non-removable.
if os.path.basename(device).startswith("mmcblk"):
return True

path = _get_device_removable_file_path(device)

if not path:
return False

contents = None
try:
with open(path, "r") as f:
contents = f.readline()
except IOError:
return False

if contents.strip() == "1":
return True
return False


def _get_device_removable_file_path(device):
"""
Get a device's "removable" file path.

This function figures out the C{/sys/block/<device>/removable} path
associated with the given device. The file at that path contains either
a "0" if the device is not removable, or a "1" if it is.

@param device: File system path of the device.
"""
# The device will be a symlink if the disk is mounted by uuid or by label.
if os.path.islink(device):
# Paths are in the form "/dev/disk/by-uuid/<uuid>" and symlink
# to the device file under /dev
device = os.readlink(device) # /dev/disk/by-uuid/<uuid> -> ../../sda1

[device_name] = device.split("/")[-1:] # /dev/sda1 -> sda1

matched = EXTRACT_DEVICE.match(device_name) # sda1 -> sda

if not matched:
return None

device_name = matched.groups()[0]

removable_file = os.path.join("/sys/block/", device_name, "removable")
return removable_file