|
| 1 | +// Copyright 2025 Lablup Inc. and Jeongkyu Shin |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +//! bssh-keygen binary - SSH key pair generation tool |
| 16 | +//! |
| 17 | +//! This binary provides a command-line interface for generating SSH key pairs |
| 18 | +//! in OpenSSH format, supporting Ed25519 (recommended) and RSA algorithms. |
| 19 | +//! |
| 20 | +//! # Usage |
| 21 | +//! |
| 22 | +//! ```bash |
| 23 | +//! # Generate Ed25519 key (default, recommended) |
| 24 | +//! bssh-keygen |
| 25 | +//! |
| 26 | +//! # Generate Ed25519 key with custom output path |
| 27 | +//! bssh-keygen -f ~/.ssh/my_key |
| 28 | +//! |
| 29 | +//! # Generate RSA key with 4096 bits |
| 30 | +//! bssh-keygen -t rsa -b 4096 |
| 31 | +//! |
| 32 | +//! # Generate key with custom comment |
| 33 | +//! bssh-keygen -C "user@hostname" |
| 34 | +//! ``` |
| 35 | +
|
| 36 | +use anyhow::{Context, Result}; |
| 37 | +use bssh::keygen; |
| 38 | +use bssh::utils::logging; |
| 39 | +use clap::{ArgAction, Parser}; |
| 40 | +use std::io::{self, Write}; |
| 41 | +use std::path::PathBuf; |
| 42 | + |
| 43 | +/// Backend.AI SSH Key Generator - Generate SSH key pairs in OpenSSH format |
| 44 | +#[derive(Parser, Debug)] |
| 45 | +#[command(name = "bssh-keygen")] |
| 46 | +#[command(version)] |
| 47 | +#[command(about = "Generate SSH key pairs in OpenSSH format", long_about = None)] |
| 48 | +struct Cli { |
| 49 | + /// Key type: ed25519 (recommended) or rsa |
| 50 | + #[arg( |
| 51 | + short = 't', |
| 52 | + long = "type", |
| 53 | + default_value = "ed25519", |
| 54 | + value_name = "TYPE" |
| 55 | + )] |
| 56 | + key_type: String, |
| 57 | + |
| 58 | + /// Output file path (default: ~/.ssh/id_<type>) |
| 59 | + #[arg(short = 'f', long = "file", value_name = "FILE")] |
| 60 | + output: Option<PathBuf>, |
| 61 | + |
| 62 | + /// RSA key bits (only for RSA, minimum 2048) |
| 63 | + #[arg( |
| 64 | + short = 'b', |
| 65 | + long = "bits", |
| 66 | + default_value = "4096", |
| 67 | + value_name = "BITS" |
| 68 | + )] |
| 69 | + bits: u32, |
| 70 | + |
| 71 | + /// Comment for the key |
| 72 | + #[arg(short = 'C', long = "comment", value_name = "COMMENT")] |
| 73 | + comment: Option<String>, |
| 74 | + |
| 75 | + /// Overwrite existing files without prompting |
| 76 | + #[arg(short = 'y', long = "yes")] |
| 77 | + yes: bool, |
| 78 | + |
| 79 | + /// Quiet mode (no output except errors) |
| 80 | + #[arg(short = 'q', long = "quiet")] |
| 81 | + quiet: bool, |
| 82 | + |
| 83 | + /// Verbosity level (-v, -vv, -vvv) |
| 84 | + #[arg(short, long, action = ArgAction::Count)] |
| 85 | + verbose: u8, |
| 86 | +} |
| 87 | + |
| 88 | +fn main() -> Result<()> { |
| 89 | + let cli = Cli::parse(); |
| 90 | + |
| 91 | + // Initialize logging based on verbosity (only if not quiet) |
| 92 | + if !cli.quiet { |
| 93 | + logging::init_logging_console_only(cli.verbose); |
| 94 | + } |
| 95 | + |
| 96 | + // Validate key type early |
| 97 | + let key_type = cli.key_type.to_lowercase(); |
| 98 | + if !matches!(key_type.as_str(), "ed25519" | "rsa") { |
| 99 | + anyhow::bail!( |
| 100 | + "Unknown key type: '{}'. Supported types: ed25519 (recommended), rsa", |
| 101 | + cli.key_type |
| 102 | + ); |
| 103 | + } |
| 104 | + |
| 105 | + // Determine output path |
| 106 | + let output = match cli.output { |
| 107 | + Some(path) => path, |
| 108 | + None => { |
| 109 | + let home = dirs::home_dir().context("Cannot determine home directory")?; |
| 110 | + let ssh_dir = home.join(".ssh"); |
| 111 | + |
| 112 | + // Ensure .ssh directory exists with proper permissions |
| 113 | + if !ssh_dir.exists() { |
| 114 | + create_ssh_directory(&ssh_dir)?; |
| 115 | + } |
| 116 | + |
| 117 | + match key_type.as_str() { |
| 118 | + "ed25519" => ssh_dir.join("id_ed25519"), |
| 119 | + "rsa" => ssh_dir.join("id_rsa"), |
| 120 | + _ => unreachable!(), |
| 121 | + } |
| 122 | + } |
| 123 | + }; |
| 124 | + |
| 125 | + // Ensure parent directory exists |
| 126 | + if let Some(parent) = output.parent() { |
| 127 | + if !parent.exists() { |
| 128 | + std::fs::create_dir_all(parent) |
| 129 | + .with_context(|| format!("Failed to create directory: {}", parent.display()))?; |
| 130 | + } |
| 131 | + } |
| 132 | + |
| 133 | + // Check if file exists and prompt for overwrite |
| 134 | + if output.exists() && !cli.yes { |
| 135 | + print!("{} already exists. Overwrite? (y/n) ", output.display()); |
| 136 | + io::stdout().flush()?; |
| 137 | + |
| 138 | + let mut response = String::new(); |
| 139 | + io::stdin().read_line(&mut response)?; |
| 140 | + if !response.trim().eq_ignore_ascii_case("y") { |
| 141 | + if !cli.quiet { |
| 142 | + println!("Aborted."); |
| 143 | + } |
| 144 | + return Ok(()); |
| 145 | + } |
| 146 | + } |
| 147 | + |
| 148 | + // Generate key |
| 149 | + let result = match key_type.as_str() { |
| 150 | + "ed25519" => keygen::generate_ed25519(&output, cli.comment.as_deref())?, |
| 151 | + "rsa" => keygen::generate_rsa(&output, cli.bits, cli.comment.as_deref())?, |
| 152 | + _ => unreachable!(), |
| 153 | + }; |
| 154 | + |
| 155 | + // Display output |
| 156 | + if !cli.quiet { |
| 157 | + println!("Your identification has been saved in {}", output.display()); |
| 158 | + println!("Your public key has been saved in {}.pub", output.display()); |
| 159 | + println!("The key fingerprint is:"); |
| 160 | + println!("{}", result.fingerprint); |
| 161 | + |
| 162 | + // Display public key for convenience |
| 163 | + println!("\nThe key's randomart image is not displayed (not implemented)."); |
| 164 | + println!("\nPublic key:"); |
| 165 | + println!("{}", result.public_key_openssh); |
| 166 | + } |
| 167 | + |
| 168 | + Ok(()) |
| 169 | +} |
| 170 | + |
| 171 | +/// Create the .ssh directory with proper permissions (0700) |
| 172 | +fn create_ssh_directory(path: &PathBuf) -> Result<()> { |
| 173 | + #[cfg(unix)] |
| 174 | + { |
| 175 | + use std::fs; |
| 176 | + use std::os::unix::fs::DirBuilderExt; |
| 177 | + |
| 178 | + fs::DirBuilder::new() |
| 179 | + .mode(0o700) // drwx------ (owner only) |
| 180 | + .create(path) |
| 181 | + .with_context(|| format!("Failed to create .ssh directory: {}", path.display()))?; |
| 182 | + } |
| 183 | + |
| 184 | + #[cfg(not(unix))] |
| 185 | + { |
| 186 | + std::fs::create_dir_all(path) |
| 187 | + .with_context(|| format!("Failed to create .ssh directory: {}", path.display()))?; |
| 188 | + } |
| 189 | + |
| 190 | + Ok(()) |
| 191 | +} |
| 192 | + |
| 193 | +#[cfg(test)] |
| 194 | +mod tests { |
| 195 | + use super::*; |
| 196 | + use clap::CommandFactory; |
| 197 | + use tempfile::tempdir; |
| 198 | + |
| 199 | + #[test] |
| 200 | + fn test_cli_parsing() { |
| 201 | + // Verify CLI structure is valid |
| 202 | + Cli::command().debug_assert(); |
| 203 | + } |
| 204 | + |
| 205 | + #[test] |
| 206 | + fn test_cli_defaults() { |
| 207 | + let cli = Cli::try_parse_from(["bssh-keygen"]).unwrap(); |
| 208 | + assert_eq!(cli.key_type, "ed25519"); |
| 209 | + assert_eq!(cli.bits, 4096); |
| 210 | + assert!(cli.output.is_none()); |
| 211 | + assert!(cli.comment.is_none()); |
| 212 | + assert!(!cli.yes); |
| 213 | + assert!(!cli.quiet); |
| 214 | + } |
| 215 | + |
| 216 | + #[test] |
| 217 | + fn test_cli_ed25519() { |
| 218 | + let cli = Cli::try_parse_from(["bssh-keygen", "-t", "ed25519"]).unwrap(); |
| 219 | + assert_eq!(cli.key_type, "ed25519"); |
| 220 | + } |
| 221 | + |
| 222 | + #[test] |
| 223 | + fn test_cli_rsa() { |
| 224 | + let cli = Cli::try_parse_from(["bssh-keygen", "-t", "rsa", "-b", "2048"]).unwrap(); |
| 225 | + assert_eq!(cli.key_type, "rsa"); |
| 226 | + assert_eq!(cli.bits, 2048); |
| 227 | + } |
| 228 | + |
| 229 | + #[test] |
| 230 | + fn test_cli_output_file() { |
| 231 | + let cli = Cli::try_parse_from(["bssh-keygen", "-f", "/tmp/my_key"]).unwrap(); |
| 232 | + assert_eq!(cli.output, Some(PathBuf::from("/tmp/my_key"))); |
| 233 | + } |
| 234 | + |
| 235 | + #[test] |
| 236 | + fn test_cli_comment() { |
| 237 | + let cli = Cli::try_parse_from(["bssh-keygen", "-C", "user@host"]).unwrap(); |
| 238 | + assert_eq!(cli.comment, Some("user@host".to_string())); |
| 239 | + } |
| 240 | + |
| 241 | + #[test] |
| 242 | + fn test_cli_flags() { |
| 243 | + let cli = Cli::try_parse_from(["bssh-keygen", "-y", "-q"]).unwrap(); |
| 244 | + assert!(cli.yes); |
| 245 | + assert!(cli.quiet); |
| 246 | + } |
| 247 | + |
| 248 | + #[test] |
| 249 | + fn test_cli_verbose() { |
| 250 | + let cli = Cli::try_parse_from(["bssh-keygen", "-vvv"]).unwrap(); |
| 251 | + assert_eq!(cli.verbose, 3); |
| 252 | + } |
| 253 | + |
| 254 | + #[test] |
| 255 | + fn test_cli_full_options() { |
| 256 | + let cli = Cli::try_parse_from([ |
| 257 | + "bssh-keygen", |
| 258 | + "-t", |
| 259 | + "rsa", |
| 260 | + "-b", |
| 261 | + "4096", |
| 262 | + "-f", |
| 263 | + "/tmp/test_key", |
| 264 | + "-C", |
| 265 | + "test@example.com", |
| 266 | + "-y", |
| 267 | + "-v", |
| 268 | + ]) |
| 269 | + .unwrap(); |
| 270 | + |
| 271 | + assert_eq!(cli.key_type, "rsa"); |
| 272 | + assert_eq!(cli.bits, 4096); |
| 273 | + assert_eq!(cli.output, Some(PathBuf::from("/tmp/test_key"))); |
| 274 | + assert_eq!(cli.comment, Some("test@example.com".to_string())); |
| 275 | + assert!(cli.yes); |
| 276 | + assert!(!cli.quiet); |
| 277 | + assert_eq!(cli.verbose, 1); |
| 278 | + } |
| 279 | + |
| 280 | + #[test] |
| 281 | + fn test_cli_long_options() { |
| 282 | + let cli = Cli::try_parse_from([ |
| 283 | + "bssh-keygen", |
| 284 | + "--type", |
| 285 | + "ed25519", |
| 286 | + "--file", |
| 287 | + "/tmp/key", |
| 288 | + "--comment", |
| 289 | + "my key", |
| 290 | + "--yes", |
| 291 | + "--quiet", |
| 292 | + ]) |
| 293 | + .unwrap(); |
| 294 | + |
| 295 | + assert_eq!(cli.key_type, "ed25519"); |
| 296 | + assert_eq!(cli.output, Some(PathBuf::from("/tmp/key"))); |
| 297 | + assert_eq!(cli.comment, Some("my key".to_string())); |
| 298 | + assert!(cli.yes); |
| 299 | + assert!(cli.quiet); |
| 300 | + } |
| 301 | + |
| 302 | + #[test] |
| 303 | + fn test_create_ssh_directory() { |
| 304 | + let temp_dir = tempdir().unwrap(); |
| 305 | + let ssh_dir = temp_dir.path().join(".ssh"); |
| 306 | + |
| 307 | + let result = create_ssh_directory(&ssh_dir); |
| 308 | + assert!(result.is_ok()); |
| 309 | + assert!(ssh_dir.exists()); |
| 310 | + |
| 311 | + #[cfg(unix)] |
| 312 | + { |
| 313 | + use std::os::unix::fs::PermissionsExt; |
| 314 | + let metadata = std::fs::metadata(&ssh_dir).unwrap(); |
| 315 | + let permissions = metadata.permissions(); |
| 316 | + assert_eq!(permissions.mode() & 0o777, 0o700); |
| 317 | + } |
| 318 | + } |
| 319 | +} |
0 commit comments