#
# Copyright 2007 sense.lab e.V.
# 
# This file is part of the CryptoBox.
#
# The CryptoBox is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# The CryptoBox 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 the CryptoBox; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#

'''
These classes detect and filter available blockdevices.
'''

__revision__ = "$Id$"


"""
TODO:
    - implement some caching
    - find the devnodes for each device (e.g. /dev/hda)
    - detect luks devices
    - detect LVM
"""


import os

class Blockdevices:
    
    def __init__(self, sysblock_dir='/sys/block', devnode_dir='/dev'):
        self.sysblock_dir = sysblock_dir
        self.devnode_dir = devnode_dir
        self.devices = []
        for devdir in find_blockdevices(self.sysblock_dir):
            blockdevice = Blockdevice(devdir, self.devnode_dir)
            if not blockdevice is None:
                self.devices.append(blockdevice)


    def get_devices(self):
        """return a copy of the device list
        """
        return self.devices[:]




class Blockdevice:

    def __init__(self, dev, devnode_dir='/dev', sysblock_dir='/sys/block'):
        self.devnode_dir = devnode_dir
        self.sysblock_dir = sysblock_dir
        if os.path.isabs(dev):
            self.devdir = dev
        else:
            self.devdir = self.__find_relative_device(dev)
            if self.devdir is None:
                self = None
                return
        self.name = os.path.basename(self.devdir)
        self.devnum = self.__get_major_minor()
        ## check valid devnum
        try:
            major, minor = self.devnum
        except TypeError:
            self = None
            return
        self.size = self.__get_size()
        self.range = self.__get_device_range()
        self.slaves = self.__get_dev_related("slaves")
        self.holders = self.__get_dev_related("holders")
        self.children = self.__get_children()


    def isstorage(self):
        if self.range > 1:
            ## partitionable blockdevice
            return False
        if self.size < 20:
            ## extended partition, unused loop device
            return False
        if self.devnum[0] == 1:
            ## ram device
            return False
        if self.children:
            ## a parent blockdevice
            return False
        return True


    def ispartitionable(self):
        if self.range > 1:
            return True
        

    def __find_relative_device(self, devname):
        for devdir in find_blockdevices(self.sysblock_dir):
            if os.path.basename(devdir) == devname:
                return devdir
        return None


    def __get_dev_related(self, subdir):
        """return the content of sub directories (e.g. 'holders' or 'slaves')
        """
        try:
            return os.listdir(os.path.join(self.devdir, subdir))
        except OSError:
            return []


    def __get_size(self):
        default = 0
        try:
            return int(file(os.path.join(self.devdir, 'size')).read())
        except OSError:
            return default
        except ValueError:
            return default


    def __get_major_minor(self):
        """return the major and minor of the device"""
        default = (0 ,0)
        try:
            content = file(os.path.join(self.devdir, "dev")).read()
        except IOError:
            return default
        major, minor = content.split(":", 2)
        try:
            return int(major), int(minor)
        except ValueError:
            return default


    def __get_device_range(self):
        """number of possible subdevices

        partitionable blockdevices have a range > 1
        """
        default = 1
        try:
            content = file(os.path.join(self.devdir, "range")).read()
        except IOError:
            return default
        try:
            return int(content)
        except ValueError:
            return default

    
    def __get_children(self):
        """return all devices depending on the current one
        
        all holders, subdevices and children of subdevices
        """
        direct_children = [Blockdevice(child).name
                for child in find_blockdevices(self.devdir)]
        direct_children.extend(self.holders[:])
        children = direct_children[:]
        for dchild in direct_children:
            children.extend(Blockdevice(dchild).children)
        return children


    def __str__(self):
        return self.name


    def info(self):
        output = "%s:\n" % self.name
        output += "\t%s:\t%s\n" % ("blockdir", self.devdir)
        output += "\t%s:\t%s\n" % ("major/minor", self.devnum)
        output += "\t%s:\t\t%s\n" % ("range", self.range)
        output += "\t%s:\t\t%s\n" % ("size", self.size)
        output += "\t%s:\t\t%s\n" % ("slaves", self.slaves)
        output += "\t%s:\t%s\n" % ("holders", self.holders)
        output += "\t%s:\t%s\n" % ("children", self.children)
        return output



def find_blockdevices(top_dir):

    dev_dirs = []
    
    def look4dev_dirs(arg, dirname, fnames):
        ## ignore the top level directory to avoid infinite recursion for
        ## get_children
        if os.path.samefile(dirname, top_dir):
            return
        ## add directories containing the file 'dev' to the list
        if (arg in fnames) and os.path.isfile(os.path.join(dirname, arg)):
            dev_dirs.append(dirname)
        for fname in fnames:
            ## remove symlinks and non-directories
            fullname = os.path.join(dirname, fname)
            if os.path.islink(fullname) or (not os.path.isdir(fullname)):
                fnames.remove(fname)
    
    os.path.walk(top_dir, look4dev_dirs, 'dev')
    return dev_dirs


if __name__ == '__main__':
    blocks = Blockdevices()
    for dev in blocks.get_devices():
        print dev.info()
    print
    print "Usable storage devices:"
    for dev in blocks.get_devices():
        if dev.isstorage():
            print dev
    print
    print "Partitionable devices:"
    for dev in blocks.get_devices():
        if dev.ispartitionable():
            print dev