@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};
88use std:: time:: Duration ;
99
1010const PATH_DEVICES : & str = "/api/devices" ;
11+ const PATH_DEVICES_SEARCH : & str = "/api/devices/search" ;
1112
1213#[ derive( Debug ) ]
1314pub struct NetshotClient {
@@ -53,6 +54,27 @@ pub struct NewDeviceCreatedPayload {
5354 pub status : String ,
5455}
5556
57+ #[ derive( Debug , Serialize , Deserialize ) ]
58+ struct UpdateDevicePayload {
59+ enabled : bool ,
60+ }
61+
62+ #[ derive( Debug , Serialize , Deserialize ) ]
63+ pub struct DeviceUpdatedPayload {
64+ pub status : String ,
65+ }
66+
67+ #[ derive( Debug , Serialize , Deserialize ) ]
68+ struct DeviceSearchQueryPayload {
69+ query : String ,
70+ }
71+
72+ #[ derive( Debug , Serialize , Deserialize ) ]
73+ pub struct DeviceSearchResultPayload {
74+ pub query : String ,
75+ pub devices : Vec < Device > ,
76+ }
77+
5678impl NetshotClient {
5779 /// Create a client with the given authentication token
5880 pub fn new (
@@ -111,7 +133,7 @@ impl NetshotClient {
111133 /// Register a given IP into Netshot and return the corresponding device
112134 pub fn register_device (
113135 & self ,
114- ip_address : & String ,
136+ ip_address : String ,
115137 domain_id : u32 ,
116138 ) -> Result < NewDeviceCreatedPayload , Error > {
117139 log:: info!( "Registering new device with IP {}" , ip_address) ;
@@ -143,6 +165,106 @@ impl NetshotClient {
143165
144166 Ok ( device_registration)
145167 }
168+
169+ /// Search for a device
170+ pub fn search_device ( & self , query_string : String ) -> Result < DeviceSearchResultPayload , Error > {
171+ let url = format ! ( "{}{}" , self . url, PATH_DEVICES_SEARCH ) ;
172+
173+ let query = DeviceSearchQueryPayload {
174+ query : query_string. clone ( ) ,
175+ } ;
176+
177+ let response = self . client . post ( url) . json ( & query) . send ( ) ?;
178+
179+ if !response. status ( ) . is_success ( ) {
180+ log:: warn!(
181+ "Failed to search for device with query `{}`: {}" ,
182+ query_string. clone( ) ,
183+ response. status( ) . to_string( )
184+ ) ;
185+ return Err ( anyhow ! (
186+ "Failed to search for device with query: {}" ,
187+ query_string
188+ ) ) ;
189+ }
190+
191+ let search_result: DeviceSearchResultPayload = response. json ( ) ?;
192+ log:: debug!(
193+ "Found {} devices with the given search" ,
194+ search_result. devices. len( ) ,
195+ ) ;
196+
197+ Ok ( search_result)
198+ }
199+
200+ /// Set the given device to a given state (enabled/disabled)
201+ fn set_device_enabled (
202+ & self ,
203+ ip_address : String ,
204+ enabled : bool ,
205+ ) -> Result < Option < DeviceUpdatedPayload > , Error > {
206+ log:: info!(
207+ "Setting device with IP {} to enabled={}" ,
208+ ip_address,
209+ enabled
210+ ) ;
211+
212+ let state = UpdateDevicePayload { enabled : enabled } ;
213+
214+ // Search for the device ID
215+ let response = self . search_device ( format ! ( "[IP] IS {}" , ip_address) ) ?;
216+ let device = response. devices . first ( ) . unwrap ( ) ;
217+
218+ if !enabled && device. status == "DISABLED" {
219+ log:: warn!(
220+ "Device {}({}) is already disabled, skipping" ,
221+ device. name,
222+ ip_address
223+ ) ;
224+ return Ok ( Option :: None ) ;
225+ } else if enabled && device. status != "DISABLED" {
226+ log:: warn!(
227+ "Device {}({}) is already enabled, skipping" ,
228+ device. name,
229+ ip_address
230+ ) ;
231+ return Ok ( Option :: None ) ;
232+ }
233+
234+ let url = format ! ( "{}{}/{}" , self . url, PATH_DEVICES , device. id) ;
235+ let response = self . client . put ( url) . json ( & state) . send ( ) ?;
236+
237+ if !response. status ( ) . is_success ( ) {
238+ log:: warn!(
239+ "Failed to update state for device {}, got status {}" ,
240+ ip_address,
241+ response. status( ) . to_string( )
242+ ) ;
243+ return Err ( anyhow ! (
244+ "Failed to update state for device {}, got status {}" ,
245+ ip_address,
246+ response. status( ) . to_string( )
247+ ) ) ;
248+ }
249+
250+ let device_update: DeviceUpdatedPayload = response. json ( ) ?;
251+ log:: debug!( "Device state of {} set to enabled={}" , ip_address, enabled) ;
252+
253+ Ok ( Option :: Some ( device_update) )
254+ }
255+
256+ /// Disable a given device
257+ pub fn disable_device (
258+ & self ,
259+ ip_address : String ,
260+ ) -> Result < Option < DeviceUpdatedPayload > , Error > {
261+ self . set_device_enabled ( ip_address, false )
262+ }
263+
264+ /// Enable a given device
265+ pub fn enable_device ( & self , ip_address : String ) -> Result < Option < DeviceUpdatedPayload > , Error > {
266+ self . set_device_enabled ( ip_address, true )
267+ }
146268}
147269
148270#[ cfg( test) ]
@@ -191,9 +313,50 @@ mod tests {
191313 . create ( ) ;
192314
193315 let client = NetshotClient :: new ( url. clone ( ) , String :: new ( ) , None , None , None ) . unwrap ( ) ;
194- let registration = client. register_device ( & String :: from ( "1.2.3.4" ) , 2 ) . unwrap ( ) ;
316+ let registration = client. register_device ( String :: from ( "1.2.3.4" ) , 2 ) . unwrap ( ) ;
195317
196318 assert_eq ! ( registration. task_id, 504 ) ;
197319 assert_eq ! ( registration. status, "SCHEDULED" ) ;
198320 }
321+
322+ #[ test]
323+ fn search_devices ( ) {
324+ let url = mockito:: server_url ( ) ;
325+
326+ let _mock = mockito:: mock ( "POST" , PATH_DEVICES_SEARCH )
327+ . match_query ( mockito:: Matcher :: Any )
328+ . match_body ( r#"{"query":"[IP] IS 1.2.3.4"}"# )
329+ . with_body_from_file ( "tests/data/netshot/search.json" )
330+ . create ( ) ;
331+
332+ let client = NetshotClient :: new ( url. clone ( ) , String :: new ( ) , None , None , None ) . unwrap ( ) ;
333+ let result = client
334+ . search_device ( String :: from ( "[IP] IS 1.2.3.4" ) )
335+ . unwrap ( ) ;
336+
337+ assert_eq ! ( result. devices. len( ) , 2 ) ;
338+ assert_eq ! ( result. query, "[IP] IS 1.2.3.4" ) ;
339+ }
340+
341+ #[ test]
342+ fn disable_device ( ) {
343+ let url = mockito:: server_url ( ) ;
344+
345+ let _mock = mockito:: mock ( "PUT" , format ! ( "{}/{}" , PATH_DEVICES , 2318 ) . as_str ( ) )
346+ . match_query ( mockito:: Matcher :: Any )
347+ . match_body ( r#"{"enabled":false}"# )
348+ . with_body_from_file ( "tests/data/netshot/disable_device.json" )
349+ . create ( ) ;
350+
351+ let _mock2 = mockito:: mock ( "POST" , PATH_DEVICES_SEARCH )
352+ . match_query ( mockito:: Matcher :: Any )
353+ . match_body ( r#"{"query":"[IP] IS 1.2.3.4"}"# )
354+ . with_body_from_file ( "tests/data/netshot/search.json" )
355+ . create ( ) ;
356+
357+ let client = NetshotClient :: new ( url. clone ( ) , String :: new ( ) , None , None , None ) . unwrap ( ) ;
358+ let registration = client. disable_device ( String :: from ( "1.2.3.4" ) ) . unwrap ( ) ;
359+
360+ assert_eq ! ( registration. unwrap( ) . status, "DISABLED" ) ;
361+ }
199362}
0 commit comments