#! /usr/bin/env python
# -*- coding: utf-8 -*-

import json
import requests
from pushover import init, Client
from elasticsearch import Elasticsearch
from elasticsearch_dsl import Search
from elasticsearch_dsl import UpdateByQuery

pushMessage = ""

# function to send message through Pushover
# -----------------------------------------
pushTitle = "GeoIP Lookup Fixes - Firewall"
def pushMessageSend(msg,title):
    "Send Message to Pushover"
    userKey = "[PUSHOVER-USER-KEY]"
    applKey = "[PUSHOVER_APPLICATION_KEY]"
    init(applKey)
    Client(userKey).send_message(msg, title=title)
    return []

# test HTTP connection to ipgeolocation.io URL
# --------------------------------------------
urlLookup = "https://api.ipgeolocation.io/ipgeo"
accountKey = "[IPGEOLOCATION-ACCOUNT-KEY]"
src_ip =  "8.8.8.8"

try:
    testConnection = requests.get(urlLookup)
except requests.exceptions.ConnectionError as error:
    pushMessage += "Connection failed for URL: " + urlLookup + "\n" + str(error) + "\n\n"
    pushMessage += "Processing stopped - quitting!"
    pushMessageSend(pushMessage,pushTitle)
    quit()

# test IP lookup using ipgeolocation.io API
# ------------------------------------------
testURL = urlLookup + "?apiKey=" + accountKey + "&ip=" + src_ip
testResponse = requests.get(testURL)
if testResponse.ok:
    pass  #Test request succeeded for test IP 'src_ip'
else:
    pushMessage += "Test request failed for: " + src_ip + "\n\n"
    pushMessage += "Service Unavailable: " + str(testResponse.status_code) + " " + testResponse.reason + " " + testResponse.content + "\n\n"
    pushMessage += "Processing stopped - quitting!"
    pushMessageSend(pushMessage,pushTitle)
    quit()

# prepare for Elasticsearch-DSL search
# ------------------------------------
ipDict = {}
idList = []
newTagsList = []

client = Elasticsearch("localhost:9200")
index = "firewall-*"

tagOK = "GeoIP"
tagFailed = "_geoip_lookup_failure"
tagFixed = "_geoip_lookup_fixed"

# Elasticsearch-DSL search for geoip lookup failures
# --------------------------------------------------
geoipFailDocs = Search(using=client, index=index).filter('term', tags=tagFailed)

# loop through query hits & build dict of unique IPs
# --------------------------------------------------
i = 0
tagNeedsUpdating = True
for hit in geoipFailDocs.scan():
    i += 1
    id = hit.id
    src_ip = hit.src_ip
    tags = hit.tags

    if ipDict.get(src_ip) is None:       #add new key & initiate id list
        idList = [id]
        tmpDict = {src_ip: idList}
        ipDict.update(tmpDict)
    else:                                #append to id list for existing ip
        idList = ipDict.get(src_ip)
        idList.append(id)
        tmpDict = {src_ip: idList}
        ipDict.update(tmpDict)

    if tagNeedsUpdating:
        for tagU in hit.tags:
            if tagU != tagFailed:
                tag = tagU.encode('utf-8') # convert to string
                newTagsList.append(tag)
        newTagsList.append(tagFixed)
        tagNeedsUpdating = False

# initial counts
# --------------
pushMessage +=  "Total events = " + str(i) +  "\n"
pushMessage +=  "Unique IPs = " + str(len(ipDict.keys())) +  "\n"

# quit if no failed GeoIP lookups
# -------------------------------
if i < 1:
    pushMessage +=  "\n"
    pushMessage += "No IPs have failed lookups...\n"
    pushMessage += "Processing stopped - quitting!"
    pushMessageSend(pushMessage,pushTitle)
    quit()

