Skip to content
Snippets Groups Projects
warden-map.py 4.52 KiB
Newer Older
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# warden-map.py
#
# Copyright (C) 2016 Cesnet z.s.p.o
# Use of this source is governed by a 3-clause BSD-style license, see LICENSE file.


def getLastEvents(events, client, key, cert, cacert, secret):

  try:
    ses = Session()
    req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events)
    pre = req.prepare()
    res = ses.send(pre, cert = (cert, key), verify=cacert)
  except requests.exceptions.RequestException as error:
    print error
    sys.exit(1)

  data = res.json()
  i = 0
  eventsList = []
  for p in data['events']:
    event = {}
    if i < events:
      for key, value in { 'event': 'Category', 'time': 'DetectTime', 'origin': 'Source', 'destination': 'Target'}.iteritems():
        if value in p:
          if (key == 'origin') or (key == 'destination'):
            event[key] = {}
            if 'IP4' in p[value][0]:
              event[key]['ip'] = p[value][0]['IP4']
            else:
              event[key] = {}
          elif (key == 'event'):
            event[key] = ', '.join(p[value])
          else:
            event[key] = p[value]
        else:
          if (key == 'origin') or (key == 'destination'):
            event[key] = {}
          else:
            event[key] = {}
      if 'ip' in event['origin']:
        eventsList.append(event)
        i += 1
    else:
      break

  return eventsList

def getGeolocation(ip):

  try:
    response = requests.get('http://freegeoip.net/json/' + str(ip[0]))
  except requests.exceptions.RequestException as error:
    print error
    sys.exit(1)

  try:
    json_data = json.loads(response.text)
  except ValueError as error:
    print error
    sys.exit(1)

  return {'latitude': json_data['latitude'], 'longitude': json_data['longitude'], 'country_name': json_data['country_name'], 'city': json_data['city']}

def main(args):

  events = args.events[0]
  client = args.client[0]
  key    = args.key[0]
  cert   = args.cert[0]
  cacert = args.cacert[0]
  secret = args.secret[0]

  if args.output is not None:
    path = args.output[0] + 'warden-map.json'
  else:
    path = 'warden-map.json'

  wardenEvents = getLastEvents(events, client, key, cert, cacert, secret)

  for p in wardenEvents:
    for target in {'origin', 'destination'}:
      geoData = {}
      if 'ip' in p[target]:
        geoData = getGeolocation(p[target]['ip'])
        for value in {'latitude', 'longitude', 'country_name', 'city'}:
          if value in geoData:
            if not geoData[value]:
              p[target][value] = "???"
            else:
              p[target][value] = geoData[value]
          else:
            p[target][value] = "???"

      else:
        p[target]['ip'] = "???"
        p[target]['country_name'] = "Czech Republic"
        p[target]['city'] = "???"
        p[target]['latitude'] = 49.743
        p[target]['longitude'] = 15.338

  try:
    with open(path, 'w') as outfile:
      json.dump(wardenEvents, outfile)
  except IOError:
    print "Error: File does not appear to exist."
    sys.exit(1)

  return 0

if __name__ == '__main__':
  import sys
  import json
  from requests import Request, Session
  import requests
  import argparse

  parser = argparse.ArgumentParser(description='Creates warden-map.json for warden-map.html frontend.',
                                  formatter_class=lambda prog: argparse.HelpFormatter(prog,max_help_position=30))

  parser.add_argument('--output', metavar='path/', type=str,
                    nargs=1, help='path where warden-map.json should be saved')

  requiredNamed = parser.add_argument_group('required arguments')

  requiredNamed.add_argument('--events', metavar='<number>', type=str, required=True,
                    nargs=1, help='count of events for a map')
  requiredNamed.add_argument('--client', metavar='<org.ex.cl>', type=str, required=True,
                    nargs=1, help='client name')
  requiredNamed.add_argument('--key', metavar='path/key.pem', type=str, required=True,
                    nargs=1, help='SSL key for a client')
  requiredNamed.add_argument('--cert', metavar='path/cert.pem', type=str, required=True,
                    nargs=1, help='SSL cert for a client')
  requiredNamed.add_argument('--cacert', metavar='path/cacert.pem', type=str, required=True,
                    nargs=1, help='SSL cacert for a client')
  requiredNamed.add_argument('--secret', metavar='<SeCreT>', type=str, required=True,
                    nargs=1, help='secret key for a client')

  args = parser.parse_args()
  main(args)