#!/usr/bin/env python

import argparse
import itertools
import json
import pydot
import random
import time
import urllib.request
from dns import resolver

parser = argparse.ArgumentParser(description='''
Draw a graph from RIS BGP data. A PNG file will be generated with name
'YYYY-MM-DD-HH-MM-route.png'.
During the generation, all the paths are printed to the console in case you want
to see prepending and such.
''')
parser.add_argument('--route', dest='route', type=str,
                    help='IPv4 or IPv6 range present in the DFZ')

args = parser.parse_args()
print(args.route)

def pairwise(iterable):
    a, b = itertools.tee(iterable)
    next(b, None)
    return zip(a, b)

paths = {}

ris_json = urllib.request.urlopen('https://stat.ripe.net/data/looking-glass/data.json?resource={}'.format(args.route))

ris_data = json.load(ris_json)
prefix = ris_data['data']['rrcs'][0]['peers'][0]['prefix']
for rrc in ris_data['data']['rrcs']:
    for peer in rrc['peers']:
        path = []
        for asn in peer['as_path'].split():
            path.append(asn)
        path.append(prefix)
        paths[peer['peer']] = path

#json_graph = json.JSONEncoder().encode(paths)
nodes = {}
edges = {}

dot_graph = pydot.Dot('BGPMAP', graph_type='digraph')

resolv = resolver.Resolver()
resolv.timeout = 0.5

for peer in paths:
    print(peer, paths[peer])

    string_prepended_path = ['AS'+v for v in paths[peer][:-1]]
    print(string_prepended_path)

    for previous_asn, asn in pairwise(paths[peer]):

        try:
            as_name = resolv.query(f"AS{previous_asn}.asn.cymru.com", "TXT")
            as_name = as_name.response.answer[0].to_text().split('|')[4]
            as_name = as_name.strip().replace("'","").replace('"','')
        except Exception as e:
            print(e, previous_asn)
            as_name = ''

        print(as_name)
        edge_tuple = (previous_asn, asn)
        color = "#%x" % random.randint(0, 16777215)

        if edge_tuple not in edges:
            edge = pydot.Edge(*edge_tuple)
            edge.set_style("dashed")
            edge.set_color(color)
            dot_graph.add_edge(edge)
            edges[edge_tuple] = edge

        if previous_asn not in nodes:
            nodes[previous_asn] = pydot.Node(previous_asn)
            label = 'AS' + previous_asn + '\n' + as_name
            nodes[previous_asn].set_label(label)
            dot_graph.add_node(nodes[previous_asn])

current_time = time.localtime()
filename = '{}-{}-{}-{}-{}-{}.png'.format(
        "{:04d}".format(current_time.tm_year),
        "{:02d}".format(current_time.tm_mon),
        "{:02d}".format(current_time.tm_mday),
        "{:02d}".format(current_time.tm_hour),
        "{:02d}".format(current_time.tm_min),
        args.route.replace('/', '-')
    )
dot_graph.write_png(filename)