Browse Source

Support both, filemagic and python-magic

tags/0.5.3
Johann Schmitz 2 years ago
parent
commit
092061c11e
2 changed files with 49 additions and 16 deletions
  1. 48
    15
      amavisvt/client.py
  2. 1
    1
      requirements.txt

+ 48
- 15
amavisvt/client.py View File

@@ -9,9 +9,9 @@ import shutil
9 9
 import tempfile
10 10
 from email.utils import parseaddr
11 11
 
12
-import magic
13 12
 import memcache
14 13
 import requests
14
+import sys
15 15
 
16 16
 from amavisvt import VERSION
17 17
 from amavisvt.db import Database
@@ -23,6 +23,7 @@ logger = logging.getLogger(__name__)
23 23
 
24 24
 BUFFER_SIZE = 4096
25 25
 
26
+
26 27
 def clean_silent(paths):
27 28
     for p in paths if isinstance(paths, list) else [paths]:
28 29
         try:
@@ -35,6 +36,39 @@ def clean_silent(paths):
35 36
             logger.exception("Could not remove %s", p)
36 37
 
37 38
 
39
+magic_identify_buffer = None
40
+
41
+
42
+def python_magic_id_buffer(buf):
43
+    import magic
44
+    t = magic.from_buffer(buf, mime=True)
45
+    if t and isinstance(t, bytes) and sys.version_info > (3,):
46
+        t = t.decode('UTF-8', 'ignore')
47
+    return t
48
+
49
+
50
+def filemagic_id_buffer(buf):
51
+    import magic
52
+    with magic.Magic(flags=magic.MAGIC_MIME_TYPE) as mgc:
53
+        return mgc.id_buffer(buf)
54
+
55
+
56
+try:
57
+    import magic
58
+
59
+    if hasattr(magic, 'from_buffer'):
60
+        # python-magic
61
+        identify_buffer = python_magic_id_buffer
62
+    elif hasattr(magic, 'MAGIC_MIME_TYPE'):
63
+        # filemagic
64
+        identify_buffer = filemagic_id_buffer
65
+    else:
66
+        raise Exception("No suitable libmagic library found. Please install either filemagic or python-magic")
67
+
68
+except ImportError:
69
+    raise
70
+
71
+
38 72
 class VTResponse(object):
39 73
     def __init__(self, virustotal_response):
40 74
         self._data = virustotal_response
@@ -229,21 +263,20 @@ class Resource(object):
229 263
         self._sha1 = sha1hasher.hexdigest()
230 264
         self._sha256 = sha256hasher.hexdigest()
231 265
 
232
-        with magic.Magic(flags=magic.MAGIC_MIME_TYPE) as m:
233
-            self._mime_type = m.id_buffer(id_buffer)
234
-            logger.debug("libmagic identified %s as: %s", self, self._mime_type)
266
+        self._mime_type = identify_buffer(id_buffer)
267
+        logger.debug("libmagic identified %s as: %s", self, self._mime_type)
235 268
 
236
-            # This is a hacky way to detect mail messages. Sometimes, when the amount of other text data exceeds the
237
-            # amount of "mail-like" data in a file (like a mail message with lots of HTML), libmagic fails to detect
238
-            # the file as a mail message.
239
-            if self._mime_type in ('text/plain', 'text/html'):
240
-                try:
241
-                    msg = email.message_from_string(id_buffer.decode('utf-8'))
242
-                    if len(msg.keys()) and 'From' in msg and 'To' in msg:
243
-                        logger.debug("Identified mail in %s when libmagic could not (said it was %s)", self.filename, self.mime_type)
244
-                        self._mime_type = MAIL_MIME_TYPE
245
-                except:
246
-                    pass
269
+        # This is a hacky way to detect mail messages. Sometimes, when the amount of other text data exceeds the
270
+        # amount of "mail-like" data in a file (like a mail message with lots of HTML), libmagic fails to detect
271
+        # the file as a mail message.
272
+        if self._mime_type in ('text/plain', 'text/html'):
273
+            try:
274
+                msg = email.message_from_string(id_buffer.decode('utf-8'))
275
+                if len(msg.keys()) and 'From' in msg and 'To' in msg:
276
+                    logger.debug("Identified mail in %s when libmagic could not (said it was %s)", self.filename, self.mime_type)
277
+                    self._mime_type = MAIL_MIME_TYPE
278
+            except:
279
+                pass
247 280
 
248 281
     def unpack(self):
249 282
         unpack_func = None

+ 1
- 1
requirements.txt View File

@@ -1,6 +1,6 @@
1 1
 # Project dependencies
2 2
 requests>=2.9.1
3
-filemagic
3
+python-magic
4 4
 python-memcached
5 5
 python-levenshtein
6 6
 setproctitle

Loading…
Cancel
Save