#!/bin/bash
# dkim-pubkey - print base64-encoded DKIM public key for a domain
#
# Usage:
#   dkim-pubkey [options] <selector> <domain> [@server]
#   dkim-pubkey -h
#
# Options:
#   -s, --server <host>  Query <host> instead of the system resolver (also @host)
#   -V, --validate       Verify the extracted key is a well-formed public key
#   -q, --quiet          Suppress dig command echo, DNS response body, and validation success info
#   -h, --help           Show help message
#
# Queries the TXT record at <selector>._domainkey.<domain>, prints the raw
# DNS response to stderr, and extracts the base64 public key (the p= value)
# to stdout

_dkim_pubkey() (
    local SCRIPT_NAME; SCRIPT_NAME="$(basename "${BASH_SOURCE[0]}")"
    case "${BASH_SOURCE[0]}" in /dev/*|/proc/*) SCRIPT_NAME="" ;; esac
    case "$SCRIPT_NAME" in ""|bash|sh|zsh|dash) SCRIPT_NAME="dkim-pubkey" ;; esac

    _error() { echo "[ERR][$SCRIPT_NAME] $*" >&2; }
    _info()  { echo "[INF][$SCRIPT_NAME] $*" >&2; }

    _show_help() {
        local s; [ -t 1 ] && s=$'\033[4m'
        local r; [ -t 1 ] && r=$'\033[24m'
        echo "NAME"
        echo "  $SCRIPT_NAME - print base64-encoded DKIM public key for a domain"
        echo "SYNOPSIS"
        echo "  $SCRIPT_NAME [${s}options${r}] ${s}selector${r} ${s}domain${r} [@${s}server${r}]"
        echo "  $SCRIPT_NAME -h"
        echo "DESCRIPTION"
        echo "  Queries DNS TXT record at ${s}selector${r}._domainkey.${s}domain${r},"
        echo "  prints the raw response to stderr, and extracts the base64"
        echo "  public key (the p= value) to stdout."
        echo "OPTIONS"
        echo "  -s, --server ${s}host${r}  Query ${s}host${r} instead of the system resolver (also @${s}host${r})"
        echo "  -V, --validate     Verify the extracted key is a well-formed public key"
        echo "  -q, --quiet        Suppress dig command echo, DNS response body, and validation success info (errors still surface)"
        echo "  -h, --help         Show this help message"
        echo "EXAMPLE"
        echo "  $SCRIPT_NAME dkim-prd gmail.com"
        echo "  $SCRIPT_NAME --validate s1 github.com"
        echo "  $SCRIPT_NAME s1 github.com @8.8.8.8           # Google DNS"
        echo "  $SCRIPT_NAME --server 1.1.1.1 s1 github.com   # Cloudflare DNS"
        echo "  $SCRIPT_NAME -q s1 github.com > key.txt       # Just the key, no stderr noise"
        echo "CAVEATS"
        echo "  Selector must be known in advance; check SPF/DMARC records,"
        echo "  email headers, or the domain's DNS provider."
        echo "EXIT STATUS"
        echo "  0  Success (key printed)"
        echo "  1  Runtime failure (DNS response empty)"
        echo "  2  Usage error (missing selector or domain)"
        echo "  3  Dependency error (dig missing, or openssl missing with --validate)"
        echo "  4  Domain-specific (record found but p= value missing)"
        echo "  5  Validation failed (with --validate)"
        echo "DEPENDENCIES"
        echo "  dig"
        echo "  openssl (required only for --validate)"
    }

    _expand_short_opts() {
        # $1 = string of short-opt letters that take a value (e.g. "nXHd"); "" for flag-only scripts
        # $2..$N = "$@"
        # Populates _EXPANDED; caller does: set -- "${_EXPANDED[@]}"; unset _EXPANDED
        local value_opts="$1"; shift
        _EXPANDED=()
        local passthru=""
        local arg
        local rest
        local c
        for arg in "$@"; do
            if [ -n "$passthru" ]; then _EXPANDED+=("$arg"); continue; fi
            case "$arg" in
                --)          passthru=1; _EXPANDED+=("$arg") ;;
                --*|-|"")    _EXPANDED+=("$arg") ;;
                -[a-zA-Z]?*)
                    rest="${arg#-}"
                    while [ -n "$rest" ]; do
                        c="${rest%"${rest#?}"}"; rest="${rest#?}"
                        _EXPANDED+=("-$c")
                        case "$value_opts" in *"$c"*)
                            [ -n "$rest" ] && _EXPANDED+=("$rest")
                            rest="" ;;
                        esac
                    done ;;
                *)        _EXPANDED+=("$arg") ;;
            esac
        done
    }

    _expand_short_opts "s" "$@"
    set -- "${_EXPANDED[@]}"; unset _EXPANDED

    local validate=""
    local server=""
    local quiet=""
    local args=()
    while [ $# -gt 0 ]; do
        case "$1" in
            -h|--help)     _show_help; return 0 ;;
            -V|--validate) validate=1; shift ;;
            -q|--quiet)    quiet=1; shift ;;
            -s|--server)
                [ -n "${2-}" ] || { _error "--server requires a value. Run \`$SCRIPT_NAME -h\` for usage"; return 2; }
                [ -z "$server" ] || { _error "Multiple --server not allowed (already set to '$server'). Run \`$SCRIPT_NAME -h\` for usage"; return 2; }
                server="$2"; shift 2 ;;
            --server=*)
                server="${1#*=}"
                [ -n "$server" ] || { _error "--server requires a value. Run \`$SCRIPT_NAME -h\` for usage"; return 2; }
                shift ;;
            @*)
                [ -z "$server" ] || { _error "Multiple --server not allowed (already set to '$server'). Run \`$SCRIPT_NAME -h\` for usage"; return 2; }
                server="${1#@}"
                [ -n "$server" ] || { _error "@<host> requires a value. Run \`$SCRIPT_NAME -h\` for usage"; return 2; }
                shift ;;
            --)            shift; while [ $# -gt 0 ]; do args+=("$1"); shift; done ;;
            -*)            _error "Unknown argument '$1'. Run \`$SCRIPT_NAME -h\` for usage"; return 2 ;;
            *)             args+=("$1"); shift ;;
        esac
    done

    if ! command -v dig >/dev/null 2>&1; then
        _error "dig is required"
        return 3
    fi

    if [ -z "${args[0]:-}" ]; then
        _error "selector is required. Run \`$SCRIPT_NAME -h\` for usage"
        return 2
    fi
    if [ -z "${args[1]:-}" ]; then
        _error "domain is required. Run \`$SCRIPT_NAME -h\` for usage"
        return 2
    fi
    local selector="${args[0]}"
    local domain="${args[1]}"
    local dns_hostname="$selector._domainkey.$domain"
    local dig_args=(+short TXT "$dns_hostname")
    [ -n "$server" ] && dig_args=("@$server" "${dig_args[@]}")
    local dns_response; dns_response="$(dig "${dig_args[@]}")"
    if [ -z "$quiet" ]; then
        if [ -n "$server" ]; then
            echo "\$ dig @$server +short TXT \"$dns_hostname\"" >&2
        else
            echo "\$ dig +short TXT \"$dns_hostname\"" >&2
        fi
    fi
    if [ -z "$dns_response" ]; then
        _error "DNS response empty"
        return 1
    fi
    if [ -z "$quiet" ]; then
        echo "$dns_response" >&2
        echo >&2
    fi
    # CNAME-fronted records (Sendgrid, Mailgun, etc.) cause `dig +short TXT` to
    # return both the CNAME target and the TXT body on separate lines. Pick out
    # only the line that actually contains p=, then strip quoting/joining
    local pubkey; pubkey="$(echo "$dns_response" | grep 'p=' | sed 's/^.*p=//; s/"$//; s/" "//g')"
    if [ -z "$pubkey" ]; then
        _error "Record found but key is empty (p= value missing)"
        return 4
    fi

    if [ "$validate" ]; then
        # Strict base64 shape: alphabet only, no whitespace, length divisible by 4
        local len=${#pubkey}
        if ! printf '%s' "$pubkey" | grep -Eq '^[A-Za-z0-9+/]+=*$' || [ $((len % 4)) -ne 0 ]; then
            echo "$pubkey"
            echo >&2
            _error "Validation failed: key is not valid base64 (alphabet/padding)"
            # dig presents control characters in TXT records using RFC 1035 zone-file
            # \DDD decimal escapes. If we see one, the literal DNS record contains a
            # stray control character -- almost always a copy/paste artifact in the
            # provider's editor. Surface a hint so the user doesn't have to decode
            # \010 themselves
            local esc
            esc="$(printf '%s' "$pubkey" | grep -oE '\\[0-9]{3}' | sort -u)"
            if [ -n "$esc" ]; then
                local seq
                local name
                while IFS= read -r seq; do
                    case "$seq" in
                        '\000') name="null byte" ;;
                        '\008') name="backspace" ;;
                        '\009') name="tab"       ;;
                        '\010') name="newline"   ;;
                        '\013') name="carriage return" ;;
                        *)      name="control character" ;;
                    esac
                    _info "Hint: record contains literal $name ($seq). Likely a stray $name in the DNS provider's editor"
                done <<< "$esc"
            fi
            return 5
        fi

        if ! command -v openssl >/dev/null 2>&1; then
            _error "openssl is required"
            return 3
        fi

        # Decode and parse as a SubjectPublicKeyInfo. -text prints algorithm + size,
        # which we use for the success summary below
        local pkey_text
        if ! pkey_text="$(printf '%s' "$pubkey" | base64 -d 2>/dev/null | openssl pkey -pubin -inform DER -text -noout 2>/dev/null)"; then
            echo "$pubkey"
            echo >&2
            _error "Validation failed: key bytes do not parse as a public key"
            return 5
        fi

        echo "$pubkey"
        if [ -z "$quiet" ]; then
            # Parse "Public-Key: (2048 bit)" / "RSA Public-Key:" / "ED25519 Public-Key:" etc.
            local key_summary
            key_summary="$(printf '%s' "$pkey_text" | grep -m1 -iE 'public-key' | sed 's/^[[:space:]]*//; s/[[:space:]]*$//')"
            [ -z "$key_summary" ] && key_summary="public key"
            echo >&2
            _info "Key valid ($key_summary)"
        fi
        return 0
    fi

    echo "$pubkey"
)

_dkim_pubkey "$@"
__dkim_pubkey_rc=$?
unset -f _dkim_pubkey
if [ -n "${BASH_SOURCE[0]}" ] && [ "${BASH_SOURCE[0]}" != "$0" ]; then
    eval "unset __dkim_pubkey_rc; return $__dkim_pubkey_rc"
fi
eval "unset __dkim_pubkey_rc; exit $__dkim_pubkey_rc"