r2smt-cli 0.1.1

//! Shared CLI support: analysis-context plumbing and provider
//! lifecycle helpers used by every subcommand module.

use std::path::Path;

use anyhow::{Context, Result};
use r2smt_common::{Address, Arch};
use r2smt_core::{
    Finding, classify_finding_with_pretty, classify_lowered_upstream, dump_program, prepare_ssa,
    reconcile_folded,
};
use r2smt_ir::BinaryProvider;
use r2smt_ir::Decompiler;
use r2smt_ir::NameHints;
use r2smt_ir::program::{Function, Program};
use r2smt_r2pipe::{AnalysisLevel, R2PipeProvider};
use r2smt_slicer::{BranchCandidate, SliceLimits, collect_branches, collect_function_branches};
use r2smt_smt::SolveOptions;

use crate::args::SolverArg;
use crate::render::truncate_on_char_boundary;

pub(crate) fn analysis_level(deep: bool) -> AnalysisLevel {
    if deep {
        AnalysisLevel::Deep
    } else {
        AnalysisLevel::Standard
    }
}

pub(crate) fn open_provider(file: &Path, deep: bool) -> Result<R2PipeProvider> {
    R2PipeProvider::open_with_analysis(file, false, analysis_level(deep))
        .with_context(|| format!("opening {} with radare2", file.display()))
}

pub(crate) fn open_provider_writable(file: &Path, deep: bool) -> Result<R2PipeProvider> {
    R2PipeProvider::open_with_analysis(file, true, analysis_level(deep))
        .with_context(|| format!("opening {} with radare2 (-w)", file.display()))
}

/// Owns the [`Program`] returned by r2 plus any synthesised functions
/// produced by the shellcode finder. Subcommands look up branches
/// against the union, so candidates created from a synthetic block
/// still resolve.
pub(crate) struct AnalysisContext {
    pub(crate) program: Program,
    pub(crate) extra_functions: Vec<Function>,
}

impl AnalysisContext {
    pub(crate) fn new(program: Program) -> Self {
        Self {
            program,
            extra_functions: Vec::new(),
        }
    }

    pub(crate) fn find_function(&self, address: Address) -> Option<&Function> {
        self.program
            .functions
            .iter()
            .find(|f| f.address == address)
            .or_else(|| self.extra_functions.iter().find(|f| f.address == address))
    }

    pub(crate) fn all_functions(&self) -> impl Iterator<Item = &Function> {
        self.program
            .functions
            .iter()
            .chain(self.extra_functions.iter())
    }
}

/// Build the list of candidate branches the user asked about, plus
/// the augmented [`AnalysisContext`] (potentially extended with a
/// synthetic block when `--at addr` points outside any analysed
/// function).
pub(crate) fn resolve_targets(
    provider: &mut R2PipeProvider,
    file: &Path,
    program: Program,
    at: Option<&str>,
    function_filter: Option<&str>,
) -> Result<(AnalysisContext, Vec<BranchCandidate>)> {
    let mut ctx = AnalysisContext::new(program);

    let candidates: Vec<BranchCandidate> = match function_filter {
        None => collect_branches(&ctx.program),
        Some(raw) => {
            let address: Address = raw
                .parse()
                .with_context(|| format!("parsing --function value '{raw}'"))?;
            let function = ctx
                .program
                .functions
                .iter()
                .find(|f| f.address == address)
                .ok_or_else(|| anyhow::anyhow!("no function at {address} in {}", file.display()))?;
            collect_function_branches(function, ctx.program.arch)
        }
    };

    let filtered: Vec<BranchCandidate> = if let Some(at_raw) = at {
        let target: Address = at_raw
            .parse()
            .with_context(|| format!("parsing --at value '{at_raw}'"))?;
        if let Some(found) = candidates.into_iter().find(|c| c.address == target) {
            vec![found]
        } else {
            // Shellcode / unanalysed region fallback: synthesise the
            // basic block around `target` and look for the branch
            // inside it.
            let func = provider
                .load_block_at(target)
                .with_context(|| format!("synthesising block at {target}"))?;
            let synth_candidates = collect_function_branches(&func, ctx.program.arch);
            let candidate = synth_candidates
                .into_iter()
                .find(|c| c.address == target)
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "no candidate at {target} (program had no match; synthetic block at {} had no conditional branch at the requested address)",
                        func.address,
                    )
                })?;
            ctx.extra_functions.push(func);
            vec![candidate]
        }
    } else {
        candidates
    };

    Ok((ctx, filtered))
}

