2 files changed, 61 insertions, 8 deletions
diff --git a/retiolum/scripts/adv_graphgen/tinc_stats/Geo.py b/retiolum/scripts/adv_graphgen/tinc_stats/Geo.py
index ba3f6738..038ca9c0 100755
--- a/retiolum/scripts/adv_graphgen/tinc_stats/Geo.py
+++ b/retiolum/scripts/adv_graphgen/tinc_stats/Geo.py
@@ -1,13 +1,54 @@
 #!/usr/bin/python3
 # -*- coding: utf8 -*-
-import sys,json
-from Graph import delete_unused_nodes
+import sys,json,os
+from Graph import delete_unused_nodes,resolve_myself
+GEODB=os.environ.get("GEOCITYDB","GeoLiteCity.dat")
 
-if __name__ == "__main__":
+def add_geo(nodes):
+  from pygeoip import GeoIP
+  gi = GeoIP(GEODB)
+
+  for k,v in nodes.iteritems():
+    try:
+      nodes[k].update(gi.record_by_addr(v["external-ip"]))
+    except Exception as e:
+      sys.stderr.write(str(e))
+      sys.stderr.write("Cannot determine GeoData for %s\n"%k)
+
+  return nodes
+def add_coords_to_edges(nodes):
   from pygeoip import GeoIP
-  gi = GeoIP("GeoLiteCity.dat")
-  for node,data in delete_unused_nodes(json.load(sys.stdin)).items():
+  gi = GeoIP(GEODB)
+
+  for k,v in nodes.iteritems():
+    for i,j in enumerate(v.get("to",[])):
+      data=gi.record_by_addr(j["addr"])
+      try:
+        j["latitude"]=data["latitude"]
+        j["longitude"]=data["longitude"]
+      except Exception as e: pass
+
+  return nodes
+
+def add_jitter(nodes):
+  from random import random
+  #add a bit of jitter to all of the coordinates
+  max_jitter=0.005
+  for k,v in nodes.iteritems():
+    jitter_lat= max_jitter -random()*max_jitter*2
+    jitter_long= max_jitter -random()*max_jitter*2
     try:
-      print ("%s in %s"%(node,gi.record_by_addr(data["external-ip"])["city"]))
-    except:
-      print ("%s myself"%node)
+      v["latitude"]= v["latitude"] + jitter_lat
+      v["longitude"]= v["longitude"] + jitter_long
+      for nodek,node in nodes.iteritems():
+        for to in node['to']:
+          if to['name'] == k:
+            to['latitude'] = v["latitude"]
+            to['longitude'] = v["longitude"]
+    except Exception as e: pass
+  return nodes
+
+if __name__ == "__main__":
+  import json
+  nodes = add_jitter(add_coords_to_edges(add_geo(resolve_myself(delete_unused_nodes(json.load(sys.stdin))))))
+  print (json.dumps(nodes))
diff --git a/retiolum/scripts/adv_graphgen/tinc_stats/Graph.py b/retiolum/scripts/adv_graphgen/tinc_stats/Graph.py
index da0ed828..18c3d545 100644
--- a/retiolum/scripts/adv_graphgen/tinc_stats/Graph.py
+++ b/retiolum/scripts/adv_graphgen/tinc_stats/Graph.py
@@ -7,6 +7,18 @@ from Availability import get_node_availability
 import sys,json
 from time import time
 DUMP_FILE = "/krebs/db/availability"
+
+
+def resolve_myself(nodes):
+  #resolve MYSELF to the real ip
+  for k,v in nodes.iteritems():
+    if v["external-ip"] == "MYSELF":
+      for nodek,node in nodes.iteritems():
+         for to in node['to']:
+           if to['name'] == k:
+             v["external-ip"] = to["addr"]
+  return nodes
+
 def dump_graph(nodes):
   from time import time
   graph = {}