spftrace 0.3.0

// spftrace – utility for tracing SPF queries
// Copyright © 2022–2023 David Bürgin <dbuergin@gluet.ch>
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or (at your option) any later
// version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License along with
// this program. If not, see <https://www.gnu.org/licenses/>.

use crate::{config::Config, model::*};
use owo_colors::{OwoColorize, Stream};
use std::{fmt::Display, io, time::Duration};
use viaspf::{record::SpfRecord, EvalError, ExplanationString, SpfResult};

// Note: Throughout, `print!` and `println!` are replaced with non-panicking,
// `Result`-based alternatives.

macro_rules! print {
    ($($arg:tt)*) => {
        {
            use ::std::io::Write;
            ::std::write!(::std::io::stdout(), $($arg)*)
        }
    };
}

macro_rules! println {
    ($($arg:tt)*) => {
        {
            use ::std::io::Write;
            ::std::writeln!(::std::io::stdout(), $($arg)*)
        }
    };
}

macro_rules! error {
    () => {
        "error".if_supports_color(::owo_colors::Stream::Stdout, |s| s.red())
    };
}

trait ToStr {
    fn to_str(&self) -> &'static str;
}

impl ToStr for SpfResult {
    fn to_str(&self) -> &'static str {
        match self {
            Self::None => "none",
            Self::Neutral => "neutral",
            Self::Pass => "pass",
            Self::Fail(_) => "fail",
            Self::Softfail => "softfail",
            Self::Temperror => "temperror",
            Self::Permerror => "permerror",
        }
    }
}

struct Context {
    line_width: usize,
    tiles: Vec<&'static str>,
    lookup_counts: LookupCounts,
    // Where a recursive query results in an error, some special handling for
    // printing the wrapping recursive term is necessary.
    omit_recursive_error: bool,
}

impl Context {
    fn new(line_width: usize) -> Self {
        Self {
            line_width,
            tiles: Default::default(),
            lookup_counts: Default::default(),
            omit_recursive_error: Default::default(),
        }
    }

    fn push_indent(&mut self, last: bool) {
        if last {
            self.tiles.push("    ");
        } else {
            self.tiles.push("│   ");
        }
    }

    fn pop_indent(&mut self) {
        self.tiles.pop();
    }

    fn print_indentation(&self) -> io::Result<()> {
        for t in &self.tiles {
            print!("{t}")?;
        }
        Ok(())
    }

    fn print_indentation_leaf(&mut self, last: bool) -> io::Result<()> {
        self.pop_indent();
        self.print_indentation()?;
        if last {
            print!("└── ")?;
        } else {
            print!("├── ")?;
        }
        self.push_indent(last);
        Ok(())
    }
}

pub fn print_evaluated_query(query: TimedEvaluatedQuery, config: &Config) -> io::Result<()> {
    let TimedEvaluatedQuery { query, duration, lookup_times } = query;

    let mut cx = Context::new(config.line_width);

    print_query(&mut cx, &query)?;

    // Print the final SPF result (note that its `Display` representation
    // includes any explanation string).
    println!("{}", query.query_result)?;

    if config.time {
        print_timings(duration, &lookup_times)?;
    }

    Ok(())
}

fn print_query(cx: &mut Context, query: &EvaluatedQuery) -> io::Result<()> {
    cx.print_indentation()?;
    println!("{}", query.domain)?;

    match &query.result {
        EvaluatedQueryResult::NoRecord => {
            cx.print_indentation()?;
            println!("no SPF record found")?;
        }
        EvaluatedQueryResult::Error(msg) => {
            cx.print_indentation()?;
            println!("{}: {}, result={}", error!(), msg, query.query_result.to_str())?;

            // After printing error cause and SPF result for this query, make
            // sure to suppress repeated printing of this information in the
            // wrapping recursive terms *include* and *redirect*.
            cx.omit_recursive_error = true;
        }
        EvaluatedQueryResult::Record(record) => {
            cx.push_indent(false);

            print_evaluated_record(cx, record)?;

            cx.pop_indent();
        }
    }

    Ok(())
}

