#! /usr/bin/env python
# coding=CP437

# diskinfo.py - view, mount and unmount drives and CDs as a normal user
version='3.6'

# Author: Henry Kroll www.thenerdshow.com
# License: GPLv3
# Dependencies: Fedora/none (/sys/dev, /etc/mtab, /proc/partitions in normal places)
# Requires 2.6(+?) series kernel: abuses sysfs in new and interesting ways (tm)
# Test code with unpriviliged test user account, not root.

# Changelog
# 2.0: run as user, not root, got rid of parted requirement :)
# 3.0: use /sys/block instead of /proc/partitions, show removables, options, -h help
# 3.2: -i interactive mode, bugfix, eject now toggles open/close
# 3.3: detect blank and empty removable drives, show symlinks
# 3.4: fix_spaces in drive and device names
# 3.5: replace unnecessary eval statements with int and float
# 3.6: use os.stat and os.readlink instead of popen4 to get /dev symlinks

import os,sys
from stat import *
drives={} 	# dictionary of drives
mounted={} 	# dictionary of mounted partitions

def fix_spaces(string):
    return string.replace("\\040", "\\ ")

def get_links(dev):
	a=os.listdir('/dev')
	for b in a:
		if S_ISLNK(os.lstat('/dev/'+b)[0]) and os.readlink('/dev/'+b)==dev:
			return ' /dev/' + b
	return ''
	
def get_drives():
	mounted.clear()	# clear the global drive dictionaries
	drives.clear()		# They are global so other routines can access them
	
	f=open('/etc/mtab', 'r')
	line=f.readline()
	while line:# quick find out which devices are mounted
		values=line.split()
		if "/dev/" in values[0]: # only interested in /dev/ devices
			mounted[values[0]]=values[1:3]
			if "/mapper/" in values[0]: # add in /dev/mapper entries
				drives[values[0]]="mounted on \033[31m"+values[1]+"\033[0m " \
				+"type "+values[2]
		line=f.readline()
	f.close()

	f=open('/proc/swaps','r')
	line=f.readline()
	if line:
		line=f.readline() # skip the file header
	while line: # read in all the swaps
		swaps=line.split()
		if '/mapper/' in line:
			drives[swaps[0]]=str(round(float(swaps[2]) / 1048576.0,1))+"G swap"
		mounted[swaps[0]]=['swap', 'swap']
		line=f.readline()
	f.close()

	devices=os.listdir('/sys/block/') # get list of block devices
	for device in devices:
		if "ram" not in device and "loop" not in device and "dm-" not in device:
			devicepath="/dev/"+device
			sysblock='/sys/block/'+device
			mntmsg=output=''
			try:
				removable=''
				v=open(sysblock+'/removable','r')
				r=v.readline()[:-1].strip()
				v.close()
				if r=='1':
					removable='\033[34m'
					if '-r' in sys.argv[1:]:# suppress removable drives -r
						continue
					#try: get symlinks, if any
						#res=os.popen4('find /dev -maxdepth 1 -lname '+device+'|head -n 1')
						#output=' ' +res[1].read().strip()
					output=get_links(device)
					if devicepath in mounted.keys(): # mounted?
						mntmsg="\033[0m mounted on \033[31m"+mounted[devicepath][0] \
						+"\033[0m type "+mounted[devicepath][1] 
				else:
					if '-n' in sys.argv[1:]:# suppress non-removable drives -n
						continue
			except:
				pass
			# get device size
			try:
				f=open(sysblock+'/size','r')
				rawsize=size=f.readline().strip()
				f.close()
			except:
				size=0

			if int(size) > 2097151: # calculate drive sizes
				size=str(round(float(size) /2048/1024.0,0))[:-2]+"G"
			else:
				size=str(round(float(size) /2048.0,1))+"M"
			if size=='0.0M':
				if 'fd' in device:
					output='Empty'+output
				else:
					output='Blank'+output
			else:
				if devicepath not in mounted.keys() and 'sr' in device:
					output='Empty'+output
				else:
					output=size+output
			try:
				for line in open(sysblock+'/device/vendor','r'):
					vendor=line.strip()
			except:
				vendor="Unknown Vendor"
				try:
					for line in open(sysblock+'/device/modalias','r'):
						vendor=line.strip()
				except:
					pass
			try:
				v=open(sysblock+'/device/model','r')
				model=v.readline()[:-1].strip()
				v.close()
			except:
				if removable:#on unknown removables, set model = devtype
					try:
						v=open(sysblock+'/uevent','r')
						for x in xrange(3):
							devtype=v.readline()
						model=devtype.split('=')[1][:-1]# strip off trailing \n
						v.close()
					except: #detect defective floppy drive hardware
						model="missing, check cable"
				else: #yield up some kind of message for debugging
					model="Unknown Model"
				
			output+=" "+removable+vendor+" "+model+mntmsg+"\033[0m"
			drives[devicepath]=[output]
			# Pull in any partitions and add them to the list
			subs=os.listdir(sysblock)
			for sub in subs:
				if device in sub: # "sda1" (sub) belongs to "sda" (device)
					partpath='/dev/'+sub
					output=''
					try: # get size
						f=open(sysblock+'/'+sub+'/size','r')
						size=f.readline()
						f.close()
					except:
						size=0
					if int(size) > 1048576: # calculate drive sizes
						output=str(round(float(size) /2048/1024.0,0))[:-2]+"G"
					else:
						output=str(round(float(size) /2048.0,0))[:-2]+"M"
					if partpath in mounted.keys(): # mounted?
						output+=" mounted on \033[31m"+mounted[partpath][0] \
						+"\033[0m type "+mounted[partpath][1] 
					else: # check if it is a dm-x device
						try:
							d=os.listdir(sysblock+'/'+sub \
							+'/holders')
							if d:
								output+=" lvm (/dev/mapper)"
								for e in d:
									output+=" "+e
								output+=" ..." #hmm how do we xref this?
						except:
							pass
					drives[devicepath].append(partpath+" "+output)