/// The functions the differential-lift harness should cross-check.
///
/// [`resolve_targets`] narrows the *branch candidates* and never touches
/// `ctx.program.functions`, so the harness used to walk the whole binary
/// even under `--at` / `--function`. Triaging one disagreement then cost
/// a full re-run of the program.
///
/// `--function` resolves against the program directly rather than
/// through `candidates`, so a function holding no conditional branch is
/// still cross-checked. `--at` reads the containing function off the
/// single candidate [`resolve_targets`] kept, which is also what carries
/// the synthesised block when the address lies outside any analysed
/// function.
pub(crate) fn difflift_scope(
    ctx: &AnalysisContext,
    at: Option<&str>,
    function_filter: Option<&str>,
    candidates: &[BranchCandidate],
) -> Vec<Function> {
    let target = function_filter
        .and_then(|raw| raw.parse::<Address>().ok())
        .or_else(|| {
            candidates
                .first()
                .map(|c| c.function)
                .filter(|_| at.is_some())
        });
    match target.and_then(|addr| ctx.find_function(addr)) {
        Some(function) => vec![function.clone()],
        None => ctx.all_functions().cloned().collect(),
    }
}

/// Per-function pseudocode byte budget. Host-Side Safety: the cache
/// is keyed by function (a finite set) and every entry is truncated
/// on a UTF-8 boundary so a pathological decompilation cannot blow
/// host RAM or the `CCu` payload.
const MAX_PSEUDOCODE_BYTES: usize = 16 * 1024;

/// Attach decompiler pseudocode to every finding, one decompile per
/// function (cached). Best-effort: a missing backend or transport
/// hiccup leaves `pseudocode` as `None` and never fails the run.
pub(crate) fn attach_pseudocode(provider: &mut R2PipeProvider, findings: &mut [Finding]) {
    let mut cache: std::collections::BTreeMap<Address, Option<String>> =
        std::collections::BTreeMap::new();
    for f in findings.iter_mut() {
        let entry = cache.entry(f.function).or_insert_with(|| {
            provider
                .pseudocode(f.function)
                .ok()
                .flatten()
                .map(|s| truncate_on_char_boundary(&s, MAX_PSEUDOCODE_BYTES))
        });
        f.pseudocode = entry.clone();
    }
}

// `clippy::too_many_arguments`: same rationale as `solve` / `annotate`
// — a CLI driver threading through independent, read-at-distinct-stages
// knobs (the `with_decompiler` opt-in is the 8th). A params struct
// would only relocate the noise.
#[allow(clippy::too_many_arguments)]
pub(crate) fn compute_findings(
    file: &Path,
    deep: bool,
    at: Option<&str>,
    function_filter: Option<&str>,
    limits: &SliceLimits,
    options: SolveOptions,
    solver: SolverArg,
    with_decompiler: bool,
    ir_pcode: bool,
) -> Result<(Arch, Vec<Finding>)> {
    let mut provider = open_provider(file, deep)?;
    provider.set_attach_pcode(ir_pcode);
    let program = dump_program(&mut provider)
        .with_context(|| format!("loading program from {}", file.display()))?;
    let arch = program.arch;

    let (ctx, filtered) = resolve_targets(&mut provider, file, program, at, function_filter)?;

    let mut hint_cache: std::collections::BTreeMap<Address, NameHints> =
        std::collections::BTreeMap::new();
    let mut findings: Vec<Finding> = Vec::with_capacity(filtered.len());
    for cand in &filtered {
        if let Some(finding) = resolve_folded_branch(
            &mut provider,
            cand,
            ctx.program.arch,
            limits,
            solver,
            options,
        )? {
            findings.push(finding);
            continue;
        }
        let Some(function) = ctx.find_function(cand.function) else {
            continue;
        };
        let ssa = prepare_ssa(function, cand, limits, ctx.program.arch);
        let (verdict, z3_pretty) = dispatch_solver(solver, &ssa, options)?;
        let hints = hint_cache
            .entry(cand.function)
            .or_insert_with(|| provider.name_hints(cand.function).unwrap_or_default());
        findings.push(classify_finding_with_pretty(
            &ssa, verdict, z3_pretty, hints,
        ));
    }
    if with_decompiler {
        attach_pseudocode(&mut provider, &mut findings);
    }
    Ok((arch, findings))
}