fn print_evaluated_record(cx: &mut Context, record: &EvaluatedRecord) -> io::Result<()> {
    print_spf_record(cx, &record.spf_record)?;

    if let [terms @ .., last_term] = &record.terms[..] {
        for term in terms {
            print_evaluated_term(cx, term, false)?;
        }
        print_evaluated_term(cx, last_term, true)?;
    }

    Ok(())
}

fn print_spf_record(cx: &Context, spf_record: &SpfRecord) -> io::Result<()> {
    let s = spf_record.to_string();

    let pieces = s.split_ascii_whitespace().collect::<Vec<_>>();
    assert!(!pieces.is_empty());

    if let [first_piece, pieces @ ..] = &pieces[..] {
        // Note that SPF records may contain " and \ (eg in a macro literal), so
        // quoting them with `"` does not necessarily give a well-formed string.

        cx.print_indentation()?;
        print!("\"{first_piece}")?;

        let last_i = pieces.len().saturating_sub(1);
        let indent = cx.tiles.first()
            .map_or(0, |t| t.chars().count() * cx.tiles.len());

        let mut n = indent + first_piece.chars().count() + 1;

        for (i, piece) in pieces.iter().enumerate() {
            // Account for the additional trailing `"` of the final piece.
            let padding = if i < last_i { 1 } else { 2 };
            let len = piece.chars().count() + padding;

            if n + len > cx.line_width {
                println!()?;
                cx.print_indentation()?;
                n = indent + len;
            } else {
                n += len;
            }

            print!(" {piece}")?;
        }

        println!("\"")?;
    }

    Ok(())
}

fn print_evaluated_term(cx: &mut Context, term: &EvaluatedTerm, last: bool) -> io::Result<()> {
    match term {
        EvaluatedTerm::All(directive) => print_all(cx, directive, last),
        EvaluatedTerm::Ip4(directive) => print_ip4(cx, directive, last),
        EvaluatedTerm::Ip6(directive) => print_ip6(cx, directive, last),
        EvaluatedTerm::A(directive) => print_a(cx, directive, last),
        EvaluatedTerm::Mx(directive) => print_mx(cx, directive, last),
        EvaluatedTerm::Ptr(directive) => print_ptr(cx, directive, last),
        EvaluatedTerm::Exists(directive) => print_exists(cx, directive, last),
        EvaluatedTerm::Include(directive) => print_include(cx, directive, last),
        EvaluatedTerm::Redirect(modifier) => print_redirect(cx, modifier, last),
        EvaluatedTerm::Exp(modifier) => print_exp(cx, modifier, last),
        EvaluatedTerm::NeutralDefault => {
            cx.print_indentation_leaf(last)?;
            println!("default result=neutral")?;
            Ok(())
        }
    }
}

