makeweather.py

#!/usr/bin/env python2
from __future__ import print_function, division, with_statement, generators

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>"
MRTG_FILE_PATH = "/var/www/mrtg/"
WEATHERMAP_CONFIG_PATH = "/etc/weathermap.conf"

WIDTH = 1750
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, zero_as_black=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(-3, 0, len(colors)))
        percentages[0] = 0
    if zero_as_black:
        percentages.insert(0, 0)
        colors.insert(0, (0, 0, 0))
    return list(zip(percentages, colors))

def convert_colorscale(datapoints):
    for min_data, max_data in zip(datapoints[:-1], datapoints[1:]):
        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="{:.1f}".format(min_value*100),
            red=int(r1*255), green=int(g1*255), blue=int(b1*255),
            max="{:.1f}".format(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)
    groups = {}
    for link in links:
        name = _make_name(link)
        groups[name] = groups.get(name, []) + [link]
    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():
        with open(MRTG_CFG, "r") as mrtg_file:
            return mrtg_file.read()
    links = []
    for match in re.findall(MRTG_PATTERN, _get_content(), re.MULTILINE):
        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))
        for node in nodes
    }

def place_circle(nodes, links):
    links_per_node = {node: 0 for node in nodes}
    for link in links:
        links_per_node[link.node_from] += 1
        links_per_node[link.node_to] += 1
    max_radius = min((WIDTH, HEIGHT)) * 0.45
    inner_radius = max_radius * 0.3
    outer_radius = max_radius * 0.9
    inner_threshold = 10
    inner_nodes = []
    outer_nodes = []
    for node in nodes:
        if links_per_node[node] >= inner_threshold:
            inner_nodes.append(node)
        else:
            outer_nodes.append(node)
    positions = {}
    def _calc_circ(angle, radius):
        return (int(math.sin(angle) * radius) + WIDTH//2,
                int(math.cos(angle) * radius) + HEIGHT//2)
    for index, node in enumerate(inner_nodes):
        angle = index / len(inner_nodes) * 2 * math.pi
        positions[node] = _calc_circ(angle, inner_radius)
    for index, node in enumerate(outer_nodes):
        angle = (index + 0.5) / len(outer_nodes) * 2 * math.pi
        positions[node] = _calc_circ(angle, outer_radius)
    return positions

def place_graphviz(nodes, links):
    import pygraphviz as gv
    graph = gv.AGraph()
    for node in nodes:
        graph.add_node(node)
    for link in links:
        graph.add_edge(link.node_from, link.node_to)
    graph.layout("neato")
    positions = {
        node.name: tuple(map(float, node.attr["pos"].split(",")))
        for node in graph.nodes()
    }
    MARGIN = 100
    xs = [position[0] for position in positions.values()]
    ys = [position[1] for position in positions.values()]
    minx, maxx = min(xs), max(xs)
    miny, maxy = min(ys), max(ys)
    scale = max(maxx-minx, maxy-miny)
    def _normalize(x, y):
        nx = int((x - minx) / scale * (WIDTH - 2 * MARGIN)) + MARGIN
        ny = int((y - miny) / scale * (HEIGHT - 2 * MARGIN)) + MARGIN
        return (nx, ny)
    return {
        key: _normalize(*positions[key]) 
        for key in positions
    }

PLACEMENT_STRATEGIES = {
    "random": place_random,
    "circle": place_circle,
    "graphviz": place_graphviz,
}

class ConfigWriter:
    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, width, height):
        self.write("TITLE Network Weathermap")
        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")
        self.write("BANDWIDTH 1G", 1)
        self.write("BWLABEL bits", 1)
        self.write("WIDTH 4", 1)
        self.write()
        self.write("NODE DEFAULT")
        self.write()

    def write_all(self, nodes, links, placement_strategy):
        positions = placement_strategy(nodes, links)
        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)
        for link_group in link_groups:
            self.write_links(link_group)

    def write_node(self, node, position):
        self.write("NODE {}".format(node))
        self.write("LABEL {}".format(node), 1)
        self.write("POSITION {} {}".format(position[0], position[1]), 1)
        self.write()

    def write_links(self, links):
        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(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):
        self.write("LINK {}".format(link.name))
        self.write("NODES {} {}".format(link.node_from, link.node_to), 1)
        # TODO: bandwidth
        self.write("TARGET {}".format(
            os.path.join(self.source_path, link.name + ".html")), 1)
        self.write()
    

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(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])