focus_dns.py

#!/usr/bin/env python
# coding: utf-8


from socketserver import BaseRequestHandler, ThreadingUDPServer
from io import BytesIO
import sys
import os
import re
from os.path import exists
import socket
from socket import timeout
import struct
import time
from datetime import datetime
import importlib

DEFAULT_FILE = os.path.join(os.path.dirname(__file__), "hosts")
FOCUS_ROOT = os.environ.get("FOCUS_ROOT")
BLACKLIST_LAST_CHECKED = 0
DEFAULT_BLACKLIST = """
def domain_news_ycombinator_com(dt):
    # return dt.hour % 2 # every other hour
    return False

def domain_reddit_com(dt):
    # return dt.hour in (12, 21) # at noon-1pm, or from 9-10pm
    return False

def domain_facebook_com(dt):
    return False


def default(domain, dt):
    # do something with regular expressions here?
    return True
""".strip()
BLACKLIST_FILE = os.path.join(FOCUS_ROOT, "focus_blacklist.py")
# these are special characters that are common to domain names but must be
# replaced with an underscore in order for the domain name to be referenced
# as a function in focus_blacklist.  for example, you cannot call
# test-site.com()...you must convert it to test_site_com()
DOMAIN_SPECIAL_CHARACTERS = "-."

sys.path.append(os.path.join(FOCUS_ROOT, "etc"))

try:
    import focus_blacklist as blacklist
except ImportError:
    blacklist = None

"""
A simple DNS proxy server, support wilcard hosts, IPv6, cache. Usage:

Edit /etc/hosts, add:
    127.0.0.1 *.local
    2404:6800:8005::62 *.blogspot.com
startup dnsproxy(here use Google DNS server as delegating server):
$ sudo python dnsproxy.py -s 8.8.8.8

Then set system dns server as 127.0.0.1, you can verify it by dig:
$ dig test.local

The result should contains 127.0.0.1.

author: marlonyao<yaolei135@gmail.com>
"""


def main():
    import optparse
    import sys

    parser = optparse.OptionParser()
    parser.add_option(
        "-f",
        "--hosts-file",
        dest="hosts_file",
        metavar="<file>",
        default=DEFAULT_FILE,
        help="specify hosts file, default %s" % DEFAULT_FILE,
    )
    parser.add_option(
        "-H",
        "--host",
        dest="host",
        default="127.0.0.1",
        help="specify the address to listen on",
    )
    parser.add_option(
        "-p",
        "--port",
        dest="port",
        default=53,
        type="int",
        help="specify the port to listen on",
    )
    parser.add_option(
        "-s",
        "--server",
        dest="dns_server",
        metavar="SERVER",
        help="Required: specify the delegating dns server",
    )
    parser.add_option(
        "-b",
        "--backup-server",
        dest="backup_server",
        metavar="BACKUPSERVER",
        help="Required: specify the delegating backup dns server",
    )
    parser.add_option(
        "-C",
        "--no-cache",
        dest="disable_cache",
        default=False,
        action="store_true",
        help="disable dns cache",
    )

    opts, args = parser.parse_args()
    if not opts.dns_server:
        parser.print_help()
        sys.exit(1)
    dnsserver = DNSProxyServer(
        opts.dns_server,
        opts.backup_server,
        disable_cache=opts.disable_cache,
        host=opts.host,
        port=opts.port,
        hosts_file=opts.hosts_file,
    )
    dnsserver.serve_forever()


class Struct(object):
    def __init__(self, **kwargs):
        for name, value in list(kwargs.items()):
            setattr(self, name, value)


def parse_dns_message(data):
    message = BytesIO(data)
    message.seek(4)  # skip id, flag
    c_qd, c_an, c_ns, c_ar = struct.unpack("!4H", message.read(8))
    # parse question
    question = parse_dns_question(message)
    for i in range(1, c_qd):  # skip other question
        parse_dns_question(message)
    records = []
    for i in range(c_an + c_ns + c_ar):
        records.append(parse_dns_record(message))
    return Struct(question=question, records=records)


def parse_dns_question(message):
    qname = parse_domain_name(message)
    qtype, qclass = struct.unpack("!HH", message.read(4))
    end_offset = message.tell()
    return Struct(name=qname, type_=qtype, class_=qclass, end_offset=end_offset)


