pioneerm/update-dns repository overview

#!/usr/bin/python3
import ipaddress
import os
import re
import shutil
import sys
from typing import List

suffix_tmp = ".tmp"
suffix_checkpoint = ".checkpoint"

# Mount /etc as /etc-host
debian_dns_path = "/etc-host/NetworkManager/system-connections/eth0.nmconnection"
debian_dns_checkpoint = "/etc-host/NetworkManager/system-connections/.eth0.nmconnection" + suffix_checkpoint
is_debian = False

rpi4_dns_path = "/etc-host/dhcpcd.conf"
rpi4_dns_checkpoint = "/etc-host/.dhcpcd.conf" + suffix_checkpoint
is_rpi4 = False

file_path = ""

def check_environment():
    global is_debian
    global is_rpi4
    global file_path

    if os.path.isfile( debian_dns_path ):
        is_debian = True
        file_path = debian_dns_path
    elif os.path.isfile( rpi4_dns_path ):
        is_rpi4 = True
        file_path = rpi4_dns_path
    else:
        if os.path.isfile( debian_dns_checkpoint ):
            print(f"'{debian_dns_path}' absents but checkpoint '{debian_dns_checkpoint}' exists, recover it from checkpoint.")
            os.replace(debian_dns_checkpoint, debian_dns_path)
        elif os.path.isfile( rpi4_dns_checkpoint ):
            print(f"'{rpi4_dns_path}' absents but checkpoint '{rpi4_dns_checkpoint}' exists, recover it from checkpoint.")
            os.replace(rpi4_dns_checkpoint, rpi4_dns_path)
        else:
            print(f"Both '{debian_dns_path}' and '{rpi4_dns_path}' do not exist!!")
            sys.exit(-1)

def write_atomic(file_path: str, data: str, mode=0o600):
    temp_list = file_path.split('/')
    new_filename = '.' + temp_list[-1]

    new_list = temp_list[:-1]
    new_list.append( new_filename )

    new_file_path = '/'.join( new_list )
    tmp_filepath = new_file_path + suffix_tmp
    checkpoint_filepath = new_file_path + suffix_checkpoint

    # Write data into buffer
    with open(tmp_filepath, 'w') as tmp_file:
        tmp_file.write(data)
        tmp_file.flush()
        os.fsync(tmp_file.fileno())  # write into disk

    # change mode. NetworkManager permits only 600
    os.chmod(tmp_filepath, mode)

    # create checkpoint
    shutil.copy2(tmp_filepath, checkpoint_filepath)

    # replace temp into real config file
    os.replace(tmp_filepath, file_path)

def is_valid_ipv4(ip: str) -> bool:
    try:
        ipaddress.IPv4Address(ip)
        return True
    except ipaddress.AddressValueError:
        return False

def update_dns(dns_ips: List[str], content_original: str):
    if is_debian:
        pattern = r"dns=[\d.]+(;[\d.]+)*"
        new_dns = "dns=" + ";".join(dns_ips)
    else:
        pattern = r"static domain_name_servers=[\d.]+( [\d.]+)*"
        new_dns = "static domain_name_servers=" + " ".join(dns_ips)

    return re.sub(pattern, new_dns, content_original)

def change_ipv6_gen_rule(content_original: str):
    '''
    Configure the generation of IPv6 link-local address based on the MAC address.
    '''
    if is_debian:
        old_line = r"addr-gen-mode=stable-privacy"
        new_line = "addr-gen-mode=eui64"
        return re.sub(old_line, new_line, content_original)
    else:
        # do nothing if it is rpi4
        return content_original

if __name__ == "__main__":
    file_name = sys.argv[0]
    if len(sys.argv) < 2:
        print(f"Usage: {file_name} [DNS_IP1] [[DNS_IP2] ...]")
        sys.exit(1)

    dns_ips: List[str] = list()

    input_ips = sys.argv[1:]
    for IPs in input_ips:
        # Kubernetes YAML might use "8.8.8.8 8.8.4.4" as an argument
        for ip in IPs.split():
            if not is_valid_ipv4(ip):
                print(f"Invalid IP '{ip}' format")
                sys.exit(1)
            else:
                dns_ips.append(ip)

    # Determine the content of "file_path", "is_debian" and "is_rpi" parameters
    check_environment()

    with open(file_path, "r") as f:
        content_original = f.read()

    # update DNS server
    content1 = update_dns( dns_ips, content_original )

    # Change IPv6 linklocal address generating method
    content_updated = change_ipv6_gen_rule( content1 )

    if content_updated != content_original:
        if is_rpi4:
            write_atomic(file_path, content_updated, 0o644)
        else:
            write_atomic(file_path, content_updated, 0o600)

        print(f"'{file_path}' updated.")
    else:
        print(f"'{file_path}' unchanged, skipping update.")