malwaredb-server 0.3.4

// SPDX-License-Identifier: Apache-2.0

use crate::db::types::FileMetadata;
use crate::State;
use malwaredb_types::KnownType;
use std::io::Cursor;

use std::sync::Arc;

use anyhow::{bail, ensure, Result};
use serde_json::Value;
use tracing::{error, instrument, trace};

pub const CART_MAGIC: [u8; 4] = [0x43, 0x41, 0x52, 0x54];

#[instrument(skip(data))]
pub async fn incoming_sample(
    state: Arc<State>,
    data: Vec<u8>,
    uid: u32,
    sid: u32,
    file_name: String,
) -> Result<()> {
    let (data, footer) = if data[0..4] == CART_MAGIC {
        let mut output = Cursor::new(vec![]);
        let bytes_input = Cursor::new(data);
        let (_header, footer) = cart_container::unpack_stream(bytes_input, &mut output, None)?;
        trace!("Received CaRT file and decoded it's contents.");
        (output.into_inner(), footer)
    } else {
        (data, None)
    };

    let known_type = match KnownType::new(&data) {
        Ok(t) => t,
        Err(e) => {
            error!("Error determining type: {e}");
            return Err(e);
        }
    };

    let meta_data = FileMetadata::new(&data, Some(&file_name));
    if let Some(footer) = footer {
        // CaRT files created by Malware DB or others will always have the SHA-256 hash
        if let Some(Value::String(cart_sha256)) = footer.get("sha256") {
            let meta_sha256 = hex::encode(&meta_data.sha256);
            ensure!(
                cart_sha256.eq(&meta_sha256),
                "CaRT SHA-256 hash doesn't match calculated hash"
            );
        } else {
            bail!("CaRT file missing SHA-256 hash");
        }
    }

    let db_file_type = match state.db_type.get_type_id_for_bytes(&data).await {
        Ok(t) => t,
        Err(e) => {
            error!("db_file_type error or is none: {e}");
            bail!("file type unknown or error: {e}");
        }
    };

    match state
        .db_type
        .add_file(&meta_data, known_type, uid, sid, db_file_type, None)
        .await
    {
        Ok(added) => {
            if added.is_new {
                trace!("Storing sample {}!", added.file_id);
                if let Err(e) = state.store_bytes(&data).await {
                    error!("Error storing sample {}: {e}", added.file_id);
                    Err(e)
                } else {
                    Ok(())
                }
            } else {
                trace!("Sample {} already exists, skipping storage.", added.file_id);
                Ok(())
            }
        }
        Err(e) => {
            error!("Error storing bytes: {e}");
            Err(e)
        }
    }
}