Update weathermap

fix bandwidths
fix portchannels
change colorscale to viridis
remove vlans
use log scale
cut borders away (with margin)

Black for exactly zero does not work, since only percent-precision is
available.

mrtg/defaults/main.yml

@@ -7,3 +7,5 @@ mrtg_switches:
 
 use_weathermap: yes
 weathermap_placement_strategy: "graphviz"
+weathermap_colorscale: "viridis"
+weathermap_colorscale_hash: "sha256:389c7a479cd64136ad5bf49daab59358437f69cf3c74cf74f958b093c7df50fd"

mrtg/files/makeweather.py

@@ -5,6 +5,7 @@ import re
 import random
 import os
 import math
+import numpy as np
 
 MRTG_CFG = "/etc/mrtg.cfg"
 MRTG_PATTERN = r"^PageTop\[(?P<name>[^\]]+)\]: <h1>(?P<description>[^\s]+)\s+(?P<to>[^\s]+)\s+--\s+(?P<from>[^\s<]+)</h1>"
@@ -17,17 +18,80 @@ HEIGHT = 1750
 def normalize_node_name(node):
     return node.split(".")[0]
 
+COLORMAP_PATTERN = r"([0-9]+) +([0-9.]+) +([0-9.]+) +([0-9.]+)"
+COLORMAP_HEX_PATTERN = r"([0-9]+) +'#([0-9A-F]+)'"
+
+def load_colorscale(filename):
+    with open(filename, "r") as file:
+        for line in file:
+            match = re.match(COLORMAP_PATTERN, line)
+            if match is not None:
+                groups = match.groups()
+                yield (1 - int(groups[0]) / 255, map(float, groups[1:]))
+            else:
+                match = re.search(COLORMAP_HEX_PATTERN, line)
+                if match is None:
+                    continue
+                groups = match.groups()
+                index = int(groups[0]) / 8
+                rgb = int(groups[1], base=16)
+                r = rgb // 2**16 / 255
+                g = rgb // 2**8 % 2**8 / 255
+                b = rgb % 2**8 / 255
+                yield index, (r, g, b)
+
+def reduce_colorscale(datapoints, max_points=15, log=True):
+    step = max(int(len(datapoints) / max_points), 1)
+    last_point = datapoints[-1]
+    colors = [point[1] for point in datapoints[::step]]
+    if last_point not in datapoints:
+        colors.append(last_point[1])
+    percentages = list(np.linspace(0, 1, len(colors)))
+    if log:
+        percentages = list(np.logspace(-2, 0, len(colors)))
+        percentages.insert(0, 0)
+        colors.insert(0, (0.5, 0.5, 0.5))
+    return list(zip(percentages, colors))
+
+def convert_colorscale(datapoints):
+    for min_data, max_data in zip(datapoints[:-1], datapoints[1:]):
+        min_data, max_data = sorted((min_data, max_data))
+        min_value, (r1, g1, b1) = min_data
+        max_value, (r2, g2, b2) = max_data
+        yield "SCALE {min} {max} {red} {green} {blue} {red2} {green2} {blue2}".format(
+            min=int(min_value*100),
+            red=int(r1*255), green=int(g1*255), blue=int(b1*255),
+            max=int(max_value*100),
+            red2=int(r2*255), green2=int(g2*255), blue2=int(b2*255))
+
+
 class Link:
     def __init__(self, name, description, node_from, node_to):
         self.name = name
         self.description = description
         self.node_from = normalize_node_name(node_from)
         self.node_to = normalize_node_name(node_to)
+        self.is_10g = self.description.startswith("TenGigabitEthernet")
+        self.is_portchannel = self.description.startswith("Port-channel")
 
     @staticmethod
     def from_match(match):
         return Link(*match)
 
+    def __repr__(self):
+        return "Link({}, {}, {}, {})".format(
+            self.name, self.description, self.node_from, self.node_to)
+
+    def get_speed(self):
+        if self.is_10g:
+            return 10
+        if self.is_portchannel:
+            return getattr(self, "speed", None)
+        return 1
+
+def get_total_speed(links):
+    return sum(link.get_speed() for link in links if link.get_speed() is not None)
+
 def group_links(links):
     def _make_name(link):
         return "{} {}".format(link.node_from, link.node_to)
@@ -35,7 +99,15 @@ def group_links(links):
     for link in links:
         name = _make_name(link)
         groups[name] = groups.get(name, []) + [link]
-    return groups.values()
+    result = []
+    for group in groups.values():
+        if any(link.is_portchannel for link in group):
+            channel = [link for link in group if link.is_portchannel][0]
+            channel.speed = max(get_total_speed(group), 1)
+            result.append([channel])
+        else:
+            result.append(group)
+    return result
 
 def create_mrtg_map():
     def _get_content():
@@ -43,13 +115,27 @@ def create_mrtg_map():
             return mrtg_file.read()
     links = []
     for match in re.findall(MRTG_PATTERN, _get_content(), re.MULTILINE):
-        links.append(Link.from_match(match))
+        link = Link.from_match(match)
+        if link.description.startswith("Vlan"):
+            continue
+        links.append(link)
     nodes = set()
     for link in links:
         nodes.add(link.node_from)
         nodes.add(link.node_to)
     return nodes, links
 
