1. Moises Henriquez
  2. vasm

Source

vasm / src / vasm / backend / devhelper.py

#!/usr/bin/env python

#    This file is part of VASM.
#
#    VASM is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License v3 as published by
#    the Free Software Foundation.
#
#    VASM is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with VASM.  If not, see <http://www.gnu.org/licenses/>.

#   devhelper.py
""" Interact with system disk drives """

#   Some constants to help identify device types
DEVICE_HARD_DISK = -1
DEVICE_OPTICAL = -2
DEVICE_FLOPPY = -3

import parted
import os

class Device(object):
    def __init__(self, devpath=None):
        """ Class representing a device drive on the system.  
        Attributes:
            path - path to the device
            size - Device size in MB
            model - Device model information
            type - Device type (one of DEVICE_HARD_DISK or DEVICE_OPTICAL)
        """
        assert devpath not in (None, ''), "Invalid device path"
        self.path = devpath
        self._pdev = self._get_parted_device()
        self.tabletype = None
        self.model = None
        self.type = []
        self.partitions = None
        if self._pdev:
            self.size = self._pdev.getSize()
            self.model = self._pdev.model
            self.type = self._get_device_type()
            if self.type in (DEVICE_HARD_DISK,) and int(self.size) > 0:
                # Reading the label will tell us if the drive is ready to 
                # accept partitions or if we need to create a partition table
                # first with ` mklabel msdos|gpt|etc ` in parted
                self.tabletype = self._get_partition_table_type()
            else:
                self.tabletype = None
        if self.type == DEVICE_HARD_DISK and self.tabletype is not None:
            dsk = parted.Disk(self._pdev)
            self.partitions = dsk.partitions
    
    def _get_partition_table_type(self):
        """ When an unpartitioned disk is manipulated in parted/gparted, 
        it needs to have a 'label' on it.  parted will throw an exception if
        the disk has no label, so we need to try it and find out. """
        ret = None
        try:
            mydisk = parted.Disk(self._pdev)
            # Set the return value to whatever .type reports.  This would be
            # 'msdos' or 'GPT' or something along those lines.
            ret = mydisk.type
        except parted.DiskLabelException:
            # This is definitely an unpartitioned drive.  We must return 
            # None to expose this condition
            ret = None
            return ret
        return ret

    def _get_parted_device(self):
        """ Return a parted device object that we can interact with """
        pdev = None
        for dev in parted.getAllDevices():
            if dev.path == self.path:
                pdev = dev
                break
        return pdev
    
    def _get_device_type(self):
        """ Check if the device is listed by the kernel as a cdrom.  If true, 
        then this method will return DEVICE_OPTICAL, otherwise, it will
        return DEVICE_HARD_DISK"""
        datafile = "/proc/sys/dev/cdrom/info"
        f = open(datafile, 'r')
        data = [line for line in f.readlines() if line.startswith('drive name:')]
        f.close()
        shortpath = os.path.split(self.path)[-1].strip()
        for line in data:
            thisdev = line.split(':')[-1].strip()
            if thisdev == shortpath:
                return DEVICE_OPTICAL
        return DEVICE_HARD_DISK

def listSystemDrives():
    """ Return a list of Device objects """
    ret = []
    for drive in parted.getAllDevices():
        dev = Device(drive.path)
        if int(dev.size) > 0:
            ret.append(dev)
    return ret

def listSystemPartitions():
    """ Return a list of all system partitions """   
    drivelist = []
    ret = []
    for dev in listSystemDrives():
        if dev.type == DEVICE_HARD_DISK and dev.tabletype not in (None, ''):
            if dev.partitions:
                for part in dev.partitions:
                    if part.type != parted.PARTITION_EXTENDED:
                        ret.append(part.path)
    return ret