implemented suggestions (now both rich_signature.version and rich_sig…

…nature.toolid have overloads)
1sn0m4d · Nov 27, 2015 · 23485b4 · 23485b4
1 parent 7d224dd
commit 23485b4
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 40 deletions.
diff --git a/docs/modules/pe.rst b/docs/modules/pe.rst
@@ -425,23 +425,25 @@ Reference
 
         Data after being decrypted by XORing it with the key.
 
-    .. c:function:: version(version)
+    .. c:function:: version(version, [toolid])
 
     .. 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:
+        signature. Provide the optional *toolid* argument to only match when both match
+        for one entry. More information can be found here:
 
         http://www.ntcore.com/files/richsign.htm
 
         *Example: pe.rich_signature.version(21005)*
 
-    .. c:function:: toolid(id)
+    .. c:function:: toolid(toolid, [version])
 
     .. 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:
+        signature. Provide the optional *version* argument to only match when both
+        match for one entry. More information can be found here:
 
         http://www.ntcore.com/files/richsign.htm
 

diff --git a/libyara/modules/pe.c b/libyara/modules/pe.c
@@ -1814,25 +1814,26 @@ define_function(language)
   return_integer(0);
 }
 
-define_function(rich_version)
+static uint64_t rich_internal(YR_OBJECT* module, uint64_t version, uint64_t toolid)
 {
-    YR_OBJECT* module = module();
-    uint64_t version = 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);
+        return UNDEFINED;
 
     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);
+        return UNDEFINED;
+
+    if (version == UNDEFINED && toolid == UNDEFINED)
+        return 0;
 
     clear_rich_signature = (PRICH_SIGNATURE)rich_string->c_string;
 
@@ -1841,45 +1842,47 @@ 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))
-            return_integer(1);
+        DWORD id_version = clear_rich_signature->versions[i].id_version;
+        bool matchVersion = version == RICH_VERSION_VERSION(id_version);
+        bool matchToolid = toolid == RICH_VERSION_ID(id_version);
+        if (version != UNDEFINED && toolid != UNDEFINED) //check version and toolid
+        {
+            if (matchVersion && matchToolid)
+                return 1;
+        }
+        else if (version != UNDEFINED) //check only version
+        {
+            if (matchVersion)
+                return 1;
+        }
+        else if (toolid != UNDEFINED) //check only toolid
+        {
+            if (matchToolid)
+                return 1;
+        }
     }
 
-    return_integer(0);
+    return 0;
 }
 
-define_function(rich_toolid)
+define_function(rich_version)
 {
-    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");
-    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);
+    return_integer(rich_internal(module(), integer_argument(1), UNDEFINED));
+}
 
-    clear_rich_signature = (PRICH_SIGNATURE)rich_string->c_string;
+define_function(rich_version_toolid)
+{
+    return_integer(rich_internal(module(), integer_argument(1), integer_argument(2)));
+}
 
-    // Loop over the versions in the rich signature
-    for (i = 0;
-        i < (rich_len - sizeof(RICH_SIGNATURE)) / sizeof(RICH_VERSION_INFO);
-        i++)
-    {
-        if (toolid == RICH_VERSION_ID(clear_rich_signature->versions[i].id_version))
-            return_integer(1);
-    }
+define_function(rich_toolid)
+{
+    return_integer(rich_internal(module(), UNDEFINED, integer_argument(1)));
+}
 
-    return_integer(0);
+define_function(rich_toolid_version)
+{
+    return_integer(rich_internal(module(), integer_argument(2), integer_argument(1)));
 }
 
 begin_declarations;
@@ -2015,7 +2018,9 @@ begin_declarations;
     declare_string("raw_data");
     declare_string("clear_data");
     declare_function("version", "i", "i", rich_version);
+    declare_function("version", "ii", "i", rich_version_toolid);
     declare_function("toolid", "i", "i", rich_toolid);
+    declare_function("toolid", "ii", "i", rich_toolid_version);
   end_struct("rich_signature");
 
   #if defined(HAVE_LIBCRYPTO)