ssh_mcp/
config.rs

1//! Configuration and CLI argument parsing for SeSSHion
2
3use clap::Parser;
4use std::ffi::OsStr;
5use std::path::{Path, PathBuf};
6
7use crate::error::{Result, SshMcpError};
8use crate::ssh::HostKeyCheckMode;
9
10/// Default timeout for command execution in milliseconds
11pub const DEFAULT_TIMEOUT_MS: u64 = 300_000; // 300 seconds
12
13/// Default max characters for command length (None = unlimited)
14pub const DEFAULT_MAX_CHARS: Option<usize> = Some(64_000);
15
16/// Connection timeout in seconds
17pub const CONNECTION_TIMEOUT_SECS: u64 = 30;
18
19/// Number of reconnect retries after the initial attempt
20pub const DEFAULT_RECONNECT_RETRIES: u64 = 3;
21
22/// Base reconnect backoff in milliseconds
23pub const DEFAULT_RECONNECT_BACKOFF_MS: u64 = 250;
24
25/// Health probe timeout in milliseconds
26pub const DEFAULT_HEALTH_PROBE_TIMEOUT_MS: u64 = 1500;
27
28/// Maximum reconnect retries allowed by configuration
29pub const MAX_RECONNECT_RETRIES: u64 = 10;
30
31/// Minimum reconnect backoff in milliseconds
32pub const MIN_RECONNECT_BACKOFF_MS: u64 = 10;
33
34/// Maximum reconnect backoff in milliseconds
35pub const MAX_RECONNECT_BACKOFF_MS: u64 = 30_000;
36
37/// Minimum health probe timeout in milliseconds
38pub const MIN_HEALTH_PROBE_TIMEOUT_MS: u64 = 100;
39
40/// Maximum health probe timeout in milliseconds
41pub const MAX_HEALTH_PROBE_TIMEOUT_MS: u64 = 30_000;
42
43/// SeSSHion CLI arguments
44#[derive(Parser, Debug, Clone)]
45#[command(name = "ssh-mcp")]
46#[command(author = "0FL01")]
47#[command(version = env!("CARGO_PKG_VERSION"))]
48#[command(about = env!("CARGO_PKG_DESCRIPTION"))]
49pub struct Args {
50    /// SSH host to connect to
51    #[arg(long, env = "SSH_MCP_HOST")]
52    pub host: String,
53
54    /// SSH port
55    #[arg(long, default_value = "22", env = "SSH_MCP_PORT")]
56    pub port: u16,
57
58    /// SSH username
59    #[arg(long, env = "SSH_MCP_USER")]
60    pub user: String,
61
62    /// SSH password (alternative to key)
63    #[arg(long, env = "SSH_MCP_PASSWORD")]
64    pub password: Option<String>,
65
66    /// Path to SSH private key file (alternative to password)
67    #[arg(long, env = "SSH_MCP_KEY")]
68    pub key: Option<PathBuf>,
69
70    /// Absolute local directory for background job logs and state
71    #[arg(long, env = "SSH_MCP_SPOOL_DIR")]
72    pub spool_dir: Option<PathBuf>,
73
74    /// Password for `su` elevation
75    #[arg(long, env = "SSH_MCP_SU_PASSWORD")]
76    pub su_password: Option<String>,
77
78    /// Password for `sudo` commands (if different from su_password)
79    #[arg(long, env = "SSH_MCP_SUDO_PASSWORD")]
80    pub sudo_password: Option<String>,
81
82    /// Command execution timeout in milliseconds
83    #[arg(long, default_value = "300000", env = "SSH_MCP_TIMEOUT")]
84    pub timeout: u64,
85
86    /// Maximum characters for command length.
87    /// Use "none", "0", or negative value to disable limit.
88    /// Default: 64000
89    #[arg(long = "maxChars", env = "SSH_MCP_MAX_CHARS")]
90    pub max_chars: Option<String>,
91
92    /// Disable the sudo_shell and sudo_apply_patch tools
93    #[arg(long, default_value = "false", env = "SSH_MCP_DISABLE_SUDO")]
94    pub disable_sudo: bool,
95
96    /// Maximum output tokens for command execution.
97    /// Use "none" or "0" to disable limit.
98    /// Supports "k" suffix (e.g., "16k" for 16000).
99    /// Default: 16000 (approximately 64KB)
100    #[arg(long = "max-output-tokens", env = "SSH_MCP_MAX_OUTPUT_TOKENS")]
101    pub max_output_tokens: Option<String>,
102
103    /// Logging level: trace, debug, info, warn, error
104    #[arg(long, default_value = "info", env = "SSH_MCP_LOG_LEVEL", value_parser = clap::builder::PossibleValuesParser::new(["trace", "debug", "info", "warn", "error"]))]
105    pub log_level: String,
106
107    /// Log file path (default: stdout only)
108    #[arg(long, env = "SSH_MCP_LOG_FILE")]
109    pub log_file: Option<PathBuf>,
110
111    /// Log format: text or json
112    #[arg(long, default_value = "text", env = "SSH_MCP_LOG_FORMAT", value_parser = clap::builder::PossibleValuesParser::new(["text", "json"]))]
113    pub log_format: String,
114
115    /// Log rotation strategy: daily, hourly, never
116    #[arg(long, default_value = "daily", env = "SSH_MCP_LOG_ROTATION", value_parser = clap::builder::PossibleValuesParser::new(["daily", "hourly", "never"]))]
117    pub log_rotation: String,
118
119    /// Keepalive interval in seconds (default: 30)
120    /// Sends keepalive packets to maintain connection like a human user
121    #[arg(long, default_value = "30", env = "SSH_MCP_KEEPALIVE_INTERVAL")]
122    pub keepalive_interval: u64,
123
124    /// Maximum keepalive failures before disconnecting (default: 3)
125    /// Total idle timeout = keepalive_interval * keepalive_max
126    #[arg(long, default_value = "3", env = "SSH_MCP_KEEPALIVE_MAX")]
127    pub keepalive_max: u64,
128
129    /// Number of reconnect retries after the initial attempt (default: 3)
130    #[arg(long, default_value = "3", env = "SSH_MCP_RECONNECT_RETRIES")]
131    pub reconnect_retries: u64,
132
133    /// Base reconnect backoff in milliseconds (default: 250)
134    #[arg(long, default_value = "250", env = "SSH_MCP_RECONNECT_BACKOFF_MS")]
135    pub reconnect_backoff_ms: u64,
136
137    /// Health probe timeout in milliseconds for active session checks (default: 1500)
138    #[arg(long, default_value = "1500", env = "SSH_MCP_HEALTH_PROBE_TIMEOUT_MS")]
139    pub health_probe_timeout_ms: u64,
140
141    /// SSH host key checking mode: yes, accept-new, or no
142    #[arg(
143        long = "strict-host-key-checking",
144        env = "SSH_MCP_STRICT_HOST_KEY_CHECKING",
145        value_enum,
146        default_value_t = HostKeyCheckMode::AcceptNew
147    )]
148    pub strict_host_key_checking: HostKeyCheckMode,
149
150    /// Path to known_hosts file (default: OpenSSH user known_hosts)
151    #[arg(long = "known-hosts", env = "SSH_MCP_KNOWN_HOSTS")]
152    pub known_hosts: Option<PathBuf>,
153}
154
155/// Parsed and validated configuration
156#[derive(Debug, Clone)]
157pub struct Config {
158    /// SSH host
159    pub host: String,
160
161    /// SSH port
162    pub port: u16,
163
164    /// SSH username
165    pub user: String,
166
167    /// SSH password
168    pub password: Option<String>,
169
170    /// Path to SSH private key
171    pub key: Option<PathBuf>,
172
173    /// Password for su elevation
174    pub su_password: Option<String>,
175
176    /// Password for sudo commands
177    pub sudo_password: Option<String>,
178
179    /// Command timeout in milliseconds
180    pub timeout_ms: u64,
181
182    /// Maximum command length (None = unlimited)
183    pub max_chars: Option<usize>,
184
185    /// Maximum output tokens for command execution (None = unlimited)
186    pub max_output_tokens: Option<usize>,
187
188    /// Whether sudo_shell and sudo_apply_patch tools are disabled
189    pub disable_sudo: bool,
190
191    /// Keepalive interval in seconds
192    pub keepalive_interval: u64,
193
194    /// Maximum keepalive failures before disconnecting
195    pub keepalive_max: u64,
196
197    /// Number of reconnect retries after the initial attempt
198    pub reconnect_retries: u64,
199
200    /// Base reconnect backoff in milliseconds
201    pub reconnect_backoff_ms: u64,
202
203    /// Health probe timeout in milliseconds for active session checks
204    pub health_probe_timeout_ms: u64,
205
206    /// SSH host key checking mode
207    pub strict_host_key_checking: HostKeyCheckMode,
208
209    /// Optional known_hosts file path
210    pub known_hosts: Option<PathBuf>,
211}
212
213impl Config {
214    /// Create Config from CLI Args
215    pub fn from_args(args: Args) -> Result<Self> {
216        let home = std::env::var_os("HOME");
217        Self::from_args_with_home(args, home.as_deref())
218    }
219
220    fn from_args_with_home(mut args: Args, home: Option<&OsStr>) -> Result<Self> {
221        args.key = args
222            .key
223            .map(|path| expand_key_path(path, home))
224            .transpose()?;
225        validate_args(&args)?;
226
227        let max_chars = parse_max_chars(args.max_chars.as_deref());
228        let max_output_tokens = parse_max_output_tokens(args.max_output_tokens.as_deref());
229
230        Ok(Config {
231            host: args.host,
232            port: args.port,
233            user: args.user,
234            password: sanitize_password(args.password),
235            key: args.key,
236            su_password: sanitize_password(args.su_password),
237            sudo_password: sanitize_password(args.sudo_password),
238            timeout_ms: args.timeout,
239            max_chars,
240            max_output_tokens,
241            disable_sudo: args.disable_sudo,
242            keepalive_interval: args.keepalive_interval,
243            keepalive_max: args.keepalive_max,
244            reconnect_retries: args.reconnect_retries,
245            reconnect_backoff_ms: args.reconnect_backoff_ms,
246            health_probe_timeout_ms: args.health_probe_timeout_ms,
247            strict_host_key_checking: args.strict_host_key_checking,
248            known_hosts: args.known_hosts,
249        })
250    }
251}
252
253fn expand_key_path(path: PathBuf, home: Option<&OsStr>) -> Result<PathBuf> {
254    if !path.as_os_str().as_encoded_bytes().starts_with(b"~/") {
255        return Ok(path);
256    }
257
258    let home = home.filter(|value| !value.is_empty()).ok_or_else(|| {
259        SshMcpError::Config(format!(
260            "Cannot expand SSH key path {}: HOME is not set",
261            path.display()
262        ))
263    })?;
264    let suffix = path
265        .strip_prefix("~")
266        .expect("leading ~/ path must have a tilde component");
267
268    Ok(Path::new(home).join(suffix))
269}
270
271/// Validate CLI arguments
272fn validate_args(args: &Args) -> Result<()> {
273    let mut errors = Vec::new();
274
275    if args.host.is_empty() {
276        errors.push("Missing required --host".to_string());
277    }
278
279    if args.user.is_empty() {
280        errors.push("Missing required --user".to_string());
281    }
282
283    // Must have either password or key
284    if args.password.is_none() && args.key.is_none() {
285        errors.push("Must provide either --password or --key".to_string());
286    }
287
288    // If key is provided, check if file exists
289    if let Some(ref key_path) = args.key
290        && !key_path.exists()
291    {
292        errors.push(format!("SSH key file not found: {}", key_path.display()));
293    }
294
295    if args.reconnect_retries > MAX_RECONNECT_RETRIES {
296        errors.push(format!(
297            "--reconnect-retries must be <= {MAX_RECONNECT_RETRIES}"
298        ));
299    }
300
301    if !(MIN_RECONNECT_BACKOFF_MS..=MAX_RECONNECT_BACKOFF_MS).contains(&args.reconnect_backoff_ms) {
302        errors.push(format!(
303            "--reconnect-backoff-ms must be between {MIN_RECONNECT_BACKOFF_MS} and {MAX_RECONNECT_BACKOFF_MS}"
304        ));
305    }
306
307    if !(MIN_HEALTH_PROBE_TIMEOUT_MS..=MAX_HEALTH_PROBE_TIMEOUT_MS)
308        .contains(&args.health_probe_timeout_ms)
309    {
310        errors.push(format!(
311            "--health-probe-timeout-ms must be between {MIN_HEALTH_PROBE_TIMEOUT_MS} and {MAX_HEALTH_PROBE_TIMEOUT_MS}"
312        ));
313    }
314
315    if !errors.is_empty() {
316        return Err(SshMcpError::Config(format!(
317            "Configuration error:\n{}",
318            errors.join("\n")
319        )));
320    }
321
322    Ok(())
323}
324
325/// Default max output tokens (16_000 ≈ 64KB)
326pub const DEFAULT_MAX_OUTPUT_TOKENS: Option<usize> = Some(16_000);
327
328/// Parse max_chars argument
329///
330/// - "none" (case-insensitive) → None (unlimited)
331/// - "0" or negative → None (unlimited)
332/// - positive integer → Some(value)
333/// - None (not provided) → DEFAULT_MAX_CHARS
334pub fn parse_max_chars(value: Option<&str>) -> Option<usize> {
335    match value {
336        None => DEFAULT_MAX_CHARS,
337        Some(s) => {
338            let lowered = s.to_lowercase();
339            if lowered == "none" {
340                return None;
341            }
342
343            match s.parse::<i64>() {
344                Ok(n) if n <= 0 => None,
345                Ok(n) => Some(n as usize),
346                Err(_) => DEFAULT_MAX_CHARS,
347            }
348        }
349    }
350}
351
352/// Parse max_output_tokens argument
353///
354/// - "none" (case-insensitive) → None (unlimited)
355/// - "0" or negative → None (unlimited)
356/// - positive integer with optional "k" suffix (e.g., "12k") → Some(value)
357/// - None (not provided) → DEFAULT_MAX_OUTPUT_TOKENS
358pub fn parse_max_output_tokens(value: Option<&str>) -> Option<usize> {
359    match value {
360        None => DEFAULT_MAX_OUTPUT_TOKENS,
361        Some(s) => {
362            let lowered = s.to_lowercase().replace(" ", "");
363            if lowered == "none" {
364                return None;
365            }
366
367            // Try to parse with k suffix
368            if lowered.ends_with('k') {
369                let num_part = &lowered[..lowered.len() - 1];
370                match num_part.parse::<i64>() {
371                    Ok(n) if n <= 0 => None,
372                    Ok(n) => Some((n as usize).saturating_mul(1_000)),
373                    Err(_) => DEFAULT_MAX_OUTPUT_TOKENS,
374                }
375            } else {
376                match lowered.parse::<i64>() {
377                    Ok(n) if n <= 0 => None,
378                    Ok(n) => Some(n as usize),
379                    Err(_) => DEFAULT_MAX_OUTPUT_TOKENS,
380                }
381            }
382        }
383    }
384}
385
386/// Sanitize password: return None if empty
387fn sanitize_password(password: Option<String>) -> Option<String> {
388    password.filter(|p| !p.is_empty())
389}
390
391#[cfg(test)]
392mod tests {
393    use super::*;
394
395    fn base_args() -> Args {
396        Args {
397            host: "localhost".to_string(),
398            port: 22,
399            user: "test".to_string(),
400            password: Some("secret".to_string()),
401            key: None,
402            spool_dir: None,
403            su_password: None,
404            sudo_password: None,
405            timeout: DEFAULT_TIMEOUT_MS,
406            max_chars: None,
407            disable_sudo: false,
408            max_output_tokens: None,
409            log_level: "info".to_string(),
410            log_file: None,
411            log_format: "text".to_string(),
412            log_rotation: "daily".to_string(),
413            keepalive_interval: 30,
414            keepalive_max: 3,
415            reconnect_retries: DEFAULT_RECONNECT_RETRIES,
416            reconnect_backoff_ms: DEFAULT_RECONNECT_BACKOFF_MS,
417            health_probe_timeout_ms: DEFAULT_HEALTH_PROBE_TIMEOUT_MS,
418            strict_host_key_checking: HostKeyCheckMode::AcceptNew,
419            known_hosts: None,
420        }
421    }
422
423    #[test]
424    fn test_parse_max_chars_none_string() {
425        assert_eq!(parse_max_chars(Some("none")), None);
426        assert_eq!(parse_max_chars(Some("None")), None);
427        assert_eq!(parse_max_chars(Some("NONE")), None);
428    }
429
430    #[test]
431    fn test_parse_max_chars_zero_or_negative() {
432        assert_eq!(parse_max_chars(Some("0")), None);
433        assert_eq!(parse_max_chars(Some("-1")), None);
434        assert_eq!(parse_max_chars(Some("-100")), None);
435    }
436
437    #[test]
438    fn test_parse_max_chars_positive() {
439        assert_eq!(parse_max_chars(Some("500")), Some(500));
440        assert_eq!(parse_max_chars(Some("2000")), Some(2000));
441    }
442
443    #[test]
444    fn test_parse_max_chars_invalid() {
445        // Invalid strings should return default
446        assert_eq!(parse_max_chars(Some("abc")), DEFAULT_MAX_CHARS);
447        assert_eq!(parse_max_chars(Some("")), DEFAULT_MAX_CHARS);
448    }
449
450    #[test]
451    fn test_parse_max_chars_not_provided() {
452        assert_eq!(parse_max_chars(None), DEFAULT_MAX_CHARS);
453    }
454
455    #[test]
456    fn test_config_from_args_uses_default_max_chars() {
457        let config = Config::from_args(base_args()).unwrap();
458
459        assert_eq!(config.max_chars, Some(64_000));
460        assert_eq!(config.strict_host_key_checking, HostKeyCheckMode::AcceptNew);
461        assert!(config.known_hosts.is_none());
462    }
463
464    #[test]
465    fn test_config_expands_tilde_key_before_validation() {
466        let home = tempfile::tempdir().unwrap();
467        let key_path = home.path().join(".ssh/id_ed25519");
468        std::fs::create_dir_all(key_path.parent().unwrap()).unwrap();
469        std::fs::write(&key_path, "test key").unwrap();
470
471        let mut args = base_args();
472        args.password = None;
473        args.key = Some(PathBuf::from("~/.ssh/id_ed25519"));
474
475        let config = Config::from_args_with_home(args, Some(home.path().as_os_str())).unwrap();
476        assert_eq!(config.key, Some(key_path));
477    }
478
479    #[test]
480    fn test_expand_key_path_only_expands_leading_home_prefix() {
481        let home = OsStr::new("/home/test");
482        let cases = [
483            ("~", "~"),
484            ("~user/key", "~user/key"),
485            ("dir/~/key", "dir/~/key"),
486            ("$HOME/key", "$HOME/key"),
487            (r"~\key", r"~\key"),
488            ("/tmp/key", "/tmp/key"),
489        ];
490
491        for (input, expected) in cases {
492            assert_eq!(
493                expand_key_path(PathBuf::from(input), Some(home)).unwrap(),
494                PathBuf::from(expected)
495            );
496        }
497        assert_eq!(
498            expand_key_path(PathBuf::from("~/.ssh/id_ed25519"), Some(home)).unwrap(),
499            PathBuf::from("/home/test/.ssh/id_ed25519")
500        );
501    }
502
503    #[test]
504    fn test_expand_key_path_requires_home() {
505        for home in [None, Some(OsStr::new(""))] {
506            let error = expand_key_path(PathBuf::from("~/.ssh/id_ed25519"), home).unwrap_err();
507            assert!(error.to_string().contains("HOME is not set"));
508        }
509    }
510
511    #[test]
512    fn test_args_parse_host_key_options() {
513        let args = Args::try_parse_from([
514            "ssh-mcp",
515            "--host",
516            "example.com",
517            "--user",
518            "alice",
519            "--password",
520            "secret",
521            "--strict-host-key-checking",
522            "yes",
523            "--known-hosts",
524            "/tmp/known_hosts",
525        ])
526        .unwrap();
527
528        assert_eq!(args.strict_host_key_checking, HostKeyCheckMode::Yes);
529        assert_eq!(args.known_hosts, Some(PathBuf::from("/tmp/known_hosts")));
530    }
531
532    #[test]
533    fn test_args_parse_spool_dir() {
534        let args = Args::try_parse_from([
535            "ssh-mcp",
536            "--host",
537            "example.com",
538            "--user",
539            "alice",
540            "--password",
541            "secret",
542            "--spool-dir",
543            "/tmp/ssh-mcp-alice",
544        ])
545        .unwrap();
546
547        assert_eq!(args.spool_dir, Some(PathBuf::from("/tmp/ssh-mcp-alice")));
548    }
549
550    #[test]
551    fn test_sanitize_password() {
552        assert_eq!(
553            sanitize_password(Some("secret".to_string())),
554            Some("secret".to_string())
555        );
556        assert_eq!(sanitize_password(Some(String::new())), None);
557        assert_eq!(sanitize_password(None), None);
558    }
559
560    #[test]
561    fn test_parse_max_output_tokens_none_string() {
562        assert_eq!(parse_max_output_tokens(Some("none")), None);
563        assert_eq!(parse_max_output_tokens(Some("None")), None);
564        assert_eq!(parse_max_output_tokens(Some("NONE")), None);
565    }
566
567    #[test]
568    fn test_parse_max_output_tokens_zero_or_negative() {
569        assert_eq!(parse_max_output_tokens(Some("0")), None);
570        assert_eq!(parse_max_output_tokens(Some("-1")), None);
571        assert_eq!(parse_max_output_tokens(Some("-100")), None);
572    }
573
574    #[test]
575    fn test_parse_max_output_tokens_positive() {
576        assert_eq!(parse_max_output_tokens(Some("500")), Some(500));
577        assert_eq!(parse_max_output_tokens(Some("12000")), Some(12_000));
578    }
579
580    #[test]
581    fn test_parse_max_output_tokens_with_k_suffix() {
582        assert_eq!(parse_max_output_tokens(Some("12k")), Some(12_000));
583        assert_eq!(parse_max_output_tokens(Some("5K")), Some(5_000));
584        assert_eq!(parse_max_output_tokens(Some("100k")), Some(100_000));
585    }
586
587    #[test]
588    fn test_parse_max_output_tokens_invalid() {
589        // Invalid strings should return default
590        assert_eq!(
591            parse_max_output_tokens(Some("abc")),
592            DEFAULT_MAX_OUTPUT_TOKENS
593        );
594        assert_eq!(parse_max_output_tokens(Some("")), DEFAULT_MAX_OUTPUT_TOKENS);
595    }
596
597    #[test]
598    fn test_parse_max_output_tokens_not_provided() {
599        assert_eq!(parse_max_output_tokens(None), DEFAULT_MAX_OUTPUT_TOKENS);
600    }
601
602    #[test]
603    fn test_validate_args_rejects_reconnect_retries_out_of_range() {
604        let mut args = base_args();
605        args.reconnect_retries = MAX_RECONNECT_RETRIES.saturating_add(1);
606
607        let result = validate_args(&args);
608        assert!(result.is_err());
609    }
610
611    #[test]
612    fn test_validate_args_rejects_reconnect_backoff_out_of_range() {
613        let mut args = base_args();
614        args.reconnect_backoff_ms = MIN_RECONNECT_BACKOFF_MS.saturating_sub(1);
615
616        let result = validate_args(&args);
617        assert!(result.is_err());
618    }
619
620    #[test]
621    fn test_validate_args_rejects_health_probe_timeout_out_of_range() {
622        let mut args = base_args();
623        args.health_probe_timeout_ms = MAX_HEALTH_PROBE_TIMEOUT_MS.saturating_add(1);
624
625        let result = validate_args(&args);
626        assert!(result.is_err());
627    }
628}
ssh_mcp/config.rs

ssh_mcp/
config.rs
