#! /usr/bin/python

# Missing country data?
#
# apt-get install python-geop 
# curl http://geolite.maxmind.com/download/geoip/database/GeoLiteCountry/GeoIP.dat.gz > GeoIP.dat.gz &&
#   gzip -d < GeoIP.dat.gz > /usr/share/GeoIP/GeoIP.dat
#


import os
try:
    import GeoIP
except ImportError:
    os.system('sh /usr/local/afribin/maintenance/geoip_install2016')
    import GeoIP

import sys,re

# Signal handling for filter application: don't panic
import signal
import socket
import time

def doupdate():
    import os
    cmd="""
        cd ;
        curl http://geolite.maxmind.com/download/geoip/database/GeoLiteCountry/GeoIP.dat.gz > GeoIP.dat.gz &&
           gzip -d < GeoIP.dat.gz > /usr/share/GeoIP/GeoIP.dat
    """
    os.system(cmd)

try: signal.signal(signal.SIGPIPE, signal.SIG_DFL)
except: pass
try: signal.signal(signal.SIGINT, signal.SIG_DFL)
except: pass

specialcountries = ['ZA', ]

from optparse import OptionParser
optionparser = OptionParser()
optionparser.add_option("-d", "--dns", dest="dns", action="store_true", default=False, help="Resolve DNS")
optionparser.add_option("--timeout", dest="timeout", action="store", type="int", default=5, help="DNS resolution timeout")
optionparser.add_option("-r", "--replace", dest="replace", action="store_true", default=False, help="Replace IP with country or __")
optionparser.add_option("-t", "--tab", dest="tab", action="store_true", default=False, help="Separate country by tab, not space")
optionparser.add_option("-u", "--update", dest="update", action="store_true", default=False, help="update geoip db here")
optionparser.add_option("-c", "--colour", dest="colour", action="store_true", default=False, help="Force colour escapes")
(options, args) = optionparser.parse_args()
tab=' '
if options.tab: tab='\t'

line_buffered = False
if options.colour or sys.stdout.isatty():
    line_buffered = True
    cyan='\x1B[36m'
    red='\x1B[31m'
    unred='\x1B[0m'
    green='\x1B[32m'
    ungreen='\x1B[0m'
    blue='\x1B[34m'
    unblue='\x1B[0m'
else:
    cyan=''
    red=''
    unred=''
    green=''
    ungreen=''
    blue=''
    unblue=''


if options.update:
    doupdate()

ipaddressre = re.compile("((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?))( {,3})")
rfc1918re = re.compile(r'^(10|192\.168|172\.(1[6-9]|2[0-9]|3[01]))\.');

gi = GeoIP.new(GeoIP.GEOIP_STANDARD)

knownips={}
knowncc={}
start = time.time()
skipdns = False

def dodns(ip):
    global options
    global start 
    global skipdns
    if options.replace: return ''
    if options.dns==0: return tab+blue+ip+unblue
    if knownips.has_key(ip):
        name=knownips[ip]
    else:
        name=''
        if not skipdns:
            try:
                name = socket.gethostbyaddr(ip)
            except:
                pass
            if time.time() > start+options.timeout:
                skipdns=True
        knownips[ip]=name
    if name:
        return tab+blue+ip+unblue+' ('+green+name[0]+ungreen+')'
    else:
        return tab+blue+ip+unblue;

def cc(match):
    global gi
    global option
    global red,unred,cyan
    ip = match.group(1).strip()
    whitespace = match.group(6);
    if rfc1918re.search(ip):
        return ip+whitespace;
    countrycode = knowncc.setdefault(ip,gi.country_code_by_addr(ip))
    if countrycode in specialcountries:
        return cyan+countrycode+unred+dodns(ip)+whitespace
    if countrycode:
        return red+countrycode+unred+dodns(ip)+whitespace
    if options.replace:
        return red+"__"+unred+dodns(ip)+whitespace
    else:
        return red+"__"+unred+dodns(ip)+whitespace  # dodns(ip) is going to frob rfc1918 addresses for nothing
    
for arg in args:
    sys.stdout.write( ipaddressre.sub(cc,arg+"\n" ))
for line in sys.stdin:
    sys.stdout.write( ipaddressre.sub(cc,line) )
    sys.stdout.flush()