def parse_dns_record(message):
    parse_domain_name(message)  # skip name
    message.seek(4, os.SEEK_CUR)  # skip type, class
    ttl_offset = message.tell()
    ttl = struct.unpack("!I", message.read(4))[0]
    rd_len = struct.unpack("!H", message.read(2))[0]
    message.seek(rd_len, os.SEEK_CUR)  # skip rd_content
    return Struct(ttl_offset=ttl_offset, ttl=ttl)


def _parse_domain_labels(message):
    labels = []
    len = ord(message.read(1))
    while len > 0:
        if len >= 64:  # domain name compression
            len = len & 0x3f
            offset = (len << 8) + ord(message.read(1))
            mesg = BytesIO(message.getvalue())
            mesg.seek(offset)
            labels.extend(_parse_domain_labels(mesg))
            return labels
        else:
            labels.append(message.read(len))
            len = ord(message.read(1))
    return labels


def parse_domain_name(message):
    msgs = (i.decode() for i in _parse_domain_labels(message))
    return ".".join(msgs)


def addr_p2n(addr):
    try:
        return socket.inet_pton(socket.AF_INET, addr)
    except socket.error:
        return socket.inet_pton(socket.AF_INET6, addr)


def refresh_blacklist():
    global BLACKLIST_LAST_CHECKED, blacklist

    # we also check for not exists because the pyc file may be left around.
    # in that case, blacklist name will exist, but the file will not
    if not blacklist or not exists(BLACKLIST_FILE):
        print(("couldn't find %s, creating a default blacklist" % BLACKLIST_FILE))
        with open(BLACKLIST_FILE, "w") as h:
            h.write(DEFAULT_BLACKLIST)
        import focus_blacklist as blacklist

    # has it changed?
    changed = os.stat(BLACKLIST_FILE).st_mtime
    if changed > BLACKLIST_LAST_CHECKED:
        print(("blacklist %s changed, reloading" % BLACKLIST_FILE))
        importlib.reload(blacklist)
        BLACKLIST_LAST_CHECKED = changed


def can_visit(domain):
    """ determine if the domain is blacklisted at this time """

    refresh_blacklist()

    # here we do a cascading lookup for the function to run.  example:
    # for the domain "herp.derp.domain.com", first we try to find the
    # following functions in the following order:
    #
    # herp_derp_domain_com()
    # derp_domain_com()
    # domain_com()
    #
    # and if one still isn't found, we go with default(), if it exists
    parts = domain.split(".")
    for i in range(len(parts) - 1):
        domain_fn_name = "domain_" + ".".join(parts[i:])
        domain_fn_name = re.sub(
            "[" + DOMAIN_SPECIAL_CHARACTERS + "]", "_", domain_fn_name
        )
        fn = getattr(blacklist, domain_fn_name, None)

        if fn:
            return fn(datetime.now())

    fn = getattr(blacklist, "default", None)
    if fn:
        return fn(domain, datetime.now())
    return True


DNS_TYPE_A = 1
DNS_TYPE_AAAA = 28
DNS_CLASS_IN = 1


