#!/usr/bin/python

"""
Written by Igor Pashev <pashev.igor@gmail.com>

The author has placed this work in the Public Domain,
thereby relinquishing all copyrights. Everyone is free
to use, modify, republish, sell or give away this work
without prior consent from anybody.
"""

from subprocess import Popen, PIPE
import re
import sys

def human_capacity(b):
    b = float(b)
    G = ['%.0f B', '%.2f KiB', '%.2f MiB', '%.2f GiB', '%.2f TiB']
    for p in G:
        if b < 1024:
            return p % b
        else:
            b /= 1024
    raise Exception('Too large hard drive ;-)')

class Partition(object):
    blocks = 0
    boot = False
    bytes = 0
    capacity = None
    end = 0
    id = 0
    logical = False
    start = 0
    system = None

class HDD(object):
    bytes = 0
    capacity = None
    description = None
    geometry = []
    name = None
    partitions = []
    raw_device = None

    # fdisk -G /dev/rdsk/c0t0d0p0
    # * Physical geometry for device /dev/rdsk/c0t0d0p0
    # * PCYL     NCYL     ACYL     BCYL     NHEAD NSECT SECSIZ
    #   2088     2088     0        0        255   63    512
    _fdiskG = re.compile('\s*(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s*')

    # fdisk.ul -l /dev/rdsk/c0t4d0p0
    # XXX note a bug: p0p1, must be p1.
    #               Device Boot      Start         End      Blocks   Id  System
    # /dev/rdsk/c0t4d0p0p1              63    15631244     7815591   bf  Solaris
    # /dev/rdsk/c0t4d0p0p2        15631245    33543719     8956237+   5  Extended
    # /dev/rdsk/c0t4d0p0p5        15631308    23438834     3903763+  82  Linux swap / Solaris
    # /dev/rdsk/c0t4d0p0p6        33479523    33543719       32098+  83  Linux


    def reread_partitions(self):
        #                                          boot?    start   end    blocks         id        system
        fdisk_ul = re.compile('^/dev/rdsk/{}\S+\s+(\*|\s)\s+(\d+)\s+(\d+)\s+(\d+)\+?\s+([a-f0-9]+)\s(.+)\s*$'.format(self.name))
        fdisk_ul_cmd = Popen(['fdisk.ul', '-l', self.raw_device], stdout=PIPE, stderr=PIPE)
        self.partitions = []
        for line in fdisk_ul_cmd.stdout:
            m = fdisk_ul.match(line)
            if m:
                p          = Partition()
                p.boot     = m.group(1) == '*'
                p.start    = int(m.group(2))
                p.end      = int(m.group(3))
                p.blocks   = int(m.group(4))
                p.id       = int(m.group(5), 16)
                p.system   = m.group(6)
                p.logical  = False # Update in the second pass
                p.bytes    = (p.end - p.start) * self.geometry['secsiz']
                p.capacity = human_capacity(p.bytes)
                self.partitions.append(p)

        ext= None
        for p in self.partitions:
            if p.id == 0x05:
                ext= p
                break
        if ext:
            for p in self.partitions:
                if p.id != 0x05 and p.start >= ext.start and p.end <= ext.end:
                    p.logical = True

    def __init_capacity(self):
        fdisk_cmd = Popen(['fdisk', '-G', self.raw_device], stdout=PIPE)
        for line in fdisk_cmd.stdout:
            m = self._fdiskG.match(line)
            if m:
                self.geometry = {
                        'pcyl'   : int(m.group(1)),
                        'ncyl'   : int(m.group(2)),
                        'acyl'   : int(m.group(3)),
                        'bcyl'   : int(m.group(4)),
                        'nhead'  : int(m.group(5)),
                        'nsect'  : int(m.group(6)),
                        'secsiz' : int(m.group(7)),
                        }


        self.bytes = ( self.geometry['ncyl'] * self.geometry['nhead'] *
                         self.geometry['nsect'] * self.geometry['secsiz'] )
        self.capacity = human_capacity(self.bytes)



    def __init__(self, name, description):
        self.name = name
        self.raw_device = '/dev/rdsk/{}p0'.format(name)
        self.description = description
        self.__init_capacity()
        self.reread_partitions()