feat: aws s3 as a file storage option

Signed-off-by: Richard Zak <[email protected]>

Author: rjzak

Cargo.lock

@@ -102,6 +102,9 @@ aes-gcm = { version = "0.10.3", default-features = false }
 anyhow = { version = "1.0", default-features = false }
 app-memory-usage-fetcher = { version = "0.2.1", default-features = false }
 argon2 = { version = "0.5.3", default-features = false }
+aws-config = { version = "1.1.7", default-features = false }
+aws-sdk-s3 = { version = "1.76.0", default-features = false }
+aws-smithy-async = { version = "1.2.4", default-features = false }
 axum = { version = "0.8.1", default-features = false }
 axum-server = { version = "0.7.2", default-features = false }
 base64 = { version = "0.22.1", default-features = false }

Cargo.toml

@@ -27,6 +27,9 @@ aes-gcm = { workspace = true, features = ["aes", "alloc", "getrandom", "std"] }
 anyhow = { workspace = true, features = ["std"] }
 app-memory-usage-fetcher = { workspace = true }
 argon2 = { workspace = true, features = ["alloc", "password-hash", "std"] }
+aws-config = { workspace = true, features = ["behavior-version-latest"] }
+aws-sdk-s3 = { workspace = true }
+aws-smithy-async = { workspace = true, features = ["rt-tokio"] }
 axum = { workspace = true, features = ["http1", "http2", "json", "macros", "tokio"] }
 axum-server = { workspace = true, features = ["tls-rustls", "rustls"] }
 base64 = { workspace = true, features = ["alloc", "std"] }

crates/server/Cargo.toml

@@ -749,6 +749,7 @@ mod tests {
                 }),
                 cert: None,
                 key: None,
+                aws: None,
             };
 
             let vt: VtUpdater = state.try_into().expect("failed to create VtUpdater");
@@ -810,6 +811,7 @@ mod tests {
                 }),
                 cert: None,
                 key: None,
+                aws: None,
             };
 
             let sqlite_second = Sqlite::new(DB_FILE)

crates/server/src/db/mod.rs

@@ -435,6 +435,7 @@ mod tests {
             vt_client: None,
             cert: None,
             key: None,
+            aws: None,
         };
 
         state
@@ -495,6 +496,7 @@ mod tests {
                 vt_client: None,
                 cert: Some("../../testdata/server_ca_cert.pem".into()),
                 key: Some("../../testdata/server_key.pem".into()),
+                aws: None,
             }
         } else {
             State {
@@ -514,6 +516,7 @@ mod tests {
                 vt_client: None,
                 cert: Some("../../testdata/server_cert.der".into()),
                 key: Some("../../testdata/server_key.der".into()),
+                aws: None,
             }
         };
 

crates/server/src/http/mod.rs

@@ -33,9 +33,17 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 use anyhow::{bail, ensure, Context, Result};
+use aws_config::{AppName, BehaviorVersion, SdkConfig};
+use aws_sdk_s3::config::{Credentials, SharedCredentialsProvider};
+use aws_sdk_s3::primitives::ByteStream;
+use aws_sdk_s3::Client;
+use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep};
 use axum_server::tls_rustls::RustlsConfig;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
 use chrono::Local;
 use chrono_humanize::{Accuracy, HumanTime, Tense};
+use constcat::concat;
 use flate2::read::GzDecoder;
 use flate2::write::GzEncoder;
 use flate2::Compression;
