Merge pull request #1 from FxQnLr/eta

Eta

4 files changed, 224 insertions, 48 deletions

diff --git a/src/error.rs b/src/error.rs
index d35991b..f15c60a 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,9 +1,11 @@
-use std::fmt::Debug;
+use std::{fmt::Debug, num::ParseIntError};
 
 pub enum CliError {
     Reqwest(reqwest::Error),
     Config(config::ConfigError),
     Serde(serde_json::Error),
+    Parse(ParseIntError),
+    WsResponse,
 }
 
 impl Debug for CliError {
@@ -12,6 +14,8 @@ impl Debug for CliError {
             Self::Reqwest(err) => { err.fmt(f) },
             Self::Config(err) => { err.fmt(f) },
             Self::Serde(err) => { err.fmt(f) },
+            Self::Parse(err) => { err.fmt(f) },
+            Self::WsResponse => { f.write_str("Error in Response") },
         }
     }
 }
diff --git a/src/main.rs b/src/main.rs
index ab7e476..afe6fac 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,10 @@
-use clap::{Parser, Subcommand};
+use std::{fmt::Display, time::Duration};
+
+use clap::{Parser, Command, CommandFactory, Subcommand};
+use clap_complete::{generate, Shell, Generator};
 use config::SETTINGS;
 use error::CliError;
+use indicatif::{ProgressBar, ProgressStyle, MultiProgress};
 use requests::{start::start, device};
 use reqwest::header::{HeaderMap, HeaderValue};
 use serde::Deserialize;
@@ -9,7 +13,15 @@ mod config;
 mod error;
 mod requests;
 
