implement daemon mode (keep alive)

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"

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<String>,
+}

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::<SrunChallenge>(raw_json).with_context(|| {

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<String> {
+    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<String>) -> Result<String, Error> {
+    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)
+}

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<u64>,
+}
+
+impl SrunDaemon {
+    pub fn new(config_path: Option<String>) -> Result<SrunDaemon> {
+        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::<SrunDaemon>(&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(())
+    }
+}

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 => {}

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;