moved rich_version to rich_signature.version and added a rich_signatu…

…re.toolid version for checking toolid's
1sn0m4d · Nov 26, 2015 · 7d224dd · 7d224dd
1 parent e6f3f9e
commit 7d224dd
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 16 deletions.
diff --git a/docs/modules/pe.rst b/docs/modules/pe.rst
@@ -425,6 +425,28 @@ Reference
 
         Data after being decrypted by XORing it with the key.
 
+    .. c:function:: version(version)
+
+    .. versionadded:: 3.5.0
+
+        Function returning true if the PE has the specified *version* in the PE's rich
+        signature. More information can be found here:
+
+        http://www.ntcore.com/files/richsign.htm
+
+        *Example: pe.rich_signature.version(21005)*
+
+    .. c:function:: toolid(id)
+
+    .. versionadded:: 3.5.0
+
+        Function returning true if the PE has the specified *id* in the PE's rich
+        signature. More information can be found here:
+
+        http://www.ntcore.com/files/richsign.htm
+
+        *Example: pe.rich_signature.toolid(222)*
+
 .. c:function:: exports(function_name)
 
     Function returning true if the PE exports *function_name* or
@@ -505,14 +527,3 @@ Reference
   *addr*. *addr* can be an offset into the file or a memory address.
 
   *Example: pe.section_index(pe.entry_point)*
-
-.. c:function:: rich_version(version)
-
- .. versionadded:: 3.5.0
-
-  Function returning true if the PE has the specified *version* in the PE's rich
-  signature. More information can be found here:
-
-  http://www.ntcore.com/files/richsign.htm
-
-  *Example: pe.rich_version(21005)*
diff --git a/libyara/modules/pe.c b/libyara/modules/pe.c
@@ -1817,7 +1817,6 @@ define_function(language)
 define_function(rich_version)
 {
     YR_OBJECT* module = module();
-    PE* pe = (PE*)module->data;
     uint64_t version = integer_argument(1);
     size_t rich_len;
     PRICH_SIGNATURE clear_rich_signature;
@@ -1828,8 +1827,38 @@ define_function(rich_version)
     if (is_undefined(module, "rich_signature.length"))
         return_integer(UNDEFINED);
 
-    // If not a PE file, return UNDEFINED
-    if (pe == NULL)
+    rich_len = get_integer(module, "rich_signature.length");
+    rich_string = get_string(module, "rich_signature.clear_data");
+
+    // If the clear_data was not set, return UNDEFINED
+    if (rich_string == NULL)
+        return_integer(UNDEFINED);
+
+    clear_rich_signature = (PRICH_SIGNATURE)rich_string->c_string;
+
+    // Loop over the versions in the rich signature
+    for (i = 0;
+        i < (rich_len - sizeof(RICH_SIGNATURE)) / sizeof(RICH_VERSION_INFO);
+        i++)
+    {
+        if (version == RICH_VERSION_VERSION(clear_rich_signature->versions[i].id_version))
+            return_integer(1);
+    }
+
+    return_integer(0);
+}
+
+define_function(rich_toolid)
+{
+    YR_OBJECT* module = module();
+    uint64_t toolid = integer_argument(1);
+    size_t rich_len;
+    PRICH_SIGNATURE clear_rich_signature;
+    SIZED_STRING* rich_string;
+    int i;
+
+    // Check if the required fields are set
+    if (is_undefined(module, "rich_signature.length"))
         return_integer(UNDEFINED);
 
     rich_len = get_integer(module, "rich_signature.length");
@@ -1846,7 +1875,7 @@ define_function(rich_version)
         i < (rich_len - sizeof(RICH_SIGNATURE)) / sizeof(RICH_VERSION_INFO);
         i++)
     {
-        if(version == RICH_VERSION_VERSION(clear_rich_signature->versions[i].id_version))
+        if (toolid == RICH_VERSION_ID(clear_rich_signature->versions[i].id_version))
             return_integer(1);
     }
 
@@ -1985,6 +2014,8 @@ begin_declarations;
     declare_integer("key");
     declare_string("raw_data");
     declare_string("clear_data");
+    declare_function("version", "i", "i", rich_version);
+    declare_function("toolid", "i", "i", rich_toolid);
   end_struct("rich_signature");
 
   #if defined(HAVE_LIBCRYPTO)
@@ -1999,7 +2030,6 @@ begin_declarations;
   declare_function("imports", "s", "i", imports_dll);
   declare_function("locale", "i", "i", locale);
   declare_function("language", "i", "i", language);
-  declare_function("rich_version", "i", "i", rich_version);
 
   declare_integer("resource_timestamp")
   begin_struct("resource_version");