#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import os
import re
import GeoIP

arguments = sys.argv

if(len(arguments) == 4 and os.path.exists(arguments[1]) and re.search("^[0-9a-z\.]+$" ,arguments[2]) and re.search("^[0-9A-Z\.]+$" ,arguments[3])):
    fileName = arguments[1]
    searchArg = arguments[2]
    resultArg = arguments[3]
else:
    print('UNKNOWN - invalid arguments.\nusage: '+arguments[0]+' geolite_database_file search_string expected_result')
    sys.exit(3)

try:
    gi = GeoIP.open(fileName, GeoIP.GEOIP_MEMORY_CACHE)
    if re.search("^[A-Z]{2}$" ,resultArg):
        response = gi.country_code_by_addr(searchArg)
    elif re.search("^AS[0-9\.]+$" ,resultArg):
        response = gi.org_by_addr(searchArg)
    else:
        print('UNKNOWN - invalid arguments.\nusage: '+arguments[0]+' geolite_database_file search_string expected_result')
        sys.exit(3)
except:
    print('CRITICAL - Error, database may be corrupted.')
    sys.exit(2)

if(response == None):
    print('UNKNOWN - response is none.')
    sys.exit(3)

result = response.split()

if not re.search('^'+resultArg+'$', result[0]):
    print('WARNING - wrong response. '+searchArg+' is '+response)
    sys.exit(1)

print('OK')