dnmapper/bin/conv.py

#!/usr/bin/env python
import argparse
import os
import sys
sys.path.append("..")
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "dnmapper.settings")
import django
django.setup()

from backend.crawler import convert_crawl, net_to_json
from bgpdata.models import CrawlRun


def _convert_crawl(crawl):
    net = convert_crawl(crawl)
    if net.nodes and net.edges:
        crawl.graph = net_to_json(net)
        crawl.save()
        print("Crawl {} updated".format(crawl.id))
    else:
        print("Crawl {} had no nodes or edges, abort".format(crawl.id))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--crawl-id", type=int)
    parser.add_argument("-a", "--all", default=False, action="store_true")
    parser.add_argument("-e", "--empty-graph-only", default=False, action="store_true")
    args = parser.parse_args()

    if args.crawl_id and args.all:
        parser.error("-c and -a don't work together")

    if args.crawl_id:
        try:
            crawl = CrawlRun.objects.get(pk=args.crawl_id)
        except CrawlRun.DoesNotExist:
            parser.error("CrawlRun with id {} does not exist".format(args.crawl_id))
        _convert_crawl(crawl)
    elif args.all:
        if args.empty_graph_only:
            crawls = CrawlRun.objects.filter(graph='')
        else:
            crawls = CrawlRun.objects.all()

        for crawl in crawls:
            _convert_crawl(crawl)
    else:
        parser.error("Either specify a crawl with -c or use -a for all")


if __name__ == '__main__':
    main()