omnyssh_core/ssh/
probe.rs

1//! Probe script generation and output parsing for service discovery.
2//!
3//! The Quick Scan probe is a single bash script that collects maximum
4//! information in one SSH invocation. Output is delimited by
5//! section markers (`===OMNYSSH:SECTION===`) for easy parsing.
6
7use std::collections::HashMap;
8
9/// Generates the Quick Scan probe bash script.
10///
11/// This script runs multiple commands and delimits their output with
12/// section markers for structured parsing. All commands use stderr
13/// redirection to /dev/null for graceful failure.
14pub fn generate_quick_scan_script() -> &'static str {
15    r#"cat << 'OMNYSSH_PROBE_EOF' | bash
16echo "===OMNYSSH:OS==="
17cat /etc/os-release 2>/dev/null | head -5
18echo "===OMNYSSH:SERVICES==="
19systemctl list-units --type=service --state=running --no-pager --no-legend 2>/dev/null | awk '{print $1}' | head -50
20echo "===OMNYSSH:DOCKER==="
21docker ps --format '{{.ID}}\t{{.Names}}\t{{.Status}}\t{{.Image}}' 2>/dev/null | head -30
22echo "===OMNYSSH:LISTEN==="
23ss -tlnp 2>/dev/null | tail -n +2 | head -30
24echo "===OMNYSSH:PROCESS==="
25ps aux --sort=-%mem 2>/dev/null | head -15
26OMNYSSH_PROBE_EOF
27"#
28}
29
30/// Parsed output from the probe script, organized by section.
31#[derive(Debug, Clone, Default)]
32pub struct ProbeOutput {
33    sections: HashMap<String, String>,
34}
35
36impl ProbeOutput {
37    /// Parse the probe script output into sections.
38    ///
39    /// Each section is delimited by `===OMNYSSH:NAME===` markers.
40    /// Returns `Ok` even for empty or malformed input (graceful degradation).
41    ///
42    /// # Errors
43    /// Never returns an error — unknown or malformed output is silently ignored.
44    pub fn parse(output: &str) -> anyhow::Result<Self> {
45        let mut sections = HashMap::new();
46        let mut current_section: Option<String> = None;
47        let mut current_content = String::new();
48
49        for line in output.lines() {
50            let trimmed = line.trim();
51
52            // Check if this is a section marker
53            if trimmed.starts_with("===OMNYSSH:") && trimmed.ends_with("===") {
54                // Save previous section if it exists
55                if let Some(section_name) = current_section.take() {
56                    sections.insert(section_name, current_content.trim().to_string());
57                    current_content.clear();
58                }
59
60                // Extract new section name (remove markers)
61                let section_name = trimmed
62                    .strip_prefix("===OMNYSSH:")
63                    .and_then(|s| s.strip_suffix("==="))
64                    .unwrap_or("")
65                    .to_string();
66
67                current_section = Some(section_name);
68            } else if current_section.is_some() {
69                // Accumulate content for current section
70                current_content.push_str(line);
71                current_content.push('\n');
72            }
73        }
74
75        // Don't forget the last section
76        if let Some(section_name) = current_section {
77            sections.insert(section_name, current_content.trim().to_string());
78        }
79
80        Ok(Self { sections })
81    }
82
83    /// Check if a specific section exists and has non-empty content.
84    pub fn has_section(&self, name: &str) -> bool {
85        self.sections
86            .get(name)
87            .map(|content| !content.is_empty())
88            .unwrap_or(false)
89    }
90
91    /// Get the content of a section, or None if it doesn't exist.
92    pub fn get_section(&self, name: &str) -> Option<&str> {
93        self.sections.get(name).map(|s| s.as_str())
94    }
95
96    /// Parse OS information from the OS section.
97    ///
98    /// Extracts OS name and version from /etc/os-release format.
99    /// Returns a formatted string like "Ubuntu 22.04 LTS" or "Debian GNU/Linux 11".
100    pub fn parse_os_info(&self) -> Option<String> {
101        let os_section = self.get_section("OS")?;
102
103        let mut name: Option<String> = None;
104        let mut version: Option<String> = None;
105        let mut pretty_name: Option<String> = None;
106
107        for line in os_section.lines() {
108            let line = line.trim();
109
110            // Try to extract PRETTY_NAME first (most user-friendly)
111            if line.starts_with("PRETTY_NAME=") {
112                pretty_name = extract_value(line, "PRETTY_NAME=");
113            } else if line.starts_with("NAME=") {
114                name = extract_value(line, "NAME=");
115            } else if line.starts_with("VERSION=") {
116                version = extract_value(line, "VERSION=");
117            }
118        }
119
120        // Prefer PRETTY_NAME if available
121        if let Some(pretty) = pretty_name {
122            return Some(pretty);
123        }
124
125        // Otherwise combine NAME and VERSION
126        match (name, version) {
127            (Some(n), Some(v)) => Some(format!("{} {}", n, v)),
128            (Some(n), None) => Some(n),
129            _ => None,
130        }
131    }
132}
133
134/// Helper function to extract value from os-release format line.
135/// Handles both quoted ("value") and unquoted (value) formats.
136fn extract_value(line: &str, prefix: &str) -> Option<String> {
137    let value = line.strip_prefix(prefix)?.trim();
138
139    // Remove surrounding quotes if present (both " and ' work the same way)
140    if (value.starts_with('"') && value.ends_with('"')
141        || value.starts_with('\'') && value.ends_with('\''))
142        && value.len() >= 2
143    {
144        Some(value[1..value.len() - 1].to_string())
145    } else {
146        Some(value.to_string())
147    }
148}
149
150// ---------------------------------------------------------------------------
151// Tests
152// ---------------------------------------------------------------------------
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn test_parse_probe_empty_output() {
160        let result = ProbeOutput::parse("").expect("should parse empty");
161        assert!(result.sections.is_empty());
162    }
163
164    #[test]
165    fn test_parse_probe_garbage_output() {
166        let result = ProbeOutput::parse("random\ngarbage\ntext").expect("should parse garbage");
167        assert!(result.sections.is_empty());
168    }
169
170    #[test]
171    fn test_parse_probe_single_section() {
172        let output = "===OMNYSSH:OS===\nUbuntu 22.04 LTS\n===OMNYSSH:SERVICES===\nsshd.service\n";
173        let result = ProbeOutput::parse(output).expect("should parse");
174        assert!(result.has_section("OS"));
175        assert!(result.has_section("SERVICES"));
176        assert_eq!(result.get_section("OS"), Some("Ubuntu 22.04 LTS"));
177        assert_eq!(result.get_section("SERVICES"), Some("sshd.service"));
178    }
179
180    #[test]
181    fn test_parse_probe_multiple_sections() {
182        let output = r#"===OMNYSSH:OS===
183NAME="Ubuntu"
184VERSION="22.04 LTS"
185===OMNYSSH:DOCKER===
186abc123	nginx-proxy	Up 2 hours	nginx:latest
187def456	db-master	Up 5 days	postgres:15
188===OMNYSSH:LISTEN===
1890.0.0.0:22	LISTEN
1900.0.0.0:80	LISTEN
191"#;
192        let result = ProbeOutput::parse(output).expect("should parse");
193        assert!(result.has_section("OS"));
194        assert!(result.has_section("DOCKER"));
195        assert!(result.has_section("LISTEN"));
196
197        let docker_section = result
198            .get_section("DOCKER")
199            .expect("docker section should exist");
200        assert!(docker_section.contains("nginx-proxy"));
201        assert!(docker_section.contains("postgres:15"));
202    }
203
204    #[test]
205    fn test_parse_probe_with_empty_sections() {
206        let output = "===OMNYSSH:OS===\nUbuntu\n===OMNYSSH:DOCKER===\n===OMNYSSH:SERVICES===\nsshd.service\n";
207        let result = ProbeOutput::parse(output).expect("should parse");
208        assert!(result.has_section("OS"));
209        assert!(!result.has_section("DOCKER")); // Empty section = false
210        assert!(result.has_section("SERVICES"));
211    }
212
213    #[test]
214    fn test_generate_quick_scan_script() {
215        let script = generate_quick_scan_script();
216        assert!(script.contains("===OMNYSSH:OS==="));
217        assert!(script.contains("===OMNYSSH:SERVICES==="));
218        assert!(script.contains("===OMNYSSH:DOCKER==="));
219        assert!(script.contains("===OMNYSSH:LISTEN==="));
220        assert!(script.contains("===OMNYSSH:PROCESS==="));
221        assert!(script.contains("/etc/os-release"));
222        assert!(script.contains("systemctl list-units"));
223        assert!(script.contains("docker ps"));
224    }
225}
omnyssh_core/ssh/probe.rs

omnyssh_core/ssh/
probe.rs
