Steven Black Hosts Updater

steven_black_hosts_updater.py

#!/usr/bin/env python3

"""
Steven Black Hosts Updater

This script automates the process of updating the system's hosts file using
the curated lists provided by Steven Black in https://github.com/StevenBlack/hosts

Requirements:
    1. A Steven Black hosts file in /etc/hosts
    2. To preserve custom host records, you need to have them in it's corresponding block:

        '# Custom host records are listed here.
        '
        '127.0.0.1	   localhost
        '127.0.0.1     rubenhortas
        '192.168.1.101 laptop1
        '192.168.1.102 laptop2
        '
        '# End of custom host records.

Usage:
    `sudo python3 steven_black_hosts_updater.py`
    or
    `sudo ./steven_black_hosts_updater.py`
"""

import os
import re
import sys
import tempfile
from urllib import request

_HOSTS_FILE = "/etc/hosts"
_SB_HOSTS_URL_PATTERN = re.compile(
    r"https://raw.githubusercontent.com/StevenBlack/hosts/master/alternates/[a-z,-]*/hosts"
)
_CUSTOM_HOSTS_BLOCK_PATTERN = re.compile(
    r"# Custom host records are listed here.(.*?)# End of custom host records.",
    re.DOTALL,
)
_ENCODING = "utf-8"

# ANSI colors
_GREEN = "\033[92m"
_YELLOW = "\033[1;33m"
_RED = "\033[91m"
_RESET = "\033[0m"

_DONE = f"{_GREEN}Done{_RESET}"
_INFO = f"{_YELLOW}INFO{_RESET}"
_ERROR = f"{_RED}ERROR{_RESET}"


def _main():
    if os.getuid() != 0:
        print(f"{_ERROR}: This script requires root privileges")
        sys.exit(1)

    try:
        _update_hosts_file()
    except Exception as e:
        print(f"{_ERROR}: An unexpected error occurred: {e}")
        sys.exit(1)


def _update_hosts_file() -> None:
    print(f"Updating {_HOSTS_FILE}...")
    current_hosts = _get_current_hosts()

    if not current_hosts:
        return

    url_match = _SB_HOSTS_URL_PATTERN.search(current_hosts)

    if not url_match:
        print(f"{_ERROR}: Steven Black's hosts URL not found in {_HOSTS_FILE}")
        return

    sb_hosts_url = url_match.group()
    custom_hosts_match = _CUSTOM_HOSTS_BLOCK_PATTERN.search(current_hosts)

    if not custom_hosts_match or not custom_hosts_match.group(1).strip():
        print(f"{_INFO}: Custom hosts not found")

    print(f"Downloading updated hosts from {sb_hosts_url}...")
    new_hosts_content = _download(sb_hosts_url)

    if not new_hosts_content:
        print(f"{_ERROR} Download failed")
        return

    updated_hosts = (
        _CUSTOM_HOSTS_BLOCK_PATTERN.sub(custom_hosts_match.group(0), new_hosts_content)
        if custom_hosts_match
        else new_hosts_content
    )

    if _write(updated_hosts):
        print(_DONE)


def _get_current_hosts() -> str:
    try:
        with open(_HOSTS_FILE, "r", encoding=_ENCODING) as f:
            return f.read()
    except FileNotFoundError:
        print(f"{_ERROR}: {_HOSTS_FILE} not found.")
        return ""


def _download(url: str) -> str:
    try:
        req = request.Request(url, headers={"User-Agent": "Mozilla/5.0"})

        with request.urlopen(req) as response:
            return response.read().decode(_ENCODING)
    except Exception as e:
        print(f"{_ERROR}: Failed to download hosts: {e}")
        return ""


def _write(hosts: str) -> bool:
    target_dir = os.path.dirname(_HOSTS_FILE)
    fd, tmp_file = tempfile.mkstemp(dir=target_dir, text=True)

    try:
        with os.fdopen(fd, "w", encoding=_ENCODING) as tmp:
            tmp.write(hosts)

        # Copy permissions
        original_stat = os.stat(_HOSTS_FILE)
        os.chmod(tmp_file, original_stat.st_mode)
        os.chown(tmp_file, original_stat.st_uid, original_stat.st_gid)
        # os.replace is atomic in Unix, which is excellent for preventing corrupted files if the script fails midway through
        os.replace(tmp_file, _HOSTS_FILE)

        return True
    except Exception as e:
        print(f"{_ERROR}: Failed to update {_HOSTS_FILE}: {e}")

        if os.path.exists(tmp_file):
            os.remove(tmp_file)

        return False


if __name__ == "__main__":
    _main()