update-blacklist-zonefile

#!/usr/bin/python3

import argparse
import email.utils as eut
import hashlib
import logging
import os
import subprocess
import textwrap
from datetime import datetime
from pathlib import Path
from random import choice

import dns.inet
import dns.name
import dns.zone
import requests
from dateutil.parser import parse as parsedate

from typing import List

from jinja2 import Template

logger = logging.getLogger(__name__)

desktop_agents = [
    'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.99 Safari/537.36',
    'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.99 Safari/537.36',
    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.99 Safari/537.36',
    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_1) AppleWebKit/602.2.14 (KHTML, like Gecko) Version/10.0.1 Safari/602.2.14',
    'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.71 Safari/537.36',
    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.98 Safari/537.36',
    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.98 Safari/537.36',
    'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.71 Safari/537.36',
    'Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/54.0.2840.99 Safari/537.36',
    'Mozilla/5.0 (Windows NT 10.0; WOW64; rv:50.0) Gecko/20100101 Firefox/50.0']

settings = {
    # Config directory
    'config_directory': '/etc/bind-adblock',
    'cache_directory': '/var/cache/bind-adblock',
    # Blocklist download request timeout
    'req_timeout_s': 50,
    # Also block *.domain.tld
    'wildcard_block': False
}


def set_logging_basic_configuration(level: int) -> None:
    logging.basicConfig(level=level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')


def set_log_level(verbosity: int) -> None:
    log_level = logging.WARNING  # default
    if verbosity == 2:
        log_level = logging.INFO
    elif verbosity >= 3:
        log_level = logging.DEBUG
    set_logging_basic_configuration(log_level)


def set_quiet_mode() -> None:
    log_level = logging.FATAL
    set_logging_basic_configuration(log_level)


def print_usage(zonefile: str, origin: str) -> None:
    current_log_level = logging.getLogger().getEffectiveLevel()

    logging.getLogger().setLevel(logging.INFO)
    path = Path(zonefile).resolve()

    logger.info(textwrap.dedent(f'''
            Zone file "{path}" created.

            Add BIND options entry:
            response-policy {{
                zone "{origin}";
            }};

            Add BIND zone entry:
            zone "{origin}" {{
                type master;
                file "{path}";
                allow-query {{ none; }};
            }};'''))

    logging.getLogger().setLevel(current_log_level)


def is_comment_line(line: str) -> bool:
    return line.startswith('#')


def is_empty_line(line: str) -> bool:
    return line == ''


def remove_comments_and_empty_lines(lst: List[str]) -> List[str]:
    return list((_ for _ in lst if not is_empty_line(_) and not is_comment_line(_)))


def is_ip_address(line: str) -> bool:
    return dns.inet.is_address(line.split('/')[0])


def get_blocklist_content(url: str) -> List[str]:
    last_modified_cache: datetime = datetime.fromtimestamp(0).astimezone()

    cache = Path(settings['cache_directory'])
    if not cache.is_dir():
        cache.mkdir(parents=True)

    cache = Path(cache, hashlib.sha1(url.encode()).hexdigest())

    if cache.is_file():
        last_modified_cache = datetime.fromtimestamp(os.path.getmtime(cache)).astimezone()

    headers = {
        'If-modified-since': eut.format_datetime(last_modified_cache),
        'User-Agent': choice(desktop_agents)
    }

    r = requests.get(url, headers=headers, timeout=settings['req_timeout_s'])

    url_time = r.headers['last-modified']
    last_modified_source = parsedate(url_time)

    if last_modified_source > last_modified_cache:
        try:
            if r.status_code == 200:
                received_lines = r.text.split('\n')

                with cache.open('w') as file:

                    len_list_original = len(received_lines)
                    data = remove_comments_and_empty_lines(r.text.split('\n'))
                    logger.info(f'''\t {len_list_original - len(data):7} empty or lines with comments removed''')

                    len_list_before_removal_of_ip_addresses = len(data)
                    data = list((x.split(' ')[0] for x in data if not is_ip_address(x)))
                    logger.info(f'''\t {len_list_before_removal_of_ip_addresses - len(data):7} lines with IP adresses removed''')

                    data.sort()

                    for element in data:
                        element.lower()  # lowering since DNS is case insensitive
                        file.write(element)
                        file.write('\n')
                    file.close()

                if 'last-modified' in r.headers:
                    last_modified = eut.parsedate_to_datetime(r.headers['last-modified']).timestamp()
                    os.utime(str(cache), times=(last_modified, last_modified))

                return data

            elif r.status_code != 304:
                logger.error(f'''Error getting list at {url}, HTTP STATUS: {r.status_code}''')
        except requests.exceptions.RequestException as e:
            logger.error(e)

    else:
        logger.info(f'''\t no update found - using cached file''')

        with cache.open() as file:
            return [i.strip() for i in file.readlines()]


def get_blocked_url_list(blocklist_urls: List) -> List[str]:
    domains: List[str] = list()
#    origin_name = dns.name.from_text(origin)

    for blocklist_url in blocklist_urls:
        logger.info(f'''Processing Blocklist: {blocklist_url}''')
        url_list = get_blocklist_content(blocklist_url)
        domains.extend(url_list)

    total_domains = len(domains)
    domains = sorted(set(domains))

    logger.info(f'''{total_domains - len(domains):7} duplicate entries removed''')
    logger.info(f'''{len(domains):7} domains on Blacklist''')
    return domains


def read_blocklists_file() -> List[str]:
    with open(f'''{settings['config_directory']}/blocklists.conf''', 'r') as file_handle:
        # convert file contents into a list
        lists = file_handle.read().splitlines()
        lists[:] = remove_comments_and_empty_lines(lists)
    return lists[:]


def reload_zone(origin):
    r = subprocess.call(['rndc', 'reload', origin])
    if r != 0:
        raise Exception(f'''rndc failed with return code {r}''')


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('zonefile', help='name of the zone file to generate', type=str)
    parser.add_argument('origin', help='name of the zone', type=str)

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-v', '--verbose', action='count',
                       help='increase verbosity (specify multiple times for more output)')

    group.add_argument('-q', '--quiet', action='store_true', help='suppress output, except fatal messages')

    parser.add_argument('--print-bind-config', action='store_true',
                        help='print necessary configuration of BIND to use the generated file')

    parser.add_argument('--reload-zone', action='store_true', help='trigger a reload of the zone after update')

    args = parser.parse_args()

    zonefile = args.zonefile
    origin = args.origin

    if args.verbose:
        set_log_level(args.verbose)
    elif args.quiet:
        set_quiet_mode()
    else:
        set_log_level(1)

    if args.print_bind_config:
        print_usage(zonefile, origin)

    blocklists_urls: List[str] = list(read_blocklists_file())

    domains: List[str] = get_blocked_url_list(blocklists_urls)

    template = Template(open('/usr/lib/bind-adblock/blocklist.zone.j2').read())
    rendered = template.render(origin=origin, now=datetime.now(), wildcard_block=settings['wildcard_block'], domains=domains)

    logger.info('Generating Zonefile ...')
    with open(zonefile, 'w') as fh:
        fh.write(rendered)
    logger.info('Zonefile generation complete')

    if args.reload_zone:
        logger.debug(f'''Reload of zone will be done now''')
        reload_zone(origin)


if __name__ == "__main__":
    main()