feat: implement is_64_bits, is_32_bits and is_dll functions in …

…PE module
VirusTotal · Nov 21, 2023 · f33dde9 · f33dde9
1 parent 76f64aa
commit f33dde9
Show file tree

Hide file tree

Showing 28 changed files with 200 additions and 56 deletions.
diff --git a/yara-x/src/modules/pe/mod.rs b/yara-x/src/modules/pe/mod.rs
@@ -34,6 +34,27 @@ fn main(input: &[u8]) -> PE {
     }
 }
 
+/// Returns true if the file is a 32-bit PE.
+#[module_export]
+fn is_32_bits(ctx: &ScanContext) -> Option<bool> {
+    let magic = ctx.module_output::<PE>()?.opthdr_magic?;
+    Some(magic.value() == OptHdrMagic::IMAGE_NT_OPTIONAL_HDR32_MAGIC as i32)
+}
+
+/// Returns true if the file is a 64-bit PE.
+#[module_export]
+fn is_64_bits(ctx: &ScanContext) -> Option<bool> {
+    let magic = ctx.module_output::<PE>()?.opthdr_magic?;
+    Some(magic.value() == OptHdrMagic::IMAGE_NT_OPTIONAL_HDR64_MAGIC as i32)
+}
+
+/// Returns true if the file is dynamic link library (DLL)
+#[module_export]
+fn is_dll(ctx: &ScanContext) -> Option<bool> {
+    let characteristics = ctx.module_output::<PE>()?.characteristics?;
+    Some(characteristics & Characteristics::FILE_DLL as u32 != 0)
+}
+
 /// Returns the PE checksum, as calculated by YARA.
 ///
 /// This is useful for comparing with the checksum appearing in the PE header

diff --git a/yara-x/src/modules/pe/parser.rs b/yara-x/src/modules/pe/parser.rs
@@ -1924,7 +1924,14 @@ impl From<PE<'_>> for pe::PE {
         result.set_pointer_to_symbol_table(pe.pe_hdr.symbol_table_offset);
         result.set_number_of_symbols(pe.pe_hdr.number_of_symbols);
         result.set_size_of_optional_header(pe.pe_hdr.size_of_optional_header.into());
-
+
+        result.opthdr_magic = pe
+            .optional_hdr
+            .magic
+            .try_into()
+            .ok()
+            .map(EnumOrUnknown::<pe::OptHdrMagic>::from_i32);
+
         result.subsystem = pe
             .optional_hdr
             .subsystem
@@ -1957,7 +1964,6 @@ impl From<PE<'_>> for pe::PE {
 
         // TODO
         // number_of_version_infos
