#!/bin/bash

# This script configures running chronyd to use NTP servers obtained from

# DHCP and _ntp._udp DNS SRV records. Files with servers from DHCP are managed

# externally (e.g. by a dhclient script). Files with servers from DNS SRV

# records are updated here using the dig utility. The script can also list

# and set static sources in the chronyd configuration file.

chronyc=/usr/bin/chronyc

chrony_conf=/etc/chrony.conf

chrony_service=chronyd.service

helper_dir=/var/run/chrony-helper

added_servers_file=$helper_dir/added_servers

network_sysconfig_file=/etc/sysconfig/network

dhclient_servers_files=/var/lib/dhclient/chrony.servers.*

dnssrv_servers_files=$helper_dir/dnssrv@*

dnssrv_timer_prefix=chrony-dnssrv@

chrony_command() {

    $chronyc -a -n -m "$1"

}

is_running() {

    chrony_command "tracking" &> /dev/null

}

is_update_needed() {

    for file in $dhclient_servers_files $dnssrv_servers_files \

            $added_servers_file; do

        [ -e "$file" ] && return 0

    done

    return 1

}

update_daemon() {

    local all_servers_with_args all_servers added_servers

    if ! is_running; then

        rm -f $added_servers_file

        return 0

    fi

    all_servers_with_args=$(

        cat $dhclient_servers_files $dnssrv_servers_files 2> /dev/null)

    all_servers=$(

        echo "$all_servers_with_args" |

            while read server serverargs; do

                echo "$server"

            done | sort -u)

    added_servers=$( (

        cat $added_servers_file 2> /dev/null

        echo "$all_servers_with_args" |

            while read server serverargs; do

                [ -z "$server" ] && continue

                chrony_command "add server $server $serverargs" &> /dev/null &&

                    echo "$server"

            done) | sort -u)

    comm -23 <(echo -n "$added_servers") <(echo -n "$all_servers") |

        while read server; do

            chrony_command "delete $server" &> /dev/null

        done

    added_servers=$(comm -12 <(echo -n "$added_servers") <(echo -n "$all_servers"))

    [ -n "$added_servers" ] && echo "$added_servers" > $added_servers_file ||

        rm -f $added_servers_file

}

get_dnssrv_servers() {

    local name=$1

    if ! command -v dig &> /dev/null; then

        echo "Missing dig (DNS lookup utility)" >&2

        return 1

    fi

    (

        . $network_sysconfig_file &> /dev/null

        output=$(dig "$name" srv +short +ndots=2 +search 2> /dev/null)

        [ $? -ne 0 ] && return 0

        echo "$output" | while read prio weight port target; do

            server=${target%.}

            [ -z "$server" ] && continue

            echo "$server port $port ${NTPSERVERARGS:-iburst}"

        done

    )

}

check_dnssrv_name() {

    local name=$1

    if [ -z "$name" ]; then

        echo "No DNS SRV name specified" >&2

        return 1

    fi

    if [ "${name:0:9}" != _ntp._udp ]; then

        echo "DNS SRV name $name doesn't start with _ntp._udp" >&2

        return 1

    fi

}

update_dnssrv_servers() {

    local name=$1

    local srv_file=$helper_dir/dnssrv@$name servers

    check_dnssrv_name "$name" || return 1

    servers=$(get_dnssrv_servers "$name")

    [ -n "$servers" ] && echo "$servers" > "$srv_file" || rm -f "$srv_file"

}

set_dnssrv_timer() {

    local state=$1 name=$2

    local srv_file=$helper_dir/dnssrv@$name servers

    local timer=$dnssrv_timer_prefix$(systemd-escape "$name").timer

    check_dnssrv_name "$name" || return 1

    if [ "$state" = enable ]; then

        systemctl enable "$timer"

        systemctl start "$timer"

    elif [ "$state" = disable ]; then

        systemctl stop "$timer"

        systemctl disable "$timer"

        rm -f "$srv_file"

    fi

}

list_dnssrv_timers() {

    systemctl --all --full -t timer list-units | grep "^$dnssrv_timer_prefix" | \

            sed "s|^$dnssrv_timer_prefix\(.*\)\.timer.*|\1|" |

        while read -r name; do

            systemd-escape --unescape "$name"

        done

}

prepare_helper_dir() {

    mkdir -p $helper_dir

    exec 100> $helper_dir/lock

    if ! flock -w 20 100; then

        echo "Failed to lock $helper_dir" >&2

        return 1

    fi

}

is_source_line() {

    local pattern="^[ \t]*(server|pool|peer|refclock)[ \t]+[^ \t]+"

    [[ "$1" =~ $pattern ]]

}

list_static_sources() {

    while read line; do

        is_source_line "$line" && echo "$line" || :

    done < $chrony_conf

}

set_static_sources() {

    local new_config tmp_conf

    new_config=$(

        sources=$(

            while read line; do

                is_source_line "$line" && echo "$line"

            done)

        while read line; do

            if ! is_source_line "$line"; then

                echo "$line"

                continue

            fi

            tmp_sources=$(

                local removed=0

                echo "$sources" | while read line2; do

                    [ "$removed" -ne 0 -o "$line" != "$line2" ] && \

                        echo "$line2" || removed=1

                done)

            [ "$sources" == "$tmp_sources" ] && continue

            sources=$tmp_sources

            echo "$line"

        done < $chrony_conf

        echo "$sources"

    )

    tmp_conf=${chrony_conf}.tmp

    cp -a $chrony_conf $tmp_conf &&

        echo "$new_config" > $tmp_conf &&

        mv $tmp_conf $chrony_conf || return 1

    systemctl try-restart $chrony_service

}

print_help() {

    echo "Usage: $0 COMMAND"

    echo

    echo "Commands:"

    echo "	update-daemon"

    echo "	update-dnssrv-servers NAME"

    echo "	enable-dnssrv NAME"

    echo "	disable-dnssrv NAME"

    echo "	list-dnssrv"

    echo "	list-static-sources"

    echo "	set-static-sources < sources.list"

    echo "	is-running"

    echo "	command CHRONYC-COMMAND"

}

case "$1" in

    update-daemon|add-dhclient-servers|remove-dhclient-servers)

        is_update_needed || exit 0

        prepare_helper_dir && update_daemon

        ;;

    update-dnssrv-servers)

        prepare_helper_dir && update_dnssrv_servers "$2" && update_daemon

        ;;

    enable-dnssrv)

        set_dnssrv_timer enable "$2"

        ;;

    disable-dnssrv)

        set_dnssrv_timer disable "$2" && prepare_helper_dir && update_daemon

        ;;

    list-dnssrv)

        list_dnssrv_timers

        ;;

    list-static-sources)

        list_static_sources

        ;;

    set-static-sources)

        set_static_sources

        ;;

    is-running)

        is_running

        ;;

    command|forced-command)

        chrony_command "$2"

        ;;

    *)

        print_help

        exit 2

esac

exit $?