-/// webol http client
+static OVERVIEW_STYLE: &str = "{spinner:.green} ({elapsed}{wide_msg}";
+static OVERVIEW_ERROR: &str = "✗ ({elapsed}) {wide_msg}";
+static OVERVIEW_DONE: &str = "✓ ({elapsed}) {wide_msg}";
+static DEFAULT_STYLE: &str = "  {spinner:.green} {wide_msg}";
+static DONE_STYLE: &str = "  ✓ {wide_msg}";
+static ERROR_STYLE: &str = "  ✗ {wide_msg}";
+static TICK_SPEED: u64 = 1000 / 16;
+
+/// webol client
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Args {
@@ -21,12 +33,17 @@ struct Args {
 enum Commands {
     Start {
         /// id of the device
-        id: String
+        id: String,
+        #[arg(short, long)]
+        ping: Option<bool>
     },
     Device {
         #[command(subcommand)]
         devicecmd: DeviceCmd,
-    }
+    },
+    CliGen {
+        id: Shell,
+    },
 }
 
 #[derive(Subcommand)]
@@ -34,7 +51,8 @@ enum DeviceCmd {
     Add {
         id: String,
         mac: String,
-        broadcast_addr: String
+        broadcast_addr: String,
+        ip: String
     },
     Get {
         id: String,
@@ -42,35 +60,46 @@ enum DeviceCmd {
     Edit {
         id: String,
         mac: String,
-        broadcast_addr: String
+        broadcast_addr: String,
+        ip: String
     },
 }
 
-fn main() -> Result<(), CliError> {
+#[tokio::main]
+async fn main() -> Result<(), CliError> {
     let cli = Args::parse();
 
     match cli.commands {
-        Commands::Start { id } => {
-            start(id)?;
+        Commands::Start { id, ping } => {
+            start(id, ping.unwrap_or(true)).await?;
         },
         Commands::Device { devicecmd } => {
             match devicecmd {
-                DeviceCmd::Add { id, mac, broadcast_addr } => {
-                    device::put(id, mac, broadcast_addr)?;
+                DeviceCmd::Add { id, mac, broadcast_addr, ip } => {
+                    device::put(id, mac, broadcast_addr, ip).await?;
                 },
                 DeviceCmd::Get { id } => {
-                    device::get(id)?;
+                    device::get(id).await?;
                 },
-                DeviceCmd::Edit { id, mac, broadcast_addr } => {
-                    device::post(id, mac, broadcast_addr)?;
+                DeviceCmd::Edit { id, mac, broadcast_addr, ip } => {
+                    device::post(id, mac, broadcast_addr, ip).await?;
                 },
             }
+        },
+        Commands::CliGen { id } => {
+            eprintln!("Generating completion file for {id:?}...");
+            let mut cmd = Args::command();
+            print_completions(id, &mut cmd)
         }
     }
 
     Ok(())
 }
 
+fn print_completions<G: Generator>(gen: G, cmd: &mut Command) {
+    generate(gen, cmd, cmd.get_name().to_string(), &mut std::io::stdout());
+}
+
 fn default_headers() -> Result<HeaderMap, CliError> {
     let mut map = HeaderMap::new();
     map.append("Accept-Content", HeaderValue::from_str("application/json").unwrap());
@@ -87,14 +116,44 @@ fn default_headers() -> Result<HeaderMap, CliError> {
     Ok(map)
 }
 
-fn format_url(path: &str) -> Result<String, CliError> {
+fn format_url(path: &str, protocol: Protocols) -> Result<String, CliError> {
     Ok(format!(
-        "{}/{}",
+        "{}://{}/{}",
+        protocol,
         SETTINGS.get_string("server").map_err(CliError::Config)?,
         path
     ))
 }
 
+fn add_pb(mp: &MultiProgress, template: &str, message: String) -> ProgressBar {
+    let pb = mp.add(ProgressBar::new(1));
+    pb.set_style(ProgressStyle::with_template(template).unwrap());
+    pb.enable_steady_tick(Duration::from_millis(TICK_SPEED));
+    pb.set_message(message);
+
+    pb
+}
+
+fn finish_pb(pb: ProgressBar, message: String, template: &str) {
+    pb.set_style(ProgressStyle::with_template(template).unwrap());
+    pb.finish_with_message(message);
+
+}
+
+enum Protocols {
+    Http,
+    Websocket,
+}
+
+impl Display for Protocols {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Http => f.write_str("http"),
+            Self::Websocket => f.write_str("ws")
+        }
+    }
+}
+
 #[derive(Debug, Deserialize)]
 struct ErrorResponse {
     error: String
diff --git a/src/requests/device.rs b/src/requests/device.rs
index 525745a..cbc838e 100644
--- a/src/requests/device.rs
+++ b/src/requests/device.rs
@@ -1,55 +1,65 @@
-use crate::{error::CliError, default_headers, format_url};
+use crate::{error::CliError, default_headers, format_url, Protocols};
 
-pub fn put(id: String, mac: String, broadcast_addr: String) -> Result<(), CliError> {
-    let res = reqwest::blocking::Client::new()
-        .put(format_url("device")?)
+pub async fn put(id: String, mac: String, broadcast_addr: String, ip: String) -> Result<(), CliError> {
+    let url = format_url("device", Protocols::Http)?;
+    println!("{}", url);
+    let res = reqwest::Client::new()
+        .put(url)
         .headers(default_headers()?)
         .body(
             format!(
-                r#"{{"id": "{}", "mac": "{}", "broadcast_addr": "{}"}}"#,
+                r#"{{"id": "{}", "mac": "{}", "broadcast_addr": "{}", "ip": "{}"}}"#,
                 id,
                 mac,
-                broadcast_addr
+                broadcast_addr,
+                ip
             )
         )
         .send()
+        .await
         .map_err(CliError::Reqwest)?
-        .text();
+        .text()
+        .await;
 
     println!("{:?}", res);
     Ok(())
 }
 
-pub fn get(id: String) -> Result<(), CliError> {
-    let res = reqwest::blocking::Client::new()
-        .get(format_url("device")?)
+pub async fn get(id: String) -> Result<(), CliError> {
+    let res = reqwest::Client::new()
+        .get(format_url("device", Protocols::Http)?)
         .headers(default_headers()?)
         .body(
             format!(r#"{{"id": "{}"}}"#, id)
         )
         .send()
+        .await
         .map_err(CliError::Reqwest)?
-        .text();
+        .text()
+        .await;
 
     println!("{:?}", res);
     Ok(())
 }
 
-pub fn post(id: String, mac: String, broadcast_addr: String) -> Result<(), CliError> {
-    let res = reqwest::blocking::Client::new()
-        .post(format_url("device")?)
+pub async fn post(id: String, mac: String, broadcast_addr: String, ip: String) -> Result<(), CliError> {
+    let res = reqwest::Client::new()
+        .post(format_url("device", Protocols::Http)?)
         .headers(default_headers()?)
         .body(
             format!(
-                r#"{{"id": "{}", "mac": "{}", "broadcast_addr": "{}"}}"#,
+                r#"{{"id": "{}", "mac": "{}", "broadcast_addr": "{}", "ip": "{}"}}"#,
                 id,
                 mac,
-                broadcast_addr
+                broadcast_addr,
+                ip
             )
         )
         .send()
+        .await
         .map_err(CliError::Reqwest)?
-        .text();
+        .text()
+        .await;
 
     println!("{:?}", res);
     Ok(())
diff --git a/src/requests/start.rs b/src/requests/start.rs
index 30f65b9..ca4ca44 100644
--- a/src/requests/start.rs
+++ b/src/requests/start.rs
@@ -1,49 +1,152 @@
+use futures_util::{StreamExt, SinkExt};
+use indicatif::{MultiProgress, ProgressBar};
 use reqwest::StatusCode;
 use serde::Deserialize;
+use tokio_tungstenite::{connect_async, tungstenite::Message};
 
-use crate::{config::SETTINGS, error::CliError, default_headers, ErrorResponse};
+use crate::{error::CliError, default_headers, ErrorResponse, format_url, Protocols, OVERVIEW_STYLE, DEFAULT_STYLE, DONE_STYLE, finish_pb, ERROR_STYLE, OVERVIEW_ERROR, OVERVIEW_DONE, add_pb};
 
-pub fn start(id: String) -> Result<(), CliError> {
-    let res = reqwest::blocking::Client::new()
-        .post(
-            format!(
-                "{}/start",
-                SETTINGS.get_string("server").map_err(CliError::Config)?
-            )
-        )
+pub async fn start(id: String, ping: bool) -> Result<(), CliError> {
+
+    let send_start = MultiProgress::new();
+    let overview = add_pb(&send_start, OVERVIEW_STYLE, format!(") start {}", id));
+
+    // TODO: calculate average start-time on server
+    let url = format_url("start", Protocols::Http)?;
+    let connect = add_pb(&send_start, DEFAULT_STYLE, format!("connect to {}", url));
+    let res = reqwest::Client::new()
+        .post(url)
         .headers(default_headers()?)
         .body(
-            format!(r#"{{"id": "{}"}}"#, id)
+            format!(r#"{{"id": "{}", "ping": {}}}"#, id, ping)
         )
         .send()
+        .await
         .map_err(CliError::Reqwest)?;
+    finish_pb(connect, "connected, got response".to_string(), DONE_STYLE);
 
+    let res_pb = add_pb(&send_start, DEFAULT_STYLE, "analyzing response".to_string());
     match res.status() {
         StatusCode::OK => {
             let body = serde_json::from_str::<StartResponse>(
-                &res.text().map_err(CliError::Reqwest)?
+                &res.text().await.map_err(CliError::Reqwest)?
             )
             .map_err(CliError::Serde)?;
 
             if body.boot {
-                println!("successfully started {}", body.id);
+                finish_pb(res_pb, "sent start packet".to_string(), DONE_STYLE);
+            }
+
+            if ping {
+                let status = status_socket(body.uuid, &send_start, &overview, id).await?;
+                if status {
+                    finish_pb(overview, format!("successfully started {}", body.id), OVERVIEW_DONE);
+                } else {
+                    finish_pb(overview, format!("error while starting {}", body.id), OVERVIEW_ERROR);
+                }
             }
         },
         _ => {
             let body = serde_json::from_str::<ErrorResponse>(
-                &res.text().map_err(CliError::Reqwest)?
+                &res.text().await.map_err(CliError::Reqwest)?
             )
             .map_err(CliError::Serde)?;
 
-            println!("got error: {}", body.error);
+            res_pb.finish_with_message(format!("✗ got error: {}", body.error));
         }
     }
 
     Ok(())
 }
 
+async fn status_socket(uuid: String, pb: &MultiProgress, overview: &ProgressBar, id: String) -> Result<bool, CliError> {
+    // TODO: Remove unwraps
+    let ws_pb = add_pb(pb, DEFAULT_STYLE, "connect to websocket".to_string());
+    let (mut ws_stream, _response) = connect_async(format_url("status", Protocols::Websocket)?)
+        .await
+        .expect("Failed to connect");
+    finish_pb(ws_pb, "connected to websocket".to_string(), DONE_STYLE);
+    
+    ws_stream.send(Message::Text(uuid.clone())).await.unwrap();
+    
+    // Get ETA
+    let eta_msg = ws_stream.next().await.unwrap().unwrap();
+    let eta = get_eta(eta_msg.into_text().unwrap(), uuid.clone())? + overview.elapsed().as_secs();
+    overview.set_message(format!("/{}) start {}", eta, id));
+
+    let msg_pb = add_pb(pb, DEFAULT_STYLE, "await message".to_string());
+    let msg = ws_stream.next().await.unwrap();
+    finish_pb(msg_pb, "received message".to_string(), DONE_STYLE);
+
+    ws_stream.close(None).await.unwrap();
+
+    let v_pb = add_pb(pb, DEFAULT_STYLE, "verify response".to_string());
+    let res = verify_response(msg.unwrap().to_string(), uuid)?;
+    match res {
+        Verified::WrongUuid => {
+            finish_pb(v_pb, "returned wrong uuid".to_string(), ERROR_STYLE);
+            Ok(false)
+        },
+        Verified::ResponseType(res_type) => {
+            match res_type {
+                ResponseType::Start => {
+                    finish_pb(v_pb, "device started".to_string(), DONE_STYLE);
+                    Ok(true)
+                },
+                ResponseType::Timeout => {
+                    finish_pb(v_pb, "ping timed out".to_string(), ERROR_STYLE);
+                    Ok(false)
+                },
+                ResponseType::NotFound => {
+                    finish_pb(v_pb, "unknown uuid".to_string(), ERROR_STYLE);
+                    Ok(false)
+                },
+            }
+        }
+    }
+}
+
+fn get_eta(msg: String, uuid: String) -> Result<u64, CliError> {
+    let spl: Vec<&str> = msg.split('_').collect();
+    if (spl[0] != "eta") || (spl[2] != uuid) { return Err(CliError::WsResponse); };
+    Ok(u64::from_str_radix(spl[1], 10).map_err(CliError::Parse)?)
+}
+
+fn verify_response(res: String, org_uuid: String) -> Result<Verified, CliError> {
+    let spl: Vec<&str> = res.split('_').collect();
+    let res_type = spl[0];
+    let uuid = spl[1];
+    
+    if uuid != org_uuid { return Ok(Verified::WrongUuid) };
+    
+    Ok(Verified::ResponseType(ResponseType::from(res_type)?))
+}
+
 #[derive(Debug, Deserialize)]
 struct StartResponse {
     boot: bool,
     id: String,
+    uuid: String,
+}
+
+enum Verified {
+    ResponseType(ResponseType),
+    WrongUuid
+}
+
+enum ResponseType {
+    Start,
+    Timeout,
+    NotFound,
+}
+
+impl ResponseType {
+    fn from(value: &str) -> Result<Self, CliError> {
+        match value {
+            "start" => Ok(ResponseType::Start),
+            "timeout" => Ok(ResponseType::Timeout),
+            "notfound" => Ok(ResponseType::NotFound),
+            _ => Err(CliError::WsResponse),
+        }
+    }
 }