warden-map.py

#!/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.

import GeoIP
import codecs
import time

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

  ses = Session()
  #req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events + '&id=0')
  req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events)
  #req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?nocat=Other&client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events)
  pre = req.prepare()
  res = ses.send(pre, cert = (cert, key), verify=cacert)

  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'][0]
            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, db):

  data = db.record_by_addr(ip)

  if not data:
    return {}
  else:
    return {
      'latitude': data['latitude'], 
      'longitude': data['longitude'], 
      'country_name': unicode(data['country_name'], "utf-8") if data['country_name'] else None, 
      'city': unicode(data['city'], "utf-8") if data['city'] else None
    }

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'

  db = GeoIP.open("GeoLiteCity.dat", GeoIP.GEOIP_MEMORY_CACHE)
  db.set_charset(GeoIP.GEOIP_CHARSET_UTF8)

  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'], db)
        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

  wardenEvents.append(int(time.time()));

  with codecs.open(path, 'w', encoding="utf-8") as outfile:
    json.dump(wardenEvents, outfile, ensure_ascii = False)

  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)