From 32bce3dc4c2b19ace1b5f06b3447c8efd08aace7 Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Fri, 31 Jan 2025 16:34:23 +0000
Subject: [PATCH] feat(cli): improvements to outputs & bind options

It is now possible to specify the host and port the SSH server listens
to using CLI arguments. Also improved host key naming and some inline
strings within the CLI version menu.
---
 Cargo.lock  |  1 +
 Cargo.toml  |  1 +
 src/cli.rs  | 21 ++++++++++++++-------
 src/main.rs | 31 ++++++++++++++++++++++++-------
 4 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index eb7f33f..68c335c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3434,6 +3434,7 @@ dependencies = [
  "directories",
  "futures",
  "human-panic",
+ "indoc",
  "json5",
  "lazy_static",
  "libc",
diff --git a/Cargo.toml b/Cargo.toml
index 6a6e09d..e377ac5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ derive_deref = "1.1.1"
 directories = "5.0.1"
 futures = "0.3.31"
 human-panic = "2.0.2"
+indoc = "2.0.5"
 json5 = "0.4.1"
 lazy_static = "1.5.0"
 libc = "0.2.161"
diff --git a/src/cli.rs b/src/cli.rs
index 8696a0a..b1cee2d 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -1,4 +1,5 @@
 use clap::Parser;
+use indoc::formatdoc;
 
 use crate::config::{get_config_dir, get_data_dir};
 
@@ -12,6 +13,13 @@ pub struct Cli {
     /// Frame rate, i.e. number of frames per second
     #[arg(short, long, value_name = "FLOAT", default_value_t = 60.0)]
     pub frame_rate: f64,
+
+    /// The host address to start the SSH server on
+    #[arg(short = 'H', long, value_name = "ADDRESS", default_value_t = String::from("127.0.0.1"))]
+    pub host: String,
+    /// The port to start the SSH server on
+    #[arg(short, long, value_name = "PORT", default_value_t = 2222)]
+    pub port: u16,
 }
 
 const VERSION_MESSAGE: &str = concat!(
@@ -30,13 +38,12 @@ pub fn version() -> String {
     let config_dir_path = get_config_dir().display().to_string();
     let data_dir_path = get_data_dir().display().to_string();
 
-    format!(
-        "\
-{VERSION_MESSAGE}
+    formatdoc! {"
+        {VERSION_MESSAGE}
 
-Authors: {author}
+        Authors: {author}
 
-Config directory: {config_dir_path}
-Data directory: {data_dir_path}"
-    )
+        Config directory: {config_dir_path}
+        Data directory: {data_dir_path}
+    "}
 }
diff --git a/src/main.rs b/src/main.rs
index 42c81fb..7885fa1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -24,13 +24,28 @@ mod ssh;
 mod tui;
 
 const SSH_KEYS: &[&[u8]] = &[
-    include_bytes!(concat!(env!("OUT_DIR"), "/ssh-rsa.pem")),
-    include_bytes!(concat!(env!("OUT_DIR"), "/ecdsa-sha2-nistp256.pem")),
-    include_bytes!(concat!(env!("OUT_DIR"), "/ssh-ed25519.pem")),
+    include_bytes!(concat!(env!("OUT_DIR"), "/rsa.pem")),
+    include_bytes!(concat!(env!("OUT_DIR"), "/ecdsa.pem")),
+    include_bytes!(concat!(env!("OUT_DIR"), "/ed25519.pem")),
 ];
 lazy_static! {
     pub(crate) static ref OPTIONS: Cli = Cli::parse();
-    pub(crate) static ref SOCKET_ADDR: SocketAddr = SocketAddr::from(([127, 0, 0, 1], 2222));
+    pub(crate) static ref SOCKET_ADDR: Option<SocketAddr> = Some(SocketAddr::from((
+        // Convert the hostname IP to a fixed size array of [u8; 4]
+        TryInto::<[u8; 4]>::try_into(
+            OPTIONS
+                .host
+                .splitn(4, ".")
+                .map(|octet_str| u8::from_str_radix(octet_str, 10)
+                    .map_err(|_| eyre!("Octet component out of range (expected u8)")))
+                .collect::<Result<Vec<u8>>>()
+                .ok()?
+        )
+        .ok()?,
+
+        // The port to listen on
+        OPTIONS.port
+    )));
 }
 
 #[tokio::main]
@@ -39,10 +54,12 @@ async fn main() -> Result<()> {
     crate::logging::init()?;
     let _ = *OPTIONS; // force clap to run by evaluating it
 
+    let socket_addr = SOCKET_ADDR.ok_or(eyre!("Invalid host IP provided"))?;
     let config = ssh_config();
-    tracing::info!("Attempting to listen on {}", *SOCKET_ADDR);
+
+    tracing::info!("Attempting to listen on {}", socket_addr);
     SshServer::default()
-        .run_on_socket(Arc::new(config), &TcpListener::bind(*SOCKET_ADDR).await?)
+        .run_on_socket(Arc::new(config), &TcpListener::bind(socket_addr).await?)
         .await
         .map_err(|err| eyre!(err))
 }
@@ -55,7 +72,7 @@ fn ssh_config() -> Config {
         .iter()
         .filter_map(|pem| PrivateKey::from_openssh(pem).ok())
         .collect();
-    
+
     tracing::trace!("SSH config: {:#?}", conf);
     conf
 }