Pushing nodes to a database and generating graph from it

web/database.py

@@ -0,0 +1,80 @@
+import MySQLdb as mdb
+from graph import Node, Edge
+import time
+
+
+class NodeDB:
+	def __init__(self, config):
+		self.con = mdb.connect(
+			config['MYSQL_DATABASE_HOST'],
+			config['MYSQL_DATABASE_USER'],
+			config['MYSQL_DATABASE_PASSWORD'],
+			config['MYSQL_DATABASE_DB'])
+		self.cur = self.con.cursor()
+
+	def __enter__(self):
+		return self
+
+	def __exit__(self, type, value, traceback):
+		self.con.commit()
+		self.con.close()
+
+
+
+	def insert_node(self, node):
+		now = int(time.time())
+		self.cur.execute('''
+			INSERT INTO nodes (ip, name, version, first_seen, last_seen)
+			VALUES (%s, %s, %s, %s, %s)
+			ON DUPLICATE KEY
+			UPDATE name = %s, version = %s, last_seen = %s''', (
+			node.ip, node.label, node.version, now, now,
+			node.label, node.version, now))
+
+	def insert_edge(self, edge):
+		now = int(time.time())
+		self.cur.execute('''
+			INSERT INTO edges (a, b, first_seen, last_seen)
+			VALUES (%s, %s, %s, %s)
+			ON DUPLICATE KEY
+			UPDATE last_seen = %s''', (
+			edge.a.ip, edge.b.ip, now, now,
+			now))
+
+	def insert_graph(self, nodes, edges):
+		for n in nodes.itervalues():
+			self.insert_node(n)
+
+		for e in edges:
+			self.insert_edge(e)
+
+
+
+	def get_nodes(self, time_limit):
+		since = int(time.time() - time_limit)
+		cur = self.con.cursor(mdb.cursors.DictCursor)
+		cur.execute("SELECT ip, version, name FROM nodes WHERE last_seen > %s", (since,))
+		db_nodes = cur.fetchall()
+
+		nodes = dict()
+		for n in db_nodes:
+			nodes[n['ip']] = Node(n['ip'], n['version'], n['name'])
+
+		return nodes
+
+	def get_edges(self, nodes, time_limit):
+		since = int(time.time() - time_limit)
+		cur = self.con.cursor(mdb.cursors.DictCursor)
+		cur.execute("SELECT a, b FROM edges WHERE last_seen > %s", (since,))
+		db_edges = cur.fetchall()
+
+		edges = []
+		for e in db_edges:
+			edges.append(Edge(nodes[e['a']], nodes[e['b']]))
+
+		return edges
+
+	def get_graph(self, time_limit):
+		nodes = self.get_nodes(time_limit)
+		edges = self.get_edges(nodes, time_limit)
+		return (nodes, edges)

web/database.sql

