sysvitals/api/server.py

#!/usr/bin/env python

import errno
import json
import os
from os import path
import struct
import warnings

warnings.filterwarnings('ignore', 'libevent')
import gevent.pywsgi

import reloader

DATA_DIR = path.expanduser('~/sysvitals_data')
handlers = None

def main():
	global handlers
	handlers = {
		'data': get_data,
		'datum': post_datum,
	}
	server = gevent.pywsgi.WSGIServer(('0.0.0.0', 8892), application)
	reloader.init(server)
	server.serve_forever()

def application(environ, start_response):
	split = environ['PATH_INFO'][1:].split('/')
	if split[0] == 'v1':
		handler = handlers.get(split[2])
		if handler:
			start_response('200 OK', [('Content-type', 'text/plain')])
			return [handler(split, environ)]
		else:
			print 'no handler for', split
	else:
		print 'split was', split

formats = {
	'cpu': 'f',
	'mem': 'q',
	'net': 'q',
	'disk': 'q',
}

def get_data(split, environ):
	group = int(split[1])
	client_id = int(split[3])
	data_path = path.join(DATA_DIR, str(group), str(client_id))
	with open(data_path, 'r') as f:
		buf = f.read()
	rval = {}
	index = 0
	while index < len(buf):
		# read the key
		key_size = ord(buf[index])
		fmt = '%dp' % (key_size + 1)
		data = struct.unpack(fmt, buf[index:index+key_size+1])
		key = data[0]

		# get the format_code
		split = key.split('.')
		stat_group = split[0]
		format_code = formats[stat_group]

		# read the data
		fmt = '%dp 1440%s' % (key_size + 1, format_code)
		size = struct.calcsize(fmt)
		data = struct.unpack(fmt, buf[index:index+size])
		dict_insert(rval, split, data[1:])
		index += size
	return json.dumps(rval)

def dict_insert(d, split, value):
	if len(split) > 1:
		d.setdefault(split[0], {})
		dict_insert(d[split[0]], split[1:], value)
	else:
		d[split[0]] = value

def post_datum(split, environ):
	group = int(split[1])
	body = json.load(environ['wsgi.input'])
	client_id = body['client_id']
	group_dir = path.join(DATA_DIR, str(group))
	try:
		os.makedirs(group_dir)
	except OSError as e:
		if e.errno != errno.EEXIST:
			raise
	with open(path.join(group_dir, str(client_id)), 'w') as f:
		for stat_group, stats in body.items():
			if stat_group == 'client_id':
				continue
			format_code = formats[stat_group]
			for stat, data in stats.items():
				key = '%s.%s' % (stat_group, stat)
				write_stat(f, key.encode('utf-8'), format_code, data)
	return '{"status": "ok"}'

def write_stat(f, key, format_code, data):
	fmt = '%dp 1440%s' % (len(key) + 1, format_code)
	array = [-1] * 1440
	array[0] = data
	f.write(struct.pack(fmt, key, *array))

main()