Skip to main content

regent_sdk/state/attribute/
shared.rs

1use crate::error::RegentError;
2use crate::hosts::command::CommandResult;
3use crate::hosts::managed_host::InternalApiCallOutcome;
4use crate::hosts::properties::{HostProperties, LinuxFlavor, OsKind};
5use crate::state::attribute::HostHandler;
6use crate::state::attribute::Privilege;
7use crate::state::attribute::container::container_spec::quoted_command;
8use serde::{Deserialize, Serialize};
9use std::collections::BTreeMap;
10use std::fmt;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(rename_all = "PascalCase")]
14pub enum PackageFamily {
15    /// `apt-get` / `dpkg` (Ubuntu, Debian)
16    Apt,
17    /// `dnf` / `rpm` (RHEL and rebuilds, Fedora, Amazon Linux)
18    Dnf,
19    /// `zypper` / `rpm` (SLES, openSUSE)
20    Zypper,
21}
22
23impl fmt::Display for PackageFamily {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            PackageFamily::Apt => write!(f, "apt"),
27            PackageFamily::Dnf => write!(f, "dnf"),
28            PackageFamily::Zypper => write!(f, "zypper"),
29        }
30    }
31}
32
33impl PackageFamily {
34    /// Directories (or files) holding the repository definitions
35    pub(crate) fn repository_locations(&self) -> &'static [&'static str] {
36        match self {
37            PackageFamily::Apt => &["/etc/apt/sources.list", "/etc/apt/sources.list.d"],
38            PackageFamily::Dnf => &["/etc/yum.repos.d"],
39            PackageFamily::Zypper => &["/etc/zypp/repos.d"],
40        }
41    }
42
43    /// Command telling (through its return code) whether a repository mentioning `url` is defined
44    pub(crate) fn repository_check_command(&self, url: &str) -> String {
45        let mut args = vec!["grep", "-rlsF", url];
46        args.extend(self.repository_locations());
47        quoted_command(&args)
48    }
49
50    /// Command listing the installed packages matching the given names or glob patterns.
51    /// Its output is parsed by [`parse_installed_packages`].
52    pub(crate) fn installed_packages_command(&self, patterns: &[&str]) -> String {
53        let mut args = match self {
54            PackageFamily::Apt => vec![
55                "dpkg-query",
56                "-W",
57                "-f=${Package} ${Version} ${db:Status-Abbrev}\\n",
58            ],
59            PackageFamily::Dnf | PackageFamily::Zypper => {
60                vec!["rpm", "-qa", "--qf", "%{NAME} %{VERSION}-%{RELEASE} ii\\n"]
61            }
62        };
63        args.extend(patterns);
64        quoted_command(&args)
65    }
66
67    /// Commands installing packages non interactively, after refreshing the apt package lists
68    /// (dnf and zypper refresh expired metadata themselves)
69    pub(crate) fn install_commands(
70        &self,
71        packages: &[String],
72        allow_downgrade: bool,
73    ) -> Vec<String> {
74        let mut commands = Vec::new();
75        if *self == PackageFamily::Apt {
76            commands.push(quoted_command(&["apt-get", "update"]));
77        }
78        commands.push(self.install_command(packages, allow_downgrade));
79        commands
80    }
81
82    pub(crate) fn install_command(&self, packages: &[String], allow_downgrade: bool) -> String {
83        let mut args: Vec<&str> = match self {
84            PackageFamily::Apt => vec!["apt-get", "install", "-y"],
85            PackageFamily::Dnf if allow_downgrade => vec!["dnf", "-y", "downgrade"],
86            PackageFamily::Dnf => vec!["dnf", "-y", "install"],
87            PackageFamily::Zypper => vec![
88                "zypper",
89                "--non-interactive",
90                "--gpg-auto-import-keys",
91                "install",
92                "--force-resolution",
93            ],
94        };
95        if allow_downgrade {
96            match self {
97                PackageFamily::Apt => args.push("--allow-downgrades"),
98                PackageFamily::Zypper => args.push("--oldpackage"),
99                PackageFamily::Dnf => {}
100            }
101        }
102        args.extend(packages.iter().map(String::as_str));
103        let command = quoted_command(&args);
104        match self {
105            PackageFamily::Apt => format!("DEBIAN_FRONTEND=noninteractive {}", command),
106            _ => command,
107        }
108    }
109
110    /// Command removing packages non interactively
111    pub(crate) fn remove_command(&self, packages: &[String]) -> String {
112        let mut args: Vec<&str> = match self {
113            PackageFamily::Apt => vec!["apt-get", "remove", "--autoremove", "--purge", "-y"],
114            PackageFamily::Dnf => vec!["dnf", "-y", "remove"],
115            PackageFamily::Zypper => vec!["zypper", "--non-interactive", "remove"],
116        };
117        args.extend(packages.iter().map(String::as_str));
118        let command = quoted_command(&args);
119        match self {
120            PackageFamily::Apt => format!("DEBIAN_FRONTEND=noninteractive {}", command),
121            _ => command,
122        }
123    }
124}
125
126/// Content of `/etc/os-release` (only the fields used to pick the NVIDIA repository)
127#[derive(Debug, Clone, Default, PartialEq, Eq)]
128pub struct OsRelease {
129    /// `ID` (e.g. `ubuntu`, `rhel`, `rocky`, `sles`)
130    pub id: String,
131    /// `ID_LIKE`, split on whitespace
132    pub id_like: Vec<String>,
133    /// `VERSION_ID` (e.g. `24.04`, `9.4`, `2023`)
134    pub version_id: String,
135    /// `PRETTY_NAME`, for messages
136    pub pretty_name: String,
137}
138
139impl OsRelease {
140    /// Major version (`VERSION_ID` up to the first dot)
141    pub fn major_version(&self) -> &str {
142        self.version_id
143            .split('.')
144            .next()
145            .unwrap_or(&self.version_id)
146    }
147
148    fn description(&self) -> String {
149        if self.pretty_name.is_empty() {
150            format!("{} {}", self.id, self.version_id)
151        } else {
152            self.pretty_name.clone()
153        }
154    }
155}
156
157/// Parse the content of `/etc/os-release` (values may be quoted or not)
158pub fn parse_os_release(content: &str) -> OsRelease {
159    let mut os_release = OsRelease::default();
160    for line in content.lines() {
161        let Some((key, value)) = line.trim().split_once('=') else {
162            continue;
163        };
164        let value = value.trim().trim_matches(|c| c == '"' || c == '\'');
165        match key {
166            "ID" => os_release.id = value.to_lowercase(),
167            "ID_LIKE" => {
168                os_release.id_like = value.split_whitespace().map(str::to_lowercase).collect()
169            }
170            "VERSION_ID" => os_release.version_id = value.to_string(),
171            "PRETTY_NAME" => os_release.pretty_name = value.to_string(),
172            _ => {}
173        }
174    }
175    os_release
176}
177
178/// Package manager of a Linux distribution, from the `ID` and `ID_LIKE` of `/etc/os-release`
179pub fn package_family(os_release: &OsRelease) -> Option<PackageFamily> {
180    std::iter::once(&os_release.id)
181        .chain(&os_release.id_like)
182        .find_map(|id| match id.as_str() {
183            "debian" | "ubuntu" => Some(PackageFamily::Apt),
184            "rhel" | "fedora" | "centos" | "rocky" | "almalinux" | "ol" | "amzn" => {
185                Some(PackageFamily::Dnf)
186            }
187            "suse" | "sles" | "sled" | "opensuse" | "opensuse-leap" | "opensuse-tumbleweed" => {
188                Some(PackageFamily::Zypper)
189            }
190            _ => None,
191        })
192}
193
194/// CPU architecture as named in the NVIDIA CUDA repositories
195#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
196#[serde(rename_all = "PascalCase")]
197pub enum CudaArch {
198    X86_64,
199    /// `sbsa`: 64-bit ARM servers (Arm SBSA, e.g. Grace). Jetson boards use JetPack instead.
200    Sbsa,
201}
202
203impl CudaArch {
204    pub fn repo_name(&self) -> &'static str {
205        match self {
206            CudaArch::X86_64 => "x86_64",
207            CudaArch::Sbsa => "sbsa",
208        }
209    }
210}
211
212/// Map the output of `uname -m` to a CUDA repository architecture
213pub fn parse_machine_arch(uname_m: &str) -> Result<CudaArch, String> {
214    match uname_m.trim() {
215        "x86_64" | "amd64" => Ok(CudaArch::X86_64),
216        "aarch64" | "arm64" => Ok(CudaArch::Sbsa),
217        other => Err(format!(
218            "CPU architecture {:?} is not supported by the NVIDIA CUDA repositories (x86_64 and aarch64 only)",
219            other
220        )),
221    }
222}
223
224/// How the NVIDIA driver packages are installed and how a driver branch is selected,
225/// following the NVIDIA Driver Installation Guide
226#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
227#[serde(rename_all = "PascalCase")]
228pub enum DriverDistro {
229    /// Ubuntu 22.04 / 24.04 / 26.04: branch pinned by installing its meta package
230    /// `nvidia-driver-<branch>[-open]`
231    Ubuntu,
232    /// Debian 12 / 13: branch pinned by installing its meta package
233    /// `nvidia-driver-<branch>[-open]`
234    Debian,
235    /// RHEL-compatible 8 and 9 (RHEL, Rocky, AlmaLinux, Oracle Linux, CentOS Stream):
236    /// branch selected with the `nvidia-driver` DNF module stream
237    RhelModular,
238    /// Amazon Linux 2023: branch selected with the `nvidia-driver` DNF module stream
239    AmazonLinux,
240    /// RHEL-compatible 10 (DNF 4): branch locked with the versionlock plugin
241    Rhel10,
242    /// Fedora (DNF 5): branch locked with `dnf versionlock`
243    Fedora,
244    /// SLES / openSUSE Leap: branch locked with a zypper lock
245    Suse,
246}
247
248impl DriverDistro {
249    pub fn package_family(&self) -> PackageFamily {
250        match self {
251            DriverDistro::Ubuntu | DriverDistro::Debian => PackageFamily::Apt,
252            DriverDistro::RhelModular
253            | DriverDistro::AmazonLinux
254            | DriverDistro::Rhel10
255            | DriverDistro::Fedora => PackageFamily::Dnf,
256            DriverDistro::Suse => PackageFamily::Zypper,
257        }
258    }
259
260    /// Kernel development packages DKMS needs to build the modules for `kernel_release` (`uname -r`)
261    pub fn kernel_headers_packages(&self, kernel_release: &str) -> Vec<String> {
262        match self {
263            DriverDistro::Ubuntu | DriverDistro::Debian => {
264                vec![format!("linux-headers-{}", kernel_release)]
265            }
266            DriverDistro::RhelModular => vec![
267                format!("kernel-devel-{}", kernel_release),
268                "kernel-headers".to_string(),
269            ],
270            DriverDistro::AmazonLinux => vec![
271                format!("kernel-devel-{}", kernel_release),
272                format!("kernel-headers-{}", kernel_release),
273            ],
274            DriverDistro::Rhel10 | DriverDistro::Fedora => vec![
275                "kernel-devel-matched".to_string(),
276                "kernel-headers".to_string(),
277            ],
278            DriverDistro::Suse => {
279                // 6.4.0-150600.23.25-default -> kernel-default-devel=6.4.0-150600.23.25
280                match kernel_release.rsplit_once('-') {
281                    Some((version, variant))
282                        if !variant.is_empty() && !variant.starts_with(char::is_numeric) =>
283                    {
284                        vec![format!("kernel-{}-devel={}", variant, version)]
285                    }
286                    _ => vec!["kernel-default-devel".to_string()],
287                }
288            }
289        }
290    }
291}
292
293/// NVIDIA CUDA network repository matching the host
294#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
295#[serde(rename_all = "PascalCase")]
296pub struct CudaRepository {
297    /// Distribution name in the repository URL (e.g. `ubuntu2404`, `rhel9`)
298    pub distro_name: String,
299    pub arch: CudaArch,
300    pub distro: DriverDistro,
301}
302
303pub(crate) const CUDA_REPOS_URL: &str = "https://developer.download.nvidia.com/compute/cuda/repos";
304
305impl CudaRepository {
306    /// `https://developer.download.nvidia.com/compute/cuda/repos/<distro>/<arch>/`
307    pub fn base_url(&self) -> String {
308        format!(
309            "{}/{}/{}/",
310            CUDA_REPOS_URL,
311            self.distro_name,
312            self.arch.repo_name()
313        )
314    }
315
316    /// Part of the URL looked for in the repository definitions
317    pub fn url_fragment(&self) -> String {
318        format!(
319            "developer.download.nvidia.com/compute/cuda/repos/{}/{}",
320            self.distro_name,
321            self.arch.repo_name()
322        )
323    }
324
325    /// Configure the repository (cuda-keyring for apt, `.repo` file for dnf and zypper)
326    pub(crate) async fn setup<Handler: HostHandler>(
327        &self,
328        host_handler: &mut Handler,
329        privilege: &Privilege,
330    ) -> Result<InternalApiCallOutcome, RegentError> {
331        let family = self.distro.package_family();
332        if family != PackageFamily::Apt {
333            return add_repo_file(
334                host_handler,
335                family,
336                &format!("{}cuda-{}.repo", self.base_url(), self.distro_name),
337                &format!("cuda-{}.repo", self.distro_name),
338                privilege,
339            )
340            .await;
341        }
342        let tmp = match make_temp_dir(host_handler, privilege).await? {
343            Ok(tmp) => tmp,
344            Err(failure) => return Ok(failure),
345        };
346        let keyring = format!("{}/cuda-keyring.deb", tmp);
347        let mut outcome = download(
348            host_handler,
349            &format!("{}cuda-keyring_1.1-1_all.deb", self.base_url()),
350            &keyring,
351            privilege,
352        )
353        .await;
354        if is_success(&outcome) {
355            outcome = run_steps(
356                host_handler,
357                &[
358                    quoted_command(&["dpkg", "-i", &keyring]),
359                    quoted_command(&["apt-get", "update"]),
360                ],
361                privilege,
362            )
363            .await;
364        }
365        remove_temp_dir(host_handler, &tmp, privilege).await;
366        Ok(outcome)
367    }
368}
369
370const SUPPORTED_DISTROS: &str = "Ubuntu 22.04/24.04/26.04, Debian 12/13, RHEL/Rocky/AlmaLinux/Oracle Linux/CentOS Stream 8/9/10, Fedora 42/43/44, Amazon Linux 2023, SLES 15/16, openSUSE Leap 15";
371
372/// Pick the NVIDIA CUDA repository for a host from its `/etc/os-release` and `uname -m`
373pub fn cuda_repository(os_release: &OsRelease, uname_m: &str) -> Result<CudaRepository, String> {
374    let arch = parse_machine_arch(uname_m)?;
375    let major = os_release.major_version();
376    let unsupported = || {
377        format!(
378            "{} is not supported by the NVIDIA CUDA network repository (supported: {})",
379            os_release.description(),
380            SUPPORTED_DISTROS
381        )
382    };
383
384    let (distro_name, distro) = match os_release.id.as_str() {
385        "ubuntu" => match os_release.version_id.as_str() {
386            "22.04" | "24.04" | "26.04" => (
387                format!("ubuntu{}", os_release.version_id.replace('.', "")),
388                DriverDistro::Ubuntu,
389            ),
390            _ => return Err(unsupported()),
391        },
392        "debian" => match major {
393            "12" | "13" => (format!("debian{}", major), DriverDistro::Debian),
394            _ => return Err(unsupported()),
395        },
396        "fedora" => match major {
397            "42" | "43" | "44" => (format!("fedora{}", major), DriverDistro::Fedora),
398            _ => return Err(unsupported()),
399        },
400        "amzn" => match major {
401            "2023" => ("amzn2023".to_string(), DriverDistro::AmazonLinux),
402            _ => return Err(unsupported()),
403        },
404        "sles" | "sled" => match major {
405            "15" | "16" => (format!("sles{}", major), DriverDistro::Suse),
406            _ => return Err(unsupported()),
407        },
408        "opensuse-leap" => match major {
409            "15" => ("opensuse15".to_string(), DriverDistro::Suse),
410            _ => return Err(unsupported()),
411        },
412        "rhel" | "rocky" | "almalinux" | "ol" | "centos" => {
413            rhel_like(major).ok_or_else(unsupported)?
414        }
415        // Other RHEL rebuilds
416        _ if os_release.id_like.iter().any(|id| id == "rhel") => {
417            rhel_like(major).ok_or_else(unsupported)?
418        }
419        _ => return Err(unsupported()),
420    };
421
422    Ok(CudaRepository {
423        distro_name,
424        arch,
425        distro,
426    })
427}
428
429fn rhel_like(major: &str) -> Option<(String, DriverDistro)> {
430    match major {
431        "8" | "9" => Some((format!("rhel{}", major), DriverDistro::RhelModular)),
432        "10" => Some(("rhel10".to_string(), DriverDistro::Rhel10)),
433        _ => None,
434    }
435}
436
437/// Parse the output of [`PackageFamily::installed_packages_command`]: installed package name -> version.
438///
439/// Lines are `<name> <version> <status>`; a package is installed when the second letter of
440/// its dpkg status is `i` (`ii`, `hi`). Other lines (warnings, "no packages found") are ignored.
441pub fn parse_installed_packages(output: &str) -> BTreeMap<String, String> {
442    output
443        .lines()
444        .filter_map(|line| {
445            let mut parts = line.split_whitespace();
446            let (name, version, status) = (parts.next()?, parts.next()?, parts.next()?);
447            let is_installed = status.chars().nth(1) == Some('i');
448            let is_version = version.starts_with(|c: char| c.is_ascii_digit());
449            (is_installed && is_version && parts.next().is_none())
450                .then(|| (name.to_string(), version.to_string()))
451        })
452        .collect()
453}
454
455/// Upstream part of a package version: `1:580.105.08-1ubuntu1` -> `580.105.08`
456pub fn upstream_version(version: &str) -> &str {
457    let without_epoch = version.split_once(':').map_or(version, |(_, rest)| rest);
458    without_epoch
459        .split_once('-')
460        .map_or(without_epoch, |(upstream, _)| upstream)
461}
462
463/// `systemctl is-enabled` of a unit
464#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
465pub enum UnitEnablement {
466    #[default]
467    NotFound,
468    Enabled,
469    Disabled,
470    /// `static`, `indirect`, `generated` or `transient`: started by other units, `systemctl
471    /// enable` and `systemctl disable` have no effect
472    Static,
473    Masked,
474}
475
476#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
477pub struct UnitState {
478    pub enablement: UnitEnablement,
479    pub active: bool,
480}
481
482impl UnitState {
483    pub fn exists(&self) -> bool {
484        self.enablement != UnitEnablement::NotFound
485    }
486}
487
488/// Parse the outputs of `systemctl is-enabled <unit>` and `systemctl is-active <unit>`. A missing
489/// unit prints an error (possibly folded into the output by the handler) instead of a state.
490pub fn parse_unit_state(is_enabled_output: &str, is_active_output: &str) -> UnitState {
491    let last_line = |output: &str| {
492        output
493            .lines()
494            .map(str::trim)
495            .rfind(|line| !line.is_empty())
496            .unwrap_or_default()
497            .to_string()
498    };
499    let enablement = match last_line(is_enabled_output).as_str() {
500        "enabled" | "enabled-runtime" | "alias" | "linked" | "linked-runtime" => {
501            UnitEnablement::Enabled
502        }
503        "disabled" => UnitEnablement::Disabled,
504        "static" | "indirect" | "generated" | "transient" => UnitEnablement::Static,
505        "masked" | "masked-runtime" => UnitEnablement::Masked,
506        _ => UnitEnablement::NotFound,
507    };
508    let active = matches!(last_line(is_active_output).as_str(), "active" | "reloading");
509    UnitState { enablement, active }
510}
511
512#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
513#[serde(rename_all = "PascalCase")]
514pub enum KernelModuleFlavor {
515    /// Open GPU kernel modules (MIT/GPLv2)
516    Open,
517    /// Proprietary (closed source) kernel modules
518    Proprietary,
519}
520
521impl fmt::Display for KernelModuleFlavor {
522    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
523        match self {
524            KernelModuleFlavor::Open => write!(f, "open"),
525            KernelModuleFlavor::Proprietary => write!(f, "proprietary"),
526        }
527    }
528}
529
530/// NVIDIA kernel module currently loaded
531#[derive(Debug, Clone, PartialEq, Eq)]
532pub struct LoadedDriver {
533    /// Driver version (e.g. `580.105.08`)
534    pub version: String,
535    pub flavor: KernelModuleFlavor,
536}
537
538/// Path of the file describing the loaded NVIDIA kernel module
539pub(crate) const PROC_DRIVER_VERSION: &str = "/proc/driver/nvidia/version";
540
541/// Parse `/proc/driver/nvidia/version`:
542///
543/// ```text
544/// NVRM version: NVIDIA UNIX Open Kernel Module for x86_64  580.105.08  Release Build  (dvs-builder@U16-I3-B03-4-3)  Thu Oct 30 2025
545/// NVRM version: NVIDIA UNIX x86_64 Kernel Module  550.54.14  Thu Feb 22 01:44:30 UTC 2024
546/// ```
547pub fn parse_loaded_driver(proc_version: &str) -> Option<LoadedDriver> {
548    let line = proc_version
549        .lines()
550        .find(|line| line.trim_start().starts_with("NVRM version:"))?;
551    let version = line.split_whitespace().find(|token| {
552        let parts: Vec<&str> = token.split('.').collect();
553        parts.len() >= 2
554            && parts
555                .iter()
556                .all(|part| !part.is_empty() && part.chars().all(|c| c.is_ascii_digit()))
557    })?;
558    Some(LoadedDriver {
559        version: version.to_string(),
560        flavor: if line.contains("Open Kernel Module") {
561            KernelModuleFlavor::Open
562        } else {
563            KernelModuleFlavor::Proprietary
564        },
565    })
566}
567
568/// Read the loaded NVIDIA driver, `None` if the module is not loaded
569pub(crate) async fn loaded_driver<Handler: HostHandler>(
570    host_handler: &mut Handler,
571    privilege: &Privilege,
572) -> Result<Option<LoadedDriver>, RegentError> {
573    let result = query(
574        host_handler,
575        &quoted_command(&["cat", PROC_DRIVER_VERSION]),
576        privilege,
577    )
578    .await?;
579    if result.return_code != 0 {
580        return Ok(None);
581    }
582    Ok(parse_loaded_driver(&result.stdout))
583}
584
585/// Linux only: other OS kinds are rejected
586pub(crate) fn check_linux_host(
587    host_properties: &HostProperties,
588    attribute: &str,
589) -> Result<(), RegentError> {
590    match host_properties.os_kind() {
591        OsKind::Linux(specifics) => match specifics.linux_flavor {
592            LinuxFlavor::Arch | LinuxFlavor::Gentoo => Err(RegentError::IncompatibleHost(format!(
593                "Host is {:?} Linux but {} only supports distributions with an NVIDIA network repository (apt, dnf or zypper based)",
594                specifics.linux_flavor, attribute
595            ))),
596            _ => Ok(()),
597        },
598        incompatible_os_kind => Err(RegentError::IncompatibleHost(format!(
599            "Host is {:?} but {} is only supported on Linux",
600            incompatible_os_kind, attribute
601        ))),
602    }
603}
604
605/// Run a read-only command during assessment. Only a failure to run the command is an error:
606/// the caller interprets the return code.
607pub(crate) async fn query<Handler: HostHandler>(
608    host_handler: &mut Handler,
609    command: &str,
610    privilege: &Privilege,
611) -> Result<CommandResult, RegentError> {
612    host_handler
613        .run_command(command, privilege)
614        .await
615        .map_err(|e| {
616            RegentError::FailedDryRunEvaluation(format!("Unable to run {:?}: {:?}", command, e))
617        })
618}
619
620/// Run a command during assessment and return its output, failing on a non-zero return code
621pub(crate) async fn query_output<Handler: HostHandler>(
622    host_handler: &mut Handler,
623    command: &str,
624    privilege: &Privilege,
625) -> Result<String, RegentError> {
626    let result = query(host_handler, command, privilege).await?;
627    if result.return_code != 0 {
628        return Err(RegentError::FailedDryRunEvaluation(format!(
629            "{:?} failed (RC {}): {}{}",
630            command, result.return_code, result.stdout, result.stderr
631        )));
632    }
633    Ok(result.stdout)
634}
635
636/// Installed packages matching the given names or glob patterns: name -> version
637pub(crate) async fn installed_packages<Handler: HostHandler>(
638    host_handler: &mut Handler,
639    family: PackageFamily,
640    patterns: &[&str],
641    privilege: &Privilege,
642) -> Result<BTreeMap<String, String>, RegentError> {
643    // dpkg-query returns 1 when a pattern matches nothing: the output is what matters
644    let result = query(
645        host_handler,
646        &family.installed_packages_command(patterns),
647        privilege,
648    )
649    .await?;
650    Ok(parse_installed_packages(&result.stdout))
651}
652
653pub(crate) async fn is_repository_configured<Handler: HostHandler>(
654    host_handler: &mut Handler,
655    family: PackageFamily,
656    url_fragment: &str,
657    privilege: &Privilege,
658) -> Result<bool, RegentError> {
659    let result = query(
660        host_handler,
661        &family.repository_check_command(url_fragment),
662        privilege,
663    )
664    .await?;
665    Ok(result.return_code == 0 && !result.stdout.trim().is_empty())
666}
667
668pub(crate) async fn unit_state<Handler: HostHandler>(
669    host_handler: &mut Handler,
670    unit: &str,
671    privilege: &Privilege,
672) -> Result<UnitState, RegentError> {
673    // Both commands return non-zero for disabled or inactive units: the output is what matters
674    let is_enabled = query(
675        host_handler,
676        &quoted_command(&["systemctl", "is-enabled", unit]),
677        privilege,
678    )
679    .await?;
680    let is_active = query(
681        host_handler,
682        &quoted_command(&["systemctl", "is-active", unit]),
683        privilege,
684    )
685    .await?;
686    Ok(parse_unit_state(&is_enabled.stdout, &is_active.stdout))
687}
688
689/// Names of the loaded kernel modules, from `/proc/modules`
690pub(crate) async fn loaded_kernel_modules<Handler: HostHandler>(
691    host_handler: &mut Handler,
692    privilege: &Privilege,
693) -> Result<Vec<String>, RegentError> {
694    let modules = query_output(
695        host_handler,
696        &quoted_command(&["cat", "/proc/modules"]),
697        privilege,
698    )
699    .await?;
700    Ok(parse_proc_modules(&modules))
701}
702
703pub fn parse_proc_modules(content: &str) -> Vec<String> {
704    content
705        .lines()
706        .filter_map(|line| line.split_whitespace().next())
707        .map(str::to_string)
708        .collect()
709}
710
711/// Package family of the host, from `/etc/os-release`
712pub(crate) async fn host_package_family<Handler: HostHandler>(
713    host_handler: &mut Handler,
714    privilege: &Privilege,
715) -> Result<Option<PackageFamily>, RegentError> {
716    let os_release = parse_os_release(
717        &query_output(
718            host_handler,
719            &quoted_command(&["cat", "/etc/os-release"]),
720            privilege,
721        )
722        .await?,
723    );
724    Ok(package_family(&os_release))
725}
726
727/// Download `url` to `destination` with curl, or wget if curl is missing
728pub(crate) async fn download<Handler: HostHandler>(
729    host_handler: &mut Handler,
730    url: &str,
731    destination: &str,
732    privilege: &Privilege,
733) -> InternalApiCallOutcome {
734    let has_curl = host_handler
735        .is_this_command_available("curl", privilege)
736        .await
737        .unwrap_or(false);
738    let command = if has_curl {
739        quoted_command(&["curl", "-fsSL", "-o", destination, url])
740    } else {
741        quoted_command(&["wget", "-q", "-O", destination, url])
742    };
743    run_step(host_handler, &command, privilege).await
744}
745
746/// Create a private temporary directory: its path, or the failure outcome
747pub(crate) async fn make_temp_dir<Handler: HostHandler>(
748    host_handler: &mut Handler,
749    privilege: &Privilege,
750) -> Result<Result<String, InternalApiCallOutcome>, RegentError> {
751    let result = query(host_handler, &quoted_command(&["mktemp", "-d"]), privilege).await?;
752    let path = result.stdout.trim();
753    if result.return_code != 0 || !path.starts_with('/') || path.contains(char::is_whitespace) {
754        return Ok(Err(InternalApiCallOutcome::Failure(format!(
755            "mktemp -d failed (RC {}): {}{}",
756            result.return_code, result.stdout, result.stderr
757        ))));
758    }
759    Ok(Ok(path.to_string()))
760}
761
762/// Best effort: a leftover temporary directory doesn't change the outcome
763pub(crate) async fn remove_temp_dir<Handler: HostHandler>(
764    host_handler: &mut Handler,
765    path: &str,
766    privilege: &Privilege,
767) {
768    let _ = host_handler
769        .run_command(&quoted_command(&["rm", "-rf", path]), privilege)
770        .await;
771}
772
773/// Add a repository described by a `.repo` file (dnf and zypper)
774pub(crate) async fn add_repo_file<Handler: HostHandler>(
775    host_handler: &mut Handler,
776    family: PackageFamily,
777    repo_file_url: &str,
778    repo_file_name: &str,
779    privilege: &Privilege,
780) -> Result<InternalApiCallOutcome, RegentError> {
781    Ok(match family {
782        PackageFamily::Apt => InternalApiCallOutcome::Failure(
783            "apt repositories are not described by a .repo file".to_string(),
784        ),
785        PackageFamily::Dnf => {
786            let outcome = download(
787                host_handler,
788                repo_file_url,
789                &format!("/etc/yum.repos.d/{}", repo_file_name),
790                privilege,
791            )
792            .await;
793            if !is_success(&outcome) {
794                return Ok(outcome);
795            }
796            run_step(
797                host_handler,
798                &quoted_command(&["dnf", "clean", "expire-cache"]),
799                privilege,
800            )
801            .await
802        }
803        PackageFamily::Zypper => {
804            run_steps(
805                host_handler,
806                &[
807                    quoted_command(&["zypper", "--non-interactive", "addrepo", repo_file_url]),
808                    quoted_command(&[
809                        "zypper",
810                        "--non-interactive",
811                        "--gpg-auto-import-keys",
812                        "refresh",
813                    ]),
814                ],
815                privilege,
816            )
817            .await
818        }
819    })
820}
821
822pub(crate) fn is_success(outcome: &InternalApiCallOutcome) -> bool {
823    !matches!(outcome, InternalApiCallOutcome::Failure(_))
824}
825
826/// Run a remediation command: a non-zero return code is a failure carrying the output
827pub(crate) async fn run_step<Handler: HostHandler>(
828    host_handler: &mut Handler,
829    command: &str,
830    privilege: &Privilege,
831) -> InternalApiCallOutcome {
832    match host_handler.run_command(command, privilege).await {
833        Ok(result) if result.return_code == 0 => InternalApiCallOutcome::Success(None),
834        Ok(result) => InternalApiCallOutcome::Failure(format!(
835            "RC: {}, STDOUT: {}, STDERR: {}",
836            result.return_code, result.stdout, result.stderr
837        )),
838        Err(e) => InternalApiCallOutcome::Failure(format!("Command execution failed: {:?}", e)),
839    }
840}
841
842/// Run remediation commands in order, stopping at the first failure
843pub(crate) async fn run_steps<Handler: HostHandler>(
844    host_handler: &mut Handler,
845    commands: &[String],
846    privilege: &Privilege,
847) -> InternalApiCallOutcome {
848    for command in commands {
849        let outcome = run_step(host_handler, command, privilege).await;
850        if let InternalApiCallOutcome::Failure(_) = outcome {
851            return outcome;
852        }
853    }
854    InternalApiCallOutcome::Success(None)
855}
856
857/// Packages of the NVIDIA Container Toolkit, never removed with the driver
858pub(crate) fn is_container_toolkit_package(name: &str) -> bool {
859    name.starts_with("nvidia-container") || name.starts_with("libnvidia-container")
860}