Updated constructor to accept the raw flag and pass it to the underlying library. Helpful for cases where unicode chars exist in match results.

rmspeers · rmspeers · commit 4514350cc363 · 2018-09-20T23:04:22.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -4,6 +4,7 @@ bin/
 deb_dist
 htmlcov/
 lib/
+__pycache__/
 python_magic.egg-info
 pip-selfcheck.json
 pyvenv.cfg
diff --git a/magic.py b/magic.py
@@ -14,7 +14,6 @@
 'PDF document, version 1.2'
 >>>
 
-
 """
 
 import sys
@@ -35,11 +34,10 @@ def __init__(self, message):
 class Magic:
     """
     Magic is a wrapper around the libmagic C library.
-
     """
 
     def __init__(self, mime=False, magic_file=None, mime_encoding=False,
-                 keep_going=False, uncompress=False):
+                 keep_going=False, uncompress=False, raw=False):
         """
         Create a new libmagic wrapper.
 
@@ -48,6 +46,7 @@ def __init__(self, mime=False, magic_file=None, mime_encoding=False,
         magic_file - use a mime database other than the system default
         keep_going - don't stop at the first match, keep going
         uncompress - Try to look inside compressed files.
+        raw - Do not try to decode "non-printable" chars.
         """
         self.flags = MAGIC_NONE
         if mime:
@@ -56,9 +55,10 @@ def __init__(self, mime=False, magic_file=None, mime_encoding=False,
             self.flags |= MAGIC_MIME_ENCODING
         if keep_going:
             self.flags |= MAGIC_CONTINUE
-
         if uncompress:
             self.flags |= MAGIC_COMPRESS
+        if raw:
+            self.flags |= MAGIC_RAW
 
         self.cookie = magic_open(self.flags)
         self.lock = threading.Lock()
@@ -190,13 +190,15 @@ def from_buffer(buffer, mime=False):
 
 magic_t = ctypes.c_void_p
 
+
 def errorcheck_null(result, func, args):
     if result is None:
         err = magic_error(args[0])
         raise MagicException(err)
     else:
         return result
 
+
 def errorcheck_negative_one(result, func, args):
     if result is -1:
         err = magic_error(args[0])
@@ -213,6 +215,7 @@ def maybe_decode(s):
     else:
         return s.decode('utf-8')
 
+
 def coerce_filename(filename):
     if filename is None:
         return None
@@ -230,6 +233,7 @@ def coerce_filename(filename):
     else:
         return filename
 
+
 magic_open = libmagic.magic_open
 magic_open.restype = magic_t
 magic_open.argtypes = [c_int]
@@ -251,6 +255,7 @@ def coerce_filename(filename):
 _magic_file.argtypes = [magic_t, c_char_p]
 _magic_file.errcheck = errorcheck_null
 
+
 def magic_file(cookie, filename):
     return _magic_file(cookie, coerce_filename(filename))
 
@@ -259,6 +264,7 @@ def magic_file(cookie, filename):
 _magic_buffer.argtypes = [magic_t, c_void_p, c_size_t]
 _magic_buffer.errcheck = errorcheck_null
 
+
 def magic_buffer(cookie, buf):
     return _magic_buffer(cookie, buf, len(buf))
 
@@ -268,6 +274,7 @@ def magic_buffer(cookie, buf):
 _magic_load.argtypes = [magic_t, c_char_p]
 _magic_load.errcheck = errorcheck_negative_one
 
+
 def magic_load(cookie, filename):
     return _magic_load(cookie, coerce_filename(filename))
 
@@ -288,6 +295,7 @@ def magic_load(cookie, filename):
 _magic_setparam.argtypes = [magic_t, c_int, POINTER(c_size_t)]
 _magic_setparam.errcheck = errorcheck_negative_one
 
+
 def magic_setparam(cookie, param, val):
     v = c_size_t(val)
     return _magic_setparam(cookie, param, byref(v))
@@ -297,11 +305,13 @@ def magic_setparam(cookie, param, val):
 _magic_getparam.argtypes = [magic_t, c_int, POINTER(c_size_t)]
 _magic_getparam.errcheck = errorcheck_negative_one
 
+
 def magic_getparam(cookie, param):
     val = c_size_t()
     _magic_getparam(cookie, param, byref(val))
     return val.value
 
+
 MAGIC_NONE = 0x000000 # No flags
 MAGIC_DEBUG = 0x000001 # Turn on debugging
 MAGIC_SYMLINK = 0x000002 # Follow symlinks
diff --git a/test/test.py b/test/test.py
@@ -84,6 +84,25 @@ def test_descriptions(self):
         finally:
             del os.environ['TZ']
 
+    def test_unicode_result_nonraw(self):
+        m = magic.Magic(raw=False)
+        src = os.path.join(MagicTest.TESTDATA_DIR, 'pgpunicode')
+        result = m.from_file(src)
+        # NOTE: This check is added as otherwise some magic files don't identify the test case as a PGP key.
+        if 'PGP' in result:
+            assert r"PGP\011Secret Sub-key -" == result
+        else:
+            raise unittest.SkipTest("Magic file doesn't return expected type.")
+
+    def test_unicode_result_raw(self):
+        m = magic.Magic(raw=True)
+        src = os.path.join(MagicTest.TESTDATA_DIR, 'pgpunicode')
+        result = m.from_file(src)
+        if 'PGP' in result:
+            assert b'PGP\tSecret Sub-key -' == result.encode('utf-8')
+        else:
+            raise unittest.SkipTest("Magic file doesn't return expected type.")
+
     def test_mime_encodings(self):
         m = magic.Magic(mime_encoding=True)
         self.assert_values(m, {
diff --git a/test/testdata/pgpunicode b/test/testdata/pgpunicode
@@ -0,0 +1 @@
+&#65533;q&#656;
