#!/bin/sh

countries="cn kr pk tw sg hk pe br ua in tr vn kz it es ru pl hu ro th uy kr ar"
listurl="http://geolite.maxmind.com/download/geoip/database/GeoIPCountryCSV.zip"

me=$(cd $(dirname "$0") && pwd)/${0##*/}
base=${0##*/}; base=${base%.*}
cmdline="$@"

setname=geolite
logprefix=GEOIPSET-LOG
settype=nethash
hashsize=199000 # default 1024
hashresize=4  # default 50

drop=1
init=0
syslog=1
stderr=0

chain=$setname
etc=/opt/etc
zipfile=$etc/geolite.ipset.zip
networks=$etc/geolite.$settype
inifile=$etc/$base.ini

config=/etc/config
fire=$config/91."$setname".fire

alias awk=/usr/bin/awk
alias wget=/usr/bin/wget

[ -e $inifile ] && source $inifile

log() {
  [ "$stderr" = "1" ] && echo $base "$@" 1>&2
  [ "$syslog" = "1" ] && echo "$@" | logger -t $base.$$
}

csv2cidr() {
awk '
function ip2int(ip) {
        for (ret=0,n=split(ip,a,"\."),x=1;x<=n;x++) ret=or(lshift(ret,8),a[x])
        return ret
}

function int2ip(ip,ret,x) {
        ret=and(ip,255)
        ip=rshift(ip,8)
        for(;x<3;ret=and(ip,255)"."ret,ip=rshift(ip,8),x++);
        return ret
}

function compl32(v) {
        ret=xor(v, 0xffffffff)
        return ret
}

function deaggregate(start,end,     base, step)
{
  base=start
  while (base <= end) {
       step = 0
       while ( or(base, lshift(1, step)) != base) {
          if ( or(base, rshift((bits, (31-step)))) > end ) {
             break;
          }
          step=step+1
       }
       print int2ip(base)"/"(32-step)
       base = base + lshift(1, step)
  }
}

BEGIN { FS="[\",]"; bits=compl32(0);}

{
deaggregate($8, $11)
}
' #end awk
}

grepstr() {
   echo \\\"$@\\\" | sed -e 's/ /\\"\|\\"/g;'
}

getfile() {
  wget -q -O $zipfile.tmp "$listurl" && {
    log  getfile: success
    mv -f $zipfile.tmp $zipfile
    return 0
  } || {
    log getfile: wget failed
    rm -f $zipfile.tmp
    return 1
  }
}

checkupdate(){
  rm -f $zipfile.tmp
  host=$(echo $listurl | awk -F"/" '{print $3}')
  path=$(echo $listurl | awk -F"/" '{print substr($0, index($0,$4))}')
  day=86400
  [ -e $zipfile ] && {
    zipdate=$(unzip -l $zipfile | awk -F" "  '/GeoIPCountryWhois\.csv/ {split($2,d,"-"); print mktime("20"d[3]" "d[1]" "d[2]" 00 00 00")}')
  } || {
    log checkupdate: file missing - download needed
    return 0
  }

  lastmodheader=$(echo -e "HEAD /$path HTTP/1.1\r\nHost: $host\r\n\r\n" | nc -w30 $host 80 | tr -d '\r' | grep "Last-Modified")
  [ "$lastmodheader" = "" ] && {
     log checkupdate: bad header
     return 1
  } || {
     lastmoddate=$(echo $lastmodheader | \
     awk '
       BEGIN{
         m=split("Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec",d,"|")
         for(o=1;o<=m;o++){
           months[d[o]]=sprintf("%02d",o)
         }
       }
       { print mktime($5" " months[$4] " " $3 " 00 00 00") }
     ' #end awk
     )
     log checkupdate: zip:$zipdate lastmod:$lastmoddate age: $(( (lastmoddate - zipdate) / day ))
     [ $lastmoddate -lt $(( zipdate - day )) -o $lastmoddate -gt $(( $(date +%s) + day )) ] && {
       log checkupdate: bad mod date
       return 1
     } || {
       [ $(( (lastmoddate - zipdate) / day)) -gt 1 ] && {
         log checkupdate: download needed
         return 0
       } || {
         log checkupdate: no download needed
         return 1
       }
     }
  }
}

updatetree() {
  log updatetree: updating
  rm -f $networks
  checkupdate && getfile
  unzip -p $zipfile | \
  awk "BEGIN {FS=\"[,\\\"]\";OFS=\"-\";IGNORECASE=1;} /$(grepstr $countries)/ {print \$2,\$5}"  > $networks
  log updatetree: update done
}

gettree() {
  [ -e $networks ] || updatetree
  { echo "-N $setname.set.tmp iptreemap "
    cat $networks | sed -e "s/^ */-A $setname.set.tmp /"
    echo COMMIT
  } | ipset --restore
  ipset --swap $setname.set.tmp $setname.set
  ipset --destroy $setname.set.tmp
}

update() {
  log update: updating
  rm -f $networks
  checkupdate && getfile
  unzip -p $zipfile | \
  grep  -Ei  "$(grepstr $countries)" | \
  csv2cidr > $networks
  log update: update done
}

gethash() {
  [ -e $networks ] || update
  { echo "-N $setname.nets.tmp nethash --resize $hashresize --hashsize $hashsize"
    cat $networks | grep -v "/32" | sed -e "s/^ */-A $setname.nets.tmp /"
    echo COMMIT
  } | ipset --restore
  ipset --swap $setname.nets.tmp $setname.nets
  ipset --destroy $setname.nets.tmp

  { echo "-N $setname.hosts.tmp iphash"
    cat $networks | grep "/32" | cut -d"/" -f1  | sed -e "s/^ */-A $setname.hosts.tmp /"
    echo COMMIT
  } | ipset --restore
  ipset --swap $setname.hosts.tmp $setname.hosts
  ipset --destroy $setname.hosts.tmp
}

getset() {
  [ "$settype" = "nethash" ] && gethash || gettree
}

start() {
  [ "$settype" = "nethash" ] && modules=" ip_set_nethash ip_set_iphash ip_set_setlist " || modules=" ip_set_iptreemap "
  for m in ip_set $modules ipt_set
  do
    lsmod | grep -q $m || insmod $m
  done
  [ "$settype" = "nethash" ] && {
    ipset --swap $setname.set $setname.set 2>&1 | grep -qi "unknown set" && ipset --create $setname.set setlist
    ipset --swap $setname.nets $setname.nets 2>&1 | grep -qi "unknown set" && ipset --create $setname.nets nethash
    ipset --swap $setname.hosts $setname.hosts 2>&1 | grep -qi "unknown set" && ipset --create $setname.hosts iphash
    ipset --list $setname.set | grep -q "$setname.nets" || ipset --add $setname.set $setname.nets
    ipset --list $setname.set | grep -q "$setname.hosts" || ipset --add $setname.set $setname.hosts
  } || {
    ipset --swap $setname.set $setname.set 2>&1 | grep -qi "unknown set" && ipset --create $setname.set iptreemap
  }
  getset
  fire
}

fire() {
  logger -t $base.$$ Executing $logprefix firewall script
  [ ! -d $config ] && mkdir $config
  [ ! -e $fire ] && ln -s $me $fire
  iptables -nL | grep -q "$chain.* set $setname.set src" || {
    inputline=$(($(iptables --line-numbers -vnL INPUT | grep "ACCEPT.*state.*RELATED,ESTABLISHED" | cut -d" " -f1) + 1))
    iptables -N $chain
    [ "$drop" = "1" ] && iptables -I $chain -j DROP
    [ $(nvram get log_in) = 1 ] && {
      limit=$(nvram get log_limit)
      [ $limit = 0 ] && { iptables -I $chain -m state --state NEW -j LOG --log-prefix "$logprefix " --log-macdecode --log-tcp-sequence --log-tcp-options --log-ip-options
      } || iptables -I $chain -m state --state NEW -m limit --limit $limit/m -j LOG --log-prefix "$logprefix " --log-macdecode --log-tcp-sequence --log-tcp-options --log-ip-option
    }
    iptables -I INPUT $inputline -m set --set $setname.set src -j $chain
    iptables -I wanin -m set --set $setname.set src -j $chain
  }
}

stop() {
  logger -t $base.$$ Executing $logprefix stop script
  iptables -nL | grep -q "$chain.* set $setname.set src" && {
    iptables -D INPUT -m set --set $setname.set src -j $chain
    iptables -D wanin -m set --set $setname.set src -j $chain

    iptables -F $chain
    iptables -X $chain
  }
  ipset --swap $setname.set $setname.set 2>&1 | grep -qi "unknown set" || ipset --destroy $setname.set
  ipset --swap $setname.nets $setname.nets 2>&1 | grep -qi "unknown set" || ipset --destroy $setname.nets
  ipset --swap $setname.hosts $setname.hosts 2>&1 | grep -qi "unknown set" || ipset --destroy $setname.hosts
  rm -f $fire
}

showstats() {
  log "Begin Stats ================================== \
      $(echo ""
        echo options: $cmdline
        echo ""
        ipset --list $setname.set 2>&1 | head
        ipset --list $setname.nets 2>&1| head -n 4
        ipset --list $setname.nets 2>&1 | wc
        echo ""
        ipset --list $setname.hosts 2>&1 | head -n 4
        ipset --list $setname.hosts 2>&1 | wc
      )"
  log "End Stats ===================================="
}



[ "$me" == "$fire" ] && {
  fire
  exit
}

for p in $@
do
  case "$p" in
    "debug" )
      stderr=1 ;;
    "cron" )
      cru a $setname  "20 02 * * *  $me checkupdate" ;;
    "start" )
      init=1 ;;
    "stop" )
      stop
      cru d $setname
      init=0
      ;;
    "rebuild" )
      rm -f $networks
      init=1
      ;;
    "force" )
      rm -f $networks $zipfile
      init=1
      ;;
    "checkupdate" )
      checkupdate && {
        rm -f $networks
        init=1
      }
      ;;
  esac
done

showstats
[ "$init" = "1" ] && {
  start
  showstats
}

exit