#!/usr/bin/env python

# Working on rocky 9
# depends on python3-geoip2.noarch and geolite2-country.noarch

import os
# Signal handling for filter application: don't panic
import signal
import socket
import time
import sys
import re
from pprint import pprint as pp

def lookup_country_code(ip):
    database_fn = '/usr/share/GeoIP/GeoLite2-Country.mmdb'
    try:
        import geoip2.database
    except ImportError:
        exit("Please install: python3-geoip2.noarch and geolite2-country.noarch")
    if not os.path.isfile(database_fn):
        exit(f"Missing databasse: {database_fn}")
    with geoip2.database.Reader(database_fn) as reader:
        try:
            r = reader.country(ip)
            return r.country.iso_code
        except:
            return None

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("-c", "--colour", dest="colour", action="store_true", default=False, help="Force colour escapes")
optionparser.add_option("-u", "--update", dest="update", action="store_true", default=False, help="update")
(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:
    exit('dnf update python3-geoip2.noarch and geolite2-country.noarch')

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]))\.');

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

def dodns(ip):
    global options
    global start 
    global skipdns
    global knownips
    if options.replace: return ''
    if options.dns==0: return tab+blue+ip+unblue
    if knownips.get(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))
    countrycode = lookup_country_code(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()