add authenticate

Cargo.toml

@@ -16,3 +16,5 @@ simple-log = "1.6.0"
 chrono = "0.4.26"
 hyper = "0.14.26"
 tungstenite = "0.19.0"
+base64 = "0.21.2"
+jwt-simple = "0.11.5"

src/authenticate/black.rs

@@ -0,0 +1,84 @@
+use std::collections::HashMap;
+use std::sync::Mutex;
+use lazy_static::lazy_static;
+use simple_log::{info, warn};
+
+lazy_static! {
+    static ref BLACKLIST: Mutex<HashMap<String, u8>> = Mutex::new(HashMap::new());
+}
+
+pub async fn is_in_blacklist(ip: String) -> bool {
+    let mut trys = 0;
+    loop {
+        trys += 1;
+        if trys > 100 {
+            return true;
+        }
+
+        if let Ok(map) = BLACKLIST.try_lock() {
+            return match map.get(ip.as_str()) {
+                None => {
+                    false
+                }
+                Some(a) => {
+                    if *a > 50 {
+                        true
+                    } else {
+                        false
+                    }
+                }
+            }
+        }
+    }
+}
+
+pub async fn clear_blacklist() {
+    loop {
+        tokio::time::sleep(std::time::Duration::from_secs(1200)).await;
+
+        let mut trys = 0;
+        loop {
+            trys += 1;
+            if trys > 100 {
+                break;
+            }
+
+            if let Ok(mut map) = BLACKLIST.try_lock() {
+                if map.len() > 0 {
+                    map.clear();
+                    info!("Clearing black list.");
+                }
+                break;
+            }
+        }
+
+    }
+}
+
+pub async fn record(ip: String) {
+    warn!("Recording {} into blacklist.", ip.as_str());
+    let mut trys = 0;
+    loop {
+        trys += 1;
+        if trys > 100 {
+            return;
+        }
+
+        if let Ok(mut map) = BLACKLIST.try_lock() {
+            let num = match map.get(ip.as_str()) {
+                None => {
+                    0
+                }
+                Some(a) => {
+                    *a
+                }
+            };
+            map.insert(ip.clone(), num + 1);
+            return;
+        }
+    }
+}
+
+pub fn init() {
+    tokio::spawn(clear_blacklist());
+}

src/authenticate/mod.rs

@@ -0,0 +1,14 @@
+pub mod token;
+pub mod black;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct UserInfo {//_id: mc uuid
+    pub _id: String,
+    pub display_name: String,
+    pub enabled: bool,
+    pub group: Vec<String>,
+    pub bind_qq: Option<i64>,
+    pub ban_reason: Option<String>
+}

src/authenticate/token.rs

@@ -0,0 +1,52 @@
+use base64::Engine;
+use jwt_simple::prelude::{HS256Key, MACLike};
+use simple_log::debug;
+use simple_log::log::error;
+use crate::authenticate::black::{is_in_blacklist, record};
+use crate::authenticate::UserInfo;
+use crate::CONFIG;
+
+impl UserInfo {
+    pub async fn from_token(token: String) -> Result<UserInfo, String> {
+        let key = match &CONFIG.authenticate.secret {
+            None => { return Err("Key is not defined.".to_string())}
+            Some(a) => {a}
+        };
+        let key = HS256Key::from_bytes(base64::engine::general_purpose::STANDARD.decode(key).unwrap().as_slice());
+
+        let claims = key.verify_token::<UserInfo>(&token, None);
+        match claims {
+            Ok(info) => Ok(info.custom),
+            Err(err) => {
+
+                error!("{}", err.to_string());
+                Err("Cannot parse token or token is valid".to_string())
+            }
+        }
+    }
+}
+
+pub async fn authenticate_token(token: String, ip: String) -> Result<(), String> {
+    if is_in_blacklist(ip.clone()).await {
+        debug!("{} is in black list.", ip);
+        return Err("In blacklist.".to_string());
+    }
+
+    let info = match UserInfo::from_token(token).await {
+        Ok(a) => {a}
+        Err(err) => {
+            record(ip.clone()).await;
+            return Err(err);
+        }
+    };
+
+    let group = info.group;
+    for permission in group {
+        if permission == "admin" {
+            return Ok(());
+        }
+    }
+
+    record(ip).await;
+    Err("No permission".to_string())
+}

src/broadcast/fire.rs

@@ -1,6 +1,6 @@
 use std::ops::{DerefMut};
-use simple_log::{info, warn};
-use tungstenite::{Message, WebSocket};
+use simple_log::{debug, info, warn};
+use tungstenite::{Message};
 use crate::broadcast::{CLIENTS};
 use crate::monitor::MachineInfo;
 
