chore: some cleanup

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

Author: rjzak

client-py/src/lib.rs

@@ -17,7 +17,6 @@ use crate::types::{Label, ServerInfo, Source, SupportedFileType, UserInfo};
 use malwaredb_client::blocking::MdbClient;
 
 use anyhow::{anyhow, Result};
-use malwaredb_client::malwaredb_api::PartialHashSearchType;
 use pyo3::prelude::*;
 
 /// MDB version
@@ -108,18 +107,12 @@ impl MalwareDBClient {
         limit: u32,
         response_hash: &str,
     ) -> Result<Vec<String>> {
-        let hash_type = match hash_type {
-            "md5" => PartialHashSearchType::MD5,
-            "sha1" => PartialHashSearchType::SHA1,
-            "sha256" => PartialHashSearchType::SHA256,
-            "sha384" => PartialHashSearchType::SHA384,
-            "sha512" => PartialHashSearchType::SHA512,
-            x => return Err(anyhow!("unknown hash type {}", x)),
-        };
+        let hash_type = hash_type.try_into().map_err(|e: String| anyhow!(e))?;
+        let response_hash = response_hash.try_into().map_err(|e: String| anyhow!(e))?;
         self.inner.partial_search(
             hash.map(|h| (hash_type, h)),
             file_name,
-            Some(response_hash),
+            response_hash,
             limit,
         )
     }

client/src/blocking.rs

@@ -5,7 +5,7 @@ use std::path::{Path, PathBuf};
 
 use malwaredb_types::exec::pe32::EXE;
 
-use anyhow::{anyhow, bail, ensure, Context};
+use anyhow::{bail, ensure, Context};
 use base64::engine::general_purpose;
 use base64::Engine;
 use fuzzyhash::FuzzyHash;
@@ -372,16 +372,13 @@ impl MdbClient {
         &self,
         partial_hash: Option<(PartialHashSearchType, String)>,
         name: Option<String>,
-        response_hash: Option<&str>,
+        response: PartialHashSearchType,
         limit: u32,
     ) -> anyhow::Result<Vec<String>> {
         let query = SearchRequest {
             partial_hash,
             file_name: name,
-            response: response_hash
-                .unwrap_or("sha256")
-                .try_into()
-                .map_err(|e: &'static str| anyhow!(e))?,
+            response,
             limit,
         };
 

client/src/cli/search.rs

@@ -2,7 +2,6 @@
 
 use anyhow::{anyhow, Result};
 use clap::Parser;
-use malwaredb_api::PartialHashSearchType;
 use malwaredb_client::MdbClient;
 use std::process::ExitCode;
 
@@ -32,26 +31,23 @@ pub struct SearchRequest {
 impl SearchRequest {
     pub async fn exec(&self, config: &MdbClient) -> Result<ExitCode> {
         let hash_search = if let Some(hash) = &self.hash {
-            let hash_type = match self.hash_type.as_str() {
-                "md5" => PartialHashSearchType::MD5,
-                "sha1" => PartialHashSearchType::SHA1,
-                "sha256" => PartialHashSearchType::SHA256,
-                "sha384" => PartialHashSearchType::SHA384,
-                "sha512" => PartialHashSearchType::SHA512,
-                x => return Err(anyhow!("unknown hash type {}", x)),
-            };
+            let hash_type = self
+                .hash_type
+                .as_str()
+                .try_into()
+                .map_err(|e: String| anyhow!(e))?;
             Some((hash_type, hash.clone()))
         } else {
             None
         };
 
+        let response = self
+            .response_type
+            .as_str()
+            .try_into()
+            .map_err(|e: String| anyhow!(e))?;
         let response = config
-            .partial_search(
-                hash_search,
-                self.file_name.clone(),
-                Some(self.response_type.as_str()),
-                self.limit,
-            )
+            .partial_search(hash_search, self.file_name.clone(), response, self.limit)
             .await?;
 
         if response.is_empty() {

client/src/lib.rs

@@ -20,7 +20,7 @@ use std::fmt::{Debug, Formatter};
 use std::io::Cursor;
 use std::path::{Path, PathBuf};
 
-use anyhow::{anyhow, bail, ensure, Context, Result};
+use anyhow::{bail, ensure, Context, Result};
 use base64::engine::general_purpose;
 use base64::Engine;
 use cart_container::JsonMap;
@@ -407,16 +407,13 @@ impl MdbClient {
         &self,
         partial_hash: Option<(PartialHashSearchType, String)>,
         name: Option<String>,
-        response_hash: Option<&str>,
+        response: PartialHashSearchType,
         limit: u32,
     ) -> Result<Vec<String>> {
         let query = SearchRequest {
             partial_hash,
             file_name: name,
-            response: response_hash
-                .unwrap_or("sha256")
-                .try_into()
-                .map_err(|e: &'static str| anyhow!(e))?,
+            response,
             limit,
         };
 

crates/api/src/lib.rs

@@ -464,7 +464,7 @@ impl Display for PartialHashSearchType {
 }
 
 impl TryInto<PartialHashSearchType> for &str {
-    type Error = &'static str;
+    type Error = String;
 
     fn try_into(self) -> Result<PartialHashSearchType, Self::Error> {
         match self {
@@ -474,7 +474,19 @@ impl TryInto<PartialHashSearchType> for &str {
             "sha256" => Ok(PartialHashSearchType::SHA256),
             "sha384" => Ok(PartialHashSearchType::SHA384),
             "sha512" => Ok(PartialHashSearchType::SHA512),
-            _ => Err("Invalid hash type"),
+            x => Err(format!("Invalid hash type {x}")),
+        }
+    }
+}
+
+impl TryInto<PartialHashSearchType> for Option<&str> {
+    type Error = String;
+
+    fn try_into(self) -> Result<PartialHashSearchType, Self::Error> {
+        if let Some(hash) = self {
+            hash.try_into()
+        } else {
+            Ok(PartialHashSearchType::SHA256)
         }
     }
 }

crates/server/src/http/mod.rs

@@ -895,7 +895,7 @@ mod tests {
             .partial_search(
                 Some((PartialHashSearchType::Any, "AAAA".into())),
                 None,
-                None,
+                PartialHashSearchType::Any,
                 10,
             )
             .await