assemblyline_client/
lib.rs

1#![warn(missing_docs, non_ascii_idents, trivial_numeric_casts,
2    unused_crate_dependencies, noop_method_call, single_use_lifetimes, trivial_casts,
3    unused_lifetimes, nonstandard_style, variant_size_differences)]
4#![deny(keyword_idents)]
5// #![warn(clippy::missing_docs_in_private_items)]
6#![allow(clippy::needless_return)]
7// #![allow(clippy::needless_return, clippy::while_let_on_iterator, clippy::collapsible_else_if)]
8
9//! A library to help access the Assemblyline API from rust
10
11mod types;
12mod connection;
13mod modules;
14
15use std::sync::Arc;
16
17use modules::file::File;
18use modules::help::Help;
19use modules::alert::Alert;
20use modules::bundle::Bundle;
21use modules::error::Error;
22use modules::search::Search;
23use modules::ingest::Ingest;
24use modules::submit::Submit;
25pub use types::{Authentication, JsonMap};
26pub use connection::{Connection, TLSSettings};
27
28
29/// A client to communicate with the Assemblyline API
30pub struct Client {
31    // Connection handler
32    // _connection: Arc<Connection>,
33
34    /// alert specific API endpoints
35    pub alert: Alert,
36    /// bundle specific API endpoints
37    pub bundle: Bundle,
38    /// error specific API endpoints
39    pub error: Error,
40    /// file specific API endpoints
41    pub file: File,
42    // self.hash_search = HashSearch(self._connection)
43    /// Help API endpoints
44    pub help: Help,
45    // self.heuristics = Heuristics(self._connection)
46    /// Ingest API endpoints
47    pub ingest: Ingest,
48    // self.live = Live(self._connection)
49    // self.ontology = Ontology(self._connection)
50    // self.replay = Replay(self._connection)
51    // self.result = Result(self._connection)
52    // self.safelist = Safelist(self._connection)
53    /// search API endpoints
54    pub search: Search,
55    // self.service = Service(self._connection)
56    // self.signature = Signature(self._connection)
57    // self.socketio = SocketIO(self._connection)
58    // self.submission = Submission(self._connection)
59    /// File submission API endpoints
60    pub submit: Submit,
61    // self.system = System(self._connection)
62    // self.user = User(self._connection)
63    // self.workflow = Workflow(self._connection)
64}
65
66impl Client {
67    /// Connect to an assemblyline system
68    pub async fn connect(server: String, auth: Authentication) -> Result<Self, types::Error> {
69        let connection = Arc::new(Connection::connect(server, auth, None, connection::TLSSettings::Native, Default::default(), None).await?);
70        Ok(Self {
71            alert: Alert::new(connection.clone()),
72            bundle: Bundle::new(connection.clone()),
73            error: Error::new(connection.clone()),
74            file: File::new(connection.clone()),
75            help: Help::new(connection.clone()),
76            ingest: Ingest::new(connection.clone()),
77            search: Search::new(connection.clone()),
78            submit: Submit::new(connection)
79            // _connection,
80        })
81    }
82
83    /// Connect to an assemblyline system
84    pub async fn from_connection(connection: Arc<Connection>) -> Result<Self, types::Error> {
85        Ok(Self {
86            alert: Alert::new(connection.clone()),
87            bundle: Bundle::new(connection.clone()),
88            error: Error::new(connection.clone()),
89            file: File::new(connection.clone()),
90            help: Help::new(connection.clone()),
91            ingest: Ingest::new(connection.clone()),
92            search: Search::new(connection.clone()),
93            submit: Submit::new(connection)
94            // _connection,
95        })
96    }
97}
98
99
100#[cfg(test)]
101mod tests {
102
103    use std::sync::Arc;
104
105    use assemblyline_models::datastore::submission::{SubmissionParams, SubmissionState, ServiceSelection};
106    use assemblyline_models::types::ClassificationString;
107    use rand::Rng;
108
109    use crate::{Authentication, Client, Connection};
110
111    fn init() {
112        let _ = env_logger::builder().is_test(true).try_init();
113    }
114
115    pub (crate) async fn prepare_client() -> Client {
116        init();
117        let url = std::env::var("ASSEMBLYLINE_URL").unwrap();
118        let username = std::env::var("ASSEMBLYLINE_USER").unwrap();
119        let key = std::env::var("ASSEMBLYLINE_KEY").unwrap();
120
121        let connection = Connection::connect(
122            url, 
123            Authentication::ApiKey { username, key }, 
124            Some(2), 
125            crate::TLSSettings::Native, 
126            Default::default(), 
127            Some(30.0)
128        ).await.unwrap();
129        Client::from_connection(Arc::new(connection)).await.unwrap()
130        // Client::connect(url, Authentication::ApiKey { username, key }).await.unwrap()
131    }
132
133    fn random_body() -> Vec<u8> {
134        let mut out = vec![];
135        let mut prng = rand::rng();
136        let length = 128 + prng.random_range(0..256);
137        while out.len() < length {
138            out.push(prng.random());
139        }
140        out
141    }
142
143    #[tokio::test]
144    async fn submit_content() {
145        let client = prepare_client().await;
146        assemblyline_models::disable_global_classification();
147
148        let result = client.submit.single()
149            .metadata_item("testbatch".to_owned(), "0".to_owned())
150            .params(SubmissionParams{ ttl: 1, ..SubmissionParams::new(ClassificationString::new_unchecked("U".to_string()))})
151            .fname("test-file".to_owned())
152            .content(random_body()).await.unwrap();
153
154        assert_eq!(result.state, SubmissionState::Submitted);
155    }
156
157
158    #[tokio::test]
159    async fn search_single_page() {
160        let client = prepare_client().await;
161        assemblyline_models::disable_global_classification();
162        let batch: u64 = rand::rng().random();
163        let batch: String = batch.to_string();
164
165        let _result = client.ingest.single()
166            .metadata_item("testbatch".to_owned(), batch.clone())
167            .notification_queue(batch.clone())
168            .params(SubmissionParams{ priority: 300, ttl: 1, services: ServiceSelection{ selected: vec!["Characterize".into()], ..Default::default()}, ..SubmissionParams::new(ClassificationString::new_unchecked("U".to_owned()))})
169            .fname("test-file".to_owned())
170            .content(random_body()).await.unwrap();
171
172        for _ in 0..100 {
173            match client.ingest.get_message(&batch).await.unwrap() {
174                Some(message) => {
175                    assert_eq!(message.submission.metadata.get("testbatch").unwrap(), &batch);
176                    break;
177                },
178                None => { tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; },
179            }
180        }
181
182        for _ in 0..10 {
183            let result = client.search.submission(format!("metadata.testbatch: {batch}"))
184                .search().await.unwrap();
185
186            if result.items.len() == 1 {
187                return
188            }
189
190            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
191        }
192        panic!()
193    }
194
195
196
197}
assemblyline_client/lib.rs

assemblyline_client/
lib.rs
