A fast hosts file generator (including compression support)

hosts.py

import itertools
import pathlib
import re
import sys
from collections import defaultdict
from typing import List

import requests
from tqdm import tqdm

from tap import Tap


class Args(Tap):
    """A piece of software for updating hosts files"""
    hosts: List[str]  # A list of urls to download
    out: pathlib.Path  # Where to write the file to
    ignore: pathlib.Path = None  # A file containing ignore matches
    compression: int = 9  # The level of compression to apply


def grouper(n, iterable, padvalue=''):
    """grouper(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"""
    return itertools.zip_longest(*[iter(iterable)] * n, fillvalue=padvalue)


def download_text_with_progress(url: str) -> str:
    """Downloads a url to a string"""
    r = requests.get(url, stream=True, allow_redirects=True)
    if r.status_code != 200:
        raise RuntimeError(f"Request to {url} returned error code {r.status_code}")

    file_size = int(r.headers.get('Content-Length', 0))
    r.raw.decode_content = True
    with tqdm.wrapattr(r.raw, "read", total=file_size, desc=url) as r_raw:
        text = r_raw.read().decode()
    return text


def hosts_to_hashmap(hosts: str):
    """Creates a dictionary keyed by the destination

    For example this would convert:
    ```
    0.0.0.0 1.27 3.13
    ```
    into
    ```
    {'0.0.0.0': ['1.27', '3.13']}
    ```

    Args:
        hosts (str): A hosts file string

    Returns:
        defaultdict: {destination: list[source]}
    """
    hashmap = defaultdict(list)
    for line in hosts.splitlines():
        line = re.sub('#.+', '', line).strip()
        if not line or line.startswith('#'):
            continue
        hosts = line.split(' ')
        reroute_host = hosts[0]
        hashmap[reroute_host].extend(hosts[1:])
    return hashmap

def ignore(original_hashmap: defaultdict, new_hashmap: defaultdict):
    """Removes items from the ignore list, warning: non-pure"""
    for dest, sources in original_hashmap.items():
        original_hashmap[dest] = list(set(sources)^set(new_hashmap[dest]))

def main() -> int:
    args = Args().parse_args()

    host_file_strings = [download_text_with_progress(hosts_url) for hosts_url in args.hosts]

    hosts_hashmap = hosts_to_hashmap('\n'.join(host_file_strings))

    if args.ignore:
        ignore_hash = hosts_to_hashmap(args.ignore.read_text())
        ignore(hosts_hashmap, ignore_hash)

    lines = []
    for reroute_host, hosts in hosts_hashmap.items():
        hosts = list(set(hosts))
        lines.extend(f'{reroute_host} {" ".join(chk)}' for chk in tqdm(grouper(args.compression, hosts)))

    args.out.write_text('\n'.join(lines))
    return 0


if __name__ == '__main__':
    sys.exit(main())