Skip to main content

virustotal_rs/
references.rs

1use crate::objects::{Collection, Object};
2use crate::{Client, Result};
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6/// Represents a Reference in `VirusTotal`
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Reference {
9    #[serde(flatten)]
10    pub object: Object<ReferenceAttributes>,
11}
12
13/// Attributes for a Reference
14#[derive(Debug, Clone, Serialize, Deserialize, Default)]
15pub struct ReferenceAttributes {
16    /// Reference title
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub title: Option<String>,
19
20    /// Reference description
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub description: Option<String>,
23
24    /// Reference URL
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub url: Option<String>,
27
28    /// Author of the reference
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub author: Option<String>,
31
32    /// Source of the reference
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub source: Option<String>,
35
36    /// Publication date (UTC timestamp)
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub publication_date: Option<i64>,
39
40    /// Creation date (UTC timestamp)
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub creation_date: Option<i64>,
43
44    /// Last modification date (UTC timestamp)
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub last_modification_date: Option<i64>,
47
48    /// Tags associated with the reference
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub tags: Option<Vec<String>>,
51
52    /// Additional attributes
53    #[serde(flatten)]
54    pub additional_attributes: HashMap<String, serde_json::Value>,
55}
56
57/// Request to create a new reference
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct CreateReferenceRequest {
60    pub data: CreateReferenceData,
61}
62
63/// Data for creating a reference
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct CreateReferenceData {
66    pub attributes: CreateReferenceAttributes,
67    #[serde(rename = "type")]
68    pub object_type: String,
69}
70
71/// Attributes for creating a reference
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct CreateReferenceAttributes {
74    pub title: String,
75    pub url: String,
76    #[serde(skip_serializing_if = "Option::is_none")]
77    pub description: Option<String>,
78    #[serde(skip_serializing_if = "Option::is_none")]
79    pub author: Option<String>,
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub source: Option<String>,
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub tags: Option<Vec<String>>,
84}
85
86/// Client for References operations
87pub struct ReferencesClient<'a> {
88    pub(crate) client: &'a Client,
89}
90
91impl<'a> ReferencesClient<'a> {
92    pub fn new(client: &'a Client) -> Self {
93        Self { client }
94    }
95
96    /// Create a new reference (requires special privileges)
97    pub async fn create(&self, request: &CreateReferenceRequest) -> Result<Reference> {
98        self.client.post("references", request).await
99    }
100
101    /// Get a reference by ID (requires special privileges)
102    pub async fn get(&self, reference_id: &str) -> Result<Reference> {
103        let url = format!("references/{}", reference_id);
104        self.client.get(&url).await
105    }
106
107    /// Delete a reference (requires special privileges)
108    pub async fn delete(&self, reference_id: &str) -> Result<()> {
109        let url = format!("references/{}", reference_id);
110        self.client.delete(&url).await
111    }
112
113    /// Get objects related to a reference
114    pub async fn get_relationship<T>(
115        &self,
116        reference_id: &str,
117        relationship: &str,
118        limit: Option<u32>,
119        cursor: Option<&str>,
120    ) -> Result<Collection<T>>
121    where
122        T: serde::de::DeserializeOwned,
123    {
124        let mut url = format!("references/{}/{}?", reference_id, relationship);
125
126        if let Some(l) = limit {
127            url.push_str(&format!("limit={}&", l));
128        }
129
130        if let Some(c) = cursor {
131            url.push_str(&format!("cursor={}&", c));
132        }
133
134        // Remove trailing '&' or '?'
135        url.pop();
136
137        self.client.get(&url).await
138    }
139
140    /// Get object descriptors related to a reference
141    pub async fn get_relationship_descriptors<T>(
142        &self,
143        reference_id: &str,
144        relationship: &str,
145        limit: Option<u32>,
146        cursor: Option<&str>,
147    ) -> Result<Collection<T>>
148    where
149        T: serde::de::DeserializeOwned,
150    {
151        let mut url = format!(
152            "references/{}/relationships/{}?",
153            reference_id, relationship
154        );
155
156        if let Some(l) = limit {
157            url.push_str(&format!("limit={}&", l));
158        }
159
160        if let Some(c) = cursor {
161            url.push_str(&format!("cursor={}&", c));
162        }
163
164        // Remove trailing '&' or '?'
165        url.pop();
166
167        self.client.get(&url).await
168    }
169}
170
171/// Helper methods for creating references
172impl CreateReferenceRequest {
173    /// Create a new reference request
174    pub fn new(title: String, url: String) -> Self {
175        Self {
176            data: CreateReferenceData {
177                attributes: CreateReferenceAttributes {
178                    title,
179                    url,
180                    description: None,
181                    author: None,
182                    source: None,
183                    tags: None,
184                },
185                object_type: "reference".to_string(),
186            },
187        }
188    }
189
190    /// Set description
191    pub fn with_description(mut self, description: String) -> Self {
192        self.data.attributes.description = Some(description);
193        self
194    }
195
196    /// Set author
197    pub fn with_author(mut self, author: String) -> Self {
198        self.data.attributes.author = Some(author);
199        self
200    }
201
202    /// Set source
203    pub fn with_source(mut self, source: String) -> Self {
204        self.data.attributes.source = Some(source);
205        self
206    }
207
208    /// Set tags
209    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
210        self.data.attributes.tags = Some(tags);
211        self
212    }
213}
214
215impl Client {
216    /// Get the References client for reference operations
217    pub fn references(&self) -> ReferencesClient<'_> {
218        ReferencesClient::new(self)
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[test]
227    fn test_reference_attributes() {
228        let attrs = ReferenceAttributes {
229            title: Some("APT1 Report".to_string()),
230            description: Some("Detailed analysis of APT1 activities".to_string()),
231            url: Some("https://example.com/apt1-report".to_string()),
232            author: Some("Security Research Team".to_string()),
233            source: Some("Example Security".to_string()),
234            tags: Some(vec![
235                "apt".to_string(),
236                "china".to_string(),
237                "espionage".to_string(),
238            ]),
239            ..Default::default()
240        };
241
242        assert_eq!(attrs.title.unwrap(), "APT1 Report");
243        assert_eq!(attrs.url.unwrap(), "https://example.com/apt1-report");
244        assert_eq!(attrs.tags.unwrap().len(), 3);
245    }
246
247    #[test]
248    fn test_create_reference_request() {
249        let request = CreateReferenceRequest::new(
250            "Malware Analysis Report".to_string(),
251            "https://example.com/report".to_string(),
252        )
253        .with_description("Comprehensive malware analysis".to_string())
254        .with_author("John Doe".to_string())
255        .with_source("Security Lab".to_string())
256        .with_tags(vec!["malware".to_string(), "analysis".to_string()]);
257
258        assert_eq!(request.data.attributes.title, "Malware Analysis Report");
259        assert_eq!(request.data.attributes.url, "https://example.com/report");
260        assert_eq!(
261            request.data.attributes.description.unwrap(),
262            "Comprehensive malware analysis"
263        );
264        assert_eq!(request.data.attributes.author.unwrap(), "John Doe");
265        assert_eq!(request.data.attributes.source.unwrap(), "Security Lab");
266        assert_eq!(request.data.attributes.tags.unwrap().len(), 2);
267        assert_eq!(request.data.object_type, "reference");
268    }
269
270    #[test]
271    fn test_reference_dates() {
272        let attrs = ReferenceAttributes {
273            publication_date: Some(1609459200),       // 2021-01-01
274            creation_date: Some(1609545600),          // 2021-01-02
275            last_modification_date: Some(1609632000), // 2021-01-03
276            ..Default::default()
277        };
278
279        assert_eq!(attrs.publication_date.unwrap(), 1609459200);
280        assert_eq!(attrs.creation_date.unwrap(), 1609545600);
281        assert_eq!(attrs.last_modification_date.unwrap(), 1609632000);
282    }
283}