@@ -0,0 +1,31 @@
+CREATE DATABASE  IF NOT EXISTS `fc00`;
+USE `fc00`;
+
+
+--
+-- Table structure for table `edges`
+--
+
+DROP TABLE IF EXISTS `edges`;
+CREATE TABLE `edges` (
+  `a` varchar(39) NOT NULL,
+  `b` varchar(39) NOT NULL,
+  `first_seen` int(11) NOT NULL,
+  `last_seen` int(11) NOT NULL,
+  PRIMARY KEY (`a`,`b`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8;
+
+
+--
+-- Table structure for table `nodes`
+--
+
+DROP TABLE IF EXISTS `nodes`;
+CREATE TABLE `nodes` (
+  `ip` varchar(39) NOT NULL,
+  `name` varchar(64) DEFAULT NULL,
+  `version` int(11) DEFAULT NULL,
+  `first_seen` int(11) NOT NULL,
+  `last_seen` int(11) NOT NULL,
+  PRIMARY KEY (`ip`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8;

web/graph.py

@@ -0,0 +1,29 @@
+class Node:
+	def __init__(self, ip, version=None, label=None):
+		self.ip = ip
+		self.version = version
+		self.label = ip[-4:] if label == None else label
+
+	def __lt__(self, b):
+		return self.ip < b.ip
+
+	def __repr__(self):
+		return 'Node(ip="%s", version=%s, label="%s")' % (
+			self.ip,
+			self.version,
+			self.label)
+
+class Edge:
+	def __init__(self, a, b):
+		self.a, self.b = sorted([a, b])
+
+	def is_in(self, edges):
+		for e in edges:
+			if e.a.ip == self.a.ip and e.b.ip == self.b.ip:
+					return True
+		return False
+
+	def __repr__(self):
+		return 'Edge(a.ip="%s", b.ip="%s")' % (
+			self.a.ip,
+			self.b.ip)

web/graphData.py

@@ -1,9 +1,35 @@
 import json
+from database import NodeDB
+from graph import Node, Edge
 
-def insert_graph_data(json_str):
+def insert_graph_data(config, json_str):
 	try:
 		graph_data = json.loads(json_str)
 	except ValueError:
 		return False
 
+	nodes = dict()
+	edges = []
+
+	if not 'nodes' in graph_data or not 'edges' in graph_data:
+		return False
+	
+
+	try:
+		for n in graph_data['nodes']:
+			node = Node(n['ip'], version=n['version'])
+			nodes[n['ip']] = node
+
+		for e in graph_data['edges']:
+			edge = Edge(nodes[e['a']], nodes[e['b']])
+			edges.append(edge)
+
+	except TypeError:
+		return False
+
+	print "Received %d nodes and %d links." % (len(nodes), len(edges))
+
+	with NodeDB(config) as db:
+		db.insert_graph(nodes, edges)
+
 	return True

web/graphPlotter.py

@@ -0,0 +1,74 @@
+import pygraphviz as pgv
+import time
+import json
+
+
+def position_nodes(nodes, edges):
+	G = pgv.AGraph(strict=True, directed=False, size='10!')
+
+	for n in nodes.values():
+		G.add_node(n.ip, label=n.label, version=n.version)
+
+	for e in edges:
+		G.add_edge(e.a.ip, e.b.ip, len=1.0)
+
+	G.layout(prog='neato', args='-Gepsilon=0.0001 -Gmaxiter=100000')
+
+	return G
+
+
+
+def get_graph_json(G):
+	max_neighbors = 1
+	for n in G.iternodes():
+		neighbors = len(G.neighbors(n))
+		if neighbors > max_neighbors:
+			max_neighbors = neighbors
+	print 'Max neighbors: %d' % max_neighbors
+
+	out_data = {
+		'created': int(time.time()),
+		'nodes': [],
+		'edges': []
+	}
+
+	for n in G.iternodes():
+		neighbor_ratio = len(G.neighbors(n)) / float(max_neighbors)
+		pos = n.attr['pos'].split(',', 1)
+
+		out_data['nodes'].append({
+			'id': n.name,
+			'label': n.attr['label'],
+			'version': n.attr['version'],
+			'x': float(pos[0]),
+			'y': float(pos[1]),
+			'color': _gradient_color(neighbor_ratio, [(100, 100, 100), (0, 0, 0)]),
+			'size': neighbor_ratio
+		})
+
+	for e in G.iteredges():
+		out_data['edges'].append({
+			'sourceID': e[0],
+			'targetID': e[1]
+		})
+
+	return json.dumps(out_data)
+
+
+
+
+
+def _gradient_color(ratio, colors):
+	jump = 1.0 / (len(colors) - 1)
+	gap_num = int(ratio / (jump + 0.0000001))
+
+	a = colors[gap_num]
+	b = colors[gap_num + 1]
+
+	ratio = (ratio - gap_num * jump) * (len(colors) - 1)
+
+	r = a[0] + (b[0] - a[0]) * ratio
+	g = a[1] + (b[1] - a[1]) * ratio
+	b = a[2] + (b[2] - a[2]) * ratio
+
+	return '#%02x%02x%02x' % (r, g, b)

web/updateGraph.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+from flask import Config
+from database import NodeDB
+import graphPlotter
+
+
+def generate_graph(time_limit=60*60*3):
+	nodes, edges = load_graph_from_db(time_limit)
+
+	graph = graphPlotter.position_nodes(nodes, edges)
+	json = graphPlotter.get_graph_json(graph)
+
+	with open('static/graph.json', 'w') as f:
+		f.write(json)
+
+
+def load_graph_from_db(time_limit):
+	config = Config('./')
+	config.from_pyfile('web_config.cfg')
+
+	with NodeDB(config) as db:
+		return db.get_graph(time_limit)
+
+
+if __name__ == '__main__':
+	generate_graph()

web/web.py

@@ -2,7 +2,7 @@ from flask import Flask, render_template, request
 from graphData import insert_graph_data
 
 app = Flask(__name__)
-app.debug = False
+app.config.from_pyfile('web_config.cfg')
 
 
 @app.context_processor
@@ -21,8 +21,10 @@ def page_about():
 
 @app.route('/sendGraph', methods=['POST'])
 def page_sendGraph():
+	print "Receiving graph from %s" % (request.remote_addr)
+	
 	data = request.form['data']
-	ret = insert_graph_data(data)
+	ret = insert_graph_data(app.config, data)
 	if ret:
 		return 'OK'
 	else:

web/web_config.example.cfg

@@ -0,0 +1,8 @@
+DEBUG = False
+
+MYSQL_DATABASE_HOST     = 'localhost'
+MYSQL_DATABASE_PORT     = 3306
+MYSQL_DATABASE_USER     = 'fc00'
+MYSQL_DATABASE_PASSWORD = 'hunter2'
+MYSQL_DATABASE_DB       = 'fc00'
+MYSQL_DATABASE_CHARSET  = 'utf8'