Traffic counter for i3blocks

trafficcount.py

#!/usr/bin/env python
import os, sys
FILE=os.path.join(os.environ.get('HOME', ""), '.cache/trafficcount')
DEVICE=os.environ.get('BLOCK_INSTANCE','wlp13s0')
FILE2=os.path.join(os.environ.get('HOME', ""), '.cache/trafficcount_limit')

def parseSize(s):
	TABLE={	'K':2**10,
			'M':2**20,
			'G':2**30,
			'T':2**40,
			'P':2**50,
			'E':2**60}
	n=s.strip().upper()
	m=1
	if n[-1] in TABLE:
		m=TABLE[n[-1]]
		n=n[:-1]
	return int(m*float(n))

def getLimit():
	try:
		with open(FILE2) as f:
			return parseSize(f.readline()) or None
	except KeyError: return None

def setLimit(n):
	print >>open(FILE2, 'wt'), int(n)

def getLastCount():
	try:
		r,t = map(int, open(FILE).read().strip().split())
		a,b = diff(getCount(), (r,t))
		if a<0 or b<0:
			setLastCount(0,0)
			return 0,0
		return r,t
	except Exception: return 0,0

def setLastCount(r,t):
	print >>open(FILE,'wt'), '%i %i'%(r,t)

def getCount():
	L=open('/proc/net/dev').read().splitlines()[2:]
	for l in L:
		l=l.split()
		if DEVICE in l[0]:
			r,t = map(int, (l[1], l[9]))
			return r,t
	return None,None

def diff((a,b),(c,d)):
	return a-c, b-d

def strsize(sz):
        assert sz>=0
        tbl=' KMGTPE'
        for p in xrange(7):
                if p!=6 and sz > 1.1*2**(10*(p+1)): continue
                fsz=sz/(2.0**(10*p))
                return "%.2f%cb"%(fsz,tbl[p])

if os.environ.get('BLOCK_BUTTON') == '3':
	r,t = getCount()
	if r is not None:
		setLastCount(r,t)
elif os.environ.get('BLOCK_BUTTON') == '1':
	if os.access(FILE2, 0):
		os.remove(FILE2)
	else:
		#open(FILE2, 'wt').close()
		if not os.system('which zenity >/dev/null 2>/dev/null'):
			try:
				f=os.popen('zenity --entry --text="Set network traffic limit:"')
				setLimit(parseSize(f.readline().strip() or '0'))
			except: pass
			finally: f.close()
		else:
			try:
				setLimit(parseSize(sys.argv[1]))
			except: pass
if os.access(FILE2, 0):
	LIMIT=getLimit()
	n = r,t = diff(getCount(), getLastCount())
	print ' '.join(map(strsize, n))
	print ' '.join(map(strsize, n))
	if LIMIT is None:
		print
	elif r+t < LIMIT:
		print '#00FF00'
	else:
		print '#FF0000'
else:
	print 'UL'
	print 'UL'
	print ''