# loop through dict of unique IPs
# -------------------------------
nSucceeded = 0
nFailed = 0
n = 0
docsIPCountSum = 0
for src_ip in ipDict.keys():
    n += 1
    idList = ipDict.get(src_ip)
    idListCount = len(idList)

    # ================================================================================================
    # Get GeoIP info from ipgeolocation.io
    # ------------------------------------
    # curl 'https://api.ipgeolocation.io/ipgeo?apiKey=a06329ec66074fbf88a96af7647126a9&ip=194.26.29.120'
    #
    # values for resp = requests.get(url)
    # -----------------------------------
    # resp.apparent_encoding  resp.cookies    resp.history               resp.iter_lines  resp.raise_for_status  resp.status_code
    # resp.close              resp.elapsed    resp.is_permanent_redirect resp.json        resp.raw               resp.text
    # resp.connection         resp.encoding   resp.is_redirect           resp.links       resp.reason            resp.url
    # resp.content            resp.headers    resp.iter_content          resp.ok          resp.request
    # ================================================================================================
    # t     geoip.city_name
    # t     geoip.continent_code
    # t     geoip.country_name
    # t     geoip.country_code2
    # t     geoip.country_code3
    # t     geoip.ip
    # #     geoip.latitude
    # GP    geoip.location
    # #     geoip.longitude
    # t     geoip.postal_code
    # t     geoip.region_code
    # t     geoip.region_name
    # t     geoip.timezone
    # ================================================================================================

    # try connecting through requests API and throw error & quit if it fails
    # -----------------------------------------------------------------------
    url = urlLookup + "?apiKey=" + accountKey + "&ip=" + src_ip
    resp = requests.get(url)
    if resp.ok:
        respJSON = json.loads(resp.text)
        city_name = respJSON['city']
        continent_code = respJSON['continent_code']
        country_name = respJSON['country_name']
        country_code2 = respJSON['country_code2']
        country_code3 = respJSON['country_code3']
        ip = respJSON['ip']
        postal_code = respJSON['zipcode']
        region_code = respJSON['zipcode']
        region_name = respJSON['state_prov']
        timezone = respJSON['time_zone'].get('name')
        latitude = float(respJSON['latitude'].encode('utf-8'))
        longitude = float(respJSON['longitude'].encode('utf-8'))
        location = {"lat": latitude, "lon": longitude}

        source = "\
        ctx._source.geoip.city_name = params.city_name; \
        ctx._source.geoip.continent_code = params.continent_code; \
        ctx._source.geoip.country_name = params.country_name; \
        ctx._source.geoip.country_code2 = params.country_code2; \
        ctx._source.geoip.country_code3 = params.country_code3; \
        ctx._source.geoip.ip = params.ip; \
        ctx._source.geoip.latitude = params.latitude; \
        ctx._source.geoip.longitude = params.longitude; \
        ctx._source.geoip.location = params.location; \
        ctx._source.geoip.region_code = params.region_code; \
        ctx._source.geoip.region_name = params.region_name; \
        ctx._source.geoip.timezone = params.timezone; \
        ctx._source.tags = params.newTags"

        params = {\
        'city_name': city_name, \
        'continent_code': continent_code, \
        'country_name': country_name, \
        'country_code2': country_code2, \
        'country_code3': country_code3, \
        'ip': ip, \
        'latitude': latitude, \
        'longitude': longitude, \
        'location': location, \
        'region_code': region_code, \
        'region_name': region_name, \
        'timezone': timezone, \
        'newTags': newTagsList}

        # use Elasticsearch-DSL UpdateByQuery to add GeoIP data
        # -----------------------------------------------------
        docsUBQ = UpdateByQuery().using(client).index(index).query("match", src_ip=src_ip).filter("term", tags=tagFailed).script(source=source, params=params)
        response = docsUBQ.execute()
        nSucceeded += 1
    else:
        nFailed += 1
        pushMessage += "Request failed for: " + src_ip + " with status: " + str(resp.status_code) + " " + resp.reason + "\n"

pushMessage +=  "Failed IPs = " + str(nFailed) + "\n"
pushMessage +=  "\n"
pushMessage += "Connection made to: " + urlLookup + "\n"

pushMessageSend(pushMessage,pushTitle)