agesafetch 1.0.0

// SPDX-FileCopyrightText: Benedikt Vollmerhaus <benedikt@vollmerhaus.org>
// SPDX-License-Identifier: MIT
use caps::{CapSet, Capability};
use clap::Parser;
use colored::{ColoredString, Colorize};
use indoc::{eprintdoc, indoc, printdoc};
use linux_memutils::agesa::{
    find_agesa_version, find_agesa_version_in_memory_region,
    get_reserved_regions_in_extended_memory, FoundVersion, SearchError, SearchResult,
};
use linux_memutils::iomem::MemoryRegion;
use std::error::Error;
use std::io;
use std::io::IsTerminal;
use std::process::ExitCode;
use std::sync::LazyLock;
use std::time::{Duration, Instant};

#[derive(Parser, Debug)]
#[command(version, about, after_help = indoc! {"
    Exit Codes:
      0: An AGESA version was found
      1: No version was found in any searched memory region
      2: /proc/iomem could not be read, e.g. due to insufficient permissions
      3: /dev/mem could not be opened, e.g. due to insufficient permissions
      4: An unhandled error occurred while reading a byte in /dev/mem
"})]
struct Cli {
    /// Print only the found version (default if not in a TTY)
    #[arg(short, long, conflicts_with = "verbose")]
    quiet: bool,
    /// Print further information and a closing search summary
    #[arg(short, long)]
    verbose: bool,
}

static STATUS_PREFIX: LazyLock<ColoredString> = LazyLock::new(|| "::".blue().bold());
static RESULT_PREFIX: LazyLock<ColoredString> = LazyLock::new(|| "->".yellow().bold());
static ERROR_PREFIX: LazyLock<ColoredString> = LazyLock::new(|| "ERR".red().bold());

fn main() -> ExitCode {
    let cli = Cli::parse();

    if !caps::has_cap(None, CapSet::Effective, Capability::CAP_SYS_ADMIN).unwrap() {
        eprintdoc! {"
            {} Missing privileges for reading a memory map from /proc/iomem.
                Please run agesafetch as root or add the SYS_ADMIN capability.
            ",
            *ERROR_PREFIX,
        }
        return ExitCode::from(2);
    }

    match find_and_print_agesa_version(&cli) {
        Err(SearchError::IomemUnreadable(_)) => {
            eprintln!(
                "{} Could not read /proc/iomem. Are its permissions correct?",
                *ERROR_PREFIX,
            );
            ExitCode::from(2)
        }
        Err(SearchError::DevMemUnopenable(_)) => {
            eprintdoc! {"
                {} Could not open /dev/mem.
                    Please run agesafetch as root or add suitable capabilities.
                ",
                *ERROR_PREFIX,
            }
            ExitCode::from(3)
        }
        Err(err @ SearchError::ByteUnreadable(_)) => {
            eprintln!(
                "{} Unhandled error while reading byte in physical memory: {}",
                *ERROR_PREFIX,
                err.source().unwrap(),
            );
            ExitCode::from(4)
        }
        Ok(()) => ExitCode::SUCCESS,
    }
}

fn find_and_print_agesa_version(cli: &Cli) -> Result<(), SearchError> {
    if !io::stdout().is_terminal() || cli.quiet {
        if let Some(found_version) = find_agesa_version()? {
            println!("{}", found_version.agesa_version);
        } else {
            eprintln!("Did not find AGESA version.");
        }

        return Ok(());
    }

    let reserved_regions =
        get_reserved_regions_in_extended_memory().map_err(SearchError::IomemUnreadable)?;

    if cli.verbose {
        println!(
            "{} Memory map lists {} regions of type {} in extended memory.",
            *STATUS_PREFIX,
            reserved_regions.len().to_string().blue(),
            "Reserved".italic(),
        );
    }

    let search_start_time = Instant::now();
    let found_version = search_regions_and_print_statuses(reserved_regions)?;
    let search_duration = search_start_time.elapsed();

    if let Some(found_version) = found_version {
        println!(
            "{} Found AGESA version: {}",
            *RESULT_PREFIX,
            found_version.agesa_version.green().bold(),
        );

        if cli.verbose {
            print_search_summary(&found_version, &search_duration);
        }
    } else {
        eprintln!("{} Did not find AGESA version.", *RESULT_PREFIX);
    }

    Ok(())
}

/// Sequentially search each of the given regions for an AGESA version.
///
/// This returns as soon as a version is found or an error has occurred.
fn search_regions_and_print_statuses(regions: Vec<MemoryRegion>) -> SearchResult {
    for (i, region) in regions.into_iter().enumerate() {
        println!(
            "{} Searching {} region {} ({} KiB)...",
            *STATUS_PREFIX,
            region.region_type.to_string().italic(),
            format!("#{}", i + 1).blue(),
            region.size() / 1024,
        );

        match find_agesa_version_in_memory_region(region) {
            Ok(None) => (),
            result => return result,
        }
    }

    Ok(None)
}

/// Print a concise summary with information about the completed search.
fn print_search_summary(found_version: &FoundVersion, search_duration: &Duration) {
    printdoc! {"
        {} Search Summary:
           * Found at Physical Address: {:#x} (in {dev_mem})
           * Surrounding Memory Region: {}
           * Region Size:               {} KiB
           * Bytes Processed in Region: {} KiB
           * Search Time:               {} ms
        ",
        *STATUS_PREFIX,
        found_version.absolute_address,
        found_version.surrounding_region,
        found_version.surrounding_region.size() / 1024,
        found_version.offset_in_region() / 1024,
        search_duration.as_millis(),
        dev_mem = "/dev/mem".dimmed(),
    }
}