# parted.py -- Python module for parted

# $Progeny: parted.py,v 1.35 2002/03/19 01:42:20 epg Exp $

# Copyright (C) 2000, 2001, 2002 Progeny Linux Systems, Inc.
# AUTHORS: Eric Gillespie, Jr. <epg@progeny.com>
#          Jeff Licquia <jlicquia@progeny.com>

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License, version 2,  as
# published by the Free Software Foundation.

# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""Python interface to parted.
"""

import sys
import types

SECTOR_SIZE = 512

###############################################################################
# Sequence

class Sequence:
    """A sequence of parted objects.
    
    This class implements the *_get_next functions by creating a
    special sequence type. This allows sequences of parted objects to
    be treated just like any other Python immutable sequence. The
    sequence avoids actually instantiating any parted objects until
    they are needed, however, and will handle growing and shrinking as
    parted objects are created or destroyed properly (for example: a
    sequence referring to all partitions on a disk will remain correct
    even as partitions are removed or added).
    """

    def __init__(self, start_item_func, *start_item_params):
        self.start_item_func = start_item_func
        if len(start_item_params) > 0:
            self.start_item_param = start_item_params[0]
        else:
            self.start_item_param = None
        self.reset()

    def reset(self):
        self.list = []
        firstitem = self.start_item_func(self.start_item_param)
        if firstitem is not None:
            self.list.append(firstitem)

    def enumerate(self):
        if len(self.list) > 0:
            next_item = self.list[-1]._get_next()
            while next_item:
                self.list.append(next_item)
                next_item = next_item._get_next()

    def __len__(self):
        self.reset()
        self.enumerate()
        return len(self.list)

    def __getitem__(self, key):
        if not isinstance(key, types.IntType):
            raise IndexError, "sequence index not an integer"

        self.reset()

        if key < 0:
            self.enumerate()
        else:
            if len(self.list) < (key + 1) and len(self.list) > 0:
                next_item = self.list[-1]._get_next()
                while next_item and len(self.list) < (key + 1):
                    self.list.append(next_item)
                    next_item = next_item._get_next()

            if len(self.list) < (key + 1):
                raise IndexError, "sequence index out of range"

        return self.list[key]

    def __getslice__(self, i, j):
        self.reset()
        if i < 0 or j == sys.maxint or j < 0:
            self.enumerate()
        else:
            # Force enumeration to the last index.
            dummy = self[j]

        return self.list[i:j]

###############################################################################
# PEDDEVICE

# Enum DeviceType
DEVICE_UNKNOWN	= 0
DEVICE_SCSI	= 1
DEVICE_IDE	= 2
DEVICE_DAC960	= 3
DEVICE_CPQARRAY	= 4
DEVICE_FILE     = 5
DEVICE_ATARAID  = 6
DEVICE_I2O      = 7


class Device:
    """Represents a storage device.
    """

    def __init__(self, obj):
        self._o = obj

    # Accessor Methods

    def _get_next(self):
        return _device_get_next(self)

    def get_model(self):
        """Get a description of the hardware manufacturer and model.
        """

        return _parted.device_get_model(self._o)

    def get_path(self):
        """Get the block device, e.g. /dev/sdb.
        """

        return _parted.device_get_path(self._o)

    def get_type(self):
        return _parted.device_get_type(self._o)

    def get_sector_size(self):
        return _parted.device_get_sector_size(self._o)

    def get_heads(self):
        return _parted.device_get_heads(self._o)

    def get_sectors(self):
        return _parted.device_get_sectors(self._o)

    def get_cylinders(self):
        return _parted.device_get_cylinders(self._o)

    def get_geom_known(self):
        return _parted.device_get_geom_known(self._o)

    def get_host(self):
        return _parted.device_get_host(self._o)

    def get_did(self):
        return _parted.device_get_did(self._o)

    def get_length(self):
        """Get the size of the device in sectors.
        """

        return _parted.device_get_length(self._o)

    def get_open_count(self):
        return _parted.device_get_open_count(self._o)

    def get_dirty(self):
        return _parted.device_get_dirty(self._o)

    # Action Methods

    def open(self):
        """Attempts to open the device.

        This is to allow use of the read, write and sync
        methods. Returns zero on failure.
        """

        return _parted.device_open(self._o)

    def close(self):
        """Closes the device.
        """

        return _parted.device_close(self._o)

    def read(self, start, count):
        """Reads count sectors from the device, beginning with sector start.
        """

        return _parted.device_read(self._o, start, count)

    def write(self, buffer, start, count):
        """Writes count sectors to the device, beginning with sector start.
        """

        return _parted.device_write(self._o, buffer, start, count)

    def sync(self):
        """Flushes the write cache.
        """

        return _parted.device_sync(self._o)

    def disk_open(self):
        """Constructs a Disk object from the device.

        Also reads the partition table.

        Returns None on failure.

        WARNING: this can modify dev->cylinders, dev->heads and
        dev->sectors, because the partition table might indicate that
        the existing values were incorrect.
        """

        r = _parted.disk_open(self._o)
        if r is None:
            return None
        else:
            return Disk(r)

    def disk_create(self, type):
        """Creates a partition table on this device.

        Returns a Disk object for it.
        """

        return Disk(_parted.disk_create(self._o, type._o))

def _device_get_next(dev=None):
    """Returns the next detected device.

    The list of detected devices comes by calling device_probe_all.

    If dev is None, returns the first device.  Returns None if dev is
    the last device.
    """

    if dev:
        r = _parted.device_get_next(dev._o)
    else:
        r =_parted.device_get_next(None)

    if r:
        return Device(r)
    else:
        return None

###############################################################################
# PEDDISK, PEDDISKTYPE

class DiskType:
    def __init__(self, obj):
        self._o = obj

    def _get_next(self):
        return _disk_type_get_next(self)

    def get_name(self):
        """Get the name of the type.
        """

        return _parted.disk_type_get_name(self._o)

    def get_ops(self):
        return _parted.disk_type_get_ops(self._o)

def _disk_type_get_next(type=None):
    if type:
        r = _parted.disk_type_get_next(type._o)
    else:
        r = _parted.disk_type_get_next(None)

    if r:
        return DiskType(r)
    else:
        return None

class Disk:
    """Represents a device + partition table.
    """

    def __init__(self, obj):
        self._o = obj

    # Accessor Methods

    def get_dev(self):
        return Device(_parted.disk_get_dev(self._o))

    def get_type(self):
        """The type of the disk label.
        """

        return DiskType(_parted.disk_get_type(self._o))

    def get_part_list(self):
        """The list of partitions on the disk.
        """

        return Sequence(_disk_next_partition, self)

    # Action Methods

    def close(self):
        """Closes the disk.
        """

        return _parted.disk_close(self._o)

    def read(self):
        """Reads the partition table from the disk.
        """

        return _parted.disk_read(self._o)

    def write(self):
        """Writes the partition table to the disk.
        """

        return _parted.disk_write(self._o)

    def add_partition(self, part):
	"""Adds "part" to "disk".

        "part"'s geometry may be changed in this process. Sorry, I
        know this is bad, but there is no other way. Some brain-dead
        partition table systems, not mentionioning any
        Micros^H^H^H^H^H^Hnames require certain alignments which can't
        be pre-calculated.  "part" is assigned a number in this
        process.
        """

        part.disk = self
        return _parted.disk_add_partition(self._o, part._o)

    def delete_partition(self, part):
        """Removes part from the disk, destroying part.
        """

        return _parted.disk_delete_partition(self._o, part._o)

    def delete_all(self):
        """Removes an destroys all partitions on the disk.
        """

        return _parted.disk_delete_all(self._o)

    def set_partition_geom(self, part, start, end):
        """Sets the geometry of part.
        """

        return _parted.disk_set_partition_geom(self._o, part._o, start, end)

    def maximize_partition(self, part):
        """Grows part's geometry to the maximum possible.
        """

        return _parted.disk_maximize_partition(self._o, part._o)

    def get_max_partition_geometry(self, part):
        """Returns the maximum geometry part can be grown to.
        """

        return Geometry(_parted.disk_get_max_paritition_geometry(self._o,
                                                                 part._o))

    def minimize_extended_partition(self):
        """Reduces the extended partition on the disk to the
        minimum possible."""

        return _parted.disk_minimize_extended_partition(self._o)

    def next_partition(self, part):
        return _disk_next_partition(self, part)

    def get_partition(self, num):
        """Returns the partition numbered num.
        """

        r = _parted.disk_get_partition(self._o, num)
        if r:
            return Partition(self, None, None, None, None, r)
        else:
            return None

    def get_partition_by_sector(self, sect):
        """Returns the partition that owns sect.
        """

        return Partition(self, None, None, None, None,
                         _parted.disk_get_partition_by_sector(self._o, sect))

    def get_extended_partition(self):
        """Returns the extended partition.
        """

        return Partition(self, None, None, None, None,
                         _parted.disk_get_extended_partition(self._o))

def _disk_next_partition(disk, part=None):
    if part is not None:
        r = _parted.disk_next_partition(disk._o, part._o)
    else:
        r = _parted.disk_next_partition(disk._o)

    if r:
        return Partition(disk, None, None, None, None, r)
    else:
        return None

###############################################################################
# PEDGEOMETRY

class Geometry:
    """Represents a continuous region on a device.
    """

    def __init__(self, obj):
        self._o = obj

    # Accessor Methods

    def get_disk(self):
        """Returns the disk.
        """

        return Disk(_parted.geometry_get_disk(self._o))

    def get_start(self):
        """Returns the start of the region in sectors.
        """

        return _parted.geometry_get_start(self._o)

    def get_length(self):
        """Returns the length of the region, in sectors.
        """

        return _parted.geometry_get_length(self._o)

    def get_end(self):
        """Returns the end of the region in sectors.
        """

        return _parted.geometry_get_end(self._o)

    # I'm not going to bind anymore of this one just yet, as it
    # appears no one is using any of the action methods.

###############################################################################
# PEDPARTITION, PEDPARTITIONTYPE

# Enum PartitionType:
PARTITION_PRIMARY		= 0x00
PARTITION_LOGICAL		= 0x01
PARTITION_EXTENDED		= 0x02
PARTITION_FREESPACE		= 0x04
PARTITION_METADATA		= 0x08

class Partition:
    """Represents a partition.

    This is basically a Geometry plus some attributes.
    """

    def __init__(self, disk, type, fs_type, start, end,
                 obj=None):
        """Creates a new Partition on disk.
        """

        # Remember disk for iteration later.

        if disk is None:
            raise ValueError, "no disk object passed in constructor"
        self.disk = disk

        if obj:
            self._o = obj
        else:
            if type == PARTITION_EXTENDED:
                r = _parted.partition_new(disk._o, type, None,
                                          start, end)
            else:
                r = _parted.partition_new(disk._o, type, fs_type._o,
                                          start, end)
            if r:
                self._o = r
            else:
                raise RuntimeError, "unable to create partition"

    # Accessor methods

    def _get_next(self):
        """Retrieve next partition on the same disk.
        """

        if self.disk is None:
            return None
        else:
            return _disk_next_partition(self.disk, self)

    def get_geom(self):
        """Geometry of the partition.
        """

        return Geometry(_parted.partition_get_geom(self._o))

    def get_num(self):
        """The partition number.
        """

        return _parted.partition_get_num(self._o)

    def get_hidden(self):
        """Whether the partition is hidden.
        """

        return _parted.partition_get_hidden(self._o)

    def get_type(self):
        """The partition type.
        """

        return _parted.partition_get_type(self._o)

    def get_part_list(self):
        """List of logical partitions in an extended partition.
        """

        return Sequence(_disk_next_partition,
                        Partition(self.disk, None, None, None, None,
                                  _parted.partition_get_part_list(self._o)))


    def get_bootable(self):
        """Whether the partition is the one that should be booted.
        """

        return _parted.partition_get_bootable(self._o)

    def set_bootable(self, value):
        """Set the bootable flag to VALUE.
        """

        return _parted.partition_set_bootable(self._o, value)
        
    def get_fs_type(self):
        """Type of filesystem on the partition.
        """

        ped_object = _parted.partition_get_fs_type(self._o)
        if ped_object is None:
            return None
        else:
            return FileSystemType(ped_object)

    def set_fs_type(self, type):
        """Set the file system type on the partition to TYPE.
        """

        return _parted.partition_set_fs_type(self._o, type._o)
    
    def is_active(self):
        """Whether the partition is active.
        """

        return _parted.partition_is_active(self._o)

    def is_busy(self):
        """Whether the partition is busy (i.e. mounted).
        """

        return _parted.partition_is_busy(self._o)

    def is_efi(self):
        """Whether the partition is an EFI boot partition.
        """

        if self.get_fs_type().get_name() == "FAT":
            return _parted.partition_is_efi(self._o)

        return FALSE

    # Member Set Methods

    def mark_efi(self):
        """Sets system type to EFI.

        Returns TRUE on success, FALSE on fail, and -1 if not a FAT
        partition.

        This writes the same bits as the set_system method, so only
        use one or the other.
        """

        if self.get_fs_type().get_name() == "FAT":
            return _parted.partition_mark_efi(self._o)

        return -1

    def set_system(self, fs_type):
        """Sets the system type.
        """

        return _parted.partition_set_system(self._o, fs_type._o)

    # Action Methods

    def destroy(self):
        """Destroys a partition.

        Should not be called on a a partition that is in the partition
        table. Use Disk.delete_partition instead.
        """

        return _parted_partition_destroy(self._o)

###############################################################################
# PEDFILESYSTEM, PEDFILESYSTEMTYPE

class FileSystemType:
    def __init__(self, obj):
        self._o = obj

    def _get_next(self):
        return _file_system_type_get_next(self)

    def get_name(self):
        """Name of the filesystem type.
        """

        return _parted.file_system_type_get_name(self._o)

    def get_ops(self):
        return FileSystemType(_parted.file_system_type_get_ops(self._o))

def _file_system_type_get_next(fs_type=None):
    if fs_type:
        r = _parted.file_system_type_get_next(fs_type._o)
    else:
        r = _parted.file_system_type_get_next(None)

    if r:
        return FileSystemType(r)
    else:
        return None

class FileSystem:
    """Represents a filesystem.

    This is associated with a Geometry, NOT a Partition.
    """

    def __init__(self, geom, type, obj=None):
        if obj:
            self._o = obj
        else:
            self._o = _parted.file_system_create(geom._o, type._o)

    # Accessor Methods

    def get_type(self):
        return FileSystemType(_parted.file_system_get_type(self._o))

    def get_geom(self):
        """Where the filesystem actually is.
        """

        return Geometry(_parted.file_system_get_geom(self._o))

    def get_resize_constraint(self):
        """Returns a constraint, that represents all of the possible
        ways the file system can be resized.
        """

        obj = _parted.file_system_get_resize_constraint(self._o)
        if obj == None:
            return None
        else:
            return Constraint(None, None, None, None, None, obj)

    # Action Methods

    def close(self):
        """Closes the filesystem.
        """

        return _parted.file_system_close(self._o)

    def check(self):
        """Checks filesystem for errors.
        """

        return _parted.file_system_check(self._o)

    def copy(self, geom):
        """Creates a copy of the same type on geom.
        """

        return _parted.file_system_copy(self._o, geom._o)

    def resize(self, geom):
        """Resize the filesystem to the new geom.
        """

        return _parted.file_system_resize(self._o, geom._o)

def file_system_probe(geom):
    """Attempts to detect a filesystem on geom and returns it.
    """

    return _parted.file_system_probe(geom._o)

def file_system_open(geom):
    """Opens a filesystem on geom.
    """

    obj = _parted.file_system_open(geom._o)
    if obj == None:
        return None
    else:
        return FileSystem(None, None, obj)

###############################################################################
# PEDALIGNMENT, PEDCONSTRAINT

class Alignment:
    """Represents an alignment.
    """

    def __init__(self, offset, grain_size, obj=None):
        """Returns an alignment object (used by PedConstraint),
        representing all PedSector's that are of the form
        "offset + X * grain_size".
        """

        if obj:
            self._o = obj
        else:
            self._o = _parted.alignment_new(offset, grain_size)

    def get_offset(self):
        """Return the offset of the alignment.
        """

        return _parted.alignment_get_offset(self._o)

    def get_grain_size(self):
        """Return the grain_size of the alignment.
        """

        return _parted.alignment_get_grain_size(self._o)

class Constraint:
    """Represents a constraint.
    """

    def __init__(self, start_align, end_align, start_range, end_range, 
                 min_size, obj=None):
        """Creates a new constraint, returns NULL on failure.
        """
        
        if obj:
            self._o = obj
        else:
            self._o = _parted.constraint_new(start_align._o, end_align._o,
                                             start_range._o, end_range._o,
                                             min_size)
            
    def get_start_align(self):
        """Return the possible alignment of the starting position of
        the constraint.
        """
        
        obj = _parted.constraint_get_start_align(self._o)

        if obj == None:
            return None
        else:
            return Alignment(None, None, obj)

    def get_end_align(self):
        """Return the possible alignment of the starting position of
        the constraint.
        """
        
        obj = _parted.constraint_get_end_align(self._o)

        if obj == None:
            return None
        else:
            return Alignment(None, None, obj)

    def get_start_range(self):
        """Return the possible alignment of the starting position of
        the constraint.
        """
        
        obj = _parted.constraint_get_start_range(self._o)

        if obj == None:
            return None
        else:
            return Geometry(obj)

    def get_end_range(self):
        """Return the possible alignment of the starting position of
        the constraint.
        """
        
        obj = _parted.constraint_get_end_range(self._o)

        if obj == None:
            return None
        else:
            return Geometry(obj)

    def get_min_size(self):
        """Return the possible alignment of the starting position of
        the constraint.
        """
        
        return(_parted.constraint_get_min_size(self._o))

###############################################################################
# Exceptions

# Enum PedExceptionType
EXCEPTION_INFORMATION=1
EXCEPTION_WARNING=2
EXCEPTION_ERROR=3
EXCEPTION_FATAL=4
EXCEPTION_BUG=5
EXCEPTION_NO_FEATURE=6

# Enum PedExceptionOption
EXCEPTION_UNHANDLED=0
EXCEPTION_FIX=1
EXCEPTION_YES=2
EXCEPTION_NO=4
EXCEPTION_OK=8
EXCEPTION_RETRY=16
EXCEPTION_IGNORE=32
EXCEPTION_CANCEL=64

# These are used by the default exception handler.
type_strings = {
    EXCEPTION_INFORMATION: "information",
    EXCEPTION_WARNING:     "warning",
    EXCEPTION_ERROR:       "error",
    EXCEPTION_FATAL:       "fatal error",
    EXCEPTION_BUG:         "parted bug",
    EXCEPTION_NO_FEATURE:  "feature not implemented"
    }

response_strings = {
    EXCEPTION_FIX:    ["Fix", "f"],
    EXCEPTION_YES:    ["Yes", "y"],
    EXCEPTION_NO:     ["No", "n"],
    EXCEPTION_OK:     ["OK", "o"],
    EXCEPTION_RETRY:  ["Retry", "r"],
    EXCEPTION_IGNORE: ["Ignore", "i"],
    EXCEPTION_CANCEL: ["Cancel", "c"]
    }

class Exception:
    """Represents parted exceptions.

    Exceptions in parted are not, strictly speaking, exceptions;
    parted functions will either invoke an error handler callback or
    return an error code, depending on the preference of the
    application.  In Python, when a parted exception is thrown, an
    instance of this class will be created to encapsulate the
    information in the exception.  This is used to provide an
    easy-to-use interface for applications to use when registering
    exception handlers.
    """

    def __init__(self, message, type, options):
        self.type = type
        self.options = options
        self.message = message
        self.iscaught = 0

        # Trigger exception if options are invalid.
        dummy = self.get_options()

    def get_options(self):
        """Return a tuple of possible options to return.
        """

        optionlist = []
        for pos_option in (EXCEPTION_FIX, EXCEPTION_YES, EXCEPTION_NO,
                           EXCEPTION_OK, EXCEPTION_RETRY,
                           EXCEPTION_IGNORE, EXCEPTION_CANCEL):
            if self.options & pos_option:
                optionlist.append(pos_option)
        if length(optionlist) < 1:
            raise AttributeError, "no valid option set"
        return optionlist

    def get_type(self):
        return self.type

    def get_message(self):
        return self.message

def _lowlevel_exception_handler(message, type, options):
    """Hook our exception handling interface into the low-level parted
    exception system."""

    if pedobject is None:
        return EXCEPTION_UNHANDLED

    try:
        ped_exception = Exception(message, type, options)
    except AttributeError:
        return EXCEPTION_UNHANDLED

    return call_handler(ped_exception)

def _default_handler(ped_exception):
    """Default parted exception handler.  It writes the error and a
    list of options to standard error, and expects the letter of one
    response to be entered on standard input.
    """
    
    sys.stderr.write("parted system: %s: %s\n" %
                     (type_strings[ped_exception.get_type()],
                      ped.exception.get_message()))

    options = ped_exception.get_options()
    for option in options:
        optionstr = response_strings[option]
        sys.stderr.write("%s (%s)", optionstr[0], optionstr[1])
    sys.stderr.write("? ")

    response_value = 0
    while not response_value:
        response = f.readline()[0]
        for option in options:
            optionstr = response_strings[option]
            if response == optionstr[1]:
                response_value = option

    return response_value

################################################################################
# Top-level functions and globals

_exception_handler = _default_handler

def init():
    global _parted
    import _parted
    return _parted.init()

def done():
    """Shut down the underlying parted subsystem.
    """

    return _parted.done()

def device_probe_all():
    """Attempts to detect all devices.
    """

    return _parted.device_probe_all()

def device_get(name):
    """Gets the device "name".

    "name" is usually the block device (eg /dev/sdb).  If the device
    wasn't detected with device_probe_all, an attempt will be made
    to detect it again.
    """

    return Device(_parted.device_get(name))

def disk_type_get(name):
    return DiskType(_parted.disk_type_get(name))

def file_system_type_get(name):
    """Returns the FileSystemType corresponding to name.
    """

    return FileSystemType(_parted.file_system_type_get(name))

def call_handler(ped_exception):
    return (_exception_handler)(ped_exception)

def register_handler(handler_func):
    _exception_handler = handler_func

def unregister_handler():
    _exception_handler = _default_handler

def get_devices():
    return Sequence(_device_get_next)

def get_disk_types():
    return Sequence(_disk_type_get_next)

def get_file_system_types():
    return Sequence(_file_system_type_get_next)