From 262601a4a46913a05516788ae7315e6fe24b7757 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Sun, 3 Dec 2023 16:32:25 +0800 Subject: [PATCH 1/3] refactor all actions as functions in main --- src/main.rs | 179 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 106 insertions(+), 73 deletions(-) diff --git a/src/main.rs b/src/main.rs index f7021ef..980694e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,6 +7,8 @@ mod xencode; use anyhow::Context; use anyhow::Result; use clap::Parser; +use cli::ClientArgs; +use cli::StatusArgs; use owo_colors::OwoColorize; use owo_colors::Stream::Stderr; use owo_colors::Stream::Stdout; @@ -41,44 +43,7 @@ async fn cli() -> Result<()> { match &args.command { // check login status Some(Commands::Status(status_args)) => { - // only verbose on args.verbose = true and not outputting json - let login_state = get_login_state(&http_client, args.verbose).await?; - - // output json - if status_args.json & !args.verbose { - let raw_json = serde_json::to_string(&login_state)?; - println!("{}", raw_json); - return Ok(()); - } - - // output human readable - match login_state.error.as_str() { - "ok" => { - println!( - "{} {} {} is online", - "bitsrun:".if_supports_color(Stdout, |t| t.bright_green()), - &login_state - .online_ip - .to_string() - .if_supports_color(Stdout, |t| t.underline()), - format!("({})", login_state.user_name.clone().unwrap_or_default()) - .if_supports_color(Stdout, |t| t.dimmed()) - ); - - // print status table - print_login_state(login_state); - } - _ => { - println!( - "{} {} is offline", - "bitsrun:".if_supports_color(Stdout, |t| t.blue()), - login_state - .online_ip - .to_string() - .if_supports_color(Stdout, |t| t.underline()) - ); - } - } + srun_status(http_client, status_args, args.verbose).await? } // login or logout @@ -101,43 +66,15 @@ async fn cli() -> Result<()> { ) .await?; - if matches!(args.command, Some(Commands::Login(_))) { - let resp = srun_client.login(client_args.force, args.verbose).await?; - match resp.error.as_str() { - "ok" => println!( - "{} {} {} logged in", - "bitsrun:".if_supports_color(Stdout, |t| t.bright_green()), - resp.online_ip - .to_string() - .if_supports_color(Stdout, |t| t.underline()), - format!("({})", resp.username.clone().unwrap_or_default()) - .if_supports_color(Stdout, |t| t.dimmed()) - ), - _ => println!( - "{} failed to login, {} {}", - "bitsrun:".if_supports_color(Stdout, |t| t.red()), - resp.error, - format!("({})", resp.error_msg).if_supports_color(Stdout, |t| t.dimmed()) - ), + match &args.command { + Some(Commands::Login(_)) => { + srun_login(&srun_client, client_args, args.verbose).await? } - } else if matches!(args.command, Some(Commands::Logout(_))) { - let resp = srun_client.logout(client_args.force, args.verbose).await?; - match resp.error.as_str() { - "ok" | "logout_ok" => println!( - "{} {} logged out", - "bitsrun:".if_supports_color(Stdout, |t| t.green()), - resp.online_ip - .to_string() - .if_supports_color(Stdout, |t| t.underline()) - ), - _ => println!( - "{} failed to logout, {} {}", - "bitsrun:".if_supports_color(Stdout, |t| t.red()), - resp.error, - format!("({})", resp.error_msg).if_supports_color(Stdout, |t| t.dimmed()) - ), + Some(Commands::Logout(_)) => { + srun_logout(&srun_client, client_args, args.verbose).await? } - } + _ => {} + }; } Some(Commands::ConfigPaths) => print_config_paths(), @@ -147,3 +84,99 @@ async fn cli() -> Result<()> { Ok(()) } + +async fn srun_status( + http_client: reqwest::Client, + status_args: &StatusArgs, + verbose: bool, +) -> Result<()> { + // only verbose on args.verbose = true and not outputting json + let login_state = get_login_state(&http_client, verbose).await?; + + // output json + if status_args.json & !verbose { + let raw_json = serde_json::to_string(&login_state)?; + println!("{}", raw_json); + return Ok(()); + } + + // output human readable + match login_state.error.as_str() { + "ok" => { + println!( + "{} {} {} is online", + "bitsrun:".if_supports_color(Stdout, |t| t.bright_green()), + &login_state + .online_ip + .to_string() + .if_supports_color(Stdout, |t| t.underline()), + format!("({})", login_state.user_name.clone().unwrap_or_default()) + .if_supports_color(Stdout, |t| t.dimmed()) + ); + + // print status table + print_login_state(login_state); + } + _ => { + println!( + "{} {} is offline", + "bitsrun:".if_supports_color(Stdout, |t| t.blue()), + login_state + .online_ip + .to_string() + .if_supports_color(Stdout, |t| t.underline()) + ); + } + }; + Ok(()) +} + +async fn srun_login( + srun_client: &SrunClient, + client_args: &ClientArgs, + verbose: bool, +) -> Result<()> { + let resp = srun_client.login(client_args.force, verbose).await?; + match resp.error.as_str() { + "ok" => println!( + "{} {} {} logged in", + "bitsrun:".if_supports_color(Stdout, |t| t.bright_green()), + resp.online_ip + .to_string() + .if_supports_color(Stdout, |t| t.underline()), + format!("({})", resp.username.clone().unwrap_or_default()) + .if_supports_color(Stdout, |t| t.dimmed()) + ), + _ => println!( + "{} failed to login, {} {}", + "bitsrun:".if_supports_color(Stdout, |t| t.red()), + resp.error, + format!("({})", resp.error_msg).if_supports_color(Stdout, |t| t.dimmed()) + ), + }; + Ok(()) +} + +async fn srun_logout( + srun_client: &SrunClient, + client_args: &ClientArgs, + verbose: bool, +) -> Result<()> { + let resp = srun_client.logout(client_args.force, verbose).await?; + match resp.error.as_str() { + "ok" | "logout_ok" => println!( + "{} {} logged out", + "bitsrun:".if_supports_color(Stdout, |t| t.green()), + resp.online_ip + .to_string() + .if_supports_color(Stdout, |t| t.underline()) + ), + _ => println!( + "{} failed to logout, {} {}", + "bitsrun:".if_supports_color(Stdout, |t| t.red()), + resp.error, + format!("({})", resp.error_msg).if_supports_color(Stdout, |t| t.dimmed()) + ), + }; + Ok(()) +} From 83b1232f7ca548f642b7501236fcc4506006b026 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Mon, 4 Dec 2023 13:47:56 +0800 Subject: [PATCH 2/3] implement daemon mode (keep alive) --- Cargo.toml | 2 + src/cli.rs | 10 ++++ src/client.rs | 2 +- src/config.rs | 119 +++++++++++++++++++++++++++++++++++++++++++++++ src/daemon.rs | 112 ++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 11 ++++- src/tables.rs | 3 +- src/user.rs | 126 +++----------------------------------------------- 8 files changed, 263 insertions(+), 122 deletions(-) create mode 100644 src/config.rs create mode 100644 src/daemon.rs diff --git a/Cargo.toml b/Cargo.toml index d99ee91..2e5464e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,8 @@ tabled = { version = "0.14", features = ["color"] } humansize = "2.1" chrono-humanize = "0.2" chrono = "0.4" +log = "0.4.20" +pretty_env_logger = "0.5.0" [profile.release] strip = "symbols" diff --git a/src/cli.rs b/src/cli.rs index 8313d0d..20c48d9 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -28,6 +28,9 @@ pub enum Commands { /// List all possible config file paths ConfigPaths, + + /// Poll the server with login requests to keep the session alive + KeepAlive(DaemonArgs), } #[derive(Args)] @@ -63,3 +66,10 @@ pub struct ClientArgs { #[arg(short, long)] pub force: bool, } + +#[derive(Args)] +pub struct DaemonArgs { + /// Path to the config file + #[arg(short, long)] + pub config: Option, +} diff --git a/src/client.rs b/src/client.rs index b42f303..8131ba2 100644 --- a/src/client.rs +++ b/src/client.rs @@ -445,7 +445,7 @@ impl SrunClient { } if raw_text.len() < 8 { - bail!("logout response too short: `{}`", raw_text) + bail!("challenge response too short: `{}`", raw_text) } let raw_json = &raw_text[6..raw_text.len() - 1]; let parsed_json = serde_json::from_str::(raw_json).with_context(|| { diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..08340c8 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,119 @@ +use std::env; +use std::fs; + +use anyhow::anyhow; +use anyhow::Error; +use anyhow::Result; +use owo_colors::OwoColorize; +use owo_colors::Stream::Stdout; + +/// Enumerate possible paths to user config file (platform specific) +/// +/// On Windows: +/// * `~\AppData\Roaming\bitsrun\bit-user.json` +/// +/// On Linux: +/// * `$XDG_CONFIG_HOME/bitsrun/bit-user.json` +/// * `~/.config/bitsrun/bit-user.json` +/// * `~/.config/bit-user.json` +/// +/// On macOS: +/// * `$HOME/Library/Preferences/bitsrun/bit-user.json` +/// * `$HOME/.config/bit-user.json` +/// * `$HOME/.config/bitsrun/bit-user.json` +/// +/// Additionally, `bitsrun` will search for config file in the current working directory. +pub fn enumerate_config_paths() -> Vec { + let mut paths = Vec::new(); + + // Windows + if env::consts::OS == "windows" { + if let Some(appdata) = env::var_os("APPDATA") { + paths.push(format!( + "{}\\bitsrun\\bit-user.json", + appdata.to_str().unwrap() + )); + } + } + + // Linux (and macOS) + if let Some(home) = env::var_os("XDG_CONFIG_HOME").or_else(|| env::var_os("HOME")) { + paths.push(format!("{}/.config/bit-user.json", home.to_str().unwrap())); + paths.push(format!( + "{}/.config/bitsrun/bit-user.json", + home.to_str().unwrap() + )); + } + + // macOS + if env::consts::OS == "macos" { + if let Some(home) = env::var_os("HOME") { + paths.push(format!( + "{}/Library/Preferences/bitsrun/bit-user.json", + home.to_str().unwrap() + )); + } + } + + // current working directory + paths.push("bit-user.json".into()); + paths +} + +/// Config file validation +pub fn validate_config_file(config_path: &Option) -> Result { + let mut validated_config_path = String::new(); + match &config_path { + Some(path) => validated_config_path = path.to_owned(), + None => { + for path in enumerate_config_paths() { + if fs::metadata(&path).is_ok() { + validated_config_path = path; + break; + } + } + } + } + let meta = fs::metadata(&validated_config_path)?; + if !meta.is_file() { + return Err(anyhow!( + "`{}` is not a file", + &validated_config_path.if_supports_color(Stdout, |t| t.underline()) + )); + } + // file should only be read/writeable by the owner alone, i.e., 0o600 + // note: this check is only performed on unix systems + #[cfg(unix)] + fn check_permissions(config: &String, meta: &std::fs::Metadata) -> Result<(), anyhow::Error> { + use std::os::unix::fs::MetadataExt; + if meta.mode() & 0o777 != 0o600 { + return Err(anyhow!( + "`{}` has too open permissions {}, aborting!\n\ + {}: set permissions to {} with `chmod 600 {}`", + &config.if_supports_color(Stdout, |t| t.underline()), + (meta.mode() & 0o777) + .to_string() + .if_supports_color(Stdout, |t| t.on_red()), + "tip".if_supports_color(Stdout, |t| t.green()), + "600".if_supports_color(Stdout, |t| t.on_cyan()), + &config + )); + } + Ok(()) + } + #[cfg(windows)] + #[allow(unused)] + fn check_permissions(_config: &str, _meta: &std::fs::Metadata) -> Result<(), anyhow::Error> { + // Windows doesn't support Unix-style permissions, so we'll just return Ok here. + Ok(()) + } + check_permissions(&validated_config_path, &meta)?; + if validated_config_path.is_empty() { + return Err(anyhow!( + "file `{}` not found, available paths can be found with `{}`", + "bit-user.json".if_supports_color(Stdout, |t| t.underline()), + "bitsrun config-paths".if_supports_color(Stdout, |t| t.cyan()) + )); + } + Ok(validated_config_path) +} diff --git a/src/daemon.rs b/src/daemon.rs new file mode 100644 index 0000000..97e4882 --- /dev/null +++ b/src/daemon.rs @@ -0,0 +1,112 @@ +use crate::client::SrunClient; +use crate::config; + +use std::fs; + +use anyhow::Context; +use anyhow::Result; +use log::info; +use log::warn; +use owo_colors::OwoColorize; +use owo_colors::Stream::Stdout; + +use reqwest::Client; +use serde::Deserialize; + +use tokio::signal::ctrl_c; +use tokio::time::Duration; + +#[derive(Debug, Deserialize)] +pub struct SrunDaemon { + username: String, + password: String, + dm: bool, + // polls every 1 hour by default + poll_interval: Option, +} + +impl SrunDaemon { + pub fn new(config_path: Option) -> Result { + let finalized_cfg = config::validate_config_file(&config_path)?; + + // in daemon mode, bitsrun must be able to read all required fields from the config file, + // including `username`, `password`, and `dm`. + let daemon_cfg_str = fs::read_to_string(&finalized_cfg).with_context(|| { + format!( + "failed to read config file `{}`", + &finalized_cfg.if_supports_color(Stdout, |t| t.underline()) + ) + })?; + let daemon_cfg = + serde_json::from_str::(&daemon_cfg_str).with_context(|| { + format!( + "failed to parse config file `{}`", + &finalized_cfg.if_supports_color(Stdout, |t| t.underline()) + ) + })?; + + Ok(daemon_cfg) + } + + pub async fn start(&self, http_client: Client) -> Result<()> { + // set logger to INFO level by default + pretty_env_logger::formatted_builder() + .filter_level(log::LevelFilter::Info) + .init(); + + // set default polling intervals every 1 hour + let poll_interval = self.poll_interval.unwrap_or(3600); + + // warn if polling interval is too short + if poll_interval < 60 * 10 { + warn!("polling interval is too short, please set it to at least 10 minutes (600s)"); + } + + // start daemon + let mut srun_ticker = tokio::time::interval(Duration::from_secs(poll_interval)); + let srun = SrunClient::new( + self.username.clone(), + self.password.clone(), + Some(http_client), + None, + Some(self.dm), + ) + .await?; + + info!( + "starting daemon ({}) with polling interval={}s", + self.username, poll_interval, + ); + + loop { + let tick = srun_ticker.tick(); + let login = srun.login(true, false); + + tokio::select! { + _ = tick => { + match login.await { + Ok(resp) => { + match resp.error.as_str() { + "ok" => { + info!("{} ({}): login success, {}", resp.client_ip, self.username, resp.suc_msg.unwrap_or_default()); + } + _ => { + warn!("{} ({}): login failed, {}", resp.client_ip, self.username, resp.error); + } + } + } + Err(e) => { + warn!("{}: login failed: {}", self.username, e); + } + } + } + _ = ctrl_c() => { + info!("{}: gracefully exiting", self.username); + break; + } + } + } + + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs index 980694e..882227f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,7 @@ mod cli; mod client; +mod config; +mod daemon; mod tables; mod user; mod xencode; @@ -17,6 +19,7 @@ use cli::Arguments; use cli::Commands; use client::get_login_state; use client::SrunClient; +use daemon::SrunDaemon; use tables::print_config_paths; use tables::print_login_state; @@ -48,7 +51,7 @@ async fn cli() -> Result<()> { // login or logout Some(Commands::Login(client_args)) | Some(Commands::Logout(client_args)) => { - let bit_user = user::get_bit_user( + let bit_user = user::finalize_bit_user( &client_args.username, &client_args.password, client_args.dm, @@ -77,6 +80,12 @@ async fn cli() -> Result<()> { }; } + Some(Commands::KeepAlive(daemon_args)) => { + let config_path = daemon_args.config.to_owned(); + let daemon = SrunDaemon::new(config_path)?; + daemon.start(http_client).await?; + } + Some(Commands::ConfigPaths) => print_config_paths(), None => {} diff --git a/src/tables.rs b/src/tables.rs index 2be3b35..0f99725 100644 --- a/src/tables.rs +++ b/src/tables.rs @@ -1,5 +1,6 @@ use crate::client::SrunLoginState; -use crate::user::enumerate_config_paths; +use crate::config::enumerate_config_paths; + use chrono::Duration; use chrono_humanize::Accuracy::Rough; use chrono_humanize::HumanTime; diff --git a/src/user.rs b/src/user.rs index 585c540..c2cf85d 100644 --- a/src/user.rs +++ b/src/user.rs @@ -1,7 +1,7 @@ -use std::env; +use crate::config; + use std::fs; -use anyhow::anyhow; use anyhow::Context; use anyhow::Result; use owo_colors::OwoColorize; @@ -35,121 +35,9 @@ impl BitUserPartial { } } -/// Enumerate possible paths to user config file (platform specific) -/// -/// On Windows: -/// * `~\AppData\Roaming\bitsrun\bit-user.json` -/// -/// On Linux: -/// * `$XDG_CONFIG_HOME/bitsrun/bit-user.json` -/// * `~/.config/bitsrun/bit-user.json` -/// * `~/.config/bit-user.json` -/// -/// On macOS: -/// * `$HOME/Library/Preferences/bitsrun/bit-user.json` -/// * `$HOME/.config/bit-user.json` -/// * `$HOME/.config/bitsrun/bit-user.json` -/// -/// Additionally, `bitsrun` will search for config file in the current working directory. -pub fn enumerate_config_paths() -> Vec { - let mut paths = Vec::new(); - - // Windows - if env::consts::OS == "windows" { - if let Some(appdata) = env::var_os("APPDATA") { - paths.push(format!( - "{}\\bitsrun\\bit-user.json", - appdata.to_str().unwrap() - )); - } - } - - // Linux (and macOS) - if let Some(home) = env::var_os("XDG_CONFIG_HOME").or_else(|| env::var_os("HOME")) { - paths.push(format!("{}/.config/bit-user.json", home.to_str().unwrap())); - paths.push(format!( - "{}/.config/bitsrun/bit-user.json", - home.to_str().unwrap() - )); - } - - // macOS - if env::consts::OS == "macos" { - if let Some(home) = env::var_os("HOME") { - paths.push(format!( - "{}/Library/Preferences/bitsrun/bit-user.json", - home.to_str().unwrap() - )); - } - } - - // current working directory - paths.push("bit-user.json".into()); - paths -} - -/// Parse credentials from config file -fn parse_config_file(config_path: &Option) -> Result { - let mut config = String::new(); - match &config_path { - Some(path) => config = path.to_owned(), - None => { - for path in enumerate_config_paths() { - if fs::metadata(&path).is_ok() { - config = path; - break; - } - } - } - } - - // check if file is valid (i.e., is a file and permissions are not too open) - let meta = fs::metadata(&config)?; - if !meta.is_file() { - return Err(anyhow!( - "`{}` is not a file", - &config.if_supports_color(Stdout, |t| t.underline()) - )); - } - - // file should only be read/writeable by the owner alone, i.e., 0o600 - // note: this check is only performed on unix systems - #[cfg(unix)] - fn check_permissions(config: &String, meta: &std::fs::Metadata) -> Result<(), anyhow::Error> { - use std::os::unix::fs::MetadataExt; - if meta.mode() & 0o777 != 0o600 { - return Err(anyhow!( - "`{}` has too open permissions {}, aborting!\n\ - {}: set permissions to {} with `chmod 600 {}`", - &config.if_supports_color(Stdout, |t| t.underline()), - (meta.mode() & 0o777) - .to_string() - .if_supports_color(Stdout, |t| t.on_red()), - "tip".if_supports_color(Stdout, |t| t.green()), - "600".if_supports_color(Stdout, |t| t.on_cyan()), - &config - )); - } - Ok(()) - } - - #[cfg(windows)] - #[allow(unused)] - fn check_permissions(_config: &str, _meta: &std::fs::Metadata) -> Result<(), anyhow::Error> { - // Windows doesn't support Unix-style permissions, so we'll just return Ok here. - Ok(()) - } - - check_permissions(&config, &meta)?; - - // check if file is empty - if config.is_empty() { - return Err(anyhow!( - "file `{}` not found, available paths can be found with `{}`", - "bit-user.json".if_supports_color(Stdout, |t| t.underline()), - "bitsrun config-paths".if_supports_color(Stdout, |t| t.cyan()) - )); - } +/// Parse bit user credentials from config file +fn parse_bit_user_config(config_path: &Option) -> Result { + let config = config::validate_config_file(config_path)?; let user_str_from_file = fs::read_to_string(&config).with_context(|| { format!( @@ -171,7 +59,7 @@ fn parse_config_file(config_path: &Option) -> Result { /// /// Note that when logging out, `password` is not required. /// In this case, `require_password` should be set to `false`. -pub fn get_bit_user( +pub fn finalize_bit_user( username: &Option, password: &Option, dm: bool, @@ -183,7 +71,7 @@ pub fn get_bit_user( // username and password priority: command line > config file > prompt if bit_user.username.is_none() | (require_password & bit_user.password.is_none()) { let mut user_from_file = BitUserPartial::default(); - match parse_config_file(config_path) { + match parse_bit_user_config(config_path) { Ok(value) => user_from_file = value, Err(e) => println!( "{} {}", From f456f971cb04e93c70022eea4b81bc41b4eb1b82 Mon Sep 17 00:00:00 2001 From: spencerwooo Date: Mon, 4 Dec 2023 13:51:01 +0800 Subject: [PATCH 3/3] update readme --- README.md | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 2d0890c..a9a8b23 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ bitsrun: () is online ```console $ bitsrun --help -A headless login and logout CLI app for 10.0.0.55 at BIT +A headless login and logout CLI for 10.0.0.55 at BIT Usage: bitsrun [OPTIONS] [COMMAND] @@ -67,6 +67,7 @@ Commands: logout Logout from the campus network status Check device login status config-paths List all possible config file paths + keep-alive Poll the server with login requests to keep the session alive help Print this message or the help of the given subcommand(s) Options: @@ -86,11 +87,13 @@ To save your credentials and configurations, create config file `bit-user.json` { "username": "", "password": "", - "dm": true + "dm": true, + "poll_interval": 3600 } ``` -**`dm` is for specifying whether the current device is a dumb terminal, and requires logging out through the alternative endpoint. Set to `true` (no quotes!) if the device you are working with is a dumb terminal.** +- **`dm` is for specifying whether the current device is a dumb terminal, and requires logging out through the alternative endpoint. Set to `true` (no quotes!) if the device you are working with is a dumb terminal.** +- `poll_interval` is an optional field for specifying the interval (in seconds) of polling login requests. Default is `3600` seconds (1 hour). Used by `bitsrun keep-alive` only. Available config file paths can be listed with: