I have a terribly written Reqwest based API Struct for a webserver. I wrote it while I was learning how Rust works. I'd like some pointers on how I can optimize it and make it cleaner than it is now. It is much more verbose than it needs to be.
Any/all tips are welcome!
use reqwest::header;
use reqwest::header::HeaderMap;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tokio::time::sleep;
use tokio::time::timeout;
use std::time::Duration;
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct Service {
// Maybe make this a parent class, and custom sub classes for each service type?
// Generic service that works for anythhing without a target!
pub team_id: u32,
pub team_name: String,
pub host_name: String,
pub service_name: String,
pub service_uid: u32,
pub public_ip: String,
pub port: u32,
pub targets: Vec<String>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct WorkerRes {
// For the /api/worker/list endpoint
pub worker_id: u32,
pub session_id: u32,
pub worker_host_name: String,
pub last_ping: String,
pub assigned_services: Vec<Service>, // Unknown length array of services to score
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct ConfigRes {
// For the /api/config/list endpoint
pub config_id: u32,
pub comp_name: String,
pub check_increment: u64,
pub passed_award: u32,
pub sla_count: u32,
pub sla_penalty: u32,
pub ccs_penalty: u32,
pub mode: u8,
pub last_check_increment_time: String,
}
#[derive(Debug)]
pub struct Config {
// For authentication
pub otp: String,
pub token: String,
pub front_end: String,
pub worker_name: String,
pub port: u64,
pub cookie: String, // TODO priv
pub client: Client, // TODO priv
}
impl Config {
// A data type that holds all the api endpoints for the frontend system
pub fn new(arg_otp: &String, token: &String, front_end: &String, worker_name: &String, arg_port: &u64) -> Config {
Config {
otp: (arg_otp).to_string(),
token: token.to_string(),
front_end: front_end.to_string(),
worker_name: worker_name.to_string(),
port: *arg_port,
cookie: "".to_string(),
client: reqwest::Client::builder().cookie_store(true).build().unwrap(),
}
}
pub async fn start(&mut self) {
println!("[+] Starting login process");
if self.token == "" {
if self.otp == "" {
// Get root token (to make the token we're going to use for scoring)
// self.token = self.api_auth(self.client.clone());
loop {
match self.api_auth().await {
Ok(token) => {
self.token = token;
println!("[+] /api/auth worked!");
break;
}
Err(_err) => {
print!("\r/api/auth ERROR !!! Trying until it works...");
}
}
}
// Generate OTP with the root token
loop {
match self.api_worker_otp().await {
Ok(otp) => {
self.otp = otp;
println!("[+] /api/worker/otp worked!");
break;
}
Err(_err) => {
print!("\r/api/worker/otp ERROR !!! Trying until it works... ");
}
}
}
}
// If given a OTP we jump down here
loop {
match self.api_worker_init().await {
// Token is set inside the program
Ok(_wrk) => {
println!("[+] /api/worker/init worked!");
break;
}
Err(_err) => {
print!("\r/api/worker/init ERROR !!! Trying until it works... ");
}
}
}
} else {
self.cookie = format!("X-Auth-Token={}", self.token);
println!("{:?}", self.cookie);
// println!("TODO ensure this is not the root cookie. Currently it is the root cookie!");
}
}
async fn api_auth(&mut self) -> Result<String, Box<dyn std::error::Error>> {
// Get ROOT TOKEN
let url = format!("http://{}:{}/api/auth", self.front_end, self.port);
let res = self.client.post(url).body(r#"[{"user_name": "root","user_password":"toor"}]"#).send().await?;
let auth = res.json::<HashMap<String, String>>().await?;
println!("{:?}", auth);
let auth_key: String = auth.get("token").expect("Hashmap skill issue.").to_string();
self.token = auth_key.to_string();
println!("ROOT TOKEN GENERATED : {:?}", auth_key);
self.cookie = format!("X-Auth-Token={}", self.token);
Ok(format!("X-Auth-Token={}", auth_key)) // Return Root cookie token
}
async fn api_worker_otp(&self) -> Result<String, Box<dyn std::error::Error>> {
// Generate OTP from the ROOT token we got
let mut headers = HeaderMap::new();
headers.insert(header::COOKIE, self.cookie.parse()?);
headers.insert("Content-Type", "application/json".parse()?);
let url = format!("http://{}:{}/api/worker/otp", self.front_end, self.port);
let otp_res = self.client.get(url).headers(headers).send().await?;
let h_arg_otp = otp_res.json::<HashMap<String, String>>().await?;
println!("{:?}", h_arg_otp);
Ok(
h_arg_otp.get("key").expect("Did not generate OTP correctly").to_string(), // Return OTP
)
}
async fn api_worker_init(&mut self) -> Result<(), Box<dyn std::error::Error>> {
let mut headers = HeaderMap::new();
headers.insert("Content-Type", "application/json".parse()?);
let url = format!(
"http://{}:{}/api/worker/init?otp={}&hostname={}&hmac={}",
self.front_end, self.port, self.otp, self.worker_name, "49efef5f70d47adc2db2eb397fbef5f7bc560e29"
);
println!("URL FOR /api/worker/init : {:?}", url);
let worker_req = self.client.get(url).headers(headers).send().await?;
let worker_res = worker_req.json::<HashMap<String, String>>().await?;
println!("/api/worker/init {:?}", worker_res);
// Clean this up once understanding the authenication and hostname assignment!
if !worker_res.get("error").is_none() {
// If error, show this
println!("{:?}", worker_res.get("error").expect("Hashmap skill issue."));
println!("{:?}", self.cookie);
} else if !worker_res.get("event").expect("Hashmap skill issue.").find("reauthenticated").is_none() {
// If event
// exsists
// Reauth -> steal that token given to use!
let auth_key: String = worker_res.get("token").expect("Hashmap skill issue.").to_string();
self.token = auth_key.to_string();
println!("Re-auth token extracted : {:?}", auth_key);
self.cookie = format!("X-Auth-Token={}", self.token);
println!("X-Auth-Token={}", auth_key); // Return Root cookie token
} else {
self.cookie = format!("X-Auth-Token={}", worker_res.get("token").expect("Hashmap skill issue."));
println!("COOKIE FOR LIST : {:?}", self.cookie);
}
Ok(())
}
pub async fn api_config_list(&self) -> Result<ConfigRes, Box<dyn std::error::Error>> {
let mut headers = HeaderMap::new();
headers.insert(header::COOKIE, self.cookie.parse()?);
headers.insert("Content-Type", "application/json".parse()?);
let url = format!("http://{}:{}/api/config/list", self.front_end, self.port);
println!("This sucks.");
let mut cont_check_req = None;
loop {
let h=headers.clone();
let u=url.clone();
match timeout(Duration::from_secs(5), self.client.get(u).headers(h).send()).await {
Ok(val) => {
cont_check_req = Some(val);
break;
},
Err(_) => {
tokio::time::sleep(Duration::from_secs(3)).await;
}
}
}
let cont_check_res = cont_check_req.unwrap()?.json::<Vec<ConfigRes>>().await?;
//.expect("Config config list API GET REQUEST error. Probably did not authenticate correctly or expired token.");
Ok(cont_check_res[0].clone())
}
pub async fn api_worker_list(&self) -> Result<WorkerRes, Box<dyn std::error::Error>> {
let mut headers = HeaderMap::new();
headers.insert(header::COOKIE, self.cookie.parse().unwrap());
headers.insert("Content-Type", "application/json".parse().unwrap());
let url = format!("http://{}:{}/api/worker/list", self.front_end, self.port);
let cont_check_req;
match self.client.get(url).headers(headers).send().await {
Ok(val) => {
cont_check_req = val;
}
Err(err) => return Err(format!("{}", err).into()),
}
let cont_check_res = cont_check_req.json::<Vec<WorkerRes>>().await?;
return Ok(cont_check_res[0].clone()); // crashed here because of vector skill issue?
/* for debug use... maybe make this toggle-able
//match cont_check_res {
Ok(res) => {
return Ok(res[0].clone()); // Crashing here
}
Err(e) => {
println!("Error with response from /api/worker/list");
return Err(e.into());
},
}
*/
}
}