class DNSProxyHandler(BaseRequestHandler):
    def handle(self):
        reqdata, sock = self.request
        req = parse_dns_message(reqdata)
        q = req.question
        if q.type_ in (DNS_TYPE_A, DNS_TYPE_AAAA) and (
            q.class_ == DNS_CLASS_IN
        ):
            if can_visit(q.name):
                # Just continue and work normal
                pass
            else:
                fail_ip = addr_p2n("127.0.0.1")
                rspdata = self._generate_dns_response_data(reqdata, q, fail_ip)
                sock.sendto(rspdata, self.client_address)
                return

            for packed_ip, host in self.server.host_lines:
                if q.name.endswith(host):
                    rspdata = self._generate_dns_response_data(
                        reqdata, q, packed_ip
                    )
                    sock.sendto(rspdata, self.client_address)
                    return

        # lookup cache
        if not self.server.disable_cache:
            cache = self.server.cache
            cache_key = (q.name, q.type_, q.class_)
            cache_entry = cache.get(cache_key)
            if cache_entry:
                rspdata = update_ttl(reqdata, cache_entry)
                if rspdata:
                    sock.sendto(rspdata, self.client_address)
                    return

        rspdata = self._get_response(reqdata)

        if not rspdata:
            return

        if not self.server.disable_cache:
            cache[cache_key] = Struct(
                rspdata=rspdata, cache_time=int(time.time())
            )
        sock.sendto(rspdata, self.client_address)

    def _get_response(self, data):
        # socket for the remote DNS server
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sock.connect((self.server.dns_server, 53))
            sock.sendall(data)
            sock.settimeout(60)
            rspdata = sock.recv(65535)
            return rspdata
        except timeout:
            print((
                "Fetching DNS response from {} timed out".format(
                    self.server.dns_server
                )
                + "\n"
                + "Fetching DNS response from {}".format(
                    self.server.backup_server
                )
            ))
        except socket.error as err:
            print(err)
        finally:
            sock.close()

    def _get_response_from_backup_dns_server(self, data):
        """
        To fetch response from backup dns server
        :data: DNS request data
        :returns: response from backup dns server

        """
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            sock.connect((self.server.backup_server, 53))
            sock.sendall(data)
            sock.settimeout(60)
            rspdata = sock.recv(65535)
            return rspdata
        except timeout:
            print((
                "Fetching DNS response from backup_server {} timed out".format(
                    self.server.backup_server
                )
            ))
        except socket.error as err:
            print(err)
        finally:
            sock.close()

    def _generate_dns_response_data(self, reqdata, q, packed_ip):
        # header, qd=1, an=1, ns=0, ar=0
        rspdata = reqdata[:2] + b"\x81\x80\x00\x01\x00\x01\x00\x00\x00\x00"
        rspdata += reqdata[12 : q.end_offset]
        # answer
        rspdata += b"\xc0\x0c"  # pointer to domain name
        # type, 1 for ip4, 28 for ip6
        if len(packed_ip) == 4:
            rspdata += b"\x00\x01"  # 1 for ip4
        else:
            rspdata += b"\x00\x1c"  # 28 for ip6
        # class: 1, ttl: 2000(0x000007d0)
        rspdata += b"\x00\x01\x00\x00\x07\xd0"
        rspdata += b"\x00" + chr(len(packed_ip)).encode()  # rd_len
        rspdata += packed_ip
        return rspdata


def update_ttl(reqdata, cache_entry):
    rspdata, cache_time = cache_entry.rspdata, cache_entry.cache_time
    rspbytes = bytearray(rspdata)
    rspbytes[:2] = reqdata[:2]  # update id
    current_time = int(time.time())
    time_interval = current_time - cache_time
    rsp = parse_dns_message(rspdata)
    for record in rsp.records:
        if record.ttl <= time_interval:
            return None
        rspbytes[record.ttl_offset : record.ttl_offset + 4] = struct.pack(
            "!I", record.ttl - time_interval
        )
    return rspbytes


def load_hosts(hosts_file):
    "load hosts config, only extract config line contains wildcard domain name"

    def parts_of_line(line):
        parts = line.strip().split()[:2]
        return parts


    def wildcard_line(line):
        parts = parts_of_line(line)
        if len(parts) < 2:
            return False
        if not parts[1].startswith("*"):
            return False
        try:
            packed_ip = addr_p2n(parts[0])
            return packed_ip, parts[1][1:]
        except socket.error:
            return None

    def entry_line(line):
        if line.startswith('#'): # It is a comment, can be ignored
            return None

        parts = parts_of_line(line)
        if len(parts) < 2:
            return False
        try:
            packed_ip = addr_p2n(parts[0])
            return packed_ip, parts[1]
        except socket.error:
            return None


    with open(hosts_file) as hosts_in:
        hostlines = []
        for line in hosts_in:
            hostline = wildcard_line(line) or entry_line(line)
            if hostline:
                hostlines.append(hostline)
        return hostlines


class DNSProxyServer(ThreadingUDPServer):
    def __init__(
        self,
        dns_server,
        backup_server=None,
        disable_cache=False,
        host="127.0.0.1",
        port=53,
        hosts_file=DEFAULT_FILE,
    ):
        self.dns_server = dns_server
        self.backup_server = backup_server
        self.hosts_file = hosts_file
        self.host_lines = load_hosts(hosts_file)
        self.disable_cache = disable_cache
        self.cache = {}
        ThreadingUDPServer.__init__(self, (host, port), DNSProxyHandler)


if __name__ == "__main__":
    main()