#! /usr/bin/python

import sys, os, string
from statvfs import *

#some default definitions
colours = {	
			'none'		:	"",
			'black'		:	"\033[30m", 
			'red'		:	"\033[31m",
			'green'		:	"\033[32m",
			'yellow'	:	"\033[33m",
			'blue'		:	"\033[34m",
			'magenta'	:	"\033[35m",
			'purple'	:	"\033[35m",
			'cyan'		:	"\033[36m",
			'white'		:	"\033[37m",
			'darkgray'	:	"\033[30m"	
			}

normal_colour = 'green'
header_colour = 'yellow'
local_fs_colour = 'green'
remote_fs_colour = 'purple'
special_fs_colour = 'blue'

filled_fs_colour = 'red'
full_fs_colour = 'magenta'

sizeformat = "-h"

FILL_THRESH = 85.0
FULL_THRESH = 99.0

format = [	
			('fs', 18, "l"), ('size', 6, "r"), 
			('used', 6, "r"), ('avail', 6, "r"), ('perc', 5, "r"),
			('bar', 10, "l"), ('on', 20, "l")
		]

barchar = '#'

#end of default definitions

header =	{		
			'fs'		:	"Filesystem",
			'size'		:	"Size",
			'used'		:	"Used",
			'avail'		:	"Avail",
			'on'		:	"Mounted on",
			'fstype'	:	"Type",
			'perc'		:	"Use%",
			'bar'		:	""
			}


def hfnum(size, power):
	"human readable number"
	if size == 0:
		return "0"
	units = ["k", "M", "G", "T"]
	r = `size`[-1:] + "b"
	for i in range(len(units)):
		if size/power > 9:
			size = size/power
			r = `size`[:-1] + units[i]
	return r


def myformat(number, sizeformat, blocksize):
	"format number as file size"
	size = long(number)*blocksize
	units = ["k", "M", "G", "T"]
	if sizeformat == "-k":
		return `size/1024`[:-1]
	elif sizeformat == "-m":
		return `size/(1024*1024)`[:-1]
	elif sizeformat == "-h":
		return hfnum(size, 1024)
	elif sizeformat == "-H":
		return hfnum(size, 1000)
	elif sizeformat == "--blocks":
		return `number`

def fsline(mtab, sizeformat):
	"return list full of desired information about filesystem fs"
	mtabline = string.split(mtab)
	status = os.statvfs(mtabline[1])
	blocksize = status[F_FRSIZE]
	line = {
			'fs'		:	mtabline[0],
			'size'		:	myformat(long(status[F_BLOCKS]), sizeformat, blocksize),
			'used'		:	myformat(long(status[F_BLOCKS]-status[F_BFREE]), sizeformat, blocksize),
			'avail'		:	myformat(long(status[F_BAVAIL]), sizeformat, blocksize),
			'on'		:	mtabline[1],
			'fstype'	:	mtabline[2]
			}
	try:
		line['perc'] = `round((1.-float(status[F_BAVAIL])/status[F_BLOCKS])*100.,1)`
	except ZeroDivisionError:
		line['perc'] = "0"
	line['bar'] = "[###############]"
	return line

def myatof(s):
	"like atof, but be friendly to non-numerical values"
	try:
		return string.atof(s)
	except ValueError:
		return 0


def manglestring(s, l, pos):
	"cut string to fit exactly into l chars"
	if pos == "r":
		ns = string.rjust(s, l)
	elif pos == "l":
		ns = string.ljust(s, l)
	elif pos == "c":
		ns = string.center(s, l)
	else:
		raise ValueError, 'Error in manglestring'
	if len(ns) > l:
		ns = ns[:l/2] + "~" + ns[-l/2+1:]
	return ns

def help():
	print "Usage: pydf [OPTIONS]"
	print "Show information about mounted filesystems"
	print
	print "-h, --human-readable  print sizes in human readable format (e.g., 1K 234M 2G)"
	print "-H, --si              likewise, but use powers of 1000 not 1024"
	print "-k, --kilobytes       use 1024-byte blocks"
	print "-m, --megabytes       use 1048576-byte blocks"
	print "    --blocks          use filesystem native block size"
	print "    --bw              do not use colours"
	print "    --help            display this help and exit"
	print "    --version         output version information and exit"
	sys.exit()

def version():
	print "0.2"
	sys.exit()


# the fun begins here
for i in ["/etc/pydfrc", os.environ['HOME']+"/.pydfrc"]:
	if os.path.isfile(i):
		execfile(i)
	

for i in sys.argv[1:]:
	if i in ["--human-readable", "-h"]:
		sizeformat = "-h"
	elif i in ["--si", "-H"]:
		sizeformat = "-H"
	elif i in ["--kilobytes", "-k"]:
		sizeformat = "-k"
	elif i in ["--megabytes", "-m"]:
		sizeformat = "-m"
	elif i == "--blocks":
		sizeformat = "--blocks"
	elif i == "--bw":
		normal_colour = header_colour = local_fs_colour = remote_fs_colour = special_fs_colour = filled_fs_colour = full_fs_colour = 'none'
	elif i == "--version":
		version()
	else:
		help()

f = open("/etc/mtab","r")

lines = [header]
for mountline in f.readlines():
	mountline_list = string.split(mountline)
	lines.append(fsline(mountline, sizeformat))
	
for i in lines:
	if i['fs'] == "Filesystem":
		current_colour = header_colour
	elif i['size'] == "0":
		current_colour = special_fs_colour
	elif i['fstype'] in [ "nfs", "smbfs", "ncpfs", "afs", "coda" ]:
		current_colour = remote_fs_colour
	else:
		current_colour = local_fs_colour
	for j in format:
		print colours[current_colour],
		sys.stdout.write('')
		if j[0] in ['perc', 'avail', 'bar'] and i['fstype'] <> "iso9660":
			if myatof(i['perc']) > FILL_THRESH:
				print colours[filled_fs_colour],
				sys.stdout.write('')
			if myatof(i['perc']) > FULL_THRESH:
				print colours[full_fs_colour],
				sys.stdout.write('')
		if j[0] == 'bar':
			print '[',
			sys.stdout.write('')
			barsize = int(j[1]*myatof(i['perc'])/100.+0.5)
			print manglestring(barsize*barchar, j[1], j[2]),
			sys.stdout.write('')
			print ']',
		else:
			print manglestring(i[j[0]], j[1], j[2]),
	print
