#!/usr/bin/env python

#
# NetmonScreenlet - <http://ashysoft.blogspot.com/>
#
#	Copyright 2007-2008 Paul Ashton
#
#	This file is part of NetmonScreenlet.
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    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, see <http://www.gnu.org/licenses/>.
#
#
# INFO:
# - Displays a nice bandwidth monitor
#
# TODO:
# - bandwidth stats with hourly, daily, weekly, monthly etc.
# - SNMP support
# - Alarm when hit a certain amount of data
# - automatic detection of available interfaces
# - smooth scaling changes
#
# DONE:
# - Line graph
# - screenlet resizing
# - Graph height set to specific amount, 10meg etc.
#



import gtk, screenlets
from screenlets.options import BoolOption, ColorOption, StringOption, IntOption, FloatOption
from screenlets.options import create_option_from_node
import cairo, pango, gobject, re, os, math, datetime



def niceSpeed( bytes, mbit=False ):
	if mbit:
		bits = bytes * 8
		if bits < 1000: return "%s b" % bits
		elif bits < 1000000: return "%.1f Kb" % (bits / 1000.0)
		elif bits < 1000000000: return "%.2f Mb" % (bits / 1000000.0)
		elif bits < 1000000000000: return "%.3f Gb" % (bits / 1000000000.0)
		else: return "%.3f Tb" % (bits / 1000000000000.0)
	else:
		if bytes < 1024: return "%s B" % bytes
		elif bytes < (1024 * 1024): return "%.1f KB" % (bytes / 1024.0)
		elif bytes < (1024 * 1024 * 1024): return "%.2f MB" % (bytes / 1024.0 / 1024.0)
		elif bytes < (1024 * 1024 * 1024 * 1024): return "%.3f GB" % (bytes / 1024.0 / 1024.0 / 1024.0)
		else: return "%.3f TB" % (bytes / 1024.0 / 1024.0 / 1024.0 / 1024.0)


#use gettext for translation
import gettext

_ = screenlets.utils.get_translator(__file__)

def tdoc(obj):
	obj.__doc__ = _(obj.__doc__)
	return obj

