feat: implement checksum32 function in hash module

VirusTotal · Nov 22, 2023 · 6bb17b2 · 6bb17b2
1 parent 3917089
commit 6bb17b2
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 3 deletions.
diff --git a/yara-x/src/modules/hash/mod.rs b/yara-x/src/modules/hash/mod.rs
@@ -23,6 +23,9 @@ thread_local!(
 
     static CRC32_CACHE: RefCell<FxHashMap<(i64, i64), i64>> =
         RefCell::new(FxHashMap::default());
+
+    static CHECKSUM32_CACHE: RefCell<FxHashMap<(i64, i64), i64>> =
+        RefCell::new(FxHashMap::default());
 );
 
 #[module_main]
@@ -32,6 +35,7 @@ fn main(_data: &[u8]) -> Hash {
     SHA1_CACHE.with(|cache| cache.borrow_mut().clear());
     MD5_CACHE.with(|cache| cache.borrow_mut().clear());
     CRC32_CACHE.with(|cache| cache.borrow_mut().clear());
+    CHECKSUM32_CACHE.with(|cache| cache.borrow_mut().clear());
 
     Hash::new()
 }
@@ -184,3 +188,37 @@ fn crc_str(ctx: &ScanContext, s: RuntimeString) -> Option<i64> {
     let crc = crc32fast::hash(s.as_bstr(ctx));
     Some(crc.into())
 }
+
+#[module_export(name = "checksum32")]
+fn checksum_data(ctx: &ScanContext, offset: i64, size: i64) -> Option<i64> {
+    let cached = CHECKSUM32_CACHE.with(|cache| -> Option<i64> {
+        Some(*cache.borrow().get(&(offset, size))?)
+    });
+
+    if cached.is_some() {
+        return cached;
+    }
+
+    let range = offset.try_into().ok()?..(offset + size).try_into().ok()?;
+    let data = ctx.scanned_data().get(range)?;
+    let mut checksum = 0_u32;
+
+    for byte in data {
+        checksum = checksum.wrapping_add(*byte as u32)
+    }
+
+    CHECKSUM32_CACHE.with(|cache| {
+        cache.borrow_mut().insert((offset, size), checksum.into());
+    });
+
+    Some(checksum.into())
+}
+
+#[module_export(name = "checksum32")]
+fn checksum_str(ctx: &ScanContext, s: RuntimeString) -> Option<i64> {
+    let mut checksum = 0_u32;
+    for byte in s.as_bstr(ctx).as_bytes() {
+        checksum = checksum.wrapping_add(*byte as u32)
+    }
+    Some(checksum.into())
+}
diff --git a/yara-x/src/modules/hash/tests/mod.rs b/yara-x/src/modules/hash/tests/mod.rs
@@ -12,7 +12,7 @@ fn test_hash_module() {
             hash.md5(0, filesize) == "6df23dc03f9b54cc38a0fc1483df6e21" and
             hash.md5(3, 3) == "37b51d194a7513e45b56f6524f2d51f2" and
             hash.md5(0, filesize) == hash.md5("foobarbaz") and
-            hash.md5(3, 3) ==  hash.md5("bar")
+            hash.md5(3, 3) == hash.md5("bar")
         }
         "#,
         b"foobarbaz"
@@ -26,7 +26,7 @@ fn test_hash_module() {
             hash.sha1(0, filesize) == "5f5513f8822fdbe5145af33b64d8d970dcf95c6e" and
             hash.sha1(3, 3) == "62cdb7020ff920e5aa642c3d4066950dd1f01f4d" and
             hash.sha1(0, filesize) == hash.sha1("foobarbaz") and
-            hash.sha1(3, 3) ==  hash.sha1("bar")
+            hash.sha1(3, 3) == hash.sha1("bar")
         }
         "#,
         b"foobarbaz"
@@ -40,7 +40,7 @@ fn test_hash_module() {
             hash.sha256(0, filesize) == "97df3588b5a3f24babc3851b372f0ba71a9dcdded43b14b9d06961bfc1707d9d" and
             hash.sha256(3, 3) == "fcde2b2edba56bf408601fb721fe9b5c338d10ee429ea04fae5511b68fbf8fb9" and
             hash.sha256(0, filesize) == hash.sha256("foobarbaz") and
-            hash.sha256(3, 3) ==  hash.sha256("bar")
+            hash.sha256(3, 3) == hash.sha256("bar")
         }
         "#,
         b"foobarbaz"
@@ -59,4 +59,15 @@ fn test_hash_module() {
         "#,
         b"foobarbaz"
     );
+
+    rule_true!(
+        r#"
+        import "hash"
+        rule test {
+          condition:
+            hash.checksum32("TEST STRING") == 0x337
+        }
+        "#,
+        b"foobarbaz"
+    );
 }