+def calculate_size(nodes, margin_x=70, margin_y=15):
+    xs = [pos[0] for pos in nodes.values()]
+    ys = [pos[1] for pos in nodes.values()]
+    min_xs = max(0, min(xs) - margin_x)
+    max_xs = min(WIDTH, max(xs) + margin_x)
+    min_ys = max(0, min(ys) - margin_y)
+    max_ys = min(HEIGHT, max(ys) + margin_y)
+    for key in nodes:
+        nodes[key] = (nodes[key][0] - min_xs, nodes[key][1] - min_ys)
+    return max_xs - min_xs, max_ys - min_ys
+
 def place_random(nodes, links):
     return {
         node: (random.randint(0, WIDTH-1), random.randint(0, HEIGHT-1))
@@ -105,7 +191,6 @@ def place_graphviz(nodes, links):
     def _normalize(x, y):
         nx = int((x - minx) / scale * (WIDTH - 2 * MARGIN)) + MARGIN
         ny = int((y - miny) / scale * (HEIGHT - 2 * MARGIN)) + MARGIN
-        print(nx, ny)
         return (nx, ny)
     return {
         key: _normalize(*positions[key]) 
@@ -119,18 +204,25 @@ PLACEMENT_STRATEGIES = {
 }
 
 class ConfigWriter:
-    def __init__(self, filehandler):
+    def __init__(self, filehandler, source_config=MRTG_CFG,
+            source_path=MRTG_FILE_PATH, colorscale=None):
         self.filehandler = filehandler
+        self.source_config = source_config
+        self.source_path = source_path
+        self.colorscale = colorscale
 
     def write(self, line=None, indent=0):
         line = line or ""
         self.filehandler.write("\t" * indent + line + "\n")
 
-    def write_header(self):
+    def write_header(self, width, height):
         self.write("TITLE Network Weathermap")
-        self.write("WIDTH {}".format(WIDTH))
-        self.write("HEIGHT {}".format(HEIGHT))
+        self.write("WIDTH {}".format(width))
+        self.write("HEIGHT {}".format(height))
         self.write("KEYPOS 0 0")
+        if self.colorscale is not None:
+            for line in self.colorscale:
+                self.write(line, 0)
         # BACKGROUND HERE
         self.write()
         self.write("LINK DEFAULT")
@@ -143,7 +235,8 @@ class ConfigWriter:
 
     def write_all(self, nodes, links, placement_strategy):
         positions = placement_strategy(nodes, links)
-        self.write_header()
+        width, height = calculate_size(positions)
+        self.write_header(width, height)
         for node in nodes:
             self.write_node(node, positions[node])
         link_groups = group_links(links)
@@ -160,8 +253,11 @@ class ConfigWriter:
         self.write("LINK {}".format(",".join(link.name for link in links)))
         self.write("NODES {} {}".format(links[0].node_from, links[0].node_to), 1)
         self.write("TARGET {}".format(" ".join(
-            os.path.join(MRTG_FILE_PATH, link.name + ".html")
+            os.path.join(self.source_path, link.name + ".html")
             for link in links)), 1)
+        total_speed = get_total_speed(links)
+        self.write("BANDWIDTH {}G".format(total_speed), 1)
+        self.write("WIDTH {}".format(6 if total_speed == 2 else (8 if total_speed == 20 else 4)), 1)
         self.write()
 
     def write_link(self, link):
@@ -169,7 +265,7 @@ class ConfigWriter:
         self.write("NODES {} {}".format(link.node_from, link.node_to), 1)
         # TODO: bandwidth
         self.write("TARGET {}".format(
-            os.path.join(MRTG_FILE_PATH, link.name + ".html")), 1)
+            os.path.join(self.source_path, link.name + ".html")), 1)
         self.write()
     
 
@@ -177,9 +273,23 @@ if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser()
     parser.add_argument("strategy", default="circle", choices=PLACEMENT_STRATEGIES.keys())
+    parser.add_argument("--source-config", default=MRTG_CFG)
+    parser.add_argument("--source-path", default=MRTG_FILE_PATH)
+    parser.add_argument("--target-config", default=WEATHERMAP_CONFIG_PATH)
+    parser.add_argument("--colorscale")
     arguments = parser.parse_args()
 
     nodes, links = create_mrtg_map()
-    with open(WEATHERMAP_CONFIG_PATH, "w") as config_file:
-        config_writer = ConfigWriter(config_file)
+    with open(arguments.target_config, "w") as config_file:
+        colorscale = arguments.colorscale
+        if colorscale is not None:
+            try:
+                colorscale = convert_colorscale(
+                    reduce_colorscale(list(load_colorscale(colorscale))))
+            except IOError:
+                colorscale = None
+        config_writer = ConfigWriter(config_file,
+            source_config=arguments.source_config,
+            source_path=arguments.source_path,
+            colorscale=colorscale)
         config_writer.write_all(nodes, links, PLACEMENT_STRATEGIES[arguments.strategy])

mrtg/tasks/weathermap.yml

@@ -5,6 +5,7 @@
   apt: name="{{item}}" state=present
   with_items:
     - python-pygraphviz
+    - python-numpy
   tags: weathermap
 
 - name: enable unpacking zip files and executing rotten php stuff
@@ -44,8 +45,20 @@
     line: '$rrdtool="/usr/sbin/nologin";'
   tags: weathermap
 
+- name: upload the weathermap script
+  copy: src=makeweather.py dest=/root/makeweather.py
+  tags: weathermap
+
+- name: get the weathermap colorscale
+  get_url:
+    dest: /root/colorscale.pal
+    url: "https://raw.githubusercontent.com/Gnuplotting/gnuplot-palettes/master/{{weathermap_colorscale}}.pal"
+    checksum: "{{weathermap_colorscale_hash}}"
+  when: weathermap_colorscale is not none
+  tags: weathermap
+
 - name: create the weathermap config
-  script: "makeweather.py {{weathermap_placement_strategy}}"
+  script: "makeweather.py {{weathermap_placement_strategy}} --colorscale /root/colorscale.pal"
   tags: weathermap
 
 - name: create the weathermap regularly