botnet_core/
lib.rs

1pub mod database;
2pub mod eval;
3pub mod task;
4pub mod utils;
5
6pub use crate::database::{Database, DatabaseKey};
7pub use async_std::sync::{Arc, Mutex};
8pub use botnet_utils::type_id;
9pub use bytes::Bytes;
10pub use nom::AsBytes;
11use serde::{Deserialize, Serialize};
12pub use serde_json::Value as SerdeValue;
13use std::{
14    collections::HashMap,
15    hash::{Hash, Hasher},
16};
17use thiserror::Error;
18pub use url::Url;
19
20pub struct Input(pub Bytes);
21
22impl Input {
23    pub fn new(s: &'static str) -> Self {
24        Self(Bytes::from(s.as_bytes()))
25    }
26}
27
28pub type BotnetResult<T> = Result<T, BotnetError>;
29
30pub mod prelude {
31    pub use super::{
32        eval::Evaluator, task::Task, type_id, utils::values_to_bytes, Arc, AsBytes,
33        BotnetResult, Bytes, Database, DatabaseKey, Extractor, Extractors, Field,
34        FieldMetadata, FieldType, Input, Key, KeyType, Metadata, Mutex, SerdeValue, Url,
35    };
36}
37
38impl AsRef<str> for Input {
39    fn as_ref(&self) -> &str {
40        std::str::from_utf8(self.0.as_bytes()).expect("Bad input.")
41    }
42}
43
44impl From<&'static str> for Input {
45    fn from(value: &'static str) -> Self {
46        Self::new(value)
47    }
48}
49
50pub trait AsValue {
51    fn as_value(&self) -> Bytes;
52}
53
54impl AsValue for bool {
55    fn as_value(&self) -> Bytes {
56        match self {
57            true => Bytes::from("1"),
58            false => Bytes::from("0"),
59        }
60    }
61}
62
63impl AsValue for u64 {
64    fn as_value(&self) -> Bytes {
65        Bytes::from(u64::to_le_bytes(*self).to_vec())
66    }
67}
68
69impl AsValue for &'static str {
70    fn as_value(&self) -> Bytes {
71        Bytes::from(self.to_string())
72    }
73}
74
75#[derive(Debug, Error)]
76pub enum BotnetError {
77    #[error("ParseError: {0:#?}")]
78    ParseError(#[from] url::ParseError),
79    #[error("BincodeError: {0:#?}")]
80    BincodeError(#[from] bincode::Error),
81    #[error("Utf8Error: {0:#?}")]
82    Utf8Error(#[from] std::str::Utf8Error),
83    #[cfg(feature = "redisdb")]
84    #[error("RedisError: {0:#?}")]
85    RedisError(#[from] redis::RedisError),
86}
87
88#[derive(Debug, Serialize, Deserialize, Default, Clone, Eq)]
89pub struct ValueMap {
90    items: HashMap<usize, Bytes>,
91}
92
93impl ValueMap {
94    pub fn from_values(v: Vec<Bytes>) -> Self {
95        Self {
96            items: HashMap::from_iter(
97                v.iter().map(|v| (type_id(v.as_bytes()), v.clone())),
98            ),
99        }
100    }
101}
102
103impl Hash for ValueMap {
104    fn hash<H: Hasher>(&self, state: &mut H) {
105        for (k, v) in self.items.iter() {
106            k.hash(state);
107            v.hash(state);
108        }
109    }
110}
111
112impl PartialEq for ValueMap {
113    fn eq(&self, other: &Self) -> bool {
114        for k in self.items.keys() {
115            if other.items.contains_key(k) {
116                return true;
117            }
118        }
119        false
120    }
121}
122
123pub enum KeyType {
124    Compressed,
125    Decompressed,
126}
127
128pub enum FieldType {
129    Compressed,
130    Decompressed,
131}
132
133#[derive(Debug, Serialize, Deserialize, Default, Clone, Eq, PartialEq, Hash)]
134pub struct FieldMetadata {
135    name: String,
136    key: String,
137    value_map: ValueMap,
138    type_id: usize,
139    description: String,
140}
141
142impl FieldMetadata {
143    pub fn new(name: &str, key: &str, values: Vec<Bytes>, description: &str) -> Self {
144        Self {
145            name: name.to_string(),
146            key: key.to_string(),
147            value_map: ValueMap::from_values(values),
148            type_id: type_id(key),
149            description: description.to_string(),
150        }
151    }
152}
153
154#[derive(Debug, Serialize, Deserialize, Default, Clone, Eq)]
155pub struct KeyMetadata {
156    field_meta: HashMap<String, FieldMetadata>,
157    type_id: usize,
158}
159
160impl Hash for KeyMetadata {
161    fn hash<H: Hasher>(&self, state: &mut H) {
162        self.type_id.hash(state);
163        for (k, v) in self.field_meta.iter() {
164            k.hash(state);
165            v.hash(state);
166        }
167    }
168}
169
170impl PartialEq for KeyMetadata {
171    fn eq(&self, other: &Self) -> bool {
172        self.type_id == other.type_id
173    }
174}
175
176impl KeyMetadata {
177    pub fn new() -> Self {
178        Self {
179            type_id: 0,
180            field_meta: HashMap::default(),
181        }
182    }
183
184    pub fn as_bytes(&self) -> BotnetResult<Bytes> {
185        Ok(Bytes::from(bincode::serialize(&self)?))
186    }
187
188    pub fn field(&mut self, f: FieldMetadata) -> &mut Self {
189        self.field_meta.insert(f.name.clone(), f);
190        self
191    }
192
193    pub fn build(&self) -> Self {
194        self.clone()
195    }
196
197    pub fn with_key_code(meta: Self, type_id: usize) -> Self {
198        Self {
199            field_meta: meta.field_meta,
200            type_id,
201        }
202    }
203}
204
205#[derive(Debug, Default, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
206pub struct Field {
207    pub type_id: usize,
208    pub field_name: String,
209    pub value: Bytes,
210}
211
212impl Field {
213    pub fn new(field_name: &str, value: Bytes) -> Self {
214        Self {
215            type_id: type_id(field_name),
216            field_name: field_name.to_string(),
217            value,
218        }
219    }
220}
221
222pub trait Key {
223    const NAME: &'static str;
224    const TYPE_ID: usize;
225
226    type Item;
227    type Metadata;
228
229    fn build(&self) -> Self;
230    fn field(&mut self, f: Self::Item) -> &mut Self;
231    fn flatten(&self) -> Bytes;
232    fn get_metadata(&self) -> Self::Metadata;
233    fn metadata(&mut self, meta: KeyMetadata) -> &mut Self;
234    fn builder() -> Self;
235    fn type_id(&self) -> usize;
236    fn name(&self) -> &'static str;
237}
238
239pub struct Extractor {
240    #[allow(unused)]
241    key: String,
242    func: fn(&Input) -> BotnetResult<Field>,
243}
244
245impl Default for Extractor {
246    fn default() -> Self {
247        fn default_func(_input: &Input) -> BotnetResult<Field> {
248            Ok(Field::default())
249        }
250
251        Self {
252            key: String::default(),
253            func: default_func,
254        }
255    }
256}
257
258impl Extractor {
259    pub fn new(key: &str, func: fn(&Input) -> BotnetResult<Field>) -> Self {
260        Self {
261            key: key.to_string(),
262            func,
263        }
264    }
265
266    pub fn call(&self, input: &Input) -> BotnetResult<Field> {
267        (self.func)(input)
268    }
269}
270
271#[derive(Default)]
272pub struct Extractors {
273    pub items: HashMap<String, Extractor>,
274}
275
276impl Extractors {
277    pub fn new() -> Self {
278        Self {
279            items: HashMap::default(),
280        }
281    }
282
283    pub fn add(&mut self, key: &str, value: fn(&Input) -> BotnetResult<Field>) {
284        self.items
285            .insert(key.to_string(), Extractor::new(key, value));
286    }
287}
288
289#[derive(Default)]
290pub struct Metadata {
291    items: HashMap<usize, KeyMetadata>,
292}
293
294impl Metadata {
295    pub fn new() -> Self {
296        Self {
297            items: HashMap::default(),
298        }
299    }
300
301    pub fn insert(&mut self, ty_id: usize, meta: KeyMetadata) {
302        self.items.insert(ty_id, meta);
303    }
304
305    pub fn get(&self, ty_id: &usize) -> &KeyMetadata {
306        self.items.get(ty_id).unwrap()
307    }
308}
botnet_core/lib.rs

botnet_core/
lib.rs