@tdoc
class NetmonScreenlet(screenlets.Screenlet):
	"""Shows a nice bandwidth graph on your desktop"""
	
	# default meta-info for Screenlets
	__name__ = 'NetmonScreenlet'
	__version__ = '0.2.1'
	__author__ = 'Paul Ashton (c) 2007-2008'
	__website__ = 'http://ashysoft.blogspot.com/'
	__desc__ = __doc__

	screenletWidth = 120

	# internals
	width1 = 120
	height1 = 50
	__updateTimer = None
	__guiTimer = None
	__saveTimer = None
	p_layout = None
	lastBytesReceived = 0
	lastBytesSent = 0
	graphUp = [0 for x in range(width1/2+2)]
	graphDown = [0 for x in range(width1/2+2)]
	firstRun = True
	frame = 0
	graph_scale = 1

	peak = 0
	averageUp = 0
	averageDown = 0
	totalSessionUp = 0
	totalSessionDown = 0
	todayUp = 0
	todayDown = 0
	stats = [0,0]

	# settings
	gui_update_interval = 200 #ms
	save_interval = 1 #mins
	device_name = 'eth0'
	text_color = [1.0,1.0,1.0,1.0]
	down_color = [1.0,0.0,0.0,1.0]
	up_color = [0.0,1.0,0.0,1.0]
	back_color = [0.0,0.0,0.0,1.0]
	draw_textshadow = True
	draw_avgbars = True
	graphOverscale = 10
	graphScaleFixed = 0
	showBits = False
	lineGraph = True
	lineGraphFill = True
	lineWidth = 1
	showDownload = True
	showUpload = True
	
	#Texts...
	textTopLeft = ""
	textTopCenter = ""
	textTopRight = "%down% / %up%"
	
	textLeftTop = ""
	textLeftCenter = ""
	textLeftBottom = "%peak%"
	
	textBottomLeft = ""
	textBottomCenter = ""
	textBottomRight = ""
	
	textRightTop = ""
	textRightCenter = ""
	textRightBottom = ""




	# constructor
	def __init__(self, **keyword_args):
		screenlets.Screenlet.__init__(self, width=self.screenletWidth, height=50, uses_theme=False, **keyword_args)

		self.add_default_menuitems()
		self.add_options_group(_('Netmon Options'), _('Netmon Options'))
		#self.add_option(IntOption(_('Netmon Options'), 'screenletWidth', self.screenletWidth, _('Screenlet Width'), _('Width of the screenlet in pixels'),min=1,max=4096,increment=2), realtime=False)
		self.add_option(StringOption(_('Netmon Options'), 'device_name', self.device_name, _('Device Name'), _('Your network device name (Default: eth0)')), realtime=False)
		self.add_option(IntOption(_('Netmon Options'), 'width1', self.width1, _('Width'), _('Width of the screenlet'),min=10, max=1000),realtime=False)
		self.add_option(IntOption(_('Netmon Options'), 'height1', self.height1, _('Height'), _('Height of the screenlet'),min=10, max=1000),realtime=False)
		self.add_option(BoolOption(_('Netmon Options'), 'showDownload', self.showDownload, _('Show Download?'), _('When enabled graph will show download data')))
		self.add_option(BoolOption(_('Netmon Options'), 'showUpload', self.showUpload, _('Show Upload?'), _('When enabled graph will show upload data')))
		self.add_option(ColorOption(_('Netmon Options'), 'text_color', self.text_color, _('Text Color'), _('Color of the text')))
		self.add_option(ColorOption(_('Netmon Options'), 'down_color', self.down_color, _('Download Color'), _('Color of the download graph')))
		self.add_option(ColorOption(_('Netmon Options'), 'up_color', self.up_color, _('Upload Color'), _('Color of the upload graph')))
		self.add_option(ColorOption(_('Netmon Options'), 'back_color', self.back_color, _('Background Color'), _('Color of the graph background')))
		self.add_option(BoolOption(_('Netmon Options'), 'draw_textshadow', self.draw_textshadow, _('Text Shadow?'), _('Apply a shadow to the text')))
		self.add_option(BoolOption(_('Netmon Options'), 'draw_avgbars', self.draw_avgbars, _('Show average bars'), _('Paints lines showing average down/upload')))

		self.add_options_group(_('More Netmon Options'), _('More Netmon Options'))
		self.add_option(IntOption(_('More Netmon Options'), 'graphOverscale', self.graphOverscale, _('Graph Overscale Percent'), _('Overscale the graph by this percentage'),min=0, max=100))
		self.add_option(IntOption(_('More Netmon Options'), 'graphScaleFixed', self.graphScaleFixed, _('Graph Scale (in bytes)'), _('Scale of graph (bytes)\n125000 = 1 Megabit\n\nSet to zero for auto-scaling'), min=0, max=1000000000000, increment=125000 ))
		self.add_option(BoolOption(_('More Netmon Options'), 'showBits', self.showBits, _('Display bits'), _('When enabled show amounts in bits instead of bytes')))
		self.add_option(BoolOption(_('More Netmon Options'), 'lineGraph', self.lineGraph, _('Display as a Line Graph?'), _('When enabled show graph in the form of a line')))
		self.add_option(BoolOption(_('More Netmon Options'), 'lineGraphFill', self.lineGraphFill, _('Fill the Line Graph?'), _('When enabled fill the graph with your chosen color')))
		self.add_option(IntOption(_('More Netmon Options'), 'lineWidth', self.lineWidth, _('Line width'), _('Adjusts the thickness of the line'), min=0, max=5 ))

		self.add_options_group(_('Text Options'), _('Available Options:\n%up%, %down%, %peak%, %scale%, %avgup%,\n%avgdown%, %sessup%, %sessdown%,\n%todayup%, %todaydown%'))
		self.add_option(StringOption(_('Text Options'), 'textTopLeft', self.textTopLeft, _('Top Left'), ''))
		self.add_option(StringOption(_('Text Options'), 'textTopCenter', self.textTopCenter, _('Top Center'), _('Default: %down% / %up%')))
		self.add_option(StringOption(_('Text Options'), 'textTopRight', self.textTopRight, _('Top Right'), ''))
		self.add_option(StringOption(_('Text Options'), 'textBottomLeft', self.textBottomLeft, _('Bottom Left'), ''))
		self.add_option(StringOption(_('Text Options'), 'textBottomCenter', self.textBottomCenter, _('Bottom Center'), ''))
		self.add_option(StringOption(_('Text Options'), 'textBottomRight', self.textBottomRight, _('Bottom Right'), ''))
		self.add_option(StringOption(_('Text Options'), 'textLeftTop', self.textLeftTop, _('Left Top'), ''))
		self.add_option(StringOption(_('Text Options'), 'textLeftCenter', self.textLeftCenter, _('Left Center'), _('Default: %peak%')))
		self.add_option(StringOption(_('Text Options'), 'textLeftBottom', self.textLeftBottom, _('Left Bottom'), ''))
		self.add_option(StringOption(_('Text Options'), 'textRightTop', self.textRightTop, _('Right Top'), ''))
		self.add_option(StringOption(_('Text Options'), 'textRightCenter', self.textRightCenter, _('Right Center'), ''))
		self.add_option(StringOption(_('Text Options'), 'textRightBottom', self.textRightBottom, _('Right Bottom'), ''))

		self.gui_update_interval = self.gui_update_interval
		self.save_interval = self.save_interval



	# attribute-"setter", handles setting of attributes
	def __setattr__(self, name, value):
		# call Screenlet.__setattr__ in baseclass (ESSENTIAL!!!!)
		screenlets.Screenlet.__setattr__(self, name, value)
		# check for this Screenlet's attributes, we are interested in:
		if name == "gui_update_interval":
			self.__dict__['gui_update_interval'] = value
			if self.__guiTimer: gobject.source_remove(self.__guiTimer)
			self.__guiTimer = gobject.timeout_add(value, self.updateGUI)
		if name == "save_interval":
			self.__dict__['save_interval'] = value
			if self.__saveTimer: gobject.source_remove(self.__saveTimer)
			self.__saveTimer = gobject.timeout_add(value * 60000, self.saveStats)
		if name == "width1":
			self.width = value
			self.graphUp = [0 for x in range(self.width/2+2)]
			self.graphDown = [0 for x in range(self.width/2+2)]
			print "Changed screenlet width to %s" % value
		if name == "height1":
			self.height = value



	def updateGUI(self):
		self.redraw_canvas()
		return True



	def saveStats(self):
		date = str(datetime.datetime.now())[:10]
		filename = "%s/stats_%s.dat" % (self.get_screenlet_dir(),self.device_name)
		print "Saving stats to %s.." % filename
		FILE = None

		if not os.path.exists(filename):
			print "stats.dat does not exist, creating one.."
			FILE = open(filename,"w")
			today = [[ date, 0, 0 ]]
			FILE.write(str(today))
			FILE.close()
			print "stats.dat created."

		try:
			FILE = open(filename,"r")
		except:
			print "Couldn't open stats.dat for reading"
			return True

		if FILE:
			fileIn = FILE.read()
			FILE.close()

			lines = eval(fileIn)

			found = False
			for i in range(len(lines)):
				if date in lines[i]:
					#print "Found today: " + str(lines[i])
					lines[i][1] += self.stats[0]
					lines[i][2] += self.stats[1]
					self.todayDown = lines[i][1]
					self.todayUp = lines[i][2]
					print "Today stats: %s down, %s up" % (self.todayDown, self.todayUp)
					found = True
					break
			if not found:
				print _("Didn't find today stats, adding today..")
				lines.append([ date, self.stats[0], self.stats[1] ])
	
		FILE = open(filename,"w")
		FILE.write(str(lines))
		FILE.close()
		self.stats = [0,0]
		return True



	def addStats(self, down, up):
		self.graphDown.pop(0)
		self.graphDown.append(down)
		self.graphUp.pop(0)
		self.graphUp.append(up)

		self.totalSessionDown += down
		self.totalSessionUp += up

		self.stats[0] += down
		self.stats[1] += up

		#print "addStats() %s, %s" % (down, up)
		#print "Peak:%s Height:%s SDn:%s SUp:%s " % (self.peak, self.height, self.totalSessionDown, self.totalSessionUp)



	def updateInfo(self):
		devdata = os.popen( "cat /proc/net/dev | grep %s" % self.device_name ).readline()
		device = devdata[devdata.find(":")+1:] #get rid of anything before the colon
		device = re.findall("([\d]+)", device)

		downCounter = int(device[0])
		upCounter = int(device[8])

		if self.firstRun == True:
			self.firstRun = False
			self.lastBytesReceived = downCounter
			self.lastBytesSent = upCounter

		#Catch the bug with /proc/net/dev's 32bit byte counters *sigh*			
		#FIXME - Not accurate, loses any bytes before we hit 32bit ceiling
		if downCounter < self.lastBytesReceived:
			self.lastBytesReceived = 0
		if upCounter < self.lastBytesSent:
			self.lastBytesSent = 0

		downDiff = downCounter-self.lastBytesReceived
		upDiff = upCounter-self.lastBytesSent
		self.lastBytesReceived = downCounter
		self.lastBytesSent = upCounter

		self.addStats(downDiff,upDiff)
		#print "dC:%s uC:%s lBR:%s lBS:%s" % (downCounter, upCounter, self.lastBytesReceived, self.lastBytesSent)

		self.peak = max( max(self.graphUp), max(self.graphDown) )

		self.averageUp = int( sum(self.graphUp)/len(self.graphUp) )
		self.averageDown = int( sum(self.graphDown)/len(self.graphDown) )

		return True



	def getTextSize( self, ctx, text, size ):
		ctx.save()
		if self.p_layout == None: self.p_layout = ctx.create_layout()
		else: ctx.update_layout(self.p_layout)
		p_fdesc = pango.FontDescription()
		p_fdesc.set_family_static("Free Sans")
		p_fdesc.set_size(size*pango.SCALE)
		self.p_layout.set_font_description(p_fdesc)
		self.p_layout.set_markup(text)
		ctx.restore()
		return self.p_layout.get_pixel_size()



	def drawText( self, ctx, x, y, text, size, rgba, align=0, shadow=False ):
		ctx.save()
		if self.p_layout == None: self.p_layout = ctx.create_layout()
		else: ctx.update_layout(self.p_layout)
		p_fdesc = pango.FontDescription()
		p_fdesc.set_family_static("Free Sans")
		p_fdesc.set_size(size*pango.SCALE)
		self.p_layout.set_font_description(p_fdesc)
		self.p_layout.set_markup(text)
		textSize = self.p_layout.get_pixel_size()
		if align == 1: x = x - textSize[0]
		elif align == 2: x = x - textSize[0]/2
		if shadow:
			ctx.translate(x+0.5, y+0.5)
			ctx.set_source_rgba(0, 0, 0, rgba[3])
			ctx.show_layout(self.p_layout)
			ctx.fill()
			ctx.translate(-(x+1), -(y+1))
		ctx.translate(x, y)
		ctx.set_source_rgba(rgba[0], rgba[1], rgba[2], rgba[3])
		ctx.show_layout(self.p_layout)
		ctx.fill()
		ctx.restore()



	def on_draw(self, ctx):
		self.frame += 1
		if self.frame > 4: 
			self.frame = 0
			self.updateInfo()

		ctx.scale(self.scale, self.scale)
		ctx.set_operator(cairo.OPERATOR_OVER)

		#Scale
		if self.graphScaleFixed: tempscale = self.graphScaleFixed
		else: tempscale = self.peak
		self.graph_scale = float((tempscale+(tempscale/100*self.graphOverscale))/self.height)
		
		#Stop div0 errors..	
		if self.graph_scale == 0: self.graph_scale = 1
		
		#Draw background
		ctx.save()
		ctx.set_source_rgba(self.back_color[0],self.back_color[1],self.back_color[2],self.back_color[3])
		ctx.rectangle(0, 0, self.width, self.height)
		ctx.fill()

		ctx.set_line_width( self.lineWidth )

		offset = -0.4*self.frame
		ctx.translate( offset, 0 )

		if self.lineGraph:
			#Draw Download
			if self.showDownload:
				ctx.move_to( -2, self.height+2 ) #bottom left
				ctx.line_to( -2, self.height-int(self.graphDown[0]/self.graph_scale) )
				for i in range(len(self.graphDown)):
					height = int(self.graphDown[i]/self.graph_scale)
					ctx.line_to( (i*2), self.height-height )
				ctx.line_to( self.width+4, self.height-height )
				ctx.line_to( self.width+2, self.height+2 ) #bottom right
				ctx.set_source_rgba(self.down_color[0],self.down_color[1],self.down_color[2],self.down_color[3])
				ctx.set_operator(cairo.OPERATOR_ADD)
				if self.lineGraphFill:
					ctx.stroke_preserve()
					ctx.set_source_rgba(self.down_color[0],self.down_color[1],self.down_color[2],self.down_color[3]/2)
					ctx.fill()
				else:
					ctx.stroke()

			#Draw Upload
			if self.showUpload:
				ctx.move_to( -1, self.height )
				ctx.line_to( -1, self.height-int(self.graphUp[0]/self.graph_scale) )
				for i in range(len(self.graphUp)):
					height = int(self.graphUp[i]/self.graph_scale)
					ctx.line_to( i*2, self.height-height )
				ctx.line_to( self.width+2, self.height-height )
				ctx.line_to( self.width+2, self.height+2 ) #bottom right
				ctx.set_source_rgba(self.up_color[0],self.up_color[1],self.up_color[2],self.up_color[3])
				ctx.set_operator(cairo.OPERATOR_ADD)
				if self.lineGraphFill:
					ctx.stroke_preserve()
					ctx.set_source_rgba(self.up_color[0],self.up_color[1],self.up_color[2],self.up_color[3]/2)
					ctx.fill()
				else:
					ctx.stroke()
		else:
			if self.showDownload:
				#Draw Download
				ctx.set_operator(cairo.OPERATOR_SOURCE)
				for i in range(len(self.graphDown)):
					height = int(self.graphDown[i]/self.graph_scale)
					ctx.rectangle((i*2), self.height-height, 2, height)
				ctx.set_source_rgba(self.down_color[0],self.down_color[1],self.down_color[2],self.down_color[3])
				ctx.fill()

			if self.showUpload:
				#Draw Upload
				ctx.set_operator(cairo.OPERATOR_ADD)
				for i in range(len(self.graphUp)):
					height = int(self.graphUp[i]/self.graph_scale)
					ctx.rectangle((i*2), self.height-height, 2, height)
				ctx.set_source_rgba(self.up_color[0],self.up_color[1],self.up_color[2],self.up_color[3])
				ctx.fill()
			
		ctx.restore()

		#Draw average bars
		ctx.set_operator(cairo.OPERATOR_OVER)
		if self.draw_avgbars:
			if self.showDownload:
				ctx.set_source_rgba(self.down_color[0],self.down_color[1],self.down_color[2],self.down_color[3])
				ctx.rectangle(0, self.height-(self.averageDown/self.graph_scale), self.width, 1)
				ctx.fill()
			if self.showUpload:
				ctx.set_source_rgba(self.up_color[0],self.up_color[1],self.up_color[2],self.up_color[3])
				ctx.rectangle(0, self.height-(self.averageUp/self.graph_scale), self.width, 1)
				ctx.fill()

		#Draw top text
		if self.textTopLeft: self.drawText( ctx, 0, 0, self.fill_vars(self.textTopLeft), 8, self.text_color, 0, self.draw_textshadow )
		if self.textTopCenter: self.drawText( ctx, self.width/2, 0, self.fill_vars(self.textTopCenter), 8, self.text_color, 2, self.draw_textshadow )
		if self.textTopRight: self.drawText( ctx, self.width, 0, self.fill_vars(self.textTopRight), 8, self.text_color, 1, self.draw_textshadow )

		#Get text height
		h = self.getTextSize(ctx, "1234567890BMKGT", 8)[1]

		#Draw bottom text
		if self.textBottomLeft: self.drawText( ctx, 0, self.height-h, self.fill_vars(self.textBottomLeft), 8, self.text_color, 0, self.draw_textshadow )
		if self.textBottomCenter: self.drawText( ctx, self.width/2, self.height-h, self.fill_vars(self.textBottomCenter), 8, self.text_color, 2, self.draw_textshadow )
		if self.textBottomRight: self.drawText( ctx, self.width, self.height-h, self.fill_vars(self.textBottomRight), 8, self.text_color, 1, self.draw_textshadow )

		#Draw left text
		ctx.save()
		ctx.translate( 0, self.height )
		ctx.rotate( -1.57 ) #90 degrees ccw ;)
		if self.textLeftTop: self.drawText( ctx, self.height, 0, self.fill_vars(self.textLeftTop), 8, self.text_color, 1, self.draw_textshadow )
		if self.textLeftCenter: self.drawText( ctx, self.height/2, 0, self.fill_vars(self.textLeftCenter), 8, self.text_color, 2, self.draw_textshadow )
		if self.textLeftBottom: self.drawText( ctx, 0, 0, self.fill_vars(self.textLeftBottom), 8, self.text_color, 0, self.draw_textshadow )
		ctx.restore()

		#Draw right text
		ctx.save()
		ctx.translate( self.width, 0 )
		ctx.rotate( 1.57 ) #90 degrees cw ;)
		if self.textRightTop: self.drawText( ctx, 0, 0, self.fill_vars(self.textRightTop), 8, self.text_color, 0, self.draw_textshadow )
		if self.textRightCenter: self.drawText( ctx, self.height/2, 0, self.fill_vars(self.textRightCenter), 8, self.text_color, 2, self.draw_textshadow )
		if self.textRightBottom: self.drawText( ctx, self.height, 0, self.fill_vars(self.textRightBottom), 8, self.text_color, 1, self.draw_textshadow )
		ctx.restore()



	def on_draw_shape(self,ctx):
		ctx.scale(self.scale, self.scale)
		ctx.set_operator(cairo.OPERATOR_OVER)
		ctx.set_source_rgba(0,0,0,1)
		ctx.rectangle(0, 0, self.width, self.height)
		ctx.fill()



	def fill_vars(self, str):
		str = str.replace( "%up%", niceSpeed(self.graphUp[len(self.graphUp)-1], self.showBits) )
		str = str.replace( "%down%", niceSpeed(self.graphDown[len(self.graphDown)-1], self.showBits) )
		str = str.replace( "%peak%", niceSpeed(self.peak, self.showBits) )
		str = str.replace( "%scale%", niceSpeed(int(self.graph_scale*self.height), self.showBits) )
		str = str.replace( "%avgup%", niceSpeed(self.averageUp, self.showBits) )
		str = str.replace( "%avgdown%", niceSpeed(self.averageDown, self.showBits) )
		str = str.replace( "%sessup%", niceSpeed(self.totalSessionUp, self.showBits) )
		str = str.replace( "%sessdown%", niceSpeed(self.totalSessionDown, self.showBits) )
		str = str.replace( "%todayup%", niceSpeed(self.todayUp+self.stats[1], self.showBits) )
		str = str.replace( "%todaydown%", niceSpeed(self.todayDown+self.stats[0], self.showBits) )
		return str



# If the program is run directly or passed as an argument to the python
# interpreter then create a Screenlet instance and show it
if __name__ == "__main__":
	import screenlets.session
	screenlets.session.create_session(NetmonScreenlet)

