From 1c10fd5e50eab118fbfef87f5d4db7006089ddb9 Mon Sep 17 00:00:00 2001
From: chrisr3d <chris.studer.68@gmail.com>
Date: Wed, 31 Oct 2018 10:21:21 +0100
Subject: [PATCH] fix: Making yara query an expansion module for single
 attributes atm

---
 misp_modules/modules/expansion/yara_query.py | 32 +++++++++-----------
 1 file changed, 14 insertions(+), 18 deletions(-)

diff --git a/misp_modules/modules/expansion/yara_query.py b/misp_modules/modules/expansion/yara_query.py
index 89d2e65..a071bb4 100644
--- a/misp_modules/modules/expansion/yara_query.py
+++ b/misp_modules/modules/expansion/yara_query.py
@@ -1,5 +1,5 @@
 import json
-import sys
+import re
 
 misperrors = {'error': 'Error'}
 moduleinfo = {'version': '1', 'author': 'Christian STUDER',
@@ -7,10 +7,10 @@ moduleinfo = {'version': '1', 'author': 'Christian STUDER',
               'module-type': ['expansion', 'hover'],
               'require_standard_format': True}
 moduleconfig = []
-mispattributes = {'input': ['md5', 'sha1', 'sha256', 'filename|md5', 'filename|sha1', 'filename|sha256'], 'output': ['yara rule']}
+mispattributes = {'input': ['md5', 'sha1', 'sha256', 'filename|md5', 'filename|sha1', 'filename|sha256'], 'output': ['yara']}
 
-def hash_cond(hashtype, hashvalue):
-    condition = 'hash.{}(0, filesize) == {}'.format(hashtype, hashvalue.lower())
+def get_hash_condition(hashtype, hashvalue):
+    condition = 'hash.{}(0, filesize) == "{}"'.format(hashtype, hashvalue.lower())
     return condition, 'hash'
 
 def handler(q=False):
@@ -21,20 +21,16 @@ def handler(q=False):
     if 'event_id' in request:
         del request['event_id']
     uuid = request.pop('attribute_uuid') if 'attribute_uuid' in request else None
-    rules = []
-    types = []
-    for attribute_type, value in request.items():
-        if 'filename' in attribute_type:
-            _, attribute_type = attribute_type.split('|')
-            _, value = value.split('|')
-        condition, required_module = hash_cond(attribute_type, value)
-        condition = '\r\n\t\t'.join([condition])
-        import_section = '\r\n'.join(['import "{}"'.format(required_module)])
-        rule_start = 'rule %s {' % uuid if uuid else 'rule {'
-        condition = '\tcondition:\r\n\t\t{}'.format(condition)
-        rules.append('\r\n'.join([rule_start, condition, '}']))
-        types.append('yara')
-    return {'results': [{'types': [t], 'values': [v]} for t, v in zip(types, rules)]}
+    attribute_type, value = list(request.items())[0]
+    if 'filename' in attribute_type:
+        _, attribute_type = attribute_type.split('|')
+        _, value = value.split('|')
+    condition, required_module = get_hash_condition(attribute_type, value)
+    import_section = 'import "{}"'.format(required_module)
+    rule_start = 'import "hash" \r\nrule %s_%s {' % (attribute_type.upper(), re.sub(r'\W+', '_', uuid)) if uuid else 'import "hash"\r\nrule %s {' % attribute_type.upper()
+    condition = '\tcondition:\r\n\t\t{}'.format(condition)
+    rule = '\r\n'.join([rule_start, condition, '}'])
+    return {'results': [{'types': mispattributes['output'], 'values': [rule]}]}
 
 def introspection():
     return mispattributes