@@ -84,6 +92,9 @@ pub struct State {
 
     /// Https private key file
     key: Option<PathBuf>,
+
+    /// AWS client and S3 bucket
+    aws: Option<(Client, String)>,
 }
 
 impl State {
@@ -99,6 +110,9 @@ impl State {
         key: Option<PathBuf>,
         pg_cert: Option<PathBuf>,
         #[cfg(feature = "vt")] vt_client: Option<malwaredb_virustotal::VirusTotalClient>,
+        aws_access_key_id: Option<String>,
+        aws_secret_access_key: Option<String>,
+        aws_bucket: Option<String>,
     ) -> Result<Self> {
         if let Some(dir) = &directory {
             if !dir.exists() {
@@ -121,6 +135,34 @@ impl State {
             ensure!(key_file.exists(), "Key file {key_file:?} does not exist!");
         }
 
+        if aws_access_key_id.is_some() != aws_secret_access_key.is_some()
+            || aws_access_key_id.is_some() != aws_bucket.is_some()
+        {
+            bail!("AWS credentials and the S3 bucket must be provided, or be `None`.");
+        }
+
+        let client = if aws_access_key_id.is_none() {
+            None
+        } else {
+            // The Aws-related .unwrap() calls are okay since we know they're `Some()`
+            let creds = Credentials::new(
+                aws_access_key_id.unwrap(),
+                aws_secret_access_key.unwrap(),
+                None,
+                None,
+                "malwaredb",
+            );
+            let provider = SharedCredentialsProvider::new(creds);
+            let sleep_impl = SharedAsyncSleep::new(TokioSleep::new());
+            let config = SdkConfig::builder()
+                .credentials_provider(provider)
+                .behavior_version(BehaviorVersion::latest())
+                .sleep_impl(sleep_impl)
+                .app_name(AppName::new(concat!("malwaredb-v", MDB_VERSION)).unwrap())
+                .build();
+            Some((Client::new(&config), aws_bucket.unwrap()))
+        };
+
         let db_type = db::DatabaseType::from_string(db_string, pg_cert).await?;
         // TODO: allow config and keys to be refreshed so changes don't require a restart
         let db_config = db_type.get_config().await?;
@@ -139,15 +181,19 @@ impl State {
             started: SystemTime::now(),
             cert,
             key,
+            aws: client,
         })
     }
 
     /// Store the sample with a depth of three based on the sample's SHA-256 hash, even if compressed
     pub async fn store_bytes(&self, data: &[u8]) -> Result<bool> {
-        if let Some(dest_path) = &self.directory {
+        if self.directory.is_none() && self.aws.is_none() {
+            Ok(false)
+        } else {
             let mut hasher = Sha256::new();
             hasher.update(data);
-            let sha256 = hex::encode(hasher.finalize());
+            let sha256_bytes = hasher.finalize();
+            let sha256 = hex::encode(sha256_bytes);
 
             // Trait `HashPath` needs to be re-worked so it can work with Strings.
             // This code below ends up making the String into ASCII representations of the hash
@@ -160,16 +206,6 @@ impl State {
                 sha256
             );
 
-            // The path which has the file name included, with the storage directory prepended.
-            //let hashed_path = result.hashed_path(3);
-            let mut dest_path = dest_path.clone();
-            dest_path.push(hashed_path);
-
-            // Remove the file name so we can just have the directory path.
-            let mut just_the_dir = dest_path.clone();
-            just_the_dir.pop();
-            std::fs::create_dir_all(just_the_dir)?;
-
             let data = if self.db_config.compression {
                 let mut compressor = GzEncoder::new(Vec::new(), Compression::default());
                 compressor.write_all(data)?;
@@ -190,30 +226,79 @@ impl State {
                 data
             };
 
-            std::fs::write(dest_path, data)?;
+            if let Some(dest_path) = &self.directory {
+                // The path which has the file name included, with the storage directory prepended.
+                //let hashed_path = result.hashed_path(3);
+                let mut dest_path = dest_path.clone();
+                dest_path.push(hashed_path);
+
+                // Remove the file name so we can just have the directory path.
+                let mut just_the_dir = dest_path.clone();
+                just_the_dir.pop();
+                std::fs::create_dir_all(just_the_dir)?;
+
+                std::fs::write(dest_path, data)?;
+            } else if let Some((client, bucket)) = &self.aws {
+                // Data is compressed and/or encrypted, so re-hash for AWS
+                let sha256_b64 =
+                    if self.db_config.compression || self.db_config.default_key.is_some() {
+                        let mut hasher = Sha256::new();
+                        hasher.update(&data);
+                        BASE64_STANDARD.encode(hasher.finalize())
+                    } else {
+                        BASE64_STANDARD.encode(sha256_bytes)
+                    };
+
+                let bytes = ByteStream::from(data);
+                client
+                    .put_object()
+                    .bucket(bucket)
+                    .checksum_sha256(sha256_b64)
+                    .key(hashed_path)
+                    .body(bytes)
+                    .send()
+                    .await?;
+            }
 
             Ok(true)
-        } else {
-            Ok(false)
         }
     }
 
     /// Retrieve a sample given the SHA-256 hash
     /// Assumes that MalwareDB permissions have already been checked to ensure this is permitted.
     pub async fn retrieve_bytes(&self, sha256: &String) -> Result<Vec<u8>> {
-        if let Some(dest_path) = &self.directory {
+        if self.directory.is_none() && self.aws.is_none() {
+            bail!("files are not saved")
+        } else {
+            // Trait `HashPath` needs to be re-worked so it can work with Strings.
+            // This code below ends up making the String into ASCII representations of the hash
+            // See: https://github.com/malwaredb/malwaredb-rs/issues/60
+            //let path = sha256.as_bytes().iter().hashed_path(3);
+
             let path = format!(
                 "{}/{}/{}/{}",
                 &sha256[0..2],
                 &sha256[2..4],
                 &sha256[4..6],
                 sha256
             );
-            // Trait `HashPath` needs to be re-worked so it can work with Strings.
-            // This code below ends up making the String into ASCII representations of the hash
-            // See: https://github.com/malwaredb/malwaredb-rs/issues/60
-            //let path = sha256.as_bytes().iter().hashed_path(3);
-            let contents = std::fs::read(dest_path.join(path))?;
+
+            let contents = if let Some(dest_path) = &self.directory {
+                std::fs::read(dest_path.join(path))?
+            } else if let Some((client, bucket)) = &self.aws {
+                client
+                    .get_object()
+                    .bucket(bucket)
+                    .key(path)
+                    .send()
+                    .await?
+                    .body
+                    .collect()
+                    .await?
+                    .to_vec()
+            } else {
+                bail!("No AWS config and no local directory, shouldn't happen!")
+            };
 
             let contents = if !self.keys.is_empty() {
                 let (key_id, nonce) = self.db_type.get_file_encryption_key_id(sha256).await?;
@@ -241,8 +326,6 @@ impl State {
             } else {
                 Ok(contents)
             }
-        } else {
-            bail!("files are not saved")
         }
     }
 

crates/server/src/lib.rs

@@ -80,6 +80,21 @@ pub struct Config {
     #[cfg(feature = "vt")]
     #[serde(default)]
     pub vt_client: Option<malwaredb_virustotal::VirusTotalClient>,
+
+    /// AWS Access Key
+    #[arg(short, long)]
+    #[serde(default)]
+    pub access_key_id: Option<String>,
+
+    /// AWS Secret Access Key
+    #[arg(short, long)]
+    #[serde(default)]
+    pub secret_access_key: Option<String>,
+
+    /// AWS S3 bucket
+    #[arg(short, long)]
+    #[serde(default)]
+    pub bucket: Option<String>,
 }
 
 const fn default_port() -> u16 {
@@ -108,6 +123,9 @@ impl Config {
             self.pg_cert,
             #[cfg(feature = "vt")]
             self.vt_client,
+            self.access_key_id,
+            self.secret_access_key,
+            self.bucket,
         )
         .await
     }
@@ -174,6 +192,9 @@ impl Default for Config {
             cert: None,
             key: None,
             pg_cert: None,
+            access_key_id: None,
+            secret_access_key: None,
+            bucket: None,
         }
     }
 }

src/cli/config.rs

@@ -39,6 +39,9 @@ impl Run {
             cfg.pg_cert,
             #[cfg(feature = "vt")]
             cfg.vt_client,
+            cfg.access_key_id,
+            cfg.secret_access_key,
+            cfg.bucket,
         )
         .await
     }
@@ -76,6 +79,7 @@ impl Load {
     }
 }
 
+#[allow(clippy::large_enum_variant)]
 #[derive(Subcommand, Clone, Debug)]
 pub enum Subcommands {
     Load(Load),
@@ -105,6 +109,9 @@ mod tests {
             None,
             #[cfg(feature = "vt")]
             None,
+            None,
+            None,
+            None,
         )
         .await
         .unwrap();

src/cli/run.rs

@@ -35,6 +35,7 @@ impl VtOption {
 }
 
 /// VirusTotal config either from disk or command line for VT operations outside the normal running of MalwareDB
+#[allow(clippy::large_enum_variant)]
 #[derive(Subcommand, Clone, Debug)]
 enum VtConfig {
     /// Load VirusTotal config from a file

src/cli/vt/mod.rs

@@ -34,6 +34,9 @@ async fn main() -> anyhow::Result<ExitCode> {
             "malwaredb_types",
             "deadpool_postgres",
             "postgres",
+            "aws_config",
+            "aws_sdk_s3",
+            "aws_smithy_async",
             #[cfg(feature = "sqlite")]
             "rusqlite",
         ]