feat: add authentication to cli client

cli/src/main.rs

@@ -4,18 +4,19 @@ extern crate upend_db;
 use crate::common::{REQWEST_ASYNC_CLIENT, WEBUI_PATH};
 use crate::config::UpEndConfig;
 use actix_web::HttpServer;
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use clap::{Args, CommandFactory, FromArgMatches, Parser, Subcommand, ValueEnum};
 use filebuffer::FileBuffer;
 use rand::{thread_rng, Rng};
 use regex::Captures;
 use regex::Regex;
 use reqwest::Url;
-use serde_json::json;
+use serde_json::{json, Value};
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::path::Path;
 use std::path::PathBuf;
+use std::str::FromStr;
 use std::sync::{Arc, Mutex};
 use tracing::trace;
 use tracing::{debug, error, info, warn};
@@ -58,6 +59,9 @@ enum Commands {
         /// Output format
         #[arg(short, long, default_value = "tsv")]
         format: OutputFormat,
+        /// Credentials
+        #[arg(short, long)]
+        credentials: Credentials,
     },
     Get {
         /// URL of the UpEnd instance to query.
@@ -70,6 +74,9 @@ enum Commands {
         /// Output format
         #[arg(short, long, default_value = "tsv")]
         format: OutputFormat,
+        /// Credentials
+        #[arg(short, long)]
+        credentials: Credentials,
     },
     /// Insert an entry into an UpEnd server instance.
     Insert {
@@ -85,6 +92,17 @@ enum Commands {
         /// Output format
         #[arg(short, long, default_value = "tsv")]
         format: OutputFormat,
+        /// Credentials
+        #[arg(short, long)]
+        credentials: Credentials,
+    },
+    /// Get authorization token from an UpEnd server instance.
+    Authenticate {
+        /// URL of the UpEnd instance to query.
+        #[arg(short, long, default_value = "http://localhost:8093")]
+        url: Url,
+        /// Credentials
+        credentials: Credentials,
     },
     /// Get the address of a file, attribute, or URL.
     Address {
@@ -100,6 +118,30 @@ enum Commands {
     Serve(ServeArgs),
 }
 
+#[derive(Debug, Clone)]
+enum Credentials {
+    /// Use an API token to authenticate.
+    Token(String),
+    /// Use a username and password to authenticate.
+    Password(String, String),
+}
+
+impl FromStr for Credentials {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self> {
+        let parts: Vec<&str> = s.split(':').collect();
+        match parts.len() {
+            1 => Ok(Credentials::Token(parts[0].to_string())),
+            2 => Ok(Credentials::Password(
+                parts[0].to_string(),
+                parts[1].to_string(),
+            )),
+            _ => Err(anyhow::anyhow!("Invalid credentials format.")),
+        }
+    }
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Debug, ValueEnum)]
 enum OutputFormat {
     /// JSON
@@ -191,7 +233,13 @@ async fn main() -> Result<()> {
         .init();
 
     match args.command {
-        Commands::Query { url, query, format } => {
+        Commands::Query {
+            url,
+            query,
+            format,
+            credentials,
+        } => {
+            let api_token = get_api_token(&url, credentials).await?;
             let re = Regex::new(r#"@(="([^"]+)"|=([^ ]+))"#).unwrap();
 
             let query = re
@@ -210,6 +258,7 @@ async fn main() -> Result<()> {
             debug!("Querying \"{}\": {}", api_url, query);
             let response = REQWEST_ASYNC_CLIENT
                 .post(api_url)
+                .header("Authorization", format!("Bearer {}", api_token))
                 .body(query)
                 .send()
                 .await?;
@@ -225,8 +274,10 @@ async fn main() -> Result<()> {
             entity,
             attribute,
             format,
+            credentials,
         } => {
             let response = if let Some(attribute) = attribute {
+                let api_token = get_api_token(&url, credentials).await?;
                 let api_url = url.join("/api/query")?;
 
                 let entity = match entity {
@@ -242,6 +293,7 @@ async fn main() -> Result<()> {
                 debug!("Querying \"{}\": {}", api_url, query);
                 REQWEST_ASYNC_CLIENT
                     .post(api_url)
+                    .header("Authorization", format!("Bearer {}", api_token))
                     .body(query)
                     .send()
                     .await?
@@ -268,7 +320,9 @@ async fn main() -> Result<()> {
             attribute,
             value,
             format: _,
+            credentials,
         } => {
+            let api_token = get_api_token(&url, credentials).await?;
             let api_url = url.join("/api/obj")?;
 
             let entity = match entity {
@@ -286,7 +340,12 @@ async fn main() -> Result<()> {
             });
 
             debug!("Inserting {:?} at \"{}\"", body, api_url);
-            let response = REQWEST_ASYNC_CLIENT.put(api_url).json(&body).send().await?;
+            let response = REQWEST_ASYNC_CLIENT
+                .put(api_url)
+                .header("Authorization", format!("Bearer {}", api_token))
+                .json(&body)
+                .send()
+                .await?;
 
             match response.error_for_status_ref() {
                 Ok(_) => {
@@ -299,6 +358,16 @@ async fn main() -> Result<()> {
                 }
             }
         }
+        Commands::Authenticate { url, credentials } => match credentials {
+            Credentials::Token(_) => {
+                Err(anyhow!("Please specify a username and password for authentication in the format `username:password`."))
+            }
+            _ => {
+                let api_token = get_api_token(&url, credentials).await?;
+                println!("{}", api_token);
+                Ok(())
+            }
+        },
         Commands::Address {
             _type,
             input,
@@ -496,6 +565,34 @@ async fn main() -> Result<()> {
     }
 }
 
+async fn get_api_token(url: &Url, credentials: Credentials) -> Result<String> {
+    match credentials {
+        Credentials::Token(token) => Ok(token),
+        Credentials::Password(username, password) => {
+            debug!("Logging in as {}...", username);
+            let api_url = url.join("/api/auth/login?via=token").unwrap();
+            let body = json!({
+                "username": username,
+                "password": password
+            });
+
+            let response = REQWEST_ASYNC_CLIENT
+                .post(api_url)
+                .json(&body)
+                .send()
+                .await?;
+
+            response.error_for_status_ref()?;
+
+            let data: Value = response.json().await?;
+            data.get("key")
+                .and_then(|key| key.as_str())
+                .map(|key| key.to_string())
+                .ok_or_else(|| anyhow::anyhow!("No API token in response."))
+        }
+    }
+}
+
 type Entries = HashMap<String, serde_json::Value>;
 
 async fn print_response_entries(response: reqwest::Response, format: OutputFormat) -> Result<()> {