azalea/
container.rs

1use std::{fmt, fmt::Debug};
2
3use azalea_client::{
4    Client,
5    inventory::{CloseContainerEvent, ContainerClickEvent},
6    packet::game::ReceiveGamePacketEvent,
7};
8use azalea_core::position::BlockPos;
9use azalea_entity::inventory::Inventory;
10use azalea_inventory::{
11    ItemStack, Menu,
12    operations::{ClickOperation, PickupClick, QuickMoveClick},
13};
14use azalea_physics::collision::BlockWithShape;
15use azalea_protocol::packets::game::ClientboundGamePacket;
16use bevy_app::{App, Plugin, Update};
17use bevy_ecs::{component::Component, prelude::MessageReader, system::Commands};
18use derive_more::Deref;
19use futures_lite::Future;
20
21use crate::bot::BotClientExt;
22
23pub struct ContainerPlugin;
24impl Plugin for ContainerPlugin {
25    fn build(&self, app: &mut App) {
26        app.add_systems(Update, handle_menu_opened_event);
27    }
28}
29
30pub trait ContainerClientExt {
31    /// Open a container in the world, like a chest.
32    ///
33    /// Use [`Client::open_inventory`] to open your own inventory.
34    ///
35    /// This function times out after 5 seconds (100 ticks). Use
36    /// [`Self::open_container_at_with_timeout_ticks`] if you would like to
37    /// configure this.
38    ///
39    /// ```
40    /// # use azalea::{prelude::*, registry::builtin::BlockKind};
41    /// # async fn example(mut bot: azalea::Client) {
42    /// let target_pos = bot
43    ///     .world()
44    ///     .read()
45    ///     .find_block(bot.position(), &BlockKind::Chest.into());
46    /// let Some(target_pos) = target_pos else {
47    ///     bot.chat("no chest found");
48    ///     return;
49    /// };
50    /// let container = bot.open_container_at(target_pos).await;
51    /// # }
52    /// ```
53    fn open_container_at(
54        &self,
55        pos: BlockPos,
56    ) -> impl Future<Output = Option<ContainerHandle>> + Send;
57
58    /// Open a container in the world, or time out after a specified amount of
59    /// ticks.
60    ///
61    /// See [`Self::open_container_at`] for more information. That function
62    /// defaults to a timeout of 5 seconds (100 ticks), which is usually good
63    /// enough. However to detect failures faster or to account for server
64    /// lag, you may find it useful to adjust the timeout to a different
65    /// value.
66    ///
67    /// The timeout is measured in game ticks (on the client, not the server),
68    /// i.e. 1/20th of a second.
69    fn open_container_at_with_timeout_ticks(
70        &self,
71        pos: BlockPos,
72        timeout_ticks: Option<usize>,
73    ) -> impl Future<Output = Option<ContainerHandle>> + Send;
74
75    /// Wait until a container is open, up to the specified number of ticks.
76    ///
77    /// Returns `None` if the container was immediately opened and closed, or if
78    /// the timeout expired.
79    ///
80    /// If `timeout_ticks` is None, there will be no timeout.
81    fn wait_for_container_open(
82        &self,
83        timeout_ticks: Option<usize>,
84    ) -> impl Future<Output = Option<ContainerHandle>> + Send;
85
86    /// Open the player's inventory.
87    ///
88    /// This will return None if another container is open.
89    ///
90    /// Note that this will send a packet to the server once it's dropped. Also,
91    /// due to how it's implemented, you could call this function multiple times
92    /// while another inventory handle already exists (but you shouldn't).
93    ///
94    /// If you just want to get the items in the player's inventory without
95    /// sending any packets, use [`Client::menu`], [`Menu::player_slots_range`],
96    /// and [`Menu::slots`].
97    fn open_inventory(&self) -> Option<ContainerHandle>;
98    /// Returns a [`ContainerHandleRef`] to the client's currently open
99    /// container, or their inventory.
100    ///
101    /// This will not send a packet to close the container when it's dropped,
102    /// which may cause anticheat compatibility issues if you modify your
103    /// inventory without closing it afterwards.
104    ///
105    /// To simulate opening your own inventory (like pressing 'e') in a way that
106    /// won't trigger anticheats, use [`Client::open_inventory`].
107    ///
108    /// To open a container in the world, use [`Client::open_container_at`].
109    fn get_inventory(&self) -> ContainerHandleRef;
110    /// Get the item in the bot's hotbar that is currently being held in its
111    /// main hand.
112    fn get_held_item(&self) -> ItemStack;
113}
114
115impl ContainerClientExt for Client {
116    async fn open_container_at(&self, pos: BlockPos) -> Option<ContainerHandle> {
117        self.open_container_at_with_timeout_ticks(pos, Some(20 * 5))
118            .await
119    }
120
121    async fn open_container_at_with_timeout_ticks(
122        &self,
123        pos: BlockPos,
124        timeout_ticks: Option<usize>,
125    ) -> Option<ContainerHandle> {
126        let mut ticks = self.get_tick_broadcaster();
127        // wait until it's not air (up to 10 ticks)
128        for _ in 0..10 {
129            let block = self.world().read().get_block_state(pos).unwrap_or_default();
130            if !block.is_collision_shape_empty() {
131                break;
132            }
133            let _ = ticks.recv().await;
134        }
135
136        self.ecs
137            .lock()
138            .entity_mut(self.entity)
139            .insert(WaitingForInventoryOpen);
140        self.block_interact(pos);
141
142        self.wait_for_container_open(timeout_ticks).await
143    }
144
145    async fn wait_for_container_open(
146        &self,
147        timeout_ticks: Option<usize>,
148    ) -> Option<ContainerHandle> {
149        let mut ticks = self.get_tick_broadcaster();
150        let mut elapsed_ticks = 0;
151        while ticks.recv().await.is_ok() {
152            let ecs = self.ecs.lock();
153            if ecs.get::<WaitingForInventoryOpen>(self.entity).is_none() {
154                break;
155            }
156
157            elapsed_ticks += 1;
158            if let Some(timeout_ticks) = timeout_ticks
159                && elapsed_ticks >= timeout_ticks
160            {
161                return None;
162            }
163        }
164
165        let ecs = self.ecs.lock();
166        let inventory = ecs.get::<Inventory>(self.entity).expect("no inventory");
167        if inventory.id == 0 {
168            None
169        } else {
170            Some(ContainerHandle::new(inventory.id, self.clone()))
171        }
172    }
173
174    fn open_inventory(&self) -> Option<ContainerHandle> {
175        let ecs = self.ecs.lock();
176        let inventory = ecs.get::<Inventory>(self.entity).expect("no inventory");
177        if inventory.id == 0 {
178            Some(ContainerHandle::new(0, self.clone()))
179        } else {
180            None
181        }
182    }
183
184    fn get_inventory(&self) -> ContainerHandleRef {
185        self.query_self::<&Inventory, _>(|inv| ContainerHandleRef::new(inv.id, self.clone()))
186    }
187
188    fn get_held_item(&self) -> ItemStack {
189        self.query_self::<&Inventory, _>(|inv| inv.held_item().clone())
190    }
191}
192
193/// A handle to a container that may be open.
194///
195/// This does not close the container when it's dropped. See [`ContainerHandle`]
196/// if that behavior is desired.
197pub struct ContainerHandleRef {
198    id: i32,
199    client: Client,
200}
201impl Debug for ContainerHandleRef {
202    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
203        f.debug_struct("ContainerHandle")
204            .field("id", &self.id())
205            .finish()
206    }
207}
208impl ContainerHandleRef {
209    pub fn new(id: i32, client: Client) -> Self {
210        Self { id, client }
211    }
212
213    pub fn close(&self) {
214        self.client.ecs.lock().trigger(CloseContainerEvent {
215            entity: self.client.entity,
216            id: self.id,
217        });
218    }
219
220    /// Get the ID of the container.
221    ///
222    /// If this is 0, that means it's the player's inventory. Otherwise, the
223    /// number isn't really meaningful since only one container can be open
224    /// at a time.
225    pub fn id(&self) -> i32 {
226        self.id
227    }
228
229    /// Returns the menu of the container.
230    ///
231    /// If the container is closed, this will return `None`.
232    ///
233    /// Note that any modifications you make to the `Menu` you're given will not
234    /// actually cause any packets to be sent. If you're trying to modify your
235    /// inventory, use [`Self::click`] instead
236    pub fn menu(&self) -> Option<Menu> {
237        let ecs = self.client.ecs.lock();
238        let inventory = ecs
239            .get::<Inventory>(self.client.entity)
240            .expect("no inventory");
241
242        // this also makes sure we can't access the inventory while a container is open
243        if inventory.id == self.id {
244            if self.id == 0 {
245                Some(inventory.inventory_menu.clone())
246            } else {
247                Some(inventory.container_menu.clone().unwrap())
248            }
249        } else {
250            None
251        }
252    }
253
254    /// Returns the item slots in the container, not including the player's
255    /// inventory.
256    ///
257    /// If the container is closed, this will return `None`.
258    pub fn contents(&self) -> Option<Vec<ItemStack>> {
259        self.menu().map(|menu| menu.contents())
260    }
261
262    /// Return the contents of the menu, including the player's inventory.
263    ///
264    /// If the container is closed, this will return `None`.
265    pub fn slots(&self) -> Option<Vec<ItemStack>> {
266        self.menu().map(|menu| menu.slots())
267    }
268
269    /// A shortcut for [`Self::click`] with `PickupClick::Left`.
270    pub fn left_click(&self, slot: impl Into<usize>) {
271        self.click(PickupClick::Left {
272            slot: Some(slot.into() as u16),
273        });
274    }
275    /// A shortcut for [`Self::click`] with `QuickMoveClick::Left`.
276    pub fn shift_click(&self, slot: impl Into<usize>) {
277        self.click(QuickMoveClick::Left {
278            slot: slot.into() as u16,
279        });
280    }
281    /// A shortcut for [`Self::click`] with `PickupClick::Right`.
282    pub fn right_click(&self, slot: impl Into<usize>) {
283        self.click(PickupClick::Right {
284            slot: Some(slot.into() as u16),
285        });
286    }
287
288    /// Simulate a click in the container and send the packet to perform the
289    /// action.
290    pub fn click(&self, operation: impl Into<ClickOperation>) {
291        let operation = operation.into();
292        self.client.ecs.lock().trigger(ContainerClickEvent {
293            entity: self.client.entity,
294            window_id: self.id,
295            operation,
296        });
297    }
298}
299
300/// A handle to the open container.
301///
302/// The container will be closed once this is dropped.
303#[derive(Deref)]
304pub struct ContainerHandle(ContainerHandleRef);
305
306impl Drop for ContainerHandle {
307    fn drop(&mut self) {
308        self.0.close();
309    }
310}
311impl Debug for ContainerHandle {
312    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
313        f.debug_struct("ContainerHandle")
314            .field("id", &self.id())
315            .finish()
316    }
317}
318impl ContainerHandle {
319    fn new(id: i32, client: Client) -> Self {
320        Self(ContainerHandleRef { id, client })
321    }
322
323    /// Closes the inventory by dropping the handle.
324    pub fn close(self) {
325        // implicitly calls drop
326    }
327}
328
329#[derive(Component, Debug)]
330pub struct WaitingForInventoryOpen;
331
332pub fn handle_menu_opened_event(
333    mut commands: Commands,
334    mut events: MessageReader<ReceiveGamePacketEvent>,
335) {
336    for event in events.read() {
337        if let ClientboundGamePacket::ContainerSetContent { .. } = event.packet.as_ref() {
338            commands
339                .entity(event.entity)
340                .remove::<WaitingForInventoryOpen>();
341        }
342    }
343}