From 6c121b918555a77cb6d8a97519004e31dd7f6563 Mon Sep 17 00:00:00 2001
From: Shayne Hartford <shaybox@shaybox.com>
Date: Tue, 29 Oct 2024 03:37:44 -0400
Subject: [PATCH] 0.3.1 - AutoExit & AutoReconnect Plugins

---
 Cargo.lock                    |  2 +-
 Cargo.toml                    |  2 +-
 README.md                     |  1 +
 src/minecraft/auto_exit.rs    | 27 ++++++++++++++++
 src/minecraft/commands/mod.rs |  2 --
 src/minecraft/mod.rs          | 59 +++++++++++++++++++++++++----------
 src/minecraft/prelude.rs      |  1 +
 7 files changed, 74 insertions(+), 20 deletions(-)
 create mode 100644 src/minecraft/auto_exit.rs

diff --git a/Cargo.lock b/Cargo.lock
index 4f86cf5..321491f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2900,7 +2900,7 @@ dependencies = [
 
 [[package]]
 name = "shaysbot"
-version = "0.3.0"
+version = "0.3.1"
 dependencies = [
  "anyhow",
  "azalea",
diff --git a/Cargo.toml b/Cargo.toml
index 20f890e..dd6fcce 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "shaysbot"
-version = "0.3.0"
+version = "0.3.1"
 authors = ["Shayne Hartford <shaybox@shaybox.com>"]
 edition = "2021"
 description = "My personal Minecraft bot using Azalea"
diff --git a/README.md b/README.md
index 7aaab39..d53cd0b 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,7 @@ My personal Minecraft bot using Azalea
 - [PearlCommandPlugin](src/minecraft/commands/pearl.rs)
 - [AntiAfkPlugin](src/minecraft/anti_afk.rs)
 - [AutoEatPlugin](src/minecraft/auto_eat.rs)
+- [AutoExitPlugin](src/minecraft/auto_exit.rs)
 - [AutoLookPlugin](src/minecraft/auto_look.rs)
 - [AutoPearlPlugin](src/minecraft/auto_pearl.rs)
 - [PearlTrackerPlugin](src/minecraft/pearl_tracker.rs)
diff --git a/src/minecraft/auto_exit.rs b/src/minecraft/auto_exit.rs
new file mode 100644
index 0000000..4258a87
--- /dev/null
+++ b/src/minecraft/auto_exit.rs
@@ -0,0 +1,27 @@
+use azalea::{
+    app::{App, Plugin, Update},
+    disconnect::DisconnectEvent,
+    ecs::prelude::*,
+};
+
+pub struct AutoExitPlugin;
+
+impl Plugin for AutoExitPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_systems(Update, handle_disconnect_event);
+    }
+}
+
+pub fn handle_disconnect_event(mut events: EventReader<DisconnectEvent>) {
+    for event in events.read() {
+        let Some(reason) = &event.reason else {
+            continue;
+        };
+
+        println!("Disconnect Reason: {}", reason.to_ansi());
+        if ["AutoDisconnect"].contains(&&*reason.to_string()) {
+            eprintln!("Exiting to stay disconnected...");
+            std::process::exit(1);
+        }
+    }
+}
diff --git a/src/minecraft/commands/mod.rs b/src/minecraft/commands/mod.rs
index 54b7b0b..976f6dd 100644
--- a/src/minecraft/commands/mod.rs
+++ b/src/minecraft/commands/mod.rs
@@ -60,8 +60,6 @@ pub fn handle_chat_received_event(
     commands: Res<Commands>,
 ) {
     for event in events.read() {
-        println!("{}", event.packet.message().to_ansi());
-
         let Ok(settings) = query.get(event.entity) else {
             continue;
         };
diff --git a/src/minecraft/mod.rs b/src/minecraft/mod.rs
index 6822b4c..73433e9 100644
--- a/src/minecraft/mod.rs
+++ b/src/minecraft/mod.rs
@@ -1,6 +1,12 @@
 use std::sync::Arc;
 
-use azalea::{Account, ClientBuilder};
+use anyhow::Result;
+use azalea::{
+    ecs::prelude::*,
+    prelude::*,
+    swarm::{Swarm, SwarmBuilder, SwarmEvent},
+    Account,
+};
 use derive_config::{DeriveTomlConfig, DeriveYamlConfig};
 use parking_lot::RwLock;
 
@@ -16,18 +22,22 @@ pub mod prelude;
 
 mod anti_afk;
 mod auto_eat;
+mod auto_exit;
 mod auto_look;
 mod auto_pearl;
 mod auto_totem;
 mod commands;
 mod pearl_tracker;
 
+#[derive(Clone, Component, Default, Resource)]
+pub struct SwarmState;
+
 /// # Create and start the Minecraft bot client
 ///
 /// # Errors
 /// Will return `Err` if `ClientBuilder::start` fails.
 #[allow(clippy::future_not_send)]
-pub async fn start() -> anyhow::Result<()> {
+pub async fn start() -> Result<()> {
     let settings = Settings::load().unwrap_or_default();
     let trapdoors = Trapdoors::load().unwrap_or_default();
     let address = settings.server_address.clone();
@@ -40,18 +50,35 @@ pub async fn start() -> anyhow::Result<()> {
     settings.save()?;
     let settings = Arc::new(RwLock::new(settings));
     let trapdoors = Arc::new(RwLock::new(trapdoors));
-    let client = ClientBuilder::new().add_plugins((
-        CommandsPlugin,
-        PearlCommandPlugin,
-        AntiAfkPlugin,
-        AutoEatPlugin,
-        AutoLookPlugin,
-        AutoPearlPlugin,
-        AutoTotemPlugin,
-        PearlTrackerPlugin,
-        SettingsPlugin(settings),
-        TrapdoorsPlugin(trapdoors),
-    ));
-
-    client.start(account, address).await?
+    let client = SwarmBuilder::new()
+        .set_swarm_handler(swarm_handler)
+        .add_account(account)
+        .add_plugins((CommandsPlugin, PearlCommandPlugin))
+        .add_plugins((
+            AntiAfkPlugin,
+            AutoEatPlugin,
+            AutoExitPlugin,
+            AutoLookPlugin,
+            AutoPearlPlugin,
+            AutoTotemPlugin,
+            PearlTrackerPlugin,
+            SettingsPlugin(settings),
+            TrapdoorsPlugin(trapdoors),
+        ));
+
+    client.start(address).await?
+}
+
+/// # Errors
+/// Will return `Err` if `Swarm::add_with_opts` fails.
+pub async fn swarm_handler(mut swarm: Swarm, event: SwarmEvent, state: SwarmState) -> Result<()> {
+    match event {
+        SwarmEvent::Chat(chat_packet) => println!("{}", chat_packet.message().to_ansi()),
+        SwarmEvent::Disconnect(account, options) => {
+            swarm.add_with_opts(&account, state, &options).await?;
+        }
+        _ => {}
+    }
+
+    Ok(())
 }
diff --git a/src/minecraft/prelude.rs b/src/minecraft/prelude.rs
index 35cb04f..37b9607 100644
--- a/src/minecraft/prelude.rs
+++ b/src/minecraft/prelude.rs
@@ -1,6 +1,7 @@
 pub use super::{
     anti_afk::*,
     auto_eat::*,
+    auto_exit::*,
     auto_look::*,
     auto_pearl::*,
     auto_totem::*,