usg

#!/usr/bin/env bash
# Utilities to make managing a Unifi Security Gateway a little easier.
#
# 1. $ usg config-refresh
# 2. $ usg dump > config.json
# 3. Make changes to USG
# 4. $ usg diff config.json
# 5. $ usg patch config.json default
# 6. Force provisioning in the Unifi UI
# 7. $ usg config-refresh
# 8. $ usg diff config.json
# 9. The previous usg diff fields show show up again, in addition to
#    unifi.mgmt.cfgversion.

# ssh arguments to connect to the USG and the UniFi Controller. The arguments
# should be escaped, if necessary.
USG_CONNECTION="${USG_CONNECTION:-router}"
CONTROLLER_CONNECTION="${CONTROLLER_CONNECTION:-unifi}"

# https://github.com/evanphx/json-patch
# Currently relying on my forked version:
# https://github.com/evanphx/json-patch/pull/154
jsonpatch="${GOBIN:-"${GOPATH}/bin"}/json-patch"

# sudo apt install jq
jq="jq"

warn() {
    echo "$@" 1>&2
}

die() {
    warn "$@"
    exit 1
}

# usg::dump dumps the current, complete USG config.gateway.json to stdout.
usg::dump() {
    [[ $# -eq 0 ]] || die "dump does not take any arguments."

    ssh -o loglevel=ERROR ${USG_CONNECTION} -- \
        mca-ctrl -t dump-cfg
}

# usg::diff diffs the current, complete USG config.gateway.json with $1.
usg::diff() {
    [[ $# -eq 1 ]] || die "diff takes exactly one argument: old config"
    local old="$1"

    "${jsonpatch}" create "${old}" <(usg::dump)
}

# usg::partial-copy copies the given JSON value from the USG's
# config.gateway.json to the Controller's config.gateway.json.
#
# The JSON values are identified by paths of the form "system.host-name".
#
# This command has two forms:
#   1) <site name> <JSON path>...
#      This form copies the values from the USG's current, complete
#      config.gateway.json.
#   2) <config file> <site name> <JSON path>...
#      This form copies the values from the provided config file.
usg::partial-copy() {
    [[ $# -ge 2 ]] || die "partial-copy needs at least two arguments: <site name> <JSON path>... or <config file> <site name> <JSON path>..."
    local source_config
    local site
    local paths
    if [[ -f "$1" ]]; then
        source_config="$(cat "$1")"
        if [[ $? != 0 ]]; then
            die "Could not read $1"
        fi
        site="$2"
        paths=( "${@:3}" )
    else
        source_config="$(usg::dump)"
        if [[ $? != 0 ]]; then
            die "Could not fetch the current, complete config.gateway.json from the USG"
        fi
        site="$1"
        paths=( "${@:2}" )
    fi

    local dest_config
    dest_config="$(usg::_fetch_controller_config "${site}")"
    if [[ $? != 0 ]]; then
        die "Could not fetch the current config.gateway.json from the controller."
    fi
    local original_dest_config="${dest_config}"

    local path
    for path in "${paths[@]}"; do
        local source_value
        source_value="$(usg::_extract_partial_json "${source_config}" "${path}")"
        if [[ $? != 0 ]]; then
            die "Did not find ${path} in config."
        fi

        local dest_value
        dest_value="$(usg::_extract_partial_json "${dest_config}" "${path}")"
        if [[ $? != 0 ]]; then
            dest_value="{}"
        fi

        local patch
        patch="$("${jsonpatch}" create <(echo "${dest_value}") <(echo "${source_value}"))"
        if [[ $? != 0 ]]; then
            die "Could not calculate patch for ${path}."
        elif [[ "${patch}" == "{}" ]]; then
            warn
            warn "${path} has no diff, ignoring."
            continue
        fi

        echo
        echo "${path}"
        echo
        echo "${source_value}"

        dest_config="$("${jsonpatch}" apply --indent "  " -p <(echo "${patch}") <(echo "${dest_config}"))"
        if [[ $? != 0 ]]; then
            die "Could not update config.gateway.json."
        fi
    done

    local diff
    diff="$("${jsonpatch}" create <(echo "${original_dest_config}") <(echo "${dest_config}"))"
    if [[ $? != 0 ]]; then
        die "Could not calculate config.gateway.json diff."
    elif [[ "${diff}" == "{}" ]]; then
        die "config.gateway.json would not be updated."
    fi

    echo
    echo
    echo "Original controller config:"
    echo "${original_dest_config}"
    echo
    echo "Updated controller config:"
    echo "${dest_config}"
    echo
    echo "Diff:"
    echo "${diff}"
    echo
    if dotfiles::user_permission "Do you want to commit these changes?"; then
        usg::_set_controller_config "${site}" "${dest_config}"
        if [[ $? != 0 ]]; then
            die
        fi
    fi
}

# usg::patch determines what has changed in the USG's config by diffing its
# current, complete config.gateway.json with a locally stored copy (normally
# stored with usg::dump before modifying the USG's config), and then applies the
# diff to the Controller's config.gateway.json.
usg::patch(){
    [[ $# -eq 2 ]] || die "patch takes exactly two arguments: old config, site name"
    local old="$1"
    local site="$2"

    local patch
    patch="$(usg::diff "${old}")"
    if [[ $? != 0 ]]; then
        die "diffing failed."
    elif [[ "${patch}" == "{}" ]]; then
        die "${old} appears to be the same as what's currently uploaded to the USG. There's nothing to do!"
    fi

    local controller_config
    controller_config="$(usg::_fetch_controller_config "${site}")"
    if [[ $? != 0 ]]; then
        die "Could not fetch controller's config.gateway.json"
    fi

    echo
    echo "Patch:"
    echo
    echo "${patch}"
    echo

    echo "Current config.gateway.json:"
    echo
    echo "${controller_config}"
    echo

    local patched
    patched="$("${jsonpatch}" apply --indent "  " -p <(echo "${patch}") <(echo "${controller_config}"))"
    if [[ $? != 0 ]]; then
        die "Could not apply patch to controller's config.gateway.json."
    fi

    echo "New config.gateway.json:"
    echo
    echo "${patched}"
    echo

    local effective_diff
    effective_diff="$("${jsonpatch}" create <(echo "${controller_config}") <(echo "${patched}"))"
    if [[ $? != 0 ]]; then
        die "Could not calculate effective diff."
    elif [[ "${effective_diff}" == "{}" ]]; then
        die "It doesn't look like applying the patch to config.gateway.json would do anything. Does it already have the same values that the patch would be setting?"
    fi

    echo "Effective diff:"
    echo
    echo "${effective_diff}"
    echo


    if dotfiles::user_permission "Do you want to commit these changes?"; then
        usg::_set_controller_config "${site}" "${patched}"
        if [[ $? != 0 ]]; then
            die
        fi
    fi
}

# usg::config-refresh refreshes the USG's config.gateway.json after a force
# provision to get rid of any transient differences that might exist (e.g.
# "123" -> 123).
usg::config-refresh() {
    [[ $# == 0 ]] || die "config-refresh doesn't take any arguments."

    local dir="$(dirname "${BASH_SOURCE}")"
    ssh -o loglevel=ERROR ${USG_CONNECTION} -- \
        "vbash -s" < "${dir}/usg_scripts/config_refresh.sh"
}

# usg::rollback rollsback the UniFi controller's config.gateway.json to its most
# recent config.gateway.json.old file.
usg::rollback() {
    [[ $# == 1 ]] || die "rollback takes exactly one argument: site name"
    local site="$1"

    local controller_config_remote_path="$(usg::_controller_config_remote_path "${site}")"
    ssh ${CONTROLLER_CONNECTION} -- \
        cp "${controller_config_remote_path}" \
           "${controller_config_remote_path}.rolled_back"
    if [[ $? != 0 ]]; then
        die "Something went wrong making a backup of the existing config.gateway.json"
    fi

    ssh ${CONTROLLER_CONNECTION} -- \
        cp "${controller_config_remote_path}.old" \
           "${controller_config_remote_path}"
    if [[ $? != 0 ]]; then
        die "Something went wrong rolling back config.gateway.json"
    fi
}

usg::_controller_config_remote_path() {
    [[ $# == 1 ]] || die "_controller_config_remote_path takes exactly one argument: site name"
    local site="$1"
    echo "/srv/unifi/data/sites/${site}/config.gateway.json"
}

usg::_fetch_controller_config() {
    [[ $# == 1 ]] || die "_fetch_controller_config takes exactly one argument: site name"

    ssh ${CONTROLLER_CONNECTION} -- \
        cat "$(usg::_controller_config_remote_path "$1")"
}

usg::_set_controller_config() {
    [[ $# == 2 ]] || die "_set_controller_config takes two arguments: site name, new config"
    local site="$1"
    local new_config="$2"

    local controller_config_remote_path="$(usg::_controller_config_remote_path "${site}")"

    ssh ${CONTROLLER_CONNECTION} -- \
        cp "${controller_config_remote_path}" \
           "${controller_config_remote_path}.old"
    if [[ $? != 0 ]]; then
        warn "Failed to backup existing config.gateway.json."
        return 1
    fi

    ssh ${CONTROLLER_CONNECTION} -- \
        "cat > \"${controller_config_remote_path}\"" <<< "${new_config}"
    if  [[ $? != 0 ]]; then
        warn "Something went wrong uploading the patched config.gateway.json"
        return 1
    fi

    # TODO: Automatically force provision the USG, then diff against the
    # expected state. Rollback if something went wrong.
}

usg::_extract_partial_json() {
    [[ $# == 2 ]] || die "_extract_partial_json takes two arguments: JSON and JSON path."
    local json="$1"
    local path="$2"

    local IFS='.'
    local parts=( ${path} )
    unset IFS

    local quoted_parts=()
    local part
    for part in "${parts[@]}"; do
        quoted_parts+=( "\"${part}\"" )
    done

    local IFS='.'
    path="$(echo "${quoted_parts[*]}")"
    unset IFS

    local value
    local ret
    value=$("${jq}" ".${path}" <<< "${json}")
    ret=$?
    if [[ ${ret} != 0 ]]; then
        return $ret
    fi

    local i
    for ((i=${#parts[@]} - 1; i >= 0; i--)); do
        value="{\"${parts[$i]}\": ${value}}"
    done

    echo "${value}"
}


##########
#  Main  #
##########

if [[ "$1" = _* || "$(type -t "usg::$1")" != "function" ]]; then
    echo "Unknown command $1"
    exit 1
fi

"usg::$1" "${@:2}"