fn print_all(cx: &mut Context, directive: &AllDirective, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold("all")?;

    print!(" ")?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_ip4(cx: &mut Context, directive: &Ip4Directive, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(directive.mechanism)?;

    print!(" ")?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_ip6(cx: &mut Context, directive: &Ip6Directive, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(directive.mechanism)?;

    print!(" ")?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_a(cx: &mut Context, directive: &ADirective, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(&directive.mechanism)?;

    if let Some(target_domain) = &directive.target_domain {
        print!(" → {target_domain}")?;
    }

    print_lookup_counts(cx, &directive.lookup_counts)?;

    if let [ips @ .., last_ip] = &directive.ips[..] {
        cx.push_indent(last);

        for ip in ips {
            cx.print_indentation_leaf(false)?;
            println!("{ip}")?;
        }
        cx.print_indentation_leaf(true)?;
        println!("{last_ip}")?;

        cx.pop_indent();
    }

    cx.print_indentation()?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_mx(cx: &mut Context, directive: &MxDirective, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(&directive.mechanism)?;

    if let Some(target_domain) = &directive.target_domain {
        print!(" → {target_domain}")?;
    }

    print_lookup_counts(cx, &directive.lookup_counts)?;

    if let [mxs @ .., last_mx] = &directive.mxs[..] {
        cx.push_indent(last);

        for mx in mxs {
            print_mx_name(cx, mx, false)?;
        }
        print_mx_name(cx, last_mx, true)?;

        cx.pop_indent();
    }

    cx.print_indentation()?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_mx_name(cx: &mut Context, mx: &MxName, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    println!("{}", mx.target_domain)?;

    if let [ips @ .., last_ip] = &mx.ips[..] {
        cx.push_indent(last);

        for ip in ips {
            cx.print_indentation_leaf(false)?;
            println!("{ip}")?;
        }
        cx.print_indentation_leaf(true)?;
        println!("{last_ip}")?;

        cx.pop_indent();
    }

    Ok(())
}

fn print_ptr(cx: &mut Context, directive: &PtrDirective, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(&directive.mechanism)?;

    if let Some(target_domain) = &directive.target_domain {
        print!(" → {target_domain}")?;
    }

    print_lookup_counts(cx, &directive.lookup_counts)?;

    if let Some(ip) = directive.ip {
        cx.print_indentation()?;
        println!("{ip}")?;
    }

    if let Some(e) = &directive.lookup_error {
        cx.print_indentation()?;
        println!("{}: {}", error!(), e)?;
    }

    if let [ptrs @ .., last_ptr] = &directive.ptrs[..] {
        cx.push_indent(last);

        for ptr in ptrs {
            print_ptr_name(cx, ptr, false)?;
        }
        print_ptr_name(cx, last_ptr, true)?;

        cx.pop_indent();
    }

    cx.print_indentation()?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_ptr_name(cx: &mut Context, ptr: &PtrName, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    println!("{}", ptr.target_domain)?;

    if let Some(e) = &ptr.error {
        cx.print_indentation()?;
        println!("{}: {}", error!(), e)?;
    }

    if let [ips @ .., last_ip] = &ptr.ips[..] {
        cx.push_indent(last);

        for ip in ips {
            cx.print_indentation_leaf(false)?;
            println!("{ip}")?;
        }
        cx.print_indentation_leaf(true)?;
        println!("{last_ip}")?;

        cx.pop_indent();
    }

    if ptr.validated {
        cx.print_indentation()?;
        println!("validated")?;
    }

    Ok(())
}

fn print_exists(cx: &mut Context, directive: &ExistsDirective, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(&directive.mechanism)?;

    if let Some(target_domain) = &directive.target_domain {
        print!(" → {target_domain}")?;
    }

    print_lookup_counts(cx, &directive.lookup_counts)?;

    cx.print_indentation()?;
    print_directive_result(&directive.result)?;

    Ok(())
}

fn print_include(cx: &mut Context, directive: &IncludeDirective, last: bool) -> io::Result<()> {
    use EvalError::*;

    cx.print_indentation_leaf(last)?;
    print_bold(&directive.mechanism)?;

    if let Some(target_domain) = &directive.target_domain {
        print!(" → {target_domain}")?;
    }

    print_lookup_counts(cx, &directive.lookup_counts)?;

    if let Some(query) = &directive.query {
        print_query(cx, query)?;
    }

    match &directive.result {
        DirectiveResult::Error(RecursiveTemperror | RecursivePermerror, spf_result) => {
            // Don’t show the useless error message when propagating a recursive
            // *temperror* or *permerror*.
            if cx.omit_recursive_error {
                cx.omit_recursive_error = false;
            } else {
                cx.print_indentation()?;
                println!("{} result={}", error!(), spf_result.to_str())?;
            }
        }
        result => {
            cx.print_indentation()?;
            print_directive_result(result)?;
        }
    }

    Ok(())
}

fn print_redirect(cx: &mut Context, modifier: &RedirectModifier, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(&modifier.modifier)?;

    if let Some(target_domain) = &modifier.target_domain {
        print!(" → {target_domain}")?;
    }

    print_lookup_counts(cx, &modifier.lookup_counts)?;

    if let Some(query) = &modifier.query {
        print_query(cx, query)?;
    }

    if let Some(e) = &modifier.error {
        cx.print_indentation()?;
        println!("{}: {}, result={}", error!(), e, modifier.result.to_str())?;
    }

    // If an error was printed above or was already printed in the query, then
    // the SPF result was printed already alongside it. Omit redundant output in
    // these cases.
    if cx.omit_recursive_error {
        cx.omit_recursive_error = false;
    } else if modifier.error.is_none() {
        cx.print_indentation()?;
        println!("result={}", modifier.result.to_str())?;
    }

    Ok(())
}

fn print_exp(cx: &mut Context, modifier: &ExpModifier, last: bool) -> io::Result<()> {
    cx.print_indentation_leaf(last)?;
    print_bold(&modifier.modifier)?;

    if let Some(target_domain) = &modifier.target_domain {
        println!(" → {target_domain}")?;
    }

    if let Some(explain_string) = &modifier.explain_string {
        cx.print_indentation()?;
        println!("{:?}", explain_string.to_string())?;
    }

    if let Some(e) = &modifier.error {
        cx.print_indentation()?;
        println!("{}: {}", error!(), e)?;
    }

    cx.print_indentation()?;
    match &modifier.explanation {
        ExplanationString::Default => println!("explanation=none")?,
        ExplanationString::External(s) => println!("explanation={s:?}")?,
    }

    Ok(())
}

fn print_bold(x: impl Display) -> io::Result<()> {
    print!("{}", x.if_supports_color(Stream::Stdout, |x| x.bold()))
}

fn print_directive_result(result: &DirectiveResult) -> io::Result<()> {
    match result {
        DirectiveResult::Match(spf_result) => println!("match result={}", spf_result.to_str()),
        DirectiveResult::NotMatch => println!("not-match"),
        DirectiveResult::Error(e, spf_result) => {
            println!("{}: {}, result={}", error!(), e, spf_result.to_str())
        }
    }
}

fn print_lookup_counts(cx: &mut Context, counts: &LookupCounts) -> io::Result<()> {
    macro_rules! color_count {
        ($max:expr, $n:expr) => {
            $n.if_supports_color(::owo_colors::Stream::Stdout, |n| {
                n.color(if $n <= $max {
                    ::owo_colors::AnsiColors::Green
                } else {
                    ::owo_colors::AnsiColors::Red
                })
            })
        };
    }

    cx.lookup_counts.lookups += counts.lookups;
    cx.lookup_counts.void += counts.void;

    print!(" (lookups: {}/10", color_count!(10, cx.lookup_counts.lookups))?;
    if counts.nested != 0 {
        print!(", nested: {}/10", color_count!(10, counts.nested))?;
    }
    if counts.void != 0 {
        print!(", void: {}/2", color_count!(2, cx.lookup_counts.void))?;
    }
    println!(")")?;

    Ok(())
}

fn print_timings(query_duration: Duration, lookup_times: &[TimedLookup]) -> io::Result<()> {
    fn print_duration(t: f32) -> io::Result<()> {
        if t < 1.0 {
            print!("{:8.2}ms", t * 1000.0)
        } else {
            print!("{:8.2}s ", t)
        }
    }

    println!()?;

    for t in lookup_times {
        print_duration(t.duration.as_secs_f32())?;
        match &t.lookup {
            LookupTarget::A(name)    => println!(" A    {name}")?,
            LookupTarget::Aaaa(name) => println!(" AAAA {name}")?,
            LookupTarget::Mx(name)   => println!(" MX   {name}")?,
            LookupTarget::Txt(name)  => println!(" TXT  {name}")?,
            LookupTarget::Ptr(ip)    => println!(" PTR  {ip}")?,
        }
    }

    print_duration(query_duration.as_secs_f32())?;
    println!(" total time")?;

    Ok(())
}