#!/usr/bin/env python
"""
Track and plot data using RRDtool.
Copyright 2006, 2008 Adam Sampson <ats@offog.org>

There are a good dozen frontends out there already, but none that do what I
want, so here's mine.

This requires the Net-SNMP Python bindings.
"""

import sys, iniparser, os, re, time, getopt, traceback, netsnmp

verbose = False

def run(args, hide_stdout = False, hide_stderr = False):
	if verbose:
		print "Running: " + " ".join(args)
	# This would use os.spawnvp(), except that "rrdtool graph" prints stuff
	# to stdout, so we need to be able to redirect its output to
	# /dev/null...
	pid = os.fork()
	if pid == 0:
		if hide_stdout or hide_stderr:
			fd = os.open("/dev/null", os.O_WRONLY)
			if hide_stdout:
				os.dup2(fd, 1)
			if hide_stderr:
				os.dup2(fd, 2)
			os.close(fd)
		os.execvp(args[0], args)
		sys._exit(42)
	(pid, status) = os.waitpid(pid, 0)
	if status >> 8 != 0:
		print "Command invocation failed: " + " ".join(args)
		sys.exit(1)

class UpdateError(Exception): pass

class Source:
	def __init__(self, name, config):
		self.name = name
		self.config = config
		self.data = None
		self.time = None

	def get_data(self):
		if self.data is not None:
			return self.data

		def read(f):
			s = f.read()
			f.close()
			return s

		if "command" in self.config:
			# command = COMMAND
			d = read(os.popen(self.config["command"]))
		elif "file" in self.config:
			# file = FILENAME
			d = read(open(self.config["file"]))
		elif "snmp" in self.config:
			# snmp = HOST COMMUNITY VAR1 VAR2 ...
			fields = self.config["snmp"].split()
			if len(fields) < 3:
				raise UpdateError("Source " + self.name + " has invalid SNMP syntax")
			results = []
			for v in fields[2:]:
				label = v

				# This is a bit of a hack: for some reason,
				# netsnmp will only match 1-digit interface
				# IDs, so we need to look out for them
				# ourselves in order to handle higher ones.
				m = re.match(r'^(.*)\.(\d+)$', v)
				if m is not None:
					(v, iid) = m.group(1, 2)
				else:
					iid = None

				var = netsnmp.Varbind(v, iid)
				r = netsnmp.snmpget(var,
				                    Version = 1,
				                    DestHost = fields[0],
				                    Community = fields[1])
				if r is None or r == [] or r[0] is None:
					raise UpdateError("Source " + self.name + " had no result for SNMP variable " + v)
				results.append(label + " " + r[0] + "\n")
			d = "".join(results)
		else:
			# Otherwise, use the name of the source as a filename.
			d = read(open(self.name))
		self.time = time.time()

		lines = [l.strip() for l in d.split("\n")]
		lines = [l for l in lines if l != ""]
		if "exclude" in self.config:
			lines = [l for l in lines if re.search(self.config["exclude"], l) is None]
		r = self.config.get("split", "\s+")
		self.data = [re.split(r, l) for l in lines]

		if self.data == []:
			raise UpdateError("Source " + self.name + " returned no data")
		return self.data

	def on_fail(self):
		if "onfail" in self.config:
			os.system(self.config["onfail"])
			return True
		else:
			return False

def match(pattern, fields):
	if pattern == "":
		return True
	(f, r) = pattern.split(None, 1)
	return re.search(r, fields[int(f)]) is not None

