From 5814c79a5c92db9917934e4b5ba99bbd41660ad4 Mon Sep 17 00:00:00 2001
From: MatejKastak <matej.kastak@gmail.com>
Date: Fri, 23 Sep 2022 18:41:57 +0200
Subject: [PATCH] feat: Add support for `entrypoint` and `filesize`

Enable this in complex expressions. This does not work in regular eval
because both are represented as `u64` and `YrValue::Integer` has `i64`.

This might result in losing some information if we just convert it from
u64 to i64. We will have to come back to this.
---
 yari-sys/src/lib.rs          |  5 +--
 yari-sys/tests/common.rs     | 21 +++++++++++++
 yari-sys/tests/tests_eval.rs | 61 ++++++++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/yari-sys/src/lib.rs b/yari-sys/src/lib.rs
index 3facd3d..03e310a 100644
--- a/yari-sys/src/lib.rs
+++ b/yari-sys/src/lib.rs
@@ -1216,10 +1216,7 @@ impl Context {
         self.iterator.context = &mut *self.block as *mut _ as *mut _;
         self.iterator.first = Some(_yr_get_first_block);
         self.iterator.next = Some(_yr_get_next_block);
-
-        // Upstream yara
-        // (*iterator).file_size = Some(_yr_get_file_size);
-        // (*iterator).last_error = ERROR_SUCCESS as i32;
+        self.iterator.file_size = Some(_yr_get_file_size);
     }
 
     unsafe fn compile_string(&mut self, rule_cstr: &CString) -> Result<(), YariError> {
diff --git a/yari-sys/tests/common.rs b/yari-sys/tests/common.rs
index a5792b8..2a5cb27 100644
--- a/yari-sys/tests/common.rs
+++ b/yari-sys/tests/common.rs
@@ -56,3 +56,24 @@ rule r {
         ),
     )
 }
+
+pub fn context_with_elf_sample_and_rule() -> Context {
+    context_with_sample(
+        "tests/assets/elf_hello_world",
+        Some(
+            "import \"elf\"
+private rule PRIVATE {
+    condition:
+        elf.number_of_sections == 4
+}
+
+rule r {
+    strings:
+        $s00 = \"Hello\"
+        $s01 = \"this is a pretty unique string that should not be found in the provided sample\"
+    condition:
+        all of them and PRIVATE
+}",
+        ),
+    )
+}
diff --git a/yari-sys/tests/tests_eval.rs b/yari-sys/tests/tests_eval.rs
index a70d487..908259d 100644
--- a/yari-sys/tests/tests_eval.rs
+++ b/yari-sys/tests/tests_eval.rs
@@ -369,3 +369,64 @@ fn test_invalid_data() {
     let val = ContextBuilder::parse_module_data_str("cuckoo=too_many=equals");
     assert_eq!(val, None);
 }
+
+#[test]
+fn test_eval_filesize() {
+    let mut context = common::context_with_pe_sample_and_rule();
+    assert_eq!(
+        context.eval("r|filesize == 8704").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|filesize == 0").unwrap(),
+        YrValue::Integer(0)
+    );
+    assert_eq!(
+        context.eval("r|filesize == 0x2200").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|filesize != 0").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|filesize <= 0x10000").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(context.eval("r|filesize > 0").unwrap(), YrValue::Integer(1));
+}
+
+#[test]
+fn test_eval_entrypoint() {
+    let mut context = common::context_with_pe_sample_and_rule();
+    assert_eq!(
+        context.eval("r|entrypoint == 2166").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|entrypoint == 0").unwrap(),
+        YrValue::Integer(0)
+    );
+    assert_eq!(
+        context.eval("r|entrypoint == 0x876").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|entrypoint != 0").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|entrypoint <= 0x10000").unwrap(),
+        YrValue::Integer(1)
+    );
+    assert_eq!(
+        context.eval("r|entrypoint > 0").unwrap(),
+        YrValue::Integer(1)
+    );
+
+    let mut context = common::context_with_elf_sample_and_rule();
+    assert_eq!(
+        context.eval("r|entrypoint == 4160").unwrap(),
+        YrValue::Integer(1)
+    );
+}