azalea/
container.rs

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