class Var:
	def __init__(self, name, config, get_source):
		self.name = name
		if len(name) + 2 > len("router-adsl-attenua"):
			raise UpdateError("Variable " + self.name + " has name too long for rrdtool")
		self.config = config
		self.get_source = get_source
		self.def_source = self.config.get("source")

	def get_fields(self):
		return self.config.get("fields", "0").split()

	def get_dsnames(self):
		nfields = len(self.get_fields())
		if nfields == 1:
			dsnames = [self.name]
		else:
			dsnames = ["%s_%d" % (self.name, i) for i in range(nfields)]
		return dsnames

	def get_labels(self):
		if "labels" in self.config:
			labels = self.config["labels"].split()
		else:
			labels = [self.name]
		if len(labels) != len(self.get_fields()):
			raise UpdateError("Variable " + self.name + " has different length fields and labels attributes")
		return labels

	def get_values(self):
		lines = []
		data_time = -1
		matches = self.config.getall("match")
		if matches == []:
			matches = [""]
		for m in matches:
			om = m
			source = self.def_source
			if m[0] == "<":
				ss = m.split(None, 1)
				source = ss[0][1:]
				if len(ss) == 2:
					m = ss[1]
				else:
					m = ""
			if source is None:
				raise UpdateError("Variable " + self.name + " match " + om + " has no source")
			so = self.get_source(source)

			for attempt in (0, 1):
				try:
					data = so.get_data()
					break
				except UpdateError, e:
					if attempt == 0 and so.on_fail():
						print >>sys.stderr, "Source " + source + " failed; retrying"
						pass
					else:
						raise e

			if so.time > data_time:
				data_time = so.time

			line = None
			for d in data:
				if match(m, d):
					line = d
					break
			if line is None:
				raise UpdateError("Variable " + self.name + " match " + om + " matched no data from " + repr(data))
			lines.append(line)

		results = []
		for exp in self.get_fields():
			def calc_value(f):
				fs = f.split(":")
				if len(fs) == 1:
					r = lines[0][int(fs[0])]
				elif len(fs) == 2:
					r = lines[int(fs[0])][int(fs[1])]
				else:
					raise UpdateError("Variable " + self.name + " has bad field spec: " + f)
				return r

			parts = re.split(r'(i\+|i-|f\+|f-)', exp)
			r = calc_value(parts[0])
			for i in range(1, len(parts), 2):
				op = parts[i]
				v = calc_value(parts[i + 1])
				if op == "i+":
					r = str(int(r) + int(v))
				elif op == "i-":
					r = str(int(r) - int(v))
				elif op == "f+":
					r = str(float(r) + float(v))
				elif op == "f-":
					r = str(float(r) - float(v))
			results.append(r)

		return (results, data_time)

def write_html(title, body, fn):
	title = "rrdpage: " + title
	fnn = fn + ".new"
	f = open(fnn, "w")
	f.write("""<!DOCTYPE HTML PUBLIC "-//W3C//DTD HTML 4.01//EN"
   "http://www.w3.org/TR/html4/strict.dtd">
<html lang="en">
<head>
    <meta http-equiv="Content-Type" content="text/html; charset=ISO-8859-1">
    <title>%s</title>
</head>
<body>
<div id="head">
<h1>%s</h1>
</div>
<div id="content">
%s
</div>
<div id="footer">
<p>Generated by <a href="http://offog.org/code/misccode.html">rrdpage</a>
by <a href="http://offog.org/">Adam Sampson</a>;
graphs by <a href="http://oss.oetiker.ch/rrdtool/">RRDtool</a></p>
</div>
</body>
</html>
""" % (title, title, body))
	f.close()
	os.rename(fnn, fn)

def usage():
	print "Usage: rrdpage [-v] [-u] [-w] config-file [var ...]"