def ppdl(): # pretty print drive list
	message="Drive\t Size\t Description"
	if '-r' not in sys.argv[1:]:
		message+="\t (Removables are highlighted in blue)"
	print message
	drivelist=sorted(drives.keys())
	mappers=[] # print /dev/mapper entries last
	for drive in drivelist:
		if "/mapper/" in drive:
			mappers.append("... "+drive+"\033[0m "+drives[drive])
			continue
		else: # print drives in bold
			icon="\033[4m\033[1m"
		print icon+drive+"\033[0m "+fix_spaces(drives[drive][0])
		lastpart=len(drives[drive][1:])
		for p in xrange(1,lastpart+1):# print partition tree underneath
			if p<lastpart:
				icon=" ├─"
			else:
				icon=" └─"
			print icon+fix_spaces(drives[drive][p])
		print #blank line after each drive/partition tree
	if '-m' not in sys.argv[1:] and '-n' not in sys.argv[1:]:
		for m in mappers:
			print fix_spaces(m)

if '-h' in sys.argv[1:]:# print help messages -h
	name=sys.argv[0].split('/')[-1]
	print "Usage: "+name+" [options]"
	print " -i interactive mode"
	print " -r hide removable drives"
	print " -m hide /dev/mapper devices"
	print " -n hide non-removable drives, implies -m"
	print " -v print version information and exit"
	print " -h this help message"
	exit(0)
if '-v' in sys.argv[1:]:# print version and exit -v
	print version
	exit(0)
	
os.system('clear');
if '-i' in sys.argv[1:]:# Interactive Mode
	m='valid'
	while 1:
		get_drives()
		ppdl()
		print "Interactive mode: Enter Q to quit. [Enter] to refresh."
		m=raw_input("dev/partition# to (un)mount/eject? (highlight and middle-click) ")
		if m=='q' or m=='Q':
			break
		try: # eject cd roms
			if '\033[34m' in drives[m][0] and 'sr' in m:
				os.system('eject -T '+m)
				continue
		except:
			pass
		if m in mounted.keys(): # mount / unmount whatever is left
			os.system('gnome-mount -dutv '+m)
		else:
			os.system('gnome-mount -dtv '+m)
		os.system('clear');

else:
	get_drives()
	ppdl() # pretty print drive list