Skip to content

Commit 456b89f

Browse files
tgauthtgauth
committed
add default shell get and set
1 parent a4ef666 commit 456b89f

File tree

8 files changed

+376
-117
lines changed

8 files changed

+376
-117
lines changed

‎sshdconfig/src/args.rs

Lines changed: 12 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22
// Licensed under the MIT License.
33

44
use clap::{Parser, Subcommand};
5-
6-
use crate::metadata::RepeatableKeyword;
5+
use serde::{Deserialize, Serialize};
76

87
#[derive(Parser)]
98
pub struct Args {
@@ -15,36 +14,19 @@ pub struct Args {
1514
pub enum Command {
1615
/// Export sshd_config
1716
Export,
18-
/// Set a value in sshd_config (not yet implemented)
17+
/// Get default shell, eventually to be used for sshd_config and repeatable keywords
18+
Get,
19+
/// Set default shell, eventually to be used for sshd_config and repeatable keywords
1920
Set {
20-
/// The input to set
2121
#[clap(short = 'i', long, help = "input to set in sshd_config")]
22-
input: Option<String>,
23-
#[clap(subcommand, help = "set commands")]
24-
subcommand: Option<SetCommand>,
22+
input: String
2523
},
2624
}
2725

28-
#[derive(Clone, Debug, Eq, PartialEq, Subcommand)]
29-
pub enum SetCommand {
30-
/// Set the default shell
31-
DefaultShell {
32-
/// The path to the shell executable
33-
#[clap(short = 's', long, help = "path to the shell executable")]
34-
shell: String,
35-
/// Additional command options
36-
///
37-
#[clap(short = 'c', long, help = "additional command options", default_value = "-c")]
38-
cmd_option: Option<String>,
39-
#[clap(short = 'e', long, help = "skip escaping arguments", default_value = "false")]
40-
escape_arguments: bool,
41-
#[clap(short = 'a', long, help = "additional shell arguments")]
42-
shell_arguments: Option<Vec<String>>,
43-
},
44-
/// Set repeatable keywords
45-
RepeatableKeyworld {
46-
keyword: RepeatableKeyword,
47-
name: String,
48-
value: Option<String>,
49-
}
50-
}
26+
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
27+
pub struct DefaultShell {
28+
pub shell: Option<String>,
29+
pub cmd_option: Option<String>,
30+
pub escape_arguments: Option<bool>,
31+
pub shell_arguments: Option<Vec<String>>,
32+
}

‎sshdconfig/src/error.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@ use thiserror::Error;
77
pub enum SshdConfigError {
88
#[error("Command: {0}")]
99
CommandError(String),
10-
#[error("Invalid Input: {0}")]
11-
IoError(String),
1210
#[error("IO: {0}")]
1311
InvalidInput(String),
1412
#[error("JSON: {0}")]

‎sshdconfig/src/export.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@ use crate::util::invoke_sshd_config_validation;
1010
/// # Errors
1111
///
1212
/// This function will return an error if the command cannot invoke sshd -T, parse the return, or convert it to json.
13-
pub fn invoke_export() -> Result<String, SshdConfigError> {
13+
pub fn invoke_export() -> Result<(), SshdConfigError> {
1414
let sshd_config_text = invoke_sshd_config_validation()?;
1515
let sshd_config: serde_json::Map<String, serde_json::Value> = parse_text_to_map(&sshd_config_text)?;
1616
let json = serde_json::to_string_pretty(&sshd_config)?;
17-
Ok(json)
17+
println!("{json}");
18+
Ok(())
1819
}

‎sshdconfig/src/get.rs

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
use registry_lib::{config::{Registry, RegistryValueData}, RegistryHelper};
5+
6+
use crate::args::DefaultShell;
7+
use crate::error::SshdConfigError;
8+
9+
/// Invoke the get command.
10+
///
11+
/// # Errors
12+
///
13+
/// This function will return an error if the desired settings cannot be retrieved.
14+
pub fn invoke_get() -> Result<(), SshdConfigError> {
15+
// TODO: distinguish between get commands for default shell, repeatable keywords, and sshd_config
16+
get_default_shell()?;
17+
Ok(())
18+
}
19+
20+
fn get_default_shell() -> Result<(), SshdConfigError> {
21+
let registry_helper = RegistryHelper::new("HKLM\\SOFTWARE\\OpenSSH", Some("DefaultShell".to_string()), None)?;
22+
let default_shell: Registry = registry_helper.get()?;
23+
let mut shell = None;
24+
let mut shell_arguments = None;
25+
if let Some(value) = default_shell.value_data {
26+
match value {
27+
RegistryValueData::String(s) => {
28+
let parts: Vec<&str> = s.split_whitespace().collect();
29+
if parts.is_empty() {
30+
return Err(SshdConfigError::InvalidInput("DefaultShell cannot be empty".to_string()));
31+
}
32+
shell = Some(parts[0].to_string());
33+
if parts.len() > 1 {
34+
shell_arguments = Some(parts[1..].iter().map(|&s| s.to_string()).collect());
35+
}
36+
}
37+
_ => return Err(SshdConfigError::InvalidInput("DefaultShell must be a string".to_string())),
38+
}
39+
}
40+
41+
let registry_helper = RegistryHelper::new("HKLM\\SOFTWARE\\OpenSSH", Some("DefaultShellCommandOption".to_string()), None)?;
42+
let option: Registry = registry_helper.get()?;
43+
let mut cmd_option = None;
44+
if let Some(value) = option.value_data {
45+
match value {
46+
RegistryValueData::String(s) => cmd_option = Some(s),
47+
_ => return Err(SshdConfigError::InvalidInput("DefaultShellCommandOption must be a string".to_string())),
48+
}
49+
}
50+
51+
let registry_helper = RegistryHelper::new("HKLM\\SOFTWARE\\OpenSSH", Some("DefaultShellEscapeArguments".to_string()), None)?;
52+
let escape_args: Registry = registry_helper.get()?;
53+
let mut escape_arguments = None;
54+
if let Some(value) = escape_args.value_data {
55+
if let RegistryValueData::DWord(b) = value {
56+
if b == 0 || b == 1 {
57+
escape_arguments = if b == 1 { Some(true) } else { Some(false) };
58+
} else {
59+
return Err(SshdConfigError::InvalidInput("DefaultShellEscapeArguments must be a boolean".to_string()));
60+
}
61+
} else {
62+
return Err(SshdConfigError::InvalidInput("DefaultShellEscapeArguments must be a boolean".to_string()));
63+
}
64+
}
65+
66+
let result = DefaultShell {
67+
shell,
68+
cmd_option,
69+
escape_arguments,
70+
shell_arguments
71+
};
72+
73+
let output = serde_json::to_string_pretty(&result)?;
74+
println!("{output}");
75+
Ok(())
76+
}
77+
78+
#[cfg(test)]
79+
mod tests {
80+
use super::*;
81+
use registry_lib::config::RegistryValueData;
82+
83+
#[test]
84+
fn test_parse_shell_command() {
85+
let (shell, args) = parse_shell_command(r#"C:\Program Files\PowerShell\pwsh.exe -NoProfile"#);
86+
assert_eq!(shell, r#"C:\Program Files\PowerShell\pwsh.exe"#);
87+
assert_eq!(args, vec!["-NoProfile"]);
88+
}
89+
90+
#[test]
91+
fn test_parse_shell_command_with_quotes() {
92+
let (shell, args) = parse_shell_command(r#""C:\Program Files\PowerShell\pwsh.exe" -NoProfile -Command"#);
93+
assert_eq!(shell, r#"C:\Program Files\PowerShell\pwsh.exe"#);
94+
assert_eq!(args, vec!["-NoProfile", "-Command"]);
95+
}
96+
97+
#[test]
98+
fn test_extract_string_value_string() {
99+
let value = RegistryValueData::String("test".to_string());
100+
let result = extract_string_value(value, "test_field").unwrap();
101+
assert_eq!(result, Some("test".to_string()));
102+
}
103+
104+
#[test]
105+
fn test_extract_string_value_multistring() {
106+
let value = RegistryValueData::MultiString(vec!["first".to_string(), "second".to_string()]);
107+
let result = extract_string_value(value, "test_field").unwrap();
108+
assert_eq!(result, Some("first".to_string()));
109+
}
110+
}

‎sshdconfig/src/main.rs

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,38 +4,29 @@
44
use args::{Args, Command};
55
use clap::{Parser};
66
use export::invoke_export;
7+
use get::invoke_get;
78
use set::invoke_set;
89

910
mod args;
1011
mod error;
1112
mod export;
13+
mod get;
1214
mod metadata;
1315
mod parser;
1416
mod set;
1517
mod util;
1618

1719
fn main() {
1820
let args = Args::parse();
19-
match &args.command {
20-
Command::Export => {
21-
match invoke_export() {
22-
Ok(result) => {
23-
println!("{result}");
24-
}
25-
Err(e) => {
26-
eprintln!("Error exporting sshd_config: {e:?}");
27-
}
28-
}
29-
}
30-
Command::Set { input , subcommand} => {
31-
match invoke_set(input, subcommand) {
32-
Ok(()) => {
33-
println!("success");
34-
}
35-
Err(e) => {
36-
eprintln!("Error setting sshd_config: {e:?}");
37-
}
38-
}
39-
}
21+
22+
let result = match &args.command {
23+
Command::Export => invoke_export(),
24+
Command::Get => invoke_get(),
25+
Command::Set { input } => invoke_set(input),
26+
};
27+
28+
if let Err(e) = result {
29+
eprintln!("{e}");
30+
std::process::exit(1);
4031
}
4132
}

‎sshdconfig/src/set.rs

Lines changed: 49 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,92 +1,80 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT License.
33

4-
use registry_lib::{config::{Registry, RegistryValueData}, RegistryHelper};
5-
use serde_json::{json, Value};
6-
use std::{path::Path, process::Command};
4+
use registry_lib::{config::RegistryValueData, RegistryHelper};
5+
use std::path::Path;
76

8-
use crate::args::SetCommand;
7+
use crate::args::DefaultShell;
98
use crate::error::SshdConfigError;
10-
use crate::metadata::RepeatableKeyword;
119

1210
/// Invoke the set command.
1311
///
1412
/// # Errors
1513
///
16-
/// This function will return an error if the sshd_config is not updated with the desired settings.
17-
pub fn invoke_set(input: &Option<String>, subcommand: &Option<SetCommand>) -> Result<(), SshdConfigError> {
18-
if let Some(subcommand) = subcommand {
19-
match subcommand {
20-
SetCommand::DefaultShell { shell, cmd_option, escape_arguments, shell_arguments } => {
21-
set_default_shell(shell, cmd_option, escape_arguments, shell_arguments)
22-
}
23-
SetCommand::RepeatableKeyworld { keyword, name, value } => {
24-
set_repeatable_keyword(keyword, name, value)
25-
}
14+
/// This function will return an error if the desired settings cannot be applied.
15+
pub fn invoke_set(input: &String) -> Result<(), SshdConfigError> {
16+
match serde_json::from_str::<DefaultShell>(input) {
17+
Ok(default_shell) => {
18+
set_default_shell(&default_shell.shell, &default_shell.cmd_option, &default_shell.escape_arguments, &default_shell.shell_arguments)
19+
},
20+
Err(e) => {
21+
// TODO: handle other commands like repeatable keywords or sshd_config modifications
22+
return Err(SshdConfigError::InvalidInput(format!("Failed to parse input as DefaultShell: {}", e)));
2623
}
27-
} else {
28-
set_regular_keywords(input)
2924
}
3025
}
3126

32-
fn set_default_shell(shell: &String, cmd_option: &Option<String>, escape_arguments: &bool, shell_arguments: &Option<Vec<String>>) -> Result<(), SshdConfigError> {
33-
let shell_path = Path::new(shell);
34-
if shell_path.is_relative() {
35-
if shell_path.components().any(|c| c == std::path::Component::ParentDir) {
36-
return Err(SshdConfigError::InvalidInput("shell path must not be relative".to_string()));
27+
fn set_default_shell(shell: &Option<String>, cmd_option: &Option<String>, escape_arguments: &Option<bool>, shell_arguments: &Option<Vec<String>>) -> Result<(), SshdConfigError> {
28+
if let Some(shell) = shell {
29+
let shell_path = Path::new(shell);
30+
if shell_path.is_relative() {
31+
if shell_path.components().any(|c| c == std::path::Component::ParentDir) {
32+
return Err(SshdConfigError::InvalidInput("shell path must not be relative".to_string()));
33+
}
34+
}
35+
if !shell_path.exists() {
36+
return Err(SshdConfigError::InvalidInput(format!("shell path does not exist: {}", shell)));
3737
}
38-
}
39-
// TODO check if binary exists when it is not an absolute path
40-
if !shell_path.exists() {
41-
return Err(SshdConfigError::InvalidInput(format!("shell path does not exist: {}", shell)));
42-
}
4338

44-
let mut shell_data = shell.clone();
45-
if let Some(shell_args) = shell_arguments {
46-
let args_str = shell_args.join(" ");
47-
shell_data = format!("{} {}", shell, args_str);
48-
}
39+
let mut shell_data = shell.clone();
40+
if let Some(shell_args) = shell_arguments {
41+
let args_str = shell_args.join(" ");
42+
shell_data = format!("{} {}", shell, args_str);
43+
}
44+
45+
set_registry("DefaultShell", RegistryValueData::String(shell_data))?;
46+
} else {
47+
remove_registry("DefaultShell")?;
48+
};
4949

50-
set_registry_default_shell("DefaultShell", RegistryValueData::String(shell_data))?;
5150

5251
if let Some(cmd_option) = cmd_option {
53-
set_registry_default_shell("DefaultShellCommandOption", RegistryValueData::String(cmd_option.clone()))?;
54-
}
55-
else {
56-
return Err(SshdConfigError::InvalidInput("cmd_option is required".to_string()));
52+
set_registry("DefaultShellCommandOption", RegistryValueData::String(cmd_option.clone()))?;
53+
} else {
54+
remove_registry("DefaultShellCommandOption")?;
5755
}
5856

59-
let mut escape_data = 0;
60-
if *escape_arguments {
61-
escape_data = 1;
57+
if let Some(escape_args) = escape_arguments {
58+
let mut escape_data = 0;
59+
if *escape_args {
60+
escape_data = 1;
61+
}
62+
set_registry("DefaultShellEscapeArguments", RegistryValueData::DWord(escape_data))?;
63+
} else {
64+
remove_registry("DefaultShellEscapeArguments")?;
6265
}
63-
set_registry_default_shell("DefaultShellEscapeArguments ", RegistryValueData::DWord(escape_data))?;
6466

6567
Ok(())
6668
}
6769

68-
pub fn set_repeatable_keyword(_keyword: &RepeatableKeyword, _name: &String, _value: &Option<String>) -> Result<(), SshdConfigError> {
69-
// TODO: handle repeat keywords like subsystem
70-
Ok(())
71-
}
72-
73-
pub fn set_regular_keywords(_input: &Option<String>) -> Result<(), SshdConfigError> {
74-
// modify or add all values in the input file
75-
// have a banner that we are managing the file and check that before modifying
76-
// if the file is managed by this tool, we can modify it
77-
// if the file is not managed by the tool, we should back it up
78-
// get existing sshd_config settings from sshd -T
79-
// get default sshd_config settings from sshd -T
80-
// save explicit settings
81-
// for each of the input keys, update it with the new value
82-
// or insert it if it doesn't exist (above any match statements)
83-
// write the updated settings back to the sshd_config file
84-
// ensure sshd -T is valid after the update?
70+
fn set_registry(name: &str, data: RegistryValueData) -> Result<(), SshdConfigError> {
71+
let registry_helper = RegistryHelper::new("HKLM\\SOFTWARE\\OpenSSH", Some(name.to_string()), Some(data))?;
72+
registry_helper.set()?;
8573
Ok(())
8674
}
8775

88-
fn set_registry_default_shell(name: &str, data: RegistryValueData) -> Result<(), SshdConfigError> {
89-
let registry_helper = RegistryHelper::new("HKLM\\SOFTWARE\\OpenSSH", Some(name.to_string()), Some(data))?;
90-
registry_helper.set()?;
76+
fn remove_registry(name: &str) -> Result<(), SshdConfigError> {
77+
let registry_helper = RegistryHelper::new("HKLM\\SOFTWARE\\OpenSSH", Some(name.to_string()), None)?;
78+
registry_helper.remove()?;
9179
Ok(())
9280
}

‎sshdconfig/sshd.dsc.resource.json

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,11 @@
1212
"set": {
1313
"executable": "sshdconfig",
1414
"args": [
15-
"set"
15+
"set",
16+
{
17+
"jsonInputArg": "--input",
18+
"mandatory": true
19+
}
1620
]
1721
},
1822
"schema": {

0 commit comments

Comments
 (0)