#! /usr/bin/python

import sys, os, string
try:
	from statvfs import *
	statvfs=os.statvfs
except ImportError:
	F_BSIZE   = 0       # Preferred file system block size
	F_FRSIZE  = 1       # Fundamental file system block size
	F_BLOCKS  = 2       # Total number of file system blocks (FRSIZE)
	F_BFREE   = 3       # Total number of free blocks
	F_BAVAIL  = 4       # Free blocks available to non-superuser
	F_FILES   = 5       # Total number of file nodes
	F_FFREE   = 6       # Total number of free file nodes
	F_FAVAIL  = 7       # Free nodes available to non-superuser
	F_FLAG    = 8       # Flags (see your local statvfs man page)
	F_NAMEMAX = 9       # Maximum file name length
	import commands
	def statvfs(fs):		
    		l = string.split(commands.getoutput(statfs_command + " " + fs))
		del l[0]
		l1 = []
		for i in l:
			try:
				l1.append(string.atoi(i))
			except ValueError:
				l1.append(0)
		l1.insert(1, l1[0])
		l1.insert(6, l1[6])
		l1.insert(8, 0)
		return(l1)


#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",
			'default'	:	"\033[0m"
			}

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

filled_fs_colour = 'red'
full_fs_colour = 'magenta'

sizeformat = "-h"

FILL_THRESH = 85.0
FULL_THRESH = 99.0

format = [	
			('fs', 15, "l"), ('size', 9, "r"), 
			('used', 9, "r"), ('avail', 9, "r"), ('perc', 5, "r"),
			('bar', 8, "l"), ('on', 16, "l")
		]

barchar = '#'

mountfile = "/etc/mtab"

statfs_command = "/usr/bin/statfs -t"


#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 = statvfs(mtabline[1])
	blocksize = status[F_BSIZE]
	used = long(status[F_BLOCKS]-status[F_BFREE])
	line = {
			'fs'		:	mtabline[0],
			'size'		:	myformat(long(status[F_BLOCKS]), sizeformat, blocksize),
			'used'		:	myformat(used, sizeformat, blocksize),
			'avail'		:	myformat(long(status[F_BAVAIL]), sizeformat, blocksize),
			'on'		:	mtabline[1],
			'fstype'	:	mtabline[2]
			}
	try:
		line['perc'] = `round(100.*used/(used + status[F_BAVAIL]))`
	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 "    --mounts=FILE     file to get mount information from."
	print "                      On normal linux system, only /etc/mtab"
	print "                      or /proc/mounts make sense."
	print "                      Use /proc/mounts when /etc/mtab is corrupted"
	print "                      or inaccesable (the output looks a bit weird"
	print "                      in this case though)"
	print "    --help            display this help and exit"
	print "    --version         output version information and exit"
	sys.exit()

def version():
	print "0.4"
	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[:9] == "--mounts=":
		mountfile = i[9:]
	elif i == "--version":
		version()
	else:
		help()

f = open(mountfile,"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['fstype'] in [ "nfs", "smbfs", "ncpfs", "afs", "coda" ]:
		current_colour = remote_fs_colour
	elif i['size'] == "0":
		current_colour = special_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
print colours['default'],

