564 lines
18 KiB
Rust
564 lines
18 KiB
Rust
/*
|
|
* Fuzzy hash implementations, matching, and signature support
|
|
*
|
|
* Copyright (C) 2022 Cisco Systems, Inc. and/or its affiliates. All rights reserved.
|
|
*
|
|
* Authors: Micah Snyder, Mickey Sola, Scott Hutton
|
|
*
|
|
* This program is free software; you can redistribute it and/or modify
|
|
* it under the terms of the GNU General Public License version 2 as
|
|
* published by the Free Software Foundation.
|
|
*
|
|
* This program is distributed in the hope that it will be useful,
|
|
* but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
* GNU General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU General Public License
|
|
* along with this program; if not, write to the Free Software
|
|
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
|
|
* MA 02110-1301, USA.
|
|
*/
|
|
|
|
use std::{
|
|
collections::HashMap,
|
|
convert::{TryFrom, TryInto},
|
|
ffi::CStr,
|
|
mem::ManuallyDrop,
|
|
os::raw::c_char,
|
|
panic,
|
|
slice,
|
|
};
|
|
|
|
use image::{imageops::FilterType::Lanczos3, DynamicImage, ImageBuffer, Luma, Pixel, Rgb};
|
|
use log::{debug, error, warn};
|
|
use num_traits::{NumCast, ToPrimitive, Zero};
|
|
use rustdct::DctPlanner;
|
|
use thiserror::Error;
|
|
use transpose::transpose;
|
|
|
|
use crate::{ffi_error, ffi_util::FFIError, rrf_call, sys, validate_str_param};
|
|
|
|
/// CdiffError enumerates all possible errors returned by this library.
|
|
#[derive(Error, Debug)]
|
|
pub enum FuzzyHashError {
|
|
#[error("Invalid format")]
|
|
Format,
|
|
|
|
#[error("Unknown algorithm: {0}")]
|
|
UnknownAlgorithm(String),
|
|
|
|
#[error("Failed to convert hamming distance to unsigned 32bit integer: {0}")]
|
|
FormatHammingDistance(String),
|
|
|
|
#[error("Invalid hamming distance: {0}")]
|
|
InvalidHammingDistance(u32),
|
|
|
|
#[error("Invalid hash: {0}")]
|
|
FormatHashBytes(String),
|
|
|
|
#[error("Failed to load image: {0}")]
|
|
ImageLoad(image::ImageError),
|
|
|
|
#[error("Failed to load image due to bug in image decoder")]
|
|
ImageLoadPanic(),
|
|
|
|
#[error("Invalid parameter: {0}")]
|
|
InvalidParameter(String),
|
|
|
|
#[error("{0} parmeter is NULL")]
|
|
NullParam(&'static str),
|
|
}
|
|
|
|
#[derive(PartialEq, Eq, Hash, Debug)]
|
|
pub struct ImageFuzzyHash {
|
|
bytes: [u8; 8],
|
|
}
|
|
|
|
#[derive(PartialEq, Eq, Hash, Debug)]
|
|
pub enum FuzzyHash {
|
|
Image(ImageFuzzyHash),
|
|
}
|
|
|
|
impl TryFrom<&str> for ImageFuzzyHash {
|
|
type Error = &'static str;
|
|
|
|
fn try_from(value: &str) -> Result<Self, Self::Error> {
|
|
if value.len() != 16 {
|
|
return Err("Image fuzzy hash must be 16 characters in length");
|
|
}
|
|
|
|
let mut hashbytes = [0; 8];
|
|
if hex::decode_to_slice(value, &mut hashbytes).is_ok() {
|
|
Ok(ImageFuzzyHash { bytes: hashbytes })
|
|
} else {
|
|
Err("Failed to decode image fuzzy hash bytes from hex to bytes")
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Display for FuzzyHash {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
|
match self {
|
|
FuzzyHash::Image(hash_bytes) => {
|
|
write!(f, "{}", hex::encode(hash_bytes.bytes))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct FuzzyHashMap {
|
|
hashmap: HashMap<FuzzyHash, Vec<FuzzyHashMeta>>,
|
|
}
|
|
|
|
#[derive(Debug, Copy, Clone)]
|
|
pub struct FuzzyHashMeta {
|
|
lsigid: u32,
|
|
subsigid: u32,
|
|
#[cfg(feature = "not_ready")]
|
|
hamming_distance: u32,
|
|
}
|
|
|
|
/// Initialize the hashmap
|
|
#[no_mangle]
|
|
pub extern "C" fn fuzzy_hashmap_new() -> sys::fuzzyhashmap_t {
|
|
Box::into_raw(Box::new(FuzzyHashMap::default())) as sys::fuzzyhashmap_t
|
|
}
|
|
|
|
/// Free the hashmap
|
|
#[no_mangle]
|
|
pub extern "C" fn fuzzy_hash_free_hashmap(fuzzy_hashmap: sys::fuzzyhashmap_t) {
|
|
if fuzzy_hashmap.is_null() {
|
|
warn!("Attempted to free a NULL hashmap pointer. Please report this at: https://github.com/Cisco-Talos/clamav/issues");
|
|
} else {
|
|
let _ = unsafe { Box::from_raw(fuzzy_hashmap as *mut FuzzyHashMap) };
|
|
}
|
|
}
|
|
|
|
/// C interface for FuzzyHashMap::check().
|
|
/// Handles all the unsafe ffi stuff.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// No parameters may be NULL
|
|
#[export_name = "fuzzy_hash_check"]
|
|
pub unsafe extern "C" fn _fuzzy_hash_check(
|
|
fuzzy_hashmap: sys::fuzzyhashmap_t,
|
|
mdata: *mut sys::cli_ac_data,
|
|
image_fuzzy_hash: sys::image_fuzzy_hash_t,
|
|
) -> bool {
|
|
let hash_bytes = image_fuzzy_hash.hash;
|
|
|
|
let hashmap = ManuallyDrop::new(Box::from_raw(fuzzy_hashmap as *mut FuzzyHashMap));
|
|
|
|
debug!(
|
|
"Checking image fuzzy hash '{}' for signature match",
|
|
hex::encode(hash_bytes)
|
|
);
|
|
|
|
if let Some(meta_vec) = hashmap.check(hash_bytes) {
|
|
for meta in meta_vec {
|
|
sys::lsig_increment_subsig_match(mdata, meta.lsigid, meta.subsigid);
|
|
}
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
/// C interface for FuzzyHashMap::load_subsignature().
|
|
/// Handles all the unsafe ffi stuff.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// `hexsig` and `err` must not be NULL
|
|
#[export_name = "fuzzy_hash_load_subsignature"]
|
|
pub unsafe extern "C" fn _fuzzy_hash_load_subsignature(
|
|
fuzzy_hashmap: sys::fuzzyhashmap_t,
|
|
hexsig: *const c_char,
|
|
lsig_id: u32,
|
|
subsig_id: u32,
|
|
err: *mut *mut FFIError,
|
|
) -> bool {
|
|
let hexsig = validate_str_param!(hexsig);
|
|
|
|
let mut hashmap = ManuallyDrop::new(Box::from_raw(fuzzy_hashmap as *mut FuzzyHashMap));
|
|
|
|
rrf_call!(
|
|
err = err,
|
|
hashmap.load_subsignature(hexsig, lsig_id, subsig_id)
|
|
)
|
|
}
|
|
|
|
/// C interface for fuzzy_hash_calculate_image().
|
|
/// Handles all the unsafe ffi stuff.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// `file_bytes` and `hash_out` must not be NULL
|
|
#[export_name = "fuzzy_hash_calculate_image"]
|
|
pub unsafe extern "C" fn _fuzzy_hash_calculate_image(
|
|
file_bytes: *const u8,
|
|
file_size: usize,
|
|
hash_out: *mut u8,
|
|
hash_out_len: usize,
|
|
err: *mut *mut FFIError,
|
|
) -> bool {
|
|
if hash_out.is_null() {
|
|
return ffi_error!(err = err, FuzzyHashError::NullParam("hash_out"));
|
|
}
|
|
|
|
let buffer = if file_bytes.is_null() {
|
|
return ffi_error!(err = err, FuzzyHashError::NullParam("file_bytes"));
|
|
} else {
|
|
slice::from_raw_parts(file_bytes, file_size)
|
|
};
|
|
|
|
let hash_result = fuzzy_hash_calculate_image(buffer);
|
|
let hash_bytes = match hash_result {
|
|
Ok(hash) => hash,
|
|
Err(error) => return ffi_error!(err = err, error),
|
|
};
|
|
|
|
if hash_out_len < hash_bytes.len() {
|
|
return ffi_error!(
|
|
err = err,
|
|
FuzzyHashError::InvalidParameter(format!(
|
|
"hash_bytes output parameter too small to hold the hash: {} < {}",
|
|
hash_out_len,
|
|
hash_bytes.len()
|
|
))
|
|
);
|
|
}
|
|
|
|
hash_out.copy_from(hash_bytes.as_ptr(), hash_bytes.len());
|
|
|
|
true
|
|
}
|
|
|
|
impl FuzzyHashMap {
|
|
/// Check for fuzzy hash matches.
|
|
///
|
|
/// In this initial version, we're just doing a simple hash lookup and the
|
|
/// hamming distance is not considered.
|
|
///
|
|
/// TODO: In a future version, replace this with an implementation that can find
|
|
/// any hashes within the signature meta.hamming_distance.
|
|
pub fn check(&self, hash: [u8; 8]) -> Option<&Vec<FuzzyHashMeta>> {
|
|
let hash = FuzzyHash::Image(ImageFuzzyHash { bytes: hash });
|
|
self.hashmap.get(&hash)
|
|
}
|
|
|
|
/// Load a fuzzy hash subsignature
|
|
/// Parse a fuzzy hash logical sig subsignature.
|
|
/// Add the fuzzy hash to the matcher so it can be matched.
|
|
pub fn load_subsignature(
|
|
&mut self,
|
|
hexsig: &str,
|
|
lsig_id: u32,
|
|
subsig_id: u32,
|
|
) -> Result<(), FuzzyHashError> {
|
|
let mut hexsig_split = hexsig.split('#');
|
|
|
|
let algorithm = match hexsig_split.next() {
|
|
Some(x) => x,
|
|
None => return Err(FuzzyHashError::Format),
|
|
};
|
|
|
|
let hash = match hexsig_split.next() {
|
|
Some(x) => x,
|
|
None => return Err(FuzzyHashError::Format),
|
|
};
|
|
|
|
let distance: u32 = match hexsig_split.next() {
|
|
Some(x) => match x.parse::<u32>() {
|
|
Ok(n) => n,
|
|
Err(_) => {
|
|
return Err(FuzzyHashError::FormatHammingDistance(x.to_string()));
|
|
}
|
|
},
|
|
None => 0,
|
|
};
|
|
|
|
// TODO: Support non-zero distance
|
|
if distance != 0 {
|
|
error!(
|
|
"Non-zero hamming distances for image fuzzy hashes are not supported in this version."
|
|
);
|
|
return Err(FuzzyHashError::InvalidHammingDistance(distance));
|
|
}
|
|
|
|
match algorithm {
|
|
"fuzzy_img" => {
|
|
// Convert the hash string to an image fuzzy hash bytes struct
|
|
let image_fuzzy_hash = hash
|
|
.try_into()
|
|
.map_err(|e| FuzzyHashError::FormatHashBytes(format!("{}: {}", e, hash)))?;
|
|
|
|
let fuzzy_hash = FuzzyHash::Image(image_fuzzy_hash);
|
|
|
|
let meta: FuzzyHashMeta = FuzzyHashMeta {
|
|
lsigid: lsig_id,
|
|
subsigid: subsig_id,
|
|
#[cfg(feature = "not_ready")]
|
|
hamming_distance: distance,
|
|
};
|
|
|
|
// If the hash key does not exist in the hashmap, insert an empty vec.
|
|
// Then add the current meta struct to the entry.
|
|
self.hashmap
|
|
.entry(fuzzy_hash)
|
|
.or_insert_with(Vec::new)
|
|
.push(meta);
|
|
|
|
Ok(())
|
|
}
|
|
_ => {
|
|
error!("Unknown fuzzy hash algorithm: {}", algorithm);
|
|
Err(FuzzyHashError::UnknownAlgorithm(algorithm.to_string()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Given a buffer and size, generate an image fuzzy hash
|
|
///
|
|
/// This algorithm attempts to reproduce the results of the `phash()` function
|
|
/// from the Python `imagehash` package.
|
|
///
|
|
/// # Notes
|
|
///
|
|
/// 1) I found that `image.grayscale() uses different RGB coefficients than
|
|
/// the python `image.convert("L"). The docs for PIL.Image.convert() state:
|
|
///
|
|
/// When translating a color image to greyscale (mode "L"),
|
|
/// the library uses the ITU-R 601-2 luma transform::
|
|
///
|
|
/// L = R * 299/1000 + G * 587/1000 + B * 114/1000
|
|
///
|
|
/// You can get near-identical** grayscale results by making a clone (or forking)
|
|
/// the image-rs crate, and changing the coefficients to match those above:
|
|
///
|
|
/// diff --git a/src/color.rs b/src/color.rs
|
|
/// index 78b5c587..92c99337 100644
|
|
/// --- a/src/color.rs
|
|
/// +++ b/src/color.rs
|
|
/// @@ -462,7 +462,7 @@ where
|
|
/// }
|
|
///
|
|
/// /// Coefficients to transform from sRGB to a CIE Y (luminance) value.
|
|
/// -const SRGB_LUMA: [f32; 3] = [0.2126, 0.7152, 0.0722];
|
|
/// +const SRGB_LUMA: [f32; 3] = [0.299, 0.587, 0.114];
|
|
///
|
|
/// #[inline]
|
|
/// fn rgb_to_luma<T: Primitive>(rgb: &[T]) -> T {
|
|
///
|
|
/// **Note that I say "near-identical" because rounding
|
|
/// appears to be slightly different and values are sometimes off-by-one.
|
|
///
|
|
/// This change doesn't appear to be required to match the phash_simple()
|
|
/// function, but to match the phash() function where the median is used instead
|
|
/// of the mean -- this change is required.
|
|
///
|
|
/// 2) scipy.fftpack.dct behaves differently on twodimensional arrays than
|
|
/// single-dimensional arrays.
|
|
/// See https://docs.scipy.org/doc/scipy/reference/generated/scipy.fftpack.dct.html:
|
|
///
|
|
/// Note the optional "axis" argument:
|
|
/// Axis along which the dct is computed; the default is over the last axis
|
|
/// (i.e., axis=-1).
|
|
///
|
|
/// For the Python `imagehash` package:
|
|
/// - The `phash_simple()` function is doing a DCT-2 transform on a 2-dimensionals
|
|
/// 32x32 array which means, just on the 2nd axis (just the rows).
|
|
/// - The `phash()` function is doing a 2D DCT-2 transform, by running the DCT-2 on
|
|
/// both X and Y axis, which is the same as transposing before or after each
|
|
/// DCT-2 call.
|
|
///
|
|
/// 3) I observed that the DCT2 results from Python are consistently 2x greater
|
|
/// than those from Rust. If I multiply every value by 2 after running the DCT,
|
|
/// then the results are the same.
|
|
///
|
|
/// 4) We need to get a subset of the 2-D array representing the lower
|
|
/// frequencies of the image, the same way the Python implementation does it.
|
|
///
|
|
/// The way the python implementation does this is with this line:
|
|
/// ```python
|
|
/// dctlowfreq = dct[:hash_size, :hash_size]
|
|
/// ```
|
|
///
|
|
/// You can't actually do that with a Python array of arrays... this is numpy
|
|
/// 2-D array manipulation magic, where you can index 2-D arrays in slices.
|
|
/// It works like this:
|
|
/// ```ipython3
|
|
/// In [0]: x = [[0, 1, 2, 3, 4], [4, 5, 6, 7, 8], [8, 9, 10, 11, 12], [12, 13, 14, 15, 16], [16, 17, 18, 19, 20]]
|
|
/// In [1]: h = 3
|
|
/// In [2]: n = np.asarray(x)
|
|
/// In [3]: lf = n[:h, 1:h+1]
|
|
/// In [4]: n
|
|
/// array([[ 0, 1, 2, 3, 4],
|
|
/// [ 4, 5, 6, 7, 8],
|
|
/// [ 8, 9, 10, 11, 12],
|
|
/// [12, 13, 14, 15, 16],
|
|
/// [16, 17, 18, 19, 20]])
|
|
///
|
|
/// In [5]: lf
|
|
/// array([[ 0, 1, 2],
|
|
/// [ 4, 5, 6],
|
|
/// [ 8, 9, 10]])
|
|
/// ```
|
|
///
|
|
/// We can do something similar, manually, to get the low-frequency selection.
|
|
///
|
|
/// param: hash_out is an output variable
|
|
/// param: hash_out_len indicates the size of the hash_out buffer
|
|
pub fn fuzzy_hash_calculate_image(buffer: &[u8]) -> Result<Vec<u8>, FuzzyHashError> {
|
|
|
|
// Load image and attempt to catch panics in case the decoders encounter unexpected issues
|
|
let result = panic::catch_unwind(|| -> Result<DynamicImage, FuzzyHashError> {
|
|
let image = image::load_from_memory(buffer).map_err(FuzzyHashError::ImageLoad)?;
|
|
Ok(image)
|
|
});
|
|
|
|
let og_image = match result {
|
|
Ok(image) => image?,
|
|
Err(_) => return Err(FuzzyHashError::ImageLoadPanic()),
|
|
};
|
|
|
|
// Drop the alpha channel (if exists).
|
|
let buff_rgb8 = og_image.to_rgb8();
|
|
|
|
// Convert image to grayscale.
|
|
let buff_luma8 = grayscale(&buff_rgb8);
|
|
|
|
// Convert back to a DynamicImage type so we can resize it.
|
|
let image_gs = DynamicImage::ImageLuma8(buff_luma8);
|
|
|
|
// Shrink to a 32x32 (1024 pixel) image.
|
|
let image_small = image::DynamicImage::resize_exact(&image_gs, 32, 32, Lanczos3);
|
|
|
|
// Convert the data to a Vec of floats.
|
|
let mut imgbuff_f32 = image_small.to_luma32f().into_raw();
|
|
|
|
//
|
|
// Compute a 2D DCT-2 in-place.
|
|
//
|
|
let dct2 = DctPlanner::new().plan_dct2(32);
|
|
|
|
// Use a scratch space so we can transpose and run DCT's without allocating any extra space.
|
|
// We'll switch back and forth between the buffer for the original small image (buffer1) and the scratch buffer (buffer2).
|
|
let buffer1: &mut [f32] = imgbuff_f32.as_mut_slice();
|
|
let buffer2: &mut [f32] = &mut [0.0; 1024];
|
|
|
|
// Transpose the image so we can run DCT on the X axis (columns) first.
|
|
transpose(buffer1, buffer2, 32, 32);
|
|
|
|
// Run DCT2 on the columns.
|
|
for (row_in, row_out) in buffer2.chunks_mut(32).zip(buffer1.chunks_mut(32)) {
|
|
dct2.process_dct2_with_scratch(row_in, row_out);
|
|
}
|
|
// Multiply each value x2, to match results from scipy.fftpack.dct() implementation.
|
|
// Note: Unsure why this is required, but it is.
|
|
buffer2.iter_mut().for_each(|f| *f *= 2.0);
|
|
|
|
// Transpose the image back so we can run DCT on the Y axis (rows).
|
|
transpose(buffer2, buffer1, 32, 32);
|
|
|
|
// Run DCT2 on the rows.
|
|
for (row_in, row_out) in buffer1.chunks_mut(32).zip(buffer2.chunks_mut(32)) {
|
|
dct2.process_dct2_with_scratch(row_in, row_out);
|
|
}
|
|
// Multiply each value x2, to match results from scipy.fftpack.dct() implementation.
|
|
// Note: Unsure why this is required, but it is.
|
|
buffer1.iter_mut().for_each(|f| *f *= 2.0);
|
|
|
|
//
|
|
// Construct a DCT low frequency vector using the top-left most 8x8 values of the 32x32 DCT array.
|
|
//
|
|
let dct_low_freq = buffer1
|
|
// 2D array is 32-elements wide.
|
|
.chunks(32)
|
|
// Grab the first 8 rows.
|
|
.take(8)
|
|
// But only take the first 8 elements (columns) from each row.
|
|
.flat_map(|chunk| chunk.chunks(8).take(1))
|
|
// Flatten the 8x8 selection down to a vector of floats.
|
|
.flatten()
|
|
.copied()
|
|
.collect::<Vec<f32>>();
|
|
|
|
// Calculate average (median) of the DCT low frequency vector.
|
|
let mut dct_low_freq_copy = dct_low_freq.clone();
|
|
dct_low_freq_copy.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
|
let median: f32 = (dct_low_freq_copy[31] + dct_low_freq_copy[32]) / 2.0;
|
|
|
|
// Construct hash vector by reducing DCT values to 1 or 0 by comparing terms vs median.
|
|
let hashvec: Vec<u64> = dct_low_freq
|
|
.into_iter()
|
|
.map(|x| if x > median { 1 } else { 0 })
|
|
.collect();
|
|
|
|
// Construct hash vec<u8> from bits.
|
|
let hash_bytes: Vec<u8> = hashvec
|
|
.chunks(8)
|
|
.map(|chunk| {
|
|
let chunk = chunk.to_owned();
|
|
chunk
|
|
.iter()
|
|
.rev()
|
|
.enumerate()
|
|
.fold(None, |accum, (n, val)| {
|
|
accum.or(Some(0)).map(|accum| accum | ((*val as u8) << n))
|
|
})
|
|
})
|
|
.take_while(|x| x.is_some())
|
|
.flatten()
|
|
.collect();
|
|
|
|
debug!("Image hash: {}", hex::encode(&hash_bytes));
|
|
|
|
Ok(hash_bytes)
|
|
}
|
|
|
|
/// Use these instead:
|
|
/// L = R * 299/1000 + G * 587/1000 + B * 114/1000
|
|
const SRGB_LUMA: [f32; 3] = [299.0 / 1000.0, 587.0 / 1000.0, 114.0 / 1000.0];
|
|
|
|
#[inline]
|
|
fn rgb_to_luma(rgb: &[u8]) -> u8 {
|
|
let l = SRGB_LUMA[0] * rgb[0].to_f32().unwrap()
|
|
+ SRGB_LUMA[1] * rgb[1].to_f32().unwrap()
|
|
+ SRGB_LUMA[2] * rgb[2].to_f32().unwrap();
|
|
NumCast::from(l.round()).unwrap()
|
|
}
|
|
|
|
/// Convert the supplied image to grayscale. Alpha channel is discarded.
|
|
///
|
|
/// This is a customized implemententation of the grayscale feature from the `image` crate.
|
|
/// This allows us to:
|
|
/// - use RGB->LUMA constants that match those used by the Python Pillow package.
|
|
/// - round the luma floating point value to the nearest integer rather than truncating.
|
|
///
|
|
/// See also: https://github.com/image-rs/image/issues/1554
|
|
fn grayscale(image: &ImageBuffer<Rgb<u8>, Vec<u8>>) -> ImageBuffer<Luma<u8>, Vec<u8>> {
|
|
let (width, height) = image.dimensions();
|
|
let mut out = ImageBuffer::new(width, height);
|
|
|
|
for y in 0..height {
|
|
for x in 0..width {
|
|
let pixel = image.get_pixel(x, y);
|
|
|
|
let mut pix = Luma([Zero::zero()]);
|
|
let gray = pix.channels_mut();
|
|
let rgb = pixel.channels();
|
|
gray[0] = rgb_to_luma(rgb);
|
|
|
|
let pixel = Luma::from_slice(gray); //.into_color(); // no-op for luma->luma
|
|
|
|
out.put_pixel(x, y, *pixel);
|
|
}
|
|
}
|
|
|
|
out
|
|
}
|