move host tool added

1 files changed, 29 insertions, 20 deletions

diff --git a/src/main.rs b/src/main.rs
index 986b7f1..c9b80ab 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,9 +3,11 @@ pub mod embed;
 pub mod error;
 pub mod info;
 pub mod state;
+pub mod tool;
+
+use anyhow::{anyhow, Context};
 use embed::*;
 use info::*;
-
 use markup::Render;
 use rocket::{
     catch, catchers,
@@ -15,24 +17,29 @@ use rocket::{
     response::{self, Responder},
     routes,
     shield::{self, Shield},
+    tokio::{self},
     Request, Response,
 };
 use state::{AdInfo, Config, Logic};
-use std::{io::Cursor, net::IpAddr};
+use std::{env::args, fs::read_to_string, io::Cursor, net::IpAddr, sync::Arc};
+use tool::tool_main;
 
-#[rocket::main]
-async fn main() {
+fn main() -> anyhow::Result<()> {
     env_logger::init_from_env("LOG");
+    let mut args = args().skip(1);
+    let config = args.next().expect("first arg needs to be the config");
+    let config = read_to_string(config).context(anyhow!("could not read config"))?;
+    let mut config: Config = toml::from_str(config.as_str()).context(anyhow!("config invalid"))?;
 
-    let config = std::env::args()
-        .nth(1)
-        .expect("first arg needs to be the config");
-    let config = rocket::tokio::fs::read_to_string(config)
-        .await
-        .expect("could not read config");
-    let mut config: Config = toml::from_str(config.as_str()).expect("config invalid");
+    if let Some(action) = args.next() {
+        return tool_main(config, &action, args.collect::<Vec<_>>());
+    }
 
-    for entry in config.ad_dir.read_dir().expect("cannot read ad directory") {
+    for entry in config
+        .ad_dir
+        .read_dir()
+        .context(anyhow!("cannot read ad directory"))?
+    {
         if let Ok(entry) = entry {
             if entry
                 .path()
@@ -45,12 +52,7 @@ async fn main() {
                 let path = entry.path();
                 let imname = path.file_stem().unwrap().to_str().unwrap();
                 let basename = imname.split_once(".").unwrap().0;
-                let info: AdInfo = toml::from_str(
-                    &rocket::tokio::fs::read_to_string(entry.path())
-                        .await
-                        .unwrap(),
-                )
-                .unwrap();
+                let info: AdInfo = toml::from_str(&read_to_string(entry.path()).unwrap()).unwrap();
                 config.ads.insert(
                     basename.to_string(),
                     AdInfo {
@@ -64,6 +66,14 @@ async fn main() {
 
     let state = Logic::new(config);
 
+    tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()?
+        .block_on(inner_main(state));
+
+    Ok(())
+}
+async fn inner_main(state: Arc<Logic>) {
     let _ = rocket::build()
         .configure(rocket::Config {
             port: state.config.port,
@@ -78,8 +88,7 @@ async fn main() {
         .mount("/", routes![r_index, r_embed, r_style, r_image, r_iptest])
         .register("/", catchers![r_catch])
         .launch()
-        .await
-        .unwrap();
+        .await;
 }
 
 #[get("/myip")]