-        // opthdr_magic
 
         result.linker_version = MessageField::some(pe::Version {
             major: Some(pe.optional_hdr.major_linker_version.into()),

diff --git a/yara-x/src/modules/pe/tests/mod.rs b/yara-x/src/modules/pe/tests/mod.rs
@@ -338,3 +338,57 @@ fn locale_and_language() {
         &pe
     );
 }
+
+#[test]
+fn is_32bits() {
+    let pe = create_binary_from_zipped_ihex(
+        "src/modules/pe/tests/testdata/0ba6042247d90a187919dd88dc2d55cd882c80e5afc511c4f7b2e0e193968f7f.in.zip",
+    );
+
+    rule_true!(
+        r#"
+        import "pe"
+        rule test {
+          condition:
+            pe.is_32_bits()
+        }
+        "#,
+        &pe
+    );
+}
+
+#[test]
+fn is_64bits() {
+    let pe = create_binary_from_zipped_ihex(
+        "src/modules/pe/tests/testdata/2e9c671b8a0411f2b397544b368c44d7f095eb395779de0ad1ac946914dfa34c.in.zip",
+    );
+
+    rule_true!(
+        r#"
+        import "pe"
+        rule test {
+          condition:
+            pe.is_64_bits()
+        }
+        "#,
+        &pe
+    );
+}
+
+#[test]
+fn is_dll() {
+    let pe = create_binary_from_zipped_ihex(
+        "src/modules/pe/tests/testdata/079a472d22290a94ebb212aa8015cdc8dd28a968c6b4d3b88acdd58ce2d3b885.in.zip",
+    );
+
+    rule_true!(
+        r#"
+        import "pe"
+        rule test {
+          condition:
+            pe.is_dll()
+        }
+        "#,
+        &pe
+    );
+}
diff --git a/...es/pe/tests/testdata/04ac6dd0c1cc33a49962ee0f3222597104f54a75683a5adee235401778279818.out b/...es/pe/tests/testdata/04ac6dd0c1cc33a49962ee0f3222597104f54a75683a5adee235401778279818.out
@@ -17,6 +17,7 @@ linker_version {
   major: 14
   minor: 29
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 258
 dll_characteristics: 33088
 timestamp: 1626863112

diff --git a/...es/pe/tests/testdata/079a472d22290a94ebb212aa8015cdc8dd28a968c6b4d3b88acdd58ce2d3b885.out b/...es/pe/tests/testdata/079a472d22290a94ebb212aa8015cdc8dd28a968c6b4d3b88acdd58ce2d3b885.out
@@ -17,6 +17,7 @@ linker_version {
   major: 10
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 8450
 dll_characteristics: 320
 timestamp: 1528213185

diff --git a/...es/pe/tests/testdata/09e7d832320e51bcc80b9aecde2a4135267a9b0156642a9596a62e85c9998cc9.out b/...es/pe/tests/testdata/09e7d832320e51bcc80b9aecde2a4135267a9b0156642a9596a62e85c9998cc9.out
@@ -17,6 +17,7 @@ linker_version {
   major: 48
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 8226
 dll_characteristics: 34112
 timestamp: 3665045795

diff --git a/...es/pe/tests/testdata/0ba6042247d90a187919dd88dc2d55cd882c80e5afc511c4f7b2e0e193968f7f.out b/...es/pe/tests/testdata/0ba6042247d90a187919dd88dc2d55cd882c80e5afc511c4f7b2e0e193968f7f.out
@@ -17,6 +17,7 @@ linker_version {
   major: 7
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 271
 dll_characteristics: 32768
 timestamp: 998081829

diff --git a/...es/pe/tests/testdata/23e72ce7e9cdbc80c0095484ebeb02f56b21e48fd67044e69e7a2ae76db631e5.out b/...es/pe/tests/testdata/23e72ce7e9cdbc80c0095484ebeb02f56b21e48fd67044e69e7a2ae76db631e5.out
@@ -17,6 +17,7 @@ linker_version {
   major: 14
   minor: 13
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 8226
 dll_characteristics: 16736
 timestamp: 1827812126

diff --git a/...es/pe/tests/testdata/2775d97f8bdb3311ace960a42eee35dbec84b9d71a6abbacb26c14e83f5897e4.out b/...es/pe/tests/testdata/2775d97f8bdb3311ace960a42eee35dbec84b9d71a6abbacb26c14e83f5897e4.out
@@ -17,6 +17,7 @@ linker_version {
   major: 6
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 271
 dll_characteristics: 0
 timestamp: 1524722207

diff --git a/...es/pe/tests/testdata/29eeeecf2c458ea3da1ce9d6d54742c0fad490cb2165f371f53b61941eedf072.out b/...es/pe/tests/testdata/29eeeecf2c458ea3da1ce9d6d54742c0fad490cb2165f371f53b61941eedf072.out
@@ -17,6 +17,7 @@ linker_version {
   major: 8
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 258
 dll_characteristics: 33088
 timestamp: 1621233906

diff --git a/...es/pe/tests/testdata/2d80c403b5c50f8bbacb65f58e7a19f272c62d1889216b7a6f1141571ec12649.out b/...es/pe/tests/testdata/2d80c403b5c50f8bbacb65f58e7a19f272c62d1889216b7a6f1141571ec12649.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 56
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 8974
 dll_characteristics: 0
 timestamp: 1274346651

diff --git a/...es/pe/tests/testdata/2e9c671b8a0411f2b397544b368c44d7f095eb395779de0ad1ac946914dfa34c.out b/...es/pe/tests/testdata/2e9c671b8a0411f2b397544b368c44d7f095eb395779de0ad1ac946914dfa34c.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 51
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 47
 dll_characteristics: 0
 timestamp: 0

diff --git a/...es/pe/tests/testdata/55cfd3bcea1aa352b4687c4d45bdcfea184ce9b58320891a2d72d4ec93766f14.out b/...es/pe/tests/testdata/55cfd3bcea1aa352b4687c4d45bdcfea184ce9b58320891a2d72d4ec93766f14.out
@@ -17,6 +17,7 @@ linker_version {
   major: 8
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 259
 dll_characteristics: 1024
 timestamp: 1162198621

diff --git a/...es/pe/tests/testdata/6c2abf4b80a87e63eee2996e5cea8f004d49ec0c1806080fa72e960529cba14c.out b/...es/pe/tests/testdata/6c2abf4b80a87e63eee2996e5cea8f004d49ec0c1806080fa72e960529cba14c.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 56
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 782
 dll_characteristics: 0
 timestamp: 0

diff --git a/...es/pe/tests/testdata/99df28014fae5f213c8decfde423b0eb69005158f981bc84560e9a5dde103d90.out b/...es/pe/tests/testdata/99df28014fae5f213c8decfde423b0eb69005158f981bc84560e9a5dde103d90.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 26
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 783
 dll_characteristics: 0
 timestamp: 1459377848

diff --git a/...es/pe/tests/testdata/9bcf79a99ffbb1bd649503ce1406dea4181f6477cfb61f2130f37d014ce21888.out b/...es/pe/tests/testdata/9bcf79a99ffbb1bd649503ce1406dea4181f6477cfb61f2130f37d014ce21888.out
@@ -17,6 +17,7 @@ linker_version {
   major: 48
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 8226
 dll_characteristics: 34112
 timestamp: 3435013737

diff --git a/...es/pe/tests/testdata/af3f20a9272489cbef4281c8c86ad42ccfb04ccedd3ada1e8c26939c726a4c8e.out b/...es/pe/tests/testdata/af3f20a9272489cbef4281c8c86ad42ccfb04ccedd3ada1e8c26939c726a4c8e.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 20
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 518
 dll_characteristics: 0
 timestamp: 0

diff --git a/...es/pe/tests/testdata/b8543d3aceec5a754292393f6602eeb966dbf9e198c94e3d74a9e9260e5f9870.out b/...es/pe/tests/testdata/b8543d3aceec5a754292393f6602eeb966dbf9e198c94e3d74a9e9260e5f9870.out
@@ -17,6 +17,7 @@ linker_version {
   major: 14
   minor: 20
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 34
 dll_characteristics: 49504
 timestamp: 1776026023

diff --git a/...es/pe/tests/testdata/bd82090d9c6e23c1e2708550f4c21bbafd719dcc7cfc28f405ad3bc2783c6a12.out b/...es/pe/tests/testdata/bd82090d9c6e23c1e2708550f4c21bbafd719dcc7cfc28f405ad3bc2783c6a12.out
@@ -17,6 +17,7 @@ linker_version {
   major: 8
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 259
 dll_characteristics: 1024
 timestamp: 1162198621

diff --git a/...es/pe/tests/testdata/c703e95b9ec0ca955b1bcf33e1318a0029d9d4e7903cd63291946b334c314831.out b/...es/pe/tests/testdata/c703e95b9ec0ca955b1bcf33e1318a0029d9d4e7903cd63291946b334c314831.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 56
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 782
 dll_characteristics: 0
 timestamp: 0

diff --git a/...es/pe/tests/testdata/c704cca0fe4c9bdee18a302952540073b860e3b4d42e081f86d27bdb1cf6ede4.out b/...es/pe/tests/testdata/c704cca0fe4c9bdee18a302952540073b860e3b4d42e081f86d27bdb1cf6ede4.out
@@ -17,6 +17,7 @@ linker_version {
   major: 2
   minor: 26
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 783
 dll_characteristics: 0
 timestamp: 1459377848

diff --git a/...es/pe/tests/testdata/d009038e4e371f9cdc3a96923700c17d29de9dfc156666c98581dfeccc07548e.out b/...es/pe/tests/testdata/d009038e4e371f9cdc3a96923700c17d29de9dfc156666c98581dfeccc07548e.out
@@ -17,6 +17,7 @@ linker_version {
   major: 9
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 258
 dll_characteristics: 1024
 timestamp: 1301987779

diff --git a/...es/pe/tests/testdata/db6a9934570fa98a93a979e7e0e218e0c9710e5a787b18c6948f2eedd9338984.out b/...es/pe/tests/testdata/db6a9934570fa98a93a979e7e0e218e0c9710e5a787b18c6948f2eedd9338984.out
@@ -17,6 +17,7 @@ linker_version {
   major: 14
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR64_MAGIC
 characteristics: 34
 dll_characteristics: 33056
 timestamp: 1629390430

diff --git a/...es/pe/tests/testdata/e3d45a2865818756068757d7e319258fef40dad54532ee4355b86bc129f27345.out b/...es/pe/tests/testdata/e3d45a2865818756068757d7e319258fef40dad54532ee4355b86bc129f27345.out
@@ -17,6 +17,7 @@ linker_version {
   major: 9
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 8450
 dll_characteristics: 0
 timestamp: 1299001425

diff --git a/...es/pe/tests/testdata/e5b038a7a8578bbe020784877bf13f8d3920141a39b502e870348ce3b254eb80.out b/...es/pe/tests/testdata/e5b038a7a8578bbe020784877bf13f8d3920141a39b502e870348ce3b254eb80.out
@@ -17,6 +17,7 @@ linker_version {
   major: 10
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 8450
 dll_characteristics: 1344
 timestamp: 1314765018

diff --git a/...es/pe/tests/testdata/eda21be1b234c84acb8257f2f77c12e0826b06cd15eb6dec9bf3df6465216dd8.out b/...es/pe/tests/testdata/eda21be1b234c84acb8257f2f77c12e0826b06cd15eb6dec9bf3df6465216dd8.out
@@ -17,6 +17,7 @@ linker_version {
   major: 0
   minor: 0
 }
+opthdr_magic: IMAGE_NT_OPTIONAL_HDR32_MAGIC
 characteristics: 270
 dll_characteristics: 0
 timestamp: 0