#!/usr/bin/env python3

import gevent
from gevent import monkey, socket
monkey.patch_all()
import requests
import sys
from dnslib import DNSRecord
from collections import OrderedDict as odict
import time

#import logging
#logging.basicConfig(level=logging.DEBUG)

port = 53
url = "https://127.0.0.1/"
cache_max = 10000
cache_refresh = 3600
status_interval = 300

sess = requests.Session()
sess.mount('https://',
           requests.adapters.HTTPAdapter(pool_connections=50,
                                         pool_maxsize=75,
                                         max_retries=3))

cache = odict()
query_cached = 0
query_total = 0

def query(url, data, sock):
    global cache
    global query_total
    global query_cached
    try:
        query_total += 1
        DNSQuery = DNSRecord.parse(data)
        qname = DNSQuery.q.qname
        #print("===request=======================================")
        #print(DNSQuery)
        #print("=================================================")
        now = time.time()
        if qname in cache and (now - cache[qname][1]) < cache_refresh:
            print('query:', qname, "cached")
            query_cached += 1
            cache.move_to_end(qname)
            resp = DNSQuery.reply()
            resp.rr = cache[qname][0]
            resp = resp.pack()
        else:
            print('query:', qname, "remote")
            headers = {'accept': 'application/dns-message',
                       'content-type': 'application/dns-message'}
            r = sess.post(url, headers = headers, data = data, allow_redirects = False)
            if r.status_code != 200:
                raise Exception
            resp = r.content
            cache[qname] = (DNSRecord.parse(resp).rr, time.time())
            if len(cache) >= cache_max:
                cache = odict(list(cache.items())[cache_max//2:])

        #print("===response======================================")
        #print(DNSRecord.parse(resp))
        #print("=================================================")
        sock.send(resp)
        #clear retrans
        while True:
            _ = sock.recv(1500)
            sock.send(resp)
    except Exception as e:
        pass
    sock.close()

def forward_query(port, url):
    while True:
        sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM, 0)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(('', port))
        data, addr = sock.recvfrom(1500)
        sock.connect(addr)
        sock.settimeout(0)
        gevent.spawn(query, url, data, sock)
def cache_status():
    global query_total
    global query_cached
    while True:
        gevent.sleep(status_interval)
        status = "{} items cached, {} times cache hit in last {} seconds, {:.2f}% "
        print(status.format(len(cache),
                            query_cached,
                            status_interval,
                            query_total == 0 and 100 or query_cached/query_total*100))
        query_cached = 0
        query_total = 0

if __name__ == "__main__":
    if len(sys.argv) == 3 and sys.argv[1] == '-f':
        exec(open(sys.argv[2]).read())
    else:
        exec(open("/etc/forwardoh.conf").read())
    version = "1.0.0"
    print("forwardoh", version)
    print("listen port: ",port, "remote url:", url)
    gevent.spawn(cache_status)
    forward_query(port, url)