@@ -22,15 +22,15 @@ pub async fn fire(data: MachineInfo) {
     for client_info in clients.deref_mut() {
         let (ip, client) = client_info.as_mut().unwrap();
         match client.write_message(Message::Text(data_str.clone())) {
-            Ok(_) => {},
+            Ok(_) => {}
             Err(_) => {
                 warn!("Removing client: {}.", ip);
                 *client_info = None;
             }
         };
     }
-    clients.retain(|x| x.is_some());
     if clients.len() > 0 {
-        info!("Send to {} clients.", clients.len());
+        clients.retain(|x| x.is_some());
+        debug!("Send to {} clients.", clients.len());
     }
 }

src/config.rs

@@ -1,11 +1,15 @@
 use std::fs::OpenOptions;
 use std::io::{Read, Write};
+use base64::Engine;
+use jwt_simple::prelude::HS256Key;
 use serde::{Deserialize, Serialize};
+use simple_log::error;
 use toml::Value;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct Config {
     pub connection: ConnectionSetting,
+    pub authenticate: AuthenticateSetting,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -14,16 +18,21 @@ pub struct ConnectionSetting {
     pub port: i32,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct AuthenticateSetting {
+    pub method: String,
+    pub secret: Option<String>,
+    pub db_url: Option<String>,
+    pub db_name: Option<String>,
+}
+
 
 pub fn get_config() -> Config {
     let mut config: String = Default::default();
     let mut file = OpenOptions::new().read(true).write(true).create(true).open("config.toml").expect("Can not open 'config.toml'");
     file.read_to_string(&mut config).expect("Can not read 'config.toml'");
-    if config == "" {
-        config = "[mongodb]\n[connect]\n[feature]".to_string();
-    }
 
-    let mut config: Value = toml::from_str(config.as_str()).expect("Cannot parse config file, you should probably generate another.");
+    let config: Value = toml::from_str(config.as_str()).expect("Cannot parse config file, you should probably generate another.");
 
     let mut edit = false;
 
@@ -32,23 +41,35 @@ pub fn get_config() -> Config {
             address: "".to_string(),
             port: 0,
         },
+        authenticate: AuthenticateSetting {
+            method: "".to_string(),
+            secret: None,
+            db_url: None,
+            db_name: None,
+        },
     };
 
     result_config.connection = if let Some(a) = config.get("connection") {
         let address = if let Some(b) = a.get("address") {
             if let Some(c) = b.as_str() {
                 c.to_string()
             } else {
+                edit = true;
                 "127.0.0.1".to_string()
             }
         } else {
             edit = true;
             "127.0.0.1".to_string()
         };
-        let port = if let Some(b) = a.get("address") {
+
+
+
+
+        let port = if let Some(b) = a.get("port") {
             if let Some(c) = b.as_integer() {
                 c as i32
             } else {
+                edit = true;
                 7890
             }
         } else {
@@ -62,15 +83,74 @@ pub fn get_config() -> Config {
             port,
         }
     } else {
+        edit = true;
         ConnectionSetting {
             address: "127.0.0.1".to_string(),
             port: 7890,
         }
     };
 
+
+    result_config.authenticate = if let Some(a) = config.get("authenticate") {
+        let method = if let Some(b) = a.get("method") {
+            if let Some(c) = b.as_str() {
+                if c == "token" || c == "key" {
+                    c.to_string()
+                } else {
+                    edit = true;
+                    "token".to_string()
+                }
+            } else {
+                edit = true;
+                "token".to_string()
+            }
+        } else {
+            edit = true;
+            "token".to_string()
+        };
+
+        let secret = if method == "token" {
+            let key = if let Some(b) = a.get("secret") {
+                if let Some(c) = b.as_str() {
+                    c.to_string()
+                } else {
+                    edit = true;
+                    let key = HS256Key::generate();
+                    base64::engine::general_purpose::STANDARD.encode(key.to_bytes())
+                }
+            } else {
+                edit = true;
+                let key = HS256Key::generate();
+                base64::engine::general_purpose::STANDARD.encode(key.to_bytes())
+            };
+
+            Some(key)
+        } else {
+            None
+        };
+
+        AuthenticateSetting {
+            method,
+            secret,
+            db_url: None,
+            db_name: None,
+        }
+    } else {
+        let key = HS256Key::generate();
+        let key = base64::engine::general_purpose::STANDARD.encode(key.to_bytes());
+        AuthenticateSetting {
+            method: "token".to_string(),
+            secret: Some(key),
+            db_url: None,
+            db_name: None,
+        }
+    };
+
     if edit {
+        error!("Saving config.");
         let mut file = OpenOptions::new().write(true).create(true).open("config.toml").expect("Can not open 'config.toml'");
-        file.write(toml::to_string(&config).unwrap().as_bytes()).expect("Cannot save config.");
+        file.write(toml::to_string(&result_config).unwrap().as_bytes()).expect("Cannot save config.");
+        panic!("Config changed")
     }
 
     result_config