def main(args):
	global verbose

	try:
		opts, args = getopt.getopt(args, "vuw")
	except getopt.GetoptError:
		usage()
		sys.exit(1)

	if len(args) < 1:
		usage()
		sys.exit(1)

	do_update = False
	do_write = False
	for (o, a) in opts:
		if o == "-v":
			verbose = True
		elif o == "-u":
			do_update = True
		elif o == "-w":
			do_write = True
	if not (do_update or do_write):
		do_update = True
		do_write = True

	ip = iniparser.Config(args[0])
	period = 60 * int(ip.default["period"])
	statedir = ip.default["statedir"]
	if not os.access(statedir, os.F_OK):
		os.makedirs(statedir)
	outdir = ip.default["outdir"]
	if not os.access(outdir, os.F_OK):
		os.makedirs(outdir)

	timeperiods = [
		("day", period, 60 * 60 * 24),
		("week", 60 * 30, 60 * 60 * 24 * 7),
		("year", 60 * 60 * 24, 60 * 60 * 24 * 365),
		("decade", 60 * 60 * 24 * 7, 60 * 60 * 24 * 365 * 10),
		]

	sources = {}
	def get_source(n):
		if not n in sources:
			opts = ip.get("source " + n, {})
			sources[n] = Source(n, opts)
		return sources[n]

	vars = {}
	for n in ip.keys():
		(t, name) = n.split(None, 1)
		if t != "var":
			continue

		vars[name] = Var(name, ip.get(n), get_source)

	if len(args) == 1:
		do_vars = vars.keys()
	else:
		do_vars = args[1:]
	do_vars.sort()

	for vn in do_vars:
		v = vars[vn]
		rrd = "%s/%s.rrd" % (statedir, v.name)
		if not os.access(rrd, os.F_OK):
			print rrd + " does not yet exist; creating it"

			rrdtype = v.config.get("type", "GAUGE")
			minvalue = v.config.get("minvalue", "U")
			maxvalue = v.config.get("maxvalue", "U")

			args = ["rrdtool", "create", rrd, "--start", "N-10"]

			for n in v.get_dsnames():
				args.append("DS:%s:%s:%d:%s:%s" % (n, rrdtype, period * 4, minvalue, maxvalue))

			for p in timeperiods:
				steps = p[1] / period
				rows = (p[2] / p[1]) + 1
				args.append("RRA:AVERAGE:0.5:%d:%d" % (steps, rows))

			run(args)

		if do_update:
			try:
				(values, data_time) = v.get_values()
				args = ["rrdtool", "update", rrd, "%d:" % (data_time,) + ":".join(values)]
				run(args)
			except UpdateError:
				print >>sys.stderr, "Exception while updating variable %s:" % (v.name,)
				traceback.print_exc(None, sys.stderr)
	if do_write:
		count = ["aaa", "aab", "aba", "abb", "baa", "bab", "bba", "bbb"]
		colours = []
		colours += ["#" + s.replace("a", "70").replace("b", "e0") for s in count]
		colours += ["#" + s.replace("a", "00").replace("b", "80") for s in count]

		now = time.time()

		index = []

		for vn in do_vars:
			v = vars[vn]
			body = []

			title = v.config.get("title", v.name)

			rrd = "%s/%s.rrd" % (statedir, v.name)
			for p in timeperiods:
				args = [
					"rrdtool", "graph",
					"%s/%s-%s.png" % (outdir, v.name, p[0]),
					"--title", "%s (last %s)" % (title, p[0]),
					"--start", "%d" % (now - p[2],),
					"--end", "%d" % (now,),
					]
				if "units" in v.config:
					args += ["--vertical-label", v.config["units"]]
				dsnames = v.get_dsnames()
				labels = v.get_labels()
				for n in dsnames:
					args.append("DEF:%s=%s:%s:AVERAGE" % (n, rrd, n))
				for i in range(len(dsnames)):
					mods = ""
					mode = "LINE1"
					if v.config.get("stacked", "no") == "yes":
						if i != 0:
							mods = ":STACK"
						mode = "AREA"
					label = labels[i]
					args.append("%s:%s%s:%s%s" % (mode, dsnames[i], colours[i], label, mods))

				# We hide both stdout (which gets the image
				# size printed to it) and stderr (which
				# occasionally gets a warning from libart).
				run(args, True, True)

				body += [
					"<h2>Last %s</h2>\n" % (p[0],),
					"""<p><img src="%s-%s.png" alt="%s, last %s"></p>\n""" % (v.name, p[0], title, p[0]),
					]

			write_html(title, "".join(body), "%s/%s.html" % (outdir, v.name))

			p = timeperiods[0]
			index += [
				"""<h2><a href="%s.html">%s</a></h2>\n""" % (v.name, title),
				"""<p><img src="%s-%s.png" alt="%s, last %s"></p>\n""" % (v.name, p[0], title, p[0]),
				]

		write_html("Index", "".join(index), "%s/index.html" % (outdir,))

if __name__ == "__main__":
	main(sys.argv[1:])