/// Build the SSA slices for the targeted branches without solving them.
/// Shares the open → dump → resolve → `prepare_ssa` path with
/// [`compute_findings`], but returns the raw `SsaLiftedSlice`s for
/// consumers (like the taint pass) that run their own analysis.
///
/// # Errors
///
/// Propagates provider / program-load failures.
pub(crate) fn compute_slices(
    file: &Path,
    deep: bool,
    at: Option<&str>,
    function_filter: Option<&str>,
    limits: &SliceLimits,
    ir_pcode: bool,
) -> Result<(Arch, Vec<r2smt_ssa::SsaLiftedSlice>)> {
    let mut provider = open_provider(file, deep)?;
    provider.set_attach_pcode(ir_pcode);
    let program = dump_program(&mut provider)
        .with_context(|| format!("loading program from {}", file.display()))?;
    let arch = program.arch;
    let (ctx, filtered) = resolve_targets(&mut provider, file, program, at, function_filter)?;
    let mut slices = Vec::with_capacity(filtered.len());
    for cand in &filtered {
        if let Some(function) = ctx.find_function(cand.function) {
            slices.push(prepare_ssa(function, cand, limits, ctx.program.arch));
        }
    }
    Ok((arch, slices))
}

/// Dispatch a single solve request to the selected backend. CVC5
/// failures (subprocess missing, garbled output) are surfaced as
/// `Err(anyhow::Error)` so the CLI can return a clear message to the
/// user; the Z3 path is infallible by contract and always succeeds.
///
/// Returns the verdict plus the C-style infix rendering of the
/// post-`aggressive_simplify` Z3 formula when the Z3 backend was
/// used; CVC5 has no Z3 AST so it returns `None`.
pub(crate) fn dispatch_solver(
    solver: SolverArg,
    slice: &r2smt_ssa::SsaLiftedSlice,
    options: SolveOptions,
) -> Result<(r2smt_common::smt::SmtResult, Option<String>)> {
    let backend = build_solver(solver);
    match backend.solve(slice, options) {
        Ok(outcome) => Ok((outcome.verdict, outcome.formula_pretty)),
        // Byte-identical to the pre-port message: the adapter's
        // `detail()` is the original text minus the backend prefix.
        Err(err) => Err(anyhow::anyhow!(
            "{} backend: {}",
            backend.name(),
            err.detail()
        )),
    }
}

/// Composition-root factory — the only place that knows concrete
/// solver adapter types. Exhaustive structural mapping over the CLI
/// enum (no domain logic): the documented exhaustive-dispatch-table
/// exception applies.
fn build_solver(solver: SolverArg) -> Box<dyn r2smt_solver_port::Solver> {
    match solver {
        SolverArg::Z3 => Box::new(r2smt_smt::Z3Solver),
        SolverArg::Cvc5 => Box::new(r2smt_smt::Cvc5Solver),
        SolverArg::Bitwuzla => Box::new(r2smt_smt::BitwuzlaSolver),
    }
}

/// Per-branch analysis-maturity fallback (always active).
///
/// Returns `Ok(None)` for branches radare2 did **not** fold (the
/// caller's normal slice → solve pipeline handles them unchanged). For
/// a branch `aaa` already collapsed to a single successor, it
/// independently re-derives a verdict: `BinaryProvider::load_block_at`
/// performs a *raw linear decode* of the containing block (its
/// shellcode-synthesis step does not apply r2's CFG folding), so the
/// genuine two-way `jcc` reappears and goes through the full
/// slice → optimize → SMT pipeline. The result is then reconciled
/// against the CFG shortcut via [`reconcile_folded`]: a sound SMT
/// proof wins; an inconclusive one falls back to r2's CFG verdict.
///
/// Cost is bounded — at most one extra block synthesis + one SMT
/// solve per folded branch — and accepted by design (always-on).
pub(crate) fn resolve_folded_branch(
    provider: &mut R2PipeProvider,
    cand: &BranchCandidate,
    arch: r2smt_common::Arch,
    limits: &SliceLimits,
    solver: SolverArg,
    options: SolveOptions,
) -> Result<Option<Finding>> {
    if cand.upstream_resolved.is_none() {
        return Ok(None);
    }
    let rederived = match provider.load_block_at(cand.address) {
        Ok(synth) => collect_function_branches(&synth, arch)
            .iter()
            .find(|c| c.address == cand.address)
            .map(|scand| {
                let ssa = prepare_ssa(&synth, scand, limits, arch);
                let (verdict, z3) = dispatch_solver(solver, &ssa, options)?;
                Ok::<_, anyhow::Error>(classify_finding_with_pretty(
                    &ssa,
                    verdict,
                    z3,
                    &NameHints::default(),
                ))
            })
            .transpose()?,
        // Synthesis failed (unmapped / undecodable): no re-derivation,
        // the CFG shortcut below is the only evidence available.
        Err(_) => None,
    };
    Ok(reconcile_folded(rederived, classify_lowered_upstream(cand)))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::panic)]

    use super::{AnalysisContext, difflift_scope};
    use r2smt_common::{Address, Arch};
    use r2smt_ir::program::{Function, Program};
    use r2smt_slicer::{BranchCandidate, BranchCondition, BranchKind};

    fn function(address: u64) -> Function {
        Function {
            address: Address::new(address),
            name: None,
            blocks: Vec::new(),
            is_thumb: false,
        }
    }

    fn ctx() -> AnalysisContext {
        AnalysisContext::new(Program {
            arch: Arch::X86_64,
            bits: 64,
            entry: None,
            functions: vec![function(0x1000), function(0x2000)],
        })
    }

    fn candidate_in(function: u64) -> BranchCandidate {
        BranchCandidate {
            address: Address::new(function + 4),
            function: Address::new(function),
            block: Address::new(function),
            kind: BranchKind::Jcc,
            mnemonic: "je".into(),
            condition: BranchCondition::Equal,
            formula: "ZF".into(),
            taken_target: None,
            fallthrough_target: None,
            compare_register: None,
            bit_index: None,
            upstream_resolved: None,
            operand_raws: Vec::new(),
            is_thumb: false,
        }
    }

    #[test]
    fn test_difflift_scope_unfiltered_covers_the_whole_program() {
        let scope = difflift_scope(&ctx(), None, None, &[]);
        assert_eq!(scope.len(), 2);
    }

    #[test]
    fn test_difflift_scope_function_filter_keeps_only_that_function() {
        let scope = difflift_scope(&ctx(), None, Some("0x2000"), &[]);
        assert_eq!(
            scope.iter().map(|f| f.address).collect::<Vec<_>>(),
            vec![Address::new(0x2000)]
        );
    }

    #[test]
    fn test_difflift_scope_function_filter_holds_without_any_candidate() {
        // A function with no conditional branch yields no candidate, and
        // must still be cross-checked rather than widening to the program.
        let scope = difflift_scope(&ctx(), None, Some("0x1000"), &[]);
        assert_eq!(scope.len(), 1);
    }

    #[test]
    fn test_difflift_scope_at_keeps_the_function_containing_the_address() {
        let scope = difflift_scope(&ctx(), Some("0x2004"), None, &[candidate_in(0x2000)]);
        assert_eq!(
            scope.iter().map(|f| f.address).collect::<Vec<_>>(),
            vec![Address::new(0x2000)]
        );
    }

    #[test]
    fn test_difflift_scope_falls_back_to_the_program_when_the_function_is_absent() {
        let scope = difflift_scope(&ctx(), None, Some("0x9999"), &[]);
        assert_eq!(scope